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