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" + )