fix(task): publish the terminal observation, not the episode's first one

CHANGELOG.md

@@ -35,6 +35,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- `RLTaskEnv.step` publishes the *terminal* observation in `info["observations"]["raw"]["obs"]` instead of the episode's first one (off-policy truncation bootstraps in clean_rl SAC/TD3 and FastTD3 read it).
 - MuJoCo: `<size memory="512M">` is reserved by default; humanoid + mesh scenes no longer die with
   `mj_stackAlloc: out of memory` (get_started/10_mount_camera.py).
 - `hf_util`: a symlinked `roboverse_data` is no longer refused as path traversal; concurrent

metasim/task/rl_task.py

@@ -148,7 +148,6 @@ def reset(self, states=None, env_ids=None, seed: int | None = None) -> tuple[tor
 
         states = self.handler.get_states(mode="tensor")
         first_obs = self._observation(states).to(self.device)
-        self._raw_observation_cache = first_obs.clone()
         priv_obs = self._privileged_observation(states)
 
         # Update observation buffers for RSL-RL compatibility
@@ -208,7 +207,11 @@ def step(
         info = {
             "privileged_observation": priv_obs,
             "episode_steps": self._episode_steps.clone(),
-            "observations": {"raw": {"obs": self._raw_observation_cache.clone()}},
+            # The terminal observation, snapshotted before the auto-reset below overwrites `obs`
+            # in place for the done envs. Off-policy learners bootstrap truncated episodes from
+            # this (V(s_T) for a time-out is a real value; V(reset state) is not), so it must be
+            # the state the episode actually ended in.
+            "observations": {"raw": {"obs": obs.clone()}},
         }
 
         done_indices = episode_done.nonzero(as_tuple=False).squeeze(-1)
@@ -217,10 +220,6 @@ def step(
             states_after = self.handler.get_states(mode="tensor")
             obs_after = self._observation(states_after).to(self.device)
             obs[done_indices] = obs_after[done_indices]
-            self._raw_observation_cache[done_indices] = obs_after[done_indices]
-        else:
-            keep_mask = (~terminated).unsqueeze(-1)
-            self._raw_observation_cache = torch.where(keep_mask, self._raw_observation_cache, obs)
 
         return obs, reward, terminated, time_out, info
 

metasim/test/test_rl_task_contract.py

@@ -211,7 +211,6 @@ def get_states(self, mode="tensor"):
     env._episode_steps = torch.zeros(1, dtype=torch.int32)
     env._action_low = torch.tensor([-10.0])
     env._action_high = torch.tensor([10.0])
-    env._raw_observation_cache = torch.zeros(1, 1)
     env._observation = lambda states: torch.zeros(1, 1)
     env._privileged_observation = lambda states: torch.zeros(1, 1)
     env._reward = lambda states: torch.zeros(1)
@@ -224,3 +223,64 @@ def get_states(self, mode="tensor"):
     assert torch.allclose(captured["applied"], torch.tensor([[3.0]])), (
         f"_process_action not applied before set_dof_targets: got {captured.get('applied')}"
     )
+
+
+@pytest.mark.general
+def test_step_publishes_the_terminal_observation_not_the_post_reset_one():
+    """``info["observations"]["raw"]["obs"]`` must be the obs the episode ended in.
+
+    ``step`` auto-resets done envs in place, so the returned ``obs`` already holds the next
+    episode's first observation for those envs. Off-policy learners bootstrap truncated
+    episodes off the raw key -- ``V(s_T)`` for a time-out is a real value, ``V(reset state)``
+    is not -- so publishing anything but the pre-reset observation silently corrupts the
+    target on every episode boundary.
+
+    This used to be served from a ``_raw_observation_cache`` that was written only in
+    ``reset()``: the update branch ran solely when *no* env was done, where ``terminated`` is
+    all-False and its ``torch.where`` was a no-op. The key therefore carried the episode's
+    *first* observation for the whole episode -- not the terminal one, and not even the
+    post-reset one.
+    """
+    step_count = {"n": 0}
+
+    class _H:
+        num_envs = 1
+
+        def set_dof_targets(self, a):
+            return None
+
+        def simulate(self):
+            step_count["n"] += 1
+
+        def get_states(self, mode="tensor"):
+            return None
+
+    env = RLTaskEnv.__new__(RLTaskEnv)
+    env.device = torch.device("cpu")
+    env.num_envs = 1
+    env.handler = _H()
+    env._episode_steps = torch.zeros(1, dtype=torch.int32)
+    env._action_low = torch.tensor([-10.0])
+    env._action_high = torch.tensor([10.0])
+    # The observation is the step counter, so the terminal obs is distinguishable from the
+    # post-reset one by value alone.
+    env._observation = lambda states: torch.full((1, 1), float(step_count["n"]))
+    env._privileged_observation = lambda states: torch.zeros(1, 1)
+    env._reward = lambda states: torch.zeros(1)
+    env._terminated = lambda states: torch.zeros(1, dtype=torch.bool)
+    env._time_out = lambda states: torch.ones(1, dtype=torch.bool)  # always truncate
+    env._process_action = lambda actions: actions
+    # reset() zeroes the counter, standing in for the state being destroyed.
+    env.reset = lambda states=None, env_ids=None, seed=None: step_count.update(n=0)
+
+    obs, _, _, time_out, info = RLTaskEnv.step(env, torch.tensor([0.0]))
+
+    assert bool(time_out[0]), "sanity: this env truncates every step"
+    raw = info["observations"]["raw"]["obs"]
+    assert torch.allclose(raw, torch.tensor([[1.0]])), (
+        f"raw obs must be the terminal observation (1.0, the state the episode ended in), got {raw.tolist()}"
+    )
+    assert torch.allclose(obs, torch.tensor([[0.0]])), (
+        "sanity: the returned obs is the post-reset one (0.0) -- which is exactly why the raw "
+        "key must not be taken from it"
+    )