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

docs/source/changelog.rst

@@ -8,6 +8,14 @@ Upcoming version (not yet released)
 Added
 ^^^^^
 
+- Added ``ManagerBasedRlEnvCfg.capture_terminal_observations``. When enabled
+  (with ``auto_reset=True``), ``step()`` stores the true terminal observation
+  (the real next state, before reset) for done environments in
+  ``extras["terminal"]``, so algorithms that need the real next state at episode
+  boundaries (value bootstrapping at truncation, AMP-style style rewards) can
+  access it without disabling auto-reset. The snapshot is read-only: it applies
+  delay and history but not observation noise, and enabling the flag does not
+  change the training trajectory (:issue:`1026`, :issue:`900`).
 - Added ``BuiltinDcMotorActuator``, a native MuJoCo ``<dcmotor>`` wrapper.
   Supports voltage / position / velocity input modes with back-EMF,
   configurable motor constants, and optional integral, slew, inductance,

src/mjlab/envs/manager_based_rl_env.py

@@ -154,6 +154,40 @@ class ManagerBasedRlEnvCfg:
   their own training loop (or a wrapper that handles the reset between steps).
   """
 
+  capture_terminal_observations: bool = False
+  """Whether to expose the true terminal observation for done environments while
+  keeping auto-reset enabled.
+
+  When True (and ``auto_reset=True``), on any step where one or more environments
+  terminate or time out, ``step()`` stores their true terminal observation (the
+  real ``s_{t+1}``, before reset) in ``extras["terminal"]``::
+
+    extras["terminal"] = {
+      "env_ids": LongTensor,                  # the done environment indices, shape [k]
+      "observations": {group: Tensor, ...},   # mirrors the obs dict, rows = [k]
+    }
+
+  The key is present only on steps where at least one environment resets, and only
+  when this flag is enabled. It is intended for algorithms that need the real next
+  state at episode boundaries (e.g. value bootstrapping at truncation, or AMP-style
+  style rewards) without giving up the standard auto-reset vectorized API. To build
+  a full ``[num_envs]`` next observation, scatter the terminal rows into the
+  returned observation::
+
+    next_obs = obs[group].clone()
+    if (term := extras.get("terminal")) is not None:
+      next_obs[term["env_ids"]] = term["observations"][group]
+
+  The terminal observation is a read-only snapshot: delay and history are applied,
+  but observation noise is *not* (it would draw from the RNG and perturb training),
+  and command resampling / step events are not re-run. For typical proprioceptive
+  observations this is exactly the terminal state; observations that depend on
+  commands reflect their value as of the previous control step. Enabling this flag
+  does not change the training trajectory; it only adds one ``forward``/``sense``
+  and an extra observation pass on steps where environments reset. This flag has no
+  effect when ``auto_reset=False`` (the returned observation is already terminal).
+  """
+
   scale_rewards_by_dt: bool = True
   """Whether to multiply rewards by the environment step duration (dt).
 
@@ -366,6 +400,7 @@ def reset(
     if seed is not None:
       self.seed(seed)
     self.extras["log"] = dict()
+    self.extras.pop("terminal", None)
     self._reset_idx(env_ids)
     self.scene.write_data_to_sim()
     self.sim.forward()
@@ -416,6 +451,7 @@ def step(self, action: torch.Tensor) -> types.VecEnvStepReturn:
       )
 
     self.extras["log"] = dict()
+    self.extras.pop("terminal", None)
     self.action_manager.process_action(action.to(self.device))
 
     for _ in range(self.cfg.decimation):
@@ -442,6 +478,14 @@ def step(self, action: torch.Tensor) -> types.VecEnvStepReturn:
 
     # Reset envs that terminated/timed-out and log the episode info.
     reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1)
+    if (
+      self.cfg.capture_terminal_observations
+      and self.cfg.auto_reset
+      and len(reset_env_ids) > 0
+    ):
+      # Snapshot the true terminal observation (real s_{t+1}) for done envs
+      # before they are reset below. Must run before _reset_idx.
+      self._store_terminal_observations(reset_env_ids)
     if self.cfg.auto_reset and len(reset_env_ids) > 0:
       self.recorder_manager.record_pre_reset(reset_env_ids)
       self._reset_idx(reset_env_ids)
@@ -589,3 +633,50 @@ def _reset_idx(self, env_ids: torch.Tensor | None = None) -> None:
     # reset the episode length buffer.
     self.episode_length_buf[env_ids] = 0
     self._manual_reset_pending[env_ids] = False
