Keep the observation an episode ended on across the auto reset.

[shubhxho]

Fixes #305.

Problem

BraxAutoResetWrapper.step replaces state.obs with the reset observation on done and keeps no copy of what the step actually produced. wrap_for_brax_training already exposes the termination/truncation distinction through info['truncation'] (set by brax's EpisodeWrapper), but computing the truncation bootstrap needs V(s_{t+1}) as well, and s_{t+1} is gone by the time the learner sees the state.

Every actor-critic built on Playground therefore bootstraps V(s_reset) at each truncation. On locomotion tasks s_reset is an upright starting pose rather than whatever the agent was doing when the time limit hit, and the bias grows with episode length and gamma.

Reproduced on CartpoleBalance with episode_length=4:

step 3: done=[1. 1.], truncation=[1. 1.]
obs after the truncating step is the reset obs: True
AutoResetWrapper_final_obs is NOT in info - terminal observation is lost

Fix

Store the pre-reset observation in info['AutoResetWrapper_final_obs']. Two details beyond the one-liner suggested in the issue:

• It is set in reset as well as step, otherwise the info pytree structure differs between the two and they cannot be used together under scan.
• Under full_reset it is written after the jax.tree.map(where_done, reset_state.info, state.info) merge, otherwise the reset state's copy wins on exactly the steps where the value matters.

The extra memory is one observation, which is what AutoResetWrapper_first_obs already keeps, so this does not change the wrapper's footprint profile.

Testing

New test_auto_reset_wrapper_keeps_final_obs, parameterised over full_reset, asserts the kept observation is the one the step produced, that state.obs is still the reset observation, and that reset and step produce the same info treedef. Both parameterisations fail without the change to wrapper.py.

$ pytest mujoco_playground/_src/wrapper_test.py
7 passed

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.