Add an option to capture true terminal observations under auto-reset.

[kevinzakka]

Some algorithms need the true terminal state at episode boundaries: value bootstrapping at truncation needs the real next state to bootstrap from (the original ask in #900), and AMP-style rewards need it because the discriminator scores (s_t, s_{t+1}) transitions, so a wrong final state means scoring a transition into the next episode's initial state instead of the actual one (#1026). Our current auto-reset returns the post-reset observation for done environments, so that terminal state is lost.

This PR adds an opt-in capture_terminal_observations flag. When it's on and an environment resets, step() stashes that environment's true terminal observation in extras["terminal"] before the reset. It's off by default and the key is only present on reset steps, so existing setups are unchanged.

The snapshot is read-only (it never touches the live history/delay buffers or the RNG), so the terminal observation is the noise-free terminal state with delay and history applied. To turn it into a full next observation:

next_obs = obs[group].clone()
if (term := extras.get("terminal")) is not None:
    next_obs[term["env_ids"]] = term["observations"][group]

Fixes #1026

[kevinzakka]

cc @realquantumcookie

[realquantumcookie]

Cool! Right now auto_reset already works for my case. Congrats on the 1.4.0 release btw :)

[zhuoqun-chen]

Hi @realquantumcookie, before this PR is landed, for auto reset env_ids, the observation is actually the initial observations in the new episodes, so will this affect your use cases especially for computing $Q(s_{t+1}, a_{t+1})$?

[realquantumcookie]

@zhuoqun-chen Before this PR lands, I disabled auto-reset so when one of the environment terminates it gives the correct $s_{t+1}$. You just have to manually reset the environment calling reset(env_idx=...) after any truncation / termination occurs