From a3c2ba03c00955b43bb1158258a287d52d47c42b Mon Sep 17 00:00:00 2001 From: Haoran Geng <71596067+geng-haoran@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:45:49 +0000 Subject: [PATCH] fix(task): publish the terminal observation, not the episode's first one RLTaskEnv.step auto-resets done envs in place, so the obs it returns already holds the next episode's first observation for those envs. info["observations"]["raw"]["obs"] exists to give off-policy learners the observation the episode actually *ended* in, which is what a truncated episode must bootstrap from: V(s_T) for a time-out is a real value, V(reset state) is not. It did not do that. _raw_observation_cache was written in exactly two places -- reset(), and the done branch, which stored the *post-reset* obs. The remaining update lived in an else branch that runs only when no env is done, where `terminated` is all-False, so its torch.where kept the old value unconditionally. Nothing ever advanced the cache during an episode. The raw key therefore carried the episode's **first** observation, from reset until the next reset -- not the terminal obs, and not even the post-reset obs. It was not a stale value; it was a constant. Verified against the real step() with a counting stub: at a time-out on step 1, the true terminal obs is 1.0 and the key published 0.0. The cache never needed to exist. At the point step() builds `info`, `obs` still holds the pre-reset observation for every env -- the auto-reset below is what clobbers it. Snapshotting it there is both correct and simpler, so the cache is removed rather than repaired. Downstream, RoboVerse's fast_td3, clean_rl/td3 and clean_rl/sac all read this key as the "true next obs" for truncation bootstrapping, and were reading a constant for every plain RLTaskEnv task. Locomotion has not obviously suffered only because LeggedRobotTask overrides step() and sets the key itself. The unused `# noqa: D401` removed alongside is a drive-by in the same file: D401 is not enabled in this repo's ruff config, so it fails RUF100 under the pinned ruff 0.14.5 today. --- CHANGELOG.md | 1 + metasim/task/rl_task.py | 11 +++-- metasim/test/test_rl_task_contract.py | 62 ++++++++++++++++++++++++++- 3 files changed, 67 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0b507a6..6e1726e 100644 --- a/CHANGELOG.md +++ b/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: `` 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 diff --git a/metasim/task/rl_task.py b/metasim/task/rl_task.py index f0b97e8..3368203 100644 --- a/metasim/task/rl_task.py +++ b/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 diff --git a/metasim/test/test_rl_task_contract.py b/metasim/test/test_rl_task_contract.py index 2643bf0..de86b8b 100644 --- a/metasim/test/test_rl_task_contract.py +++ b/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" + )