+ # Generate the lookahead_reward and state
+ ar_mask = (t % it_len == index_lookahead_reward).long() * (
+ t <= u + index_lookahead_reward
+ ).long()
+ ar(result, ar_mask)
+
+ # Generate the lookahead_reward and state
+ ar_mask = (t >= u + index_states).long() * (
+ t < u + index_states + state_len
+ ).long()