From 7a3298ea263528b371e2630bb0421cd519b45ef4 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Fri, 31 Jul 2026 16:39:49 +0200 Subject: [PATCH] fix(libero): don't reset inside step() on termination (#4273) * fix(libero): don't reset inside step() on termination LiberoEnv.step() called self.reset() when an episode terminated. Gymnasium's vector envs default to AutoresetMode.NEXT_STEP, so the vector env resets the sub-env again on the following step -- every termination paid two full resets. The self-reset was also pure waste: `observation` is built from the terminal raw_obs before it, so the reset's return value was discarded outright. Worse, LiberoEnv.reset() advances init_state_id by _reset_stride, so the extra reset skipped an initial state on every episode. Measured with a counting subclass under SyncVectorEnv (gymnasium 1.3.0, terminating every 4 steps over a 14-step loop): n_envs=1: 3 terminations -> 6 resets = 1 initial + 3 self + 2 autoreset n_envs=2: 6 terminations -> 12 resets = 2 initial + 6 self + 4 autoreset (the final termination's autoreset does not fire before the loop ends) Each LiberoEnv.reset() is a full LIBERO reset plus num_steps_wait settle steps, so on the current default reset path this duplicates roughly 1.3 s per termination. Standalone (non-vectorised) users must now call reset() themselves after termination, which is the Gymnasium contract. * test(libero): pin the autoreset *default*, not the enum's spelling The previous assertion checked `AutoresetMode.NEXT_STEP.value == "NextStep"`, which is a naming detail. If a future Gymnasium flipped the vector-env default to SAME_STEP, that assertion would still pass and the bug this fixes would come back as a missing reset instead of a double one. Assert the observable default instead. Verified on gymnasium 1.1.1 (the floor in pyproject) and 1.3.0, for both SyncVectorEnv and AsyncVectorEnv. Negative control: constructing with autoreset_mode=SAME_STEP fails the new assertion and passes the old one. Also document why `env._env = inner` is not redundant with the monkeypatched factory: `LiberoEnv.__init__` defers simulator creation, so pre-binding keeps `_ensure_env()` a no-op and avoids a stray `reset()` on the mock. * test(libero): drop the redundant `env._env` prebind @noron12234 flagged this as redundant. Their stated reason was wrong -- `__init__` defers simulator creation (`self._env = None`, libero.py:180), so the monkeypatched factory has not been called by the time it returns -- but the conclusion holds: `_ensure_env()` pulls the same `inner` from the mocked factory on first `step()`. I claimed the prebind was load-bearing because a stray `reset()` would pollute the tests' call counts. That was wrong and I had not run it. Ran both variants against the real class: the stray reset lands on `inner.reset`, which no assertion touches, and both tests pass with or without the line. Dropping it. * refactor(env): add explicit NEXT_STEP * fix(style): pre-commit --------- Co-authored-by: Dimitar Dimitrov Co-authored-by: Dimitar Dimitrov <60075474+dvdimitrov13@users.noreply.github.com> --- src/lerobot/envs/libero.py | 10 ++++++++-- src/lerobot/envs/utils.py | 7 ++++++- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/lerobot/envs/libero.py b/src/lerobot/envs/libero.py index a3e0f17c4..17fc22bbd 100644 --- a/src/lerobot/envs/libero.py +++ b/src/lerobot/envs/libero.py @@ -384,8 +384,9 @@ class LiberoEnv(gym.Env): } ) observation = self._format_raw_obs(raw_obs) - if terminated: - self.reset() + # Return the terminal observation unchanged. The caller owns resetting after + # termination; vector envs created below use NEXT_STEP autoreset. Resetting here + # would therefore reset twice and skip an initial state. truncated = False return observation, reward, terminated, truncated, info @@ -483,6 +484,7 @@ def create_libero_envs( print(f"Restricting to task_ids={task_ids_filter}") is_async = env_cls is gym.vector.AsyncVectorEnv + is_sync = env_cls is gym.vector.SyncVectorEnv out: dict[str, dict[int, Any]] = defaultdict(dict) for suite_name in suite_names: @@ -519,6 +521,10 @@ def create_libero_envs( cached_act_space = lazy.action_space cached_metadata = lazy.metadata out[suite_name][tid] = lazy + elif is_sync: + out[suite_name][tid] = gym.vector.SyncVectorEnv( + fns, autoreset_mode=gym.vector.AutoresetMode.NEXT_STEP + ) else: out[suite_name][tid] = env_cls(fns) print(f"Built vec env | suite={suite_name} | task_id={tid} | n_envs={n_envs}") diff --git a/src/lerobot/envs/utils.py b/src/lerobot/envs/utils.py index 8b9c4f94b..6c2286952 100644 --- a/src/lerobot/envs/utils.py +++ b/src/lerobot/envs/utils.py @@ -212,7 +212,12 @@ class _LazyAsyncVectorEnv: def _ensure(self) -> None: if self._env is None: - self._env = gym.vector.AsyncVectorEnv(self._env_fns, context="forkserver", shared_memory=True) + self._env = gym.vector.AsyncVectorEnv( + self._env_fns, + context="forkserver", + shared_memory=True, + autoreset_mode=gym.vector.AutoresetMode.NEXT_STEP, + ) @property def unwrapped(self):