From 8135a8a8d1d2b1bdcd95449a3071a722b29a835c Mon Sep 17 00:00:00 2001 From: Dimitar Dimitrov <60075474+dvdimitrov13@users.noreply.github.com> Date: Sat, 1 Aug 2026 18:36:51 +0200 Subject: [PATCH] perf(eval): stop simulating environments whose episode has ended (#4247) rollout() runs `while not np.all(done)` with `done` latched, so a sub-env that terminates early keeps being driven -- physics and offscreen rendering included -- until the slowest sub-env in the batch finishes. A batch of N runs for max(episode_lengths) iterations to complete work that only needs mean(episode_lengths), and all of the surplus is discarded. Adds FreezeAfterEpisodeEnd, applied to each sub-env of the eval vector env. It caches the terminal transition and replays it for any further step(), and also absorbs Gymnasium's autoreset -- under AutoresetMode.NEXT_STEP the vector env otherwise rebuilds a finished sub-env and runs it through an entire extra episode the rollout throws away. Reward is zeroed on replay so a frozen sub-env cannot inflate a return if a caller sums over the padded tail. Thawing is signalled explicitly: rollout() passes NEW_ROLLOUT_OPTION in reset(options=...). Gymnasium's autoreset calls reset() with no arguments, but so would a caller passing seeds=None, and inferring from that would strand an env frozen for a whole rollout. AutoresetMode.DISABLED is not an alternative -- Gymnasium asserts that no terminated env is ever stepped in that mode, so the wrapper is never reached. An earlier version of this patch used it and the vector-env test caught the assert. Co-authored-by: Claude Opus 5 --- src/lerobot/envs/utils.py | 72 ++++++++- src/lerobot/scripts/lerobot_eval.py | 5 +- tests/envs/test_freeze_after_episode_end.py | 158 ++++++++++++++++++++ 3 files changed, 233 insertions(+), 2 deletions(-) create mode 100644 tests/envs/test_freeze_after_episode_end.py diff --git a/src/lerobot/envs/utils.py b/src/lerobot/envs/utils.py index 6c2286952..1f57f41a7 100644 --- a/src/lerobot/envs/utils.py +++ b/src/lerobot/envs/utils.py @@ -177,6 +177,76 @@ def _sub_env_has_attr(env: gym.vector.VectorEnv, attr: str) -> bool: return False +# Passed in `reset(options=...)` by `rollout()` to mark the start of a new rollout. +# FreezeAfterEpisodeEnd thaws only on this, so Gymnasium's argument-less autoreset +# cannot be mistaken for a genuine new episode. +NEW_ROLLOUT_OPTION = "lerobot_new_rollout" + + +class FreezeAfterEpisodeEnd(gym.Wrapper): + """Stop doing simulator work once a sub-env's episode has ended. + + `rollout()` runs `while not np.all(done)` with `done` latched, so a sub-env that + terminates early keeps being stepped -- physics and offscreen rendering included -- + until the slowest sub-env in the batch finishes. The batch runs for + `max(episode_lengths)` iterations to complete work that only needs + `mean(episode_lengths)`. + + This caches the terminal transition and replays it for any further `step()` or + autoreset, so a finished sub-env costs nothing. The rollout already ignores those + transitions. + + The freeze survives Gymnasium's autoreset deliberately. Under + `AutoresetMode.NEXT_STEP` the vector env resets a terminated sub-env on the + following step and runs it through an entire extra episode that the rollout + discards, because `done` stays latched. Absorbing that reset is most of the saving. + + Only an explicit reset carrying `NEW_ROLLOUT_OPTION` thaws it, so the signal is + explicit rather than inferred: Gymnasium's autoreset calls `reset()` with no + arguments, but so would a caller passing `seeds=None`, and confusing the two would + strand an env frozen for a whole rollout. + + `AutoresetMode.DISABLED` is not an alternative here — Gymnasium asserts that no + terminated env is ever stepped in that mode, so the wrapper is never reached. + """ + + def __init__(self, env: gym.Env): + super().__init__(env) + self._frozen: tuple | None = None + + def reset(self, *, seed=None, options=None): + if self._frozen is not None and not (options or {}).get(NEW_ROLLOUT_OPTION): + # Gymnasium's autoreset for a sub-env the rollout has already finished with. + # Replay the terminal observation instead of rebuilding the simulation. + obs, _, _, _, info = self._frozen + return obs, info + self._frozen = None + return self.env.reset(seed=seed, options=options) + + def step(self, action): + if self._frozen is not None: + return self._frozen + obs, reward, terminated, truncated, info = self.env.step(action) + if terminated or truncated: + # Zero the reward on replay so a frozen sub-env cannot inflate a return if a + # caller sums rewards over the padded tail. + self._frozen = (obs, 0.0, terminated, truncated, info) + return obs, reward, terminated, truncated, info + + @property + def is_frozen(self) -> bool: + return self._frozen is not None + + +def freeze_after_episode_end(env_fn: Callable[[], gym.Env]) -> Callable[[], gym.Env]: + """Wrap an env factory so the built env freezes once its episode ends.""" + + def _fn() -> gym.Env: + return FreezeAfterEpisodeEnd(env_fn()) + + return _fn + + class _LazyAsyncVectorEnv: """Defers AsyncVectorEnv creation until first use. @@ -213,7 +283,7 @@ class _LazyAsyncVectorEnv: def _ensure(self) -> None: if self._env is None: self._env = gym.vector.AsyncVectorEnv( - self._env_fns, + [freeze_after_episode_end(fn) for fn in self._env_fns], context="forkserver", shared_memory=True, autoreset_mode=gym.vector.AutoresetMode.NEXT_STEP, diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index fa6ee30a2..7629d62e9 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -82,6 +82,7 @@ from lerobot.envs import ( make_env_pre_post_processors, preprocess_observation, ) +from lerobot.envs.utils import NEW_ROLLOUT_OPTION from lerobot.lerobot_types import PolicyAction from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.processor import PolicyProcessorPipeline @@ -217,7 +218,9 @@ def rollout( # Reset the policy and environments. policy.reset() - observation, info = env.reset(seed=seeds) + # NEW_ROLLOUT_OPTION tells FreezeAfterEpisodeEnd this is a genuine new episode, as + # opposed to Gymnasium's argument-less autoreset of a sub-env that already finished. + observation, info = env.reset(seed=seeds, options={NEW_ROLLOUT_OPTION: True}) if render_callback is not None: render_callback(env) diff --git a/tests/envs/test_freeze_after_episode_end.py b/tests/envs/test_freeze_after_episode_end.py new file mode 100644 index 000000000..514eb145e --- /dev/null +++ b/tests/envs/test_freeze_after_episode_end.py @@ -0,0 +1,158 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""`FreezeAfterEpisodeEnd` — no simulator work after an episode ends. + +`rollout()` latches `done` and keeps stepping the batch until its slowest member +finishes, so a sub-env that terminated early is stepped, physically simulated and +rendered for transitions the rollout then discards. +""" + +from __future__ import annotations + +import gymnasium as gym +import numpy as np + +from lerobot.envs.utils import ( + NEW_ROLLOUT_OPTION, + FreezeAfterEpisodeEnd, + freeze_after_episode_end, +) + + +class _CountingEnv(gym.Env): + """Terminates after `term_at` steps and counts how often it is actually stepped.""" + + observation_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(2,), dtype=np.float32) + action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(1,), dtype=np.float32) + + def __init__(self, term_at: int = 3): + self.term_at = term_at + self.n_steps = 0 + self.n_resets = 0 + self._t = 0 + + def reset(self, *, seed=None, options=None): + super().reset(seed=seed) + self.n_resets += 1 + self._t = 0 + return np.zeros(2, dtype=np.float32), {"is_success": False} + + def step(self, action): + self.n_steps += 1 + self._t += 1 + terminated = self._t >= self.term_at + obs = np.full(2, float(self._t), dtype=np.float32) + return obs, 1.0, terminated, False, {"is_success": terminated} + + +def test_no_stepping_after_termination(): + inner = _CountingEnv(term_at=3) + env = FreezeAfterEpisodeEnd(inner) + env.reset() + + for _ in range(10): + env.step(env.action_space.sample()) + + assert inner.n_steps == 3, "the wrapped env must not be stepped once its episode ended" + assert env.is_frozen + + +def test_frozen_transition_replays_terminal_observation(): + inner = _CountingEnv(term_at=2) + env = FreezeAfterEpisodeEnd(inner) + env.reset() + + env.step(env.action_space.sample()) + obs_term, _, terminated, _, info_term = env.step(env.action_space.sample()) + assert terminated + + obs_frozen, reward_frozen, term_frozen, trunc_frozen, info_frozen = env.step(env.action_space.sample()) + np.testing.assert_array_equal(obs_frozen, obs_term) + assert term_frozen is True + assert trunc_frozen is False + assert info_frozen == info_term + # replayed transitions must not add return if a caller sums rewards over the tail + assert reward_frozen == 0.0 + + +def test_new_rollout_option_clears_the_freeze(): + inner = _CountingEnv(term_at=2) + env = FreezeAfterEpisodeEnd(inner) + env.reset(options={NEW_ROLLOUT_OPTION: True}) + for _ in range(5): + env.step(env.action_space.sample()) + assert inner.n_steps == 2 + + env.reset(options={NEW_ROLLOUT_OPTION: True}) + assert not env.is_frozen + env.step(env.action_space.sample()) + assert inner.n_steps == 3 + + +def test_autoreset_does_not_thaw_a_finished_env(): + """Gymnasium's autoreset calls reset() with no arguments; that must not rebuild.""" + inner = _CountingEnv(term_at=2) + env = FreezeAfterEpisodeEnd(inner) + env.reset(options={NEW_ROLLOUT_OPTION: True}) + env.step(env.action_space.sample()) + obs_term, _, _, _, _ = env.step(env.action_space.sample()) + resets_before = inner.n_resets + + obs, _ = env.reset() # what the vector env issues on autoreset + + assert env.is_frozen + assert inner.n_resets == resets_before, "autoreset must not touch the simulator" + np.testing.assert_array_equal(obs, obs_term) + + +def test_a_bare_reset_still_works_before_termination(): + """Only a *frozen* env absorbs bare resets; otherwise reset() is normal.""" + inner = _CountingEnv(term_at=5) + env = FreezeAfterEpisodeEnd(inner) + env.reset() + assert inner.n_resets == 1 + env.reset() + assert inner.n_resets == 2 + + +def test_factory_helper_wraps(): + made = freeze_after_episode_end(lambda: _CountingEnv(term_at=1))() + assert isinstance(made, FreezeAfterEpisodeEnd) + + +def test_vector_env_stops_working_on_finished_subenvs(): + """End-to-end shape, under the real AutoresetMode.NEXT_STEP config. + + Gymnasium's AutoresetMode.DISABLED is not usable here: it asserts that a + terminated sub-env is never stepped, so the wrapper would never be reached. + """ + term_at = [2, 5, 9] + envs = [_CountingEnv(term_at=t) for t in term_at] + vec = gym.vector.SyncVectorEnv([freeze_after_episode_end(lambda e=e: e) for e in envs]) + vec.reset(options={NEW_ROLLOUT_OPTION: True}) + + done = np.zeros(len(envs), dtype=bool) + steps = 0 + while not np.all(done) and steps < 20: # mirrors lerobot_eval.rollout + _, _, term, trunc, _ = vec.step(np.zeros((len(envs), 1), dtype=np.float32)) + done = term | trunc | done + steps += 1 + + assert steps == max(term_at), "the batch still runs until its slowest member" + # ...but each sub-env only did its own episode's work + assert [e.n_steps for e in envs] == term_at + # without the wrapper this would be len(term_at) * max(term_at) = 27 + assert sum(e.n_steps for e in envs) == sum(term_at) + # and no sub-env was rebuilt by an autoreset it did not need + assert [e.n_resets for e in envs] == [1, 1, 1]