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 <noreply@anthropic.com>

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,

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)
 

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]