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 <dvdimitrov13@gmail.com>
Co-authored-by: Dimitar Dimitrov <60075474+dvdimitrov13@users.noreply.github.com>
This commit is contained in:
Steven Palma
2026-07-31 16:39:49 +02:00
committed by GitHub
parent 2d8f5f314e
commit 7a3298ea26
2 changed files with 14 additions and 3 deletions
+8 -2
View File
@@ -384,8 +384,9 @@ class LiberoEnv(gym.Env):
} }
) )
observation = self._format_raw_obs(raw_obs) observation = self._format_raw_obs(raw_obs)
if terminated: # Return the terminal observation unchanged. The caller owns resetting after
self.reset() # termination; vector envs created below use NEXT_STEP autoreset. Resetting here
# would therefore reset twice and skip an initial state.
truncated = False truncated = False
return observation, reward, terminated, truncated, info return observation, reward, terminated, truncated, info
@@ -483,6 +484,7 @@ def create_libero_envs(
print(f"Restricting to task_ids={task_ids_filter}") print(f"Restricting to task_ids={task_ids_filter}")
is_async = env_cls is gym.vector.AsyncVectorEnv is_async = env_cls is gym.vector.AsyncVectorEnv
is_sync = env_cls is gym.vector.SyncVectorEnv
out: dict[str, dict[int, Any]] = defaultdict(dict) out: dict[str, dict[int, Any]] = defaultdict(dict)
for suite_name in suite_names: for suite_name in suite_names:
@@ -519,6 +521,10 @@ def create_libero_envs(
cached_act_space = lazy.action_space cached_act_space = lazy.action_space
cached_metadata = lazy.metadata cached_metadata = lazy.metadata
out[suite_name][tid] = lazy out[suite_name][tid] = lazy
elif is_sync:
out[suite_name][tid] = gym.vector.SyncVectorEnv(
fns, autoreset_mode=gym.vector.AutoresetMode.NEXT_STEP
)
else: else:
out[suite_name][tid] = env_cls(fns) out[suite_name][tid] = env_cls(fns)
print(f"Built vec env | suite={suite_name} | task_id={tid} | n_envs={n_envs}") print(f"Built vec env | suite={suite_name} | task_id={tid} | n_envs={n_envs}")
+6 -1
View File
@@ -212,7 +212,12 @@ class _LazyAsyncVectorEnv:
def _ensure(self) -> None: def _ensure(self) -> None:
if self._env is 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 @property
def unwrapped(self): def unwrapped(self):