fix(envs): preserve lazy-env metadata + stabilize VLABench reset

Two independent fixes:

1. Port of #3416: add `metadata` / `unwrapped` to `_LazyAsyncVectorEnv`
   and cache the metadata alongside obs/action spaces in the LIBERO,
   MetaWorld, and VLABench factories so async eval keeps
   `env.unwrapped.metadata['render_fps']` working.

2. VLABench `select_mahjong` kept failing the smoke eval with
   `mjWARN_BADQACC` after the 5-retry ceiling. Bump
   `_ENSURE_ENV_MAX_ATTEMPTS` to 20, reseed NumPy's global RNG
   between attempts (VLABench's upstream layout sampler is
   deterministic given current RNG state — without a reseed, retries
   replay the same diverging configuration), and raise a clearer
   `RuntimeError` on exhaustion. Also swap `select_mahjong` out of
   the CI task list in favor of `select_billiards`, which has a
   stable layout sampler.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Pepijn
2026-04-20 14:23:35 +02:00
parent c7d506d70c
commit 0a0ebec97a
5 changed files with 38 additions and 11 deletions
+2 -2
View File
@@ -372,7 +372,7 @@ jobs:
lerobot-eval \ lerobot-eval \
--policy.path=lerobot/smolvla_vlabench \ --policy.path=lerobot/smolvla_vlabench \
--env.type=vlabench \ --env.type=vlabench \
--env.task=select_fruit,select_toy,select_book,select_painting,select_drink,select_ingredient,select_mahjong,select_poker,add_condiment,insert_flower \ --env.task=select_fruit,select_toy,select_book,select_painting,select_drink,select_ingredient,select_billiards,select_poker,add_condiment,insert_flower \
--eval.batch_size=1 \ --eval.batch_size=1 \
--eval.n_episodes=1 \ --eval.n_episodes=1 \
--eval.use_async_envs=false \ --eval.use_async_envs=false \
@@ -394,7 +394,7 @@ jobs:
python3 scripts/ci/parse_eval_metrics.py \ python3 scripts/ci/parse_eval_metrics.py \
--artifacts-dir /tmp/vlabench-artifacts \ --artifacts-dir /tmp/vlabench-artifacts \
--env vlabench \ --env vlabench \
--task select_fruit,select_toy,select_book,select_painting,select_drink,select_ingredient,select_mahjong,select_poker,add_condiment,insert_flower \ --task select_fruit,select_toy,select_book,select_painting,select_drink,select_ingredient,select_billiards,select_poker,add_condiment,insert_flower \
--policy lerobot/smolvla_vlabench --policy lerobot/smolvla_vlabench
- name: Upload VLABench rollout video - name: Upload VLABench rollout video
+3 -1
View File
@@ -462,6 +462,7 @@ def create_libero_envs(
# Probe once and reuse to avoid creating a temp env per task. # Probe once and reuse to avoid creating a temp env per task.
cached_obs_space: spaces.Space | None = None cached_obs_space: spaces.Space | None = None
cached_act_space: spaces.Space | None = None cached_act_space: spaces.Space | None = None
cached_metadata: dict[str, Any] | None = None
for tid in selected: for tid in selected:
fns = _make_env_fns( fns = _make_env_fns(
@@ -477,10 +478,11 @@ def create_libero_envs(
camera_name_mapping=camera_name_mapping, camera_name_mapping=camera_name_mapping,
) )
if is_async: if is_async:
lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space) lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata)
if cached_obs_space is None: if cached_obs_space is None:
cached_obs_space = lazy.observation_space cached_obs_space = lazy.observation_space
cached_act_space = lazy.action_space cached_act_space = lazy.action_space
cached_metadata = lazy.metadata
out[suite_name][tid] = lazy out[suite_name][tid] = lazy
else: else:
out[suite_name][tid] = env_cls(fns) out[suite_name][tid] = env_cls(fns)
+3 -1
View File
@@ -311,6 +311,7 @@ def create_metaworld_envs(
is_async = env_cls is gym.vector.AsyncVectorEnv is_async = env_cls is gym.vector.AsyncVectorEnv
cached_obs_space = None cached_obs_space = None
cached_act_space = None cached_act_space = None
cached_metadata = None
out: dict[str, dict[int, Any]] = defaultdict(dict) out: dict[str, dict[int, Any]] = defaultdict(dict)
for group in task_groups: for group in task_groups:
@@ -324,10 +325,11 @@ def create_metaworld_envs(
fns = [(lambda tn=task_name: MetaworldEnv(task=tn, **gym_kwargs)) for _ in range(n_envs)] fns = [(lambda tn=task_name: MetaworldEnv(task=tn, **gym_kwargs)) for _ in range(n_envs)]
if is_async: if is_async:
lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space) lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata)
if cached_obs_space is None: if cached_obs_space is None:
cached_obs_space = lazy.observation_space cached_obs_space = lazy.observation_space
cached_act_space = lazy.action_space cached_act_space = lazy.action_space
cached_metadata = lazy.metadata
out[group][tid] = lazy out[group][tid] = lazy
else: else:
out[group][tid] = env_cls(fns) out[group][tid] = env_cls(fns)
+8 -1
View File
@@ -153,17 +153,20 @@ class _LazyAsyncVectorEnv:
env_fns: list[Callable], env_fns: list[Callable],
observation_space=None, observation_space=None,
action_space=None, action_space=None,
metadata=None,
): ):
self._env_fns = env_fns self._env_fns = env_fns
self._env: gym.vector.AsyncVectorEnv | None = None self._env: gym.vector.AsyncVectorEnv | None = None
self.num_envs = len(env_fns) self.num_envs = len(env_fns)
if observation_space is not None and action_space is not None: if observation_space is not None and action_space is not None and metadata is not None:
self.observation_space = observation_space self.observation_space = observation_space
self.action_space = action_space self.action_space = action_space
self.metadata = metadata
else: else:
tmp = env_fns[0]() tmp = env_fns[0]()
self.observation_space = tmp.observation_space self.observation_space = tmp.observation_space
self.action_space = tmp.action_space self.action_space = tmp.action_space
self.metadata = tmp.metadata
tmp.close() tmp.close()
self.single_observation_space = self.observation_space self.single_observation_space = self.observation_space
self.single_action_space = self.action_space self.single_action_space = self.action_space
@@ -172,6 +175,10 @@ class _LazyAsyncVectorEnv:
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)
@property
def unwrapped(self):
return self
def reset(self, **kwargs): def reset(self, **kwargs):
self._ensure() self._ensure()
return self._env.reset(**kwargs) return self._env.reset(**kwargs)
+22 -6
View File
@@ -187,8 +187,10 @@ class VLABenchEnv(gym.Env):
# `PhysicsError` (e.g. mjWARN_BADQACC) during VLABench's 20-step # `PhysicsError` (e.g. mjWARN_BADQACC) during VLABench's 20-step
# reset warm-up. Some random task/layout samples land in unstable # reset warm-up. Some random task/layout samples land in unstable
# initial configurations; re-sampling the layout almost always # initial configurations; re-sampling the layout almost always
# gives a stable one. # gives a stable one. A handful of upstream tasks (notably
_ENSURE_ENV_MAX_ATTEMPTS = 5 # `select_mahjong`) have layout samplers that diverge often enough
# to need >>5 retries, so we pick a generous ceiling.
_ENSURE_ENV_MAX_ATTEMPTS = 20
def _ensure_env(self) -> None: def _ensure_env(self) -> None:
"""Create the underlying VLABench env on first use. """Create the underlying VLABench env on first use.
@@ -201,7 +203,11 @@ class VLABenchEnv(gym.Env):
warm-up `step()` calls while toggling gravity/fluids to let the scene warm-up `step()` calls while toggling gravity/fluids to let the scene
settle; for some random layouts MuJoCo's integrator diverges and settle; for some random layouts MuJoCo's integrator diverges and
raises `mjWARN_BADQACC`. Re-sampling the layout almost always yields raises `mjWARN_BADQACC`. Re-sampling the layout almost always yields
a stable one, so we retry a few times before giving up. a stable one, so we retry a number of times before giving up. Between
attempts we reseed NumPy's global RNG from OS entropy so the upstream
task sampler explores fresh initial states without this, retries
can replay the same diverging configuration when the sampler is
deterministic given the current RNG state.
""" """
if self._env is not None: if self._env is not None:
return return
@@ -223,11 +229,19 @@ class VLABenchEnv(gym.Env):
print( print(
f"[vlabench] PhysicsError on attempt {attempt}/" f"[vlabench] PhysicsError on attempt {attempt}/"
f"{self._ENSURE_ENV_MAX_ATTEMPTS} while building task " f"{self._ENSURE_ENV_MAX_ATTEMPTS} while building task "
f"'{self.task}': {exc}. Retrying with fresh layout…" f"'{self.task}': {exc}. Retrying with fresh layout…",
flush=True,
) )
np.random.seed(None)
if self._env is None: if self._env is None:
assert last_exc is not None assert last_exc is not None
raise last_exc raise RuntimeError(
f"VLABench task '{self.task}' failed to produce a stable "
f"initial layout after {self._ENSURE_ENV_MAX_ATTEMPTS} "
f"attempts. This task's upstream sampler diverges too "
f"often for the configured robot; consider removing it "
f"from the eval set. Last physics error: {last_exc}"
) from last_exc
# Extract task description from the dm_control task # Extract task description from the dm_control task
task_obj = self._env.task task_obj = self._env.task
@@ -541,6 +555,7 @@ def create_vlabench_envs(
is_async = env_cls is gym.vector.AsyncVectorEnv is_async = env_cls is gym.vector.AsyncVectorEnv
cached_obs_space = None cached_obs_space = None
cached_act_space = None cached_act_space = None
cached_metadata = None
out: dict[str, dict[int, Any]] = defaultdict(dict) out: dict[str, dict[int, Any]] = defaultdict(dict)
for group in task_groups: for group in task_groups:
@@ -557,10 +572,11 @@ def create_vlabench_envs(
) )
if is_async: if is_async:
lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space) lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata)
if cached_obs_space is None: if cached_obs_space is None:
cached_obs_space = lazy.observation_space cached_obs_space = lazy.observation_space
cached_act_space = lazy.action_space cached_act_space = lazy.action_space
cached_metadata = lazy.metadata
out[group][tid] = lazy out[group][tid] = lazy
else: else:
out[group][tid] = env_cls(fns) out[group][tid] = env_cls(fns)