[Question] Clarification on Rollout Data Format (Obs, Actions, Rewards, Resets)

[EdanToledo]

Hello, thanks for the code release.

I'm looking for some clarification on the expected data format for rollouts. After looking through the code, the interactions between the different environments (jittable vs non-jittable), the agent's actor step, and the unroll functions made it a bit unclear.

I have two specific questions:

1. Reward Alignment

Given parallel sequences of observations and actions, what is the correct temporal alignment for rewards?

• Observations: [obs_0, obs_1, obs_2, ...]
• Actions: [a_0, a_1, a_2, ...]

Is the expected reward sequence coupled in the rollout object:

• Option A: Rewards: [r_0, r_1, r_2, ...]
  • (Where r_0 is a dummy reward from the initial reset, and r_t is the reward received at the same time as obs_t i.e. obs_t-1 and a_t-1 produced r_t and obs_t so these are together in the timestep object)
• OR
• Option B: Rewards: [r_1, r_2, r_3, ...]
  • (Where r_t is the reward received after taking action a_{t-1} from obs_{t-1})

2. Auto-Reset Handling

Additionally, are auto-resets handled with a dummy action value = a_T* and consequently a dummy reward value = r_0*, or is it handled by ignoring terminal observations?

For example, is the data expected to be aligned like this across a boundary (assuming option B from the first question):

• Observations: [..., obs_T-2, obs_T-1, obs_0, obs_1, ...]
• Actions: [..., a_T-2, a_T-1, a_0, a_1, ...]
• Rewards: [..., r_T-1, r_T, r_1, r_2, ...]

OR

• Observations: [..., obs_T-2, obs_T-1, obs_T, obs_0, obs_1, ...]
• Actions: [..., a_T-2, a_T-1, a_T*, a_0, a_1, ...]
• Rewards: [..., r_T-1, r_T, r_0*, r_1, r_2, ...]

(Where obs_T is the terminal observation, r_T is its corresponding final reward, and obs_0 is the new observation from the automatic reset.)

Any clarification you could provide on this would be very helpful.

Thanks!

P.S:

Given this code in the learner step, i assume this is manually aligning the rewards to be option B?

reward = rollout.rewards[1:]
    agent_out, _ = self.unroll_net(
        learner_state.params, agent_net_state, rollout
    )

    # Compute the loss using the discovered update(s).
    # Construct inputs for the discovered update.
    eta_inputs = types.UpdateRuleInputs(
        observations=rollout.observations,  # [T, ...]
        actions=rollout.actions,  # [T, ...]
        rewards=reward,  # [T-1]
        is_terminal=rollout.discounts[1:] == 0,  # [T-1]
        behaviour_agent_out=rollout.agent_outs,  # [T, ...]
        agent_out=agent_out,  # [T, ...]
        value_out=None,
    )

If that is the case, then my only question is on the auto-reset logic.