- lookahead_rewards = self.code2lookahead_reward(seq[:, :, 0])
- states = self.code2state(seq[:, :, 1 : self.height * self.width + 1])
+ lookahead_rewards = self.code2lookahead_reward(
+ seq[:, :, self.index_lookahead_reward]
+ )
+ states = self.code2state(
+ seq[:, :, self.index_states : self.height * self.width + self.index_states]
+ )