From be146226aaf79dfe843a259de2daa200d1a73bdc Mon Sep 17 00:00:00 2001 From: Gotham-Zolio <18781106300@163.com> Date: Fri, 11 Sep 2026 11:36:57 -0400 Subject: [PATCH] fix: stop resetting the environment after the last episode rollout() reset unconditionally whenever an episode finished, including the one that satisfied num_episodes. The observation it produced was never read - the loop exits immediately afterwards. For a simulator that is a wasted rollout. For a real robot it is a pointless move back to the home pose after the run is already over. The test that proves it comes from nothingbutbut's features/robocasa branch, where this was fixed and never merged. It is the only test in this package that exercises rollout() end to end other than the two added yesterday, so it is worth having on main whatever happens to the rest of that branch. Two adaptations were needed, both because main has moved since April: * its fake agent returned an action of shape (steps, n_envs), which was only valid under the batch-axis guess that _get_action_spec made until yesterday. It now returns [H, n, *da] like the protocol specifies. * its recorder spy predates record_timing. test_cli.py::test_rollout_infer_feedback_and_partial_reset_semantics had to change with it, and this is worth being explicit about: it asserted that a partial reset happens, using num_episodes=1 as the shortest route to a done environment. Under the new behaviour that run is over before the reset would happen, so the test now uses two episodes. Its subject - that a partial reset targets only the environments that finished - is unchanged and still asserted. Co-Authored-By: nothingbutbut <2367347983@qq.com> Co-Authored-By: Claude Opus 5 (1M context) --- src/plugrl_env_client/runner/rollout.py | 13 ++- tests/test_cli.py | 12 ++- tests/test_rollout.py | 104 ++++++++++++++++++++++++ 3 files changed, 121 insertions(+), 8 deletions(-) create mode 100644 tests/test_rollout.py diff --git a/src/plugrl_env_client/runner/rollout.py b/src/plugrl_env_client/runner/rollout.py index a3cca5d..be076b8 100644 --- a/src/plugrl_env_client/runner/rollout.py +++ b/src/plugrl_env_client/runner/rollout.py @@ -211,10 +211,15 @@ def log_timing_summary(final: bool = False) -> None: if recorder is not None: recorder.on_episode_done(done_indices, obs, info) finished_episodes += done_indices.size - obs, info = env.reset(options={"reset_indices": done_indices}) - if recorder is not None: - recorder.on_reset(obs, info, reset_indices=done_indices) - step_id[done_indices] = 0 + # Only reset if the loop is going to use what comes back. A reset + # after the last episode costs a simulator a wasted rollout, and + # costs a real robot a pointless move back to its home pose - for + # an observation nothing will ever read. + if finished_episodes < num_episodes: + obs, info = env.reset(options={"reset_indices": done_indices}) + if recorder is not None: + recorder.on_reset(obs, info, reset_indices=done_indices) + step_id[done_indices] = 0 now = time.perf_counter() if now - last_timing_log_at >= 30.0: log_timing_summary(final=False) diff --git a/tests/test_cli.py b/tests/test_cli.py index d7dcad0..81453f9 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -272,20 +272,24 @@ def test_rollout_infer_feedback_and_partial_reset_semantics(): rollout( cast(gym.vector.VectorEnv, env), cast(WebSocketEnvClientAgent, agent), - num_episodes=1, + # Two episodes, so the loop continues past the first one and the + # partial reset this test is about actually happens. With one + # episode the run is already over when env 0 terminates, and + # rollout deliberately does not reset on the way out. + num_episodes=2, replan_steps=1, num_envs=num_envs, ) - assert len(agent.infer_calls) == 1 + assert len(agent.infer_calls) == 2 np.testing.assert_array_equal(agent.infer_calls[0], np.asarray([0, 1])) - assert len(env.step_calls) == 1 + assert len(env.step_calls) == 2 np.testing.assert_array_equal( env.step_calls[0], np.zeros((num_envs, action_dim), dtype=np.float32) ) - assert len(agent.feedback_calls) == 1 + assert len(agent.feedback_calls) == 2 fb = agent.feedback_calls[0] np.testing.assert_array_equal(fb["env_indices"], np.asarray([0, 1])) np.testing.assert_allclose(fb["rewards"], np.asarray([1.0, 2.0], dtype=np.float32)) diff --git a/tests/test_rollout.py b/tests/test_rollout.py new file mode 100644 index 0000000..cc9c4b5 --- /dev/null +++ b/tests/test_rollout.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import numpy as np + +from plugrl_env_client.envs.base_env import Observation +from plugrl_env_client.runner.rollout import rollout + + +class _OneStepDoneEnv: + def __init__(self) -> None: + self.num_envs = 1 + self.single_action_space = type( + "_Space", (), {"shape": (1,), "dtype": np.dtype(np.float32)} + )() + self.reset_calls: list[np.ndarray | None] = [] + self.step_calls = 0 + + def reset(self, *, seed=None, options=None): + del seed + reset_indices = None if options is None else options.get("reset_indices") + self.reset_calls.append( + None + if reset_indices is None + else np.asarray(reset_indices, dtype=np.int64).copy() + ) + obs = Observation( + images={"cam": np.zeros((1, 4, 4, 3), dtype=np.uint8)}, + states={"state": np.zeros((1, 2), dtype=np.float32)}, + text=np.asarray(["task"], dtype=np.str_), + ) + return obs, {} + + def step(self, actions): + del actions + self.step_calls += 1 + obs = Observation( + images={"cam": np.full((1, 4, 4, 3), self.step_calls, dtype=np.uint8)}, + states={"state": np.full((1, 2), self.step_calls, dtype=np.float32)}, + text=np.asarray(["task"], dtype=np.str_), + ) + reward = np.asarray([1.0], dtype=np.float32) + terminated = np.asarray([True], dtype=np.bool_) + truncated = np.asarray([False], dtype=np.bool_) + info = { + "episode": { + "r": np.asarray([1.0], dtype=np.float32), + "l": np.asarray([1], dtype=np.int32), + "s": np.asarray([True], dtype=np.bool_), + "mask": np.asarray([True], dtype=np.bool_), + "mean_success_rate": 1.0, + } + } + return obs, reward, terminated, truncated, info + + +class _OneChunkAgent: + def infer(self, obs, *, env_indices, step_ids): + del obs, env_indices, step_ids + return {"action": np.zeros((1, 1, 1), dtype=np.float32)} + + def feedback(self, **kwargs): + del kwargs + return None + + +class _SpyRecorder: + def __init__(self) -> None: + self.debug_packets_enabled = False + self.reset_calls = 0 + self.done_calls = 0 + + def on_reset(self, obs, info, *, reset_indices): + del obs, info, reset_indices + self.reset_calls += 1 + + def on_step(self, next_obs, reward, terminated, truncated, info): + del next_obs, reward, terminated, truncated, info + + def on_episode_done(self, done_indices, obs, info): + del done_indices, obs, info + self.done_calls += 1 + + def record_timing(self, timing): + del timing + + +def test_rollout_does_not_reset_after_final_episode(): + env = _OneStepDoneEnv() + agent = _OneChunkAgent() + recorder = _SpyRecorder() + + rollout( + env, + agent, + num_episodes=1, + replan_steps=1, + num_envs=1, + recorder=recorder, + ) + + assert len(env.reset_calls) == 1 + assert env.reset_calls[0] is None + assert recorder.reset_calls == 1 + assert recorder.done_calls == 1