From 8d45b2ce7265b2db090f608c92b65b536fb85220 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Fri, 17 Apr 2026 15:34:22 +0100 Subject: [PATCH] fix(vlabench): retry unstable reset, graceful PhysicsError on step VLABench's LM4ManipDMEnv.reset() runs 20 warm-up `step()` calls while toggling gravity / fluid coefficients to settle the scene. For some random layout samples MuJoCo's integrator diverges and raises `mjWARN_BADQACC` -> `dm_control.rl.control.PhysicsError` before we ever get to return from reset. In a multi-task smoke eval that kills the whole run at the next task's env construction. - `_ensure_env`: catch PhysicsError and retry `load_env()` up to five times with a fresh random layout. Almost always converges; if it really can't, re-raise the last exception. - `step`: wrap `self._env.step(ctrl)` in a PhysicsError handler. On divergence, terminate the episode gracefully (is_success=False, terminated=True, info['physics_error']=True) and drop the stale inner env so the next reset rebuilds it. Keeps the rest of the multi-task eval moving instead of propagating a hard crash. Made-with: Cursor --- src/lerobot/envs/vlabench.py | 56 ++++++++++++++++++++++++++++++------ 1 file changed, 48 insertions(+), 8 deletions(-) diff --git a/src/lerobot/envs/vlabench.py b/src/lerobot/envs/vlabench.py index a420c1d27..821b3dc22 100644 --- a/src/lerobot/envs/vlabench.py +++ b/src/lerobot/envs/vlabench.py @@ -25,6 +25,7 @@ with long-horizon reasoning, built on MuJoCo/dm_control. from __future__ import annotations +import contextlib from collections import defaultdict from collections.abc import Callable, Sequence from functools import partial @@ -181,30 +182,54 @@ class VLABenchEnv(gym.Env): low=ACTION_LOW, high=ACTION_HIGH, shape=(ACTION_DIM,), dtype=np.float32 ) + # Max attempts to rebuild the underlying env when MuJoCo throws + # `PhysicsError` (e.g. mjWARN_BADQACC) during VLABench's 20-step + # reset warm-up. Some random task/layout samples land in unstable + # initial configurations; re-sampling the layout almost always + # gives a stable one. + _ENSURE_ENV_MAX_ATTEMPTS = 5 + def _ensure_env(self) -> None: """Create the underlying VLABench env on first use. Called inside the worker subprocess after fork(), so each worker gets its own clean rendering context rather than inheriting a stale one from the parent process (which causes crashes with AsyncVectorEnv). + + Retries on `PhysicsError`: VLABench's `LM4ManipDMEnv.reset()` runs 20 + warm-up `step()` calls while toggling gravity/fluids to let the scene + settle; for some random layouts MuJoCo's integrator diverges and + raises `mjWARN_BADQACC`. Re-sampling the layout almost always yields + a stable one, so we retry a few times before giving up. """ if self._env is not None: return import VLABench.robots # noqa: F401 # type: ignore[import-untyped] import VLABench.tasks # noqa: F401 # type: ignore[import-untyped] + from dm_control.rl.control import PhysicsError # type: ignore[import-untyped] from VLABench.envs import load_env # type: ignore[import-untyped] h, w = self.render_resolution - env = load_env( - task=self.task, - robot=self.robot, - render_resolution=(h, w), - ) - self._env = env + last_exc: PhysicsError | None = None + for attempt in range(1, self._ENSURE_ENV_MAX_ATTEMPTS + 1): + try: + env = load_env(task=self.task, robot=self.robot, render_resolution=(h, w)) + self._env = env + break + except PhysicsError as exc: + last_exc = exc + print( + f"[vlabench] PhysicsError on attempt {attempt}/" + f"{self._ENSURE_ENV_MAX_ATTEMPTS} while building task " + f"'{self.task}': {exc}. Retrying with fresh layout…" + ) + if self._env is None: + assert last_exc is not None + raise last_exc # Extract task description from the dm_control task - task_obj = env.task + task_obj = self._env.task if hasattr(task_obj, "task_description"): self.task_description = task_obj.task_description elif hasattr(task_obj, "language_instruction"): @@ -374,6 +399,8 @@ class VLABenchEnv(gym.Env): return observation, info def step(self, action: np.ndarray) -> tuple[RobotObservation, float, bool, bool, dict[str, Any]]: + from dm_control.rl.control import PhysicsError # type: ignore[import-untyped] + self._ensure_env() assert self._env is not None @@ -391,7 +418,20 @@ class VLABenchEnv(gym.Env): physics = self._env.physics ctrl_dim = int(physics.data.ctrl.shape[0]) ctrl = self._build_ctrl_from_action(action, ctrl_dim) - timestep = self._env.step(ctrl) + try: + timestep = self._env.step(ctrl) + except PhysicsError as exc: + # Physics integrator diverged (e.g. mjWARN_BADQACC). Treat it as + # a graceful failed termination rather than a hard crash — the + # rest of the multi-task eval should still run. + print(f"[vlabench] PhysicsError during step on task '{self.task}': {exc}. Terminating episode.") + observation = self._get_obs() + info = {"task": self.task, "is_success": False, "physics_error": True} + # Drop the stale env so the next reset() rebuilds it cleanly. + with contextlib.suppress(Exception): + self._env.close() + self._env = None + return observation, 0.0, True, False, info # Extract reward from dm_control timestep reward = float(timestep.reward) if timestep.reward is not None else 0.0