+
+  def _store_terminal_observations(self, env_ids: torch.Tensor) -> None:
+    """Snapshot the true terminal observation for done envs into extras.
+
+    Called from ``step()`` before resetting ``env_ids``. Refreshes derived
+    quantities and sensors at the terminal (pre-reset) state so the snapshot
+    matches what ``step()`` would return with ``auto_reset=False``. Both
+    ``forward()`` and ``sense()`` are deterministic, draw no RNG, and do not
+    modify ``qpos``/``qvel``, and the peek observation pass mutates no buffers,
+    so this does not perturb the training trajectory.
+
+    This extra ``forward()`` precedes the unconditional ``forward()`` later in
+    ``step()``. That is safe because ``forward()`` is a fixed point: it writes
+    ``qacc`` into ``qacc_warmstart``, and re-solving from that warm start
+    reproduces the same ``qacc`` (verified to bit-equality on contact-rich
+    envs). So the extra call cannot change the result of the subsequent
+    ``forward()``; it is redundant compute on reset steps, not a perturbation.
+    (A deliberately under-converged solver would weaken this, at the cost of a
+    one-substep difference for done envs only.) See
+    :attr:`ManagerBasedRlEnvCfg.capture_terminal_observations`.
+    """
+    self.sim.forward()
+    self.sim.sense()
+    terminal_obs = self.observation_manager.compute(peek=True)
+    self.extras["terminal"] = {
+      "env_ids": env_ids.clone(),
+      "observations": _select_env_observations(terminal_obs, env_ids),
+    }
+
+
+def _select_env_observations(
+  obs: dict[str, torch.Tensor | dict[str, torch.Tensor]],
+  env_ids: torch.Tensor,
+) -> dict[str, torch.Tensor | dict[str, torch.Tensor]]:
+  """Extract and clone the rows for ``env_ids`` from an observation dict.
+
+  Mirrors the structure of ``obs`` (concatenated groups are tensors; non-
+  concatenated groups are term dicts), returning ``[len(env_ids), ...]`` rows.
+  Clones so the result does not alias the live observation buffers.
+  """
+  selected: dict[str, torch.Tensor | dict[str, torch.Tensor]] = {}
+  for group, value in obs.items():
+    if isinstance(value, torch.Tensor):
+      selected[group] = value[env_ids].clone()
+    else:
+      selected[group] = {k: v[env_ids].clone() for k, v in value.items()}
+  return selected

src/mjlab/managers/observation_manager.py

