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):