@@ -303,8 +303,16 @@ def _check_and_handle_nans(
     return torch.nan_to_num(tensor, nan=0.0, posinf=0.0, neginf=0.0)
 
   def compute(
-    self, update_history: bool = False
+    self, update_history: bool = False, peek: bool = False
   ) -> dict[str, torch.Tensor | dict[str, torch.Tensor]]:
+    # Peek is a read-only snapshot (see compute_group): it never reads or writes
+    # the cache, never mutates delay/history buffers, and never draws RNG.
+    if peek:
+      return {
+        group_name: self.compute_group(group_name, peek=True)
+        for group_name in self._group_obs_term_names
+      }
+
     # Return cached observations if not updating and cache exists.
     # This prevents double-pushing to delay buffers when compute() is called
     # multiple times per control step (e.g., in get_observations() after step()).
@@ -318,8 +326,17 @@ def compute(
     return obs_buffer
 
   def compute_group(
-    self, group_name: str, update_history: bool = False
+    self, group_name: str, update_history: bool = False, peek: bool = False
   ) -> torch.Tensor | dict[str, torch.Tensor]:
+    """Compute observations for a single group.
+
+    When ``peek=True``, computes a read-only snapshot of the current observation
+    *without mutating any state*: delay and history buffers are queried via their
+    ``peek_append`` methods (no pointer advance, no lag resampling), and noise is
+    skipped (it would draw from the global RNG and perturb the training
+    trajectory). This is used to capture true terminal observations for done
+    environments under auto-reset. Shapes are identical to the non-peek result.
+    """
     group_cfg = self.cfg[group_name]
     group_term_names = self._group_obs_term_names[group_name]
     group_obs: dict[str, torch.Tensor] = {}
@@ -328,10 +345,14 @@ def compute_group(
     )
     for term_name, term_cfg in obs_terms:
       obs: torch.Tensor = term_cfg.func(self._env, **term_cfg.params).clone()
-      if isinstance(term_cfg.noise, noise_cfg.NoiseCfg):
-        obs = term_cfg.noise.apply(obs)
-      elif isinstance(term_cfg.noise, noise_cfg.NoiseModelCfg):
-        obs = self._group_obs_class_instances[group_name][term_name](obs)
+      # Noise is skipped on peek: NoiseCfg draws from the global RNG (which would
+      # perturb the training trajectory) and NoiseModelCfg carries per-step state
+      # (which would be advanced twice). Terminal observations are noise-free.
+      if not peek:
+        if isinstance(term_cfg.noise, noise_cfg.NoiseCfg):
+          obs = term_cfg.noise.apply(obs)
+        elif isinstance(term_cfg.noise, noise_cfg.NoiseModelCfg):
+          obs = self._group_obs_class_instances[group_name][term_name](obs)
       if term_cfg.clip:
         obs = obs.clip_(min=term_cfg.clip[0], max=term_cfg.clip[1])
       if term_cfg.scale is not None:
@@ -347,17 +368,24 @@ def compute_group(
 
       if term_cfg.delay_max_lag > 0:
         delay_buffer = self._group_obs_term_delay_buffer[group_name][term_name]
-        delay_buffer.append(obs)
-        obs = delay_buffer.compute()
+        if peek:
+          obs = delay_buffer.peek_append(obs)
+        else:
+          delay_buffer.append(obs)
+          obs = delay_buffer.compute()
       if term_cfg.history_length > 0:
         circular_buffer = self._group_obs_term_history_buffer[group_name][term_name]
-        if update_history or not circular_buffer.is_initialized:
-          circular_buffer.append(obs)
+        if peek:
+          history = circular_buffer.peek_append(obs)
+        else:
+          if update_history or not circular_buffer.is_initialized:
+            circular_buffer.append(obs)
+          history = circular_buffer.buffer
 
         if term_cfg.flatten_history_dim:
-          group_obs[term_name] = circular_buffer.buffer.reshape(self._env.num_envs, -1)
+          group_obs[term_name] = history.reshape(self._env.num_envs, -1)
         else:
-          group_obs[term_name] = circular_buffer.buffer
+          group_obs[term_name] = history
       else:
         group_obs[term_name] = obs
 

src/mjlab/utils/buffers/circular_buffer.py

@@ -214,6 +214,49 @@ def append(self, data: torch.Tensor) -> None:
 
     self._num_pushes += 1
 
+  def peek_append(self, data: torch.Tensor) -> torch.Tensor:
+    """Return the window :meth:`buffer` would yield after ``append(data)``.
+
+    This is a pure, read-only counterpart to :meth:`append`: it computes the
+    chronological history window that *would* result from appending ``data``,
+    but mutates no state (pointer, push counts, and stored frames are all left
+    untouched).
+
+    The buffer's single pointer is shared across the batch, so a real partial
+    append cannot update a subset of rows without shifting every other row's
+    window. ``peek_append`` sidesteps that entirely by never advancing the
+    pointer, which makes it safe for computing terminal observations on a
+    subset of (done) environments without disturbing the live history that the
+    rest of the batch continues to accumulate.
+
+    Args:
+      data: Tensor of shape (batch_size, ...) to hypothetically append.
+
+    Returns:
+      Chronological window (oldest to newest) of shape (batch_size, max_len,
+      ...), where the newest frame is ``data``.
+    """
+    if data.shape[0] != self._batch_size:
+      raise ValueError(f"Expected batch size {self._batch_size}, got {data.shape[0]}")
+    data = data.to(self._device)
+
+    # Never appended: a real append would backfill every slot with this first
+    # frame, so the window is `data` repeated across the time axis. Materialize
+    # (repeat, not expand) so the result never aliases `data`.
+    if self._buffer is None:
+      return data.unsqueeze(1).repeat(1, self._max_len, *([1] * (data.ndim - 1)))
+
+    # Normal case: drop the oldest frame and append `data` as the newest.
+    old_window = self.buffer  # (batch, max_len, ...), chronological.
+    window = torch.cat([old_window[:, 1:], data.unsqueeze(1)], dim=1)
+
+    # Rows not yet pushed since a reset get backfilled with `data`, matching the
+    # first-push behavior in append().
+    is_first_push = self._num_pushes == 0
+    if torch.any(is_first_push):
+      window[is_first_push] = data[is_first_push].unsqueeze(1)
+    return window
+
   def __getitem__(self, key: torch.Tensor | int) -> torch.Tensor:
     """Retrieve lagged frames per batch (LIFO).
 

src/mjlab/utils/buffers/delay_buffer.py

@@ -226,6 +226,33 @@ def append(self, data: torch.Tensor) -> None:
     """
     self._buffer.append(data)
 
+  def peek_append(self, data: torch.Tensor) -> torch.Tensor:
+    """Return the delayed obs :meth:`compute` would yield after ``append(data)``.
+
+    Pure, read-only counterpart to ``append`` followed by ``compute``: it
+    mutates no state and, crucially, does *not* resample lags (so it draws no
+    RNG and leaves the step counter untouched). The most recently sampled
+    per-env lags are reused. This makes it safe for computing terminal
+    observations on a subset of environments without perturbing the live delay
+    process shared by the rest of the batch. See
+    :meth:`CircularBuffer.peek_append`.
+
+    Args:
+      data: Observation tensor of shape (batch_size, ...).
+
+    Returns:
+      Delayed observation of shape (batch_size, ...).
+    """
+    window = self._buffer.peek_append(data)  # (batch, buf_len, ...), chronological.
+    buf_len = window.shape[1]
+    # Valid history length after the hypothetical append, capped at capacity.
+    length_post = (self._buffer.current_length + 1).clamp_max(buf_len)
+    valid_lags = torch.minimum(self._current_lags, length_post - 1).clamp_min(0)
+    # Newest frame sits at window index buf_len - 1; lag L is L frames before it.
+    gather_idx = (buf_len - 1) - valid_lags
+    batch_idx = torch.arange(self.batch_size, device=self.device)
+    return window[batch_idx, gather_idx]
+
   def compute(self) -> torch.Tensor:
     """Compute delayed observation for current step.