fix: integrate PR #3396 review feedback (round 2)

- envs(vlabench): drop the unused `episode_index` / `n_envs`
  constructor args — they were stored but never referenced.
- envs(vlabench): replace module-level `print()` calls with the
  standard `logging.getLogger(__name__)` logger (matches the rest
  of the repo's script convention). Covers the two PhysicsError
  retry/step paths and the `create_vlabench_envs` progress lines.
- envs(vlabench): drop `_make_env_fns` and build factories
  MetaWorld-style via a lambda list comprehension — simpler, no
  extra helper.
- docker(vlabench): pin upstream clones to commit SHAs
  (`OpenMOSS/VLABench@cf588fe6`, `motion-planning/rrt-algorithms@e51d95ee`)
  so benchmark images are reproducible.
- ci(vlabench): run `scripts/ci/extract_task_descriptions.py` after
  the eval so `metrics.json` carries per-task natural-language
  labels, matching LIBERO / MetaWorld jobs. Added a VLABench
  extractor that produces cleaned-name labels keyed by `<task>_0`.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Pepijn
2026-04-20 14:33:23 +02:00
parent 0a0ebec97a
commit eff1d5a528
4 changed files with 57 additions and 46 deletions
+4
View File
@@ -379,6 +379,10 @@ jobs:
--policy.device=cuda \ --policy.device=cuda \
'--rename_map={\"observation.images.image\": \"observation.images.camera1\", \"observation.images.second_image\": \"observation.images.camera2\", \"observation.images.wrist_image\": \"observation.images.camera3\"}' \ '--rename_map={\"observation.images.image\": \"observation.images.camera1\", \"observation.images.second_image\": \"observation.images.camera2\", \"observation.images.wrist_image\": \"observation.images.camera3\"}' \
--output_dir=/tmp/eval-artifacts --output_dir=/tmp/eval-artifacts
python scripts/ci/extract_task_descriptions.py \
--env vlabench \
--task select_fruit,select_toy,select_book,select_painting,select_drink,select_ingredient,select_billiards,select_poker,add_condiment,insert_flower \
--output /tmp/eval-artifacts/task_descriptions.json
" "
- name: Copy VLABench artifacts from container - name: Copy VLABench artifacts from container
+9 -2
View File
@@ -28,8 +28,15 @@ FROM huggingface/lerobot-gpu:latest
# Patch: constant.py calls os.listdir on ~100 asset/obj/meshes/* dirs at import # Patch: constant.py calls os.listdir on ~100 asset/obj/meshes/* dirs at import
# time. Guard the call so missing dirs return [] instead of crashing (in case # time. Guard the call so missing dirs return [] instead of crashing (in case
# the asset download is partial). # the asset download is partial).
RUN git clone --depth 1 https://github.com/OpenMOSS/VLABench.git ~/VLABench && \ #
git clone --depth 1 https://github.com/motion-planning/rrt-algorithms.git ~/rrt-algorithms && \ # Pinned upstream SHAs for reproducible benchmark runs. Bump when you need
# an upstream fix; don't rely on `main`/`develop` drift.
ARG VLABENCH_SHA=cf588fe60c0c7282174fe979f5913170cfe69017
ARG RRT_ALGORITHMS_SHA=e51d95ee489a225220d6ae2a764c4111f6ba7d85
RUN git clone https://github.com/OpenMOSS/VLABench.git ~/VLABench && \
git -C ~/VLABench checkout ${VLABENCH_SHA} && \
git clone https://github.com/motion-planning/rrt-algorithms.git ~/rrt-algorithms && \
git -C ~/rrt-algorithms checkout ${RRT_ALGORITHMS_SHA} && \
python3 -c "\ python3 -c "\
import pathlib; \ import pathlib; \
p = pathlib.Path.home() / 'VLABench/VLABench/configs/constant.py'; \ p = pathlib.Path.home() / 'VLABench/VLABench/configs/constant.py'; \
+17
View File
@@ -57,6 +57,21 @@ def _metaworld_descriptions(task_name: str) -> dict[str, str]:
return {f"{task_name}_0": label} return {f"{task_name}_0": label}
def _vlabench_descriptions(task_spec: str) -> dict[str, str]:
"""For each task in the comma-separated list, emit a cleaned-name label.
VLABench tasks carry language instructions on their dm_control task
object, but pulling them requires loading the full env per task
(~seconds each). The CI smoke-eval already captures the instruction
inside its episode info; this mapping is just enough to key
`metrics.json` by `<task>_0`.
"""
out: dict[str, str] = {}
for task in (t.strip() for t in task_spec.split(",") if t.strip()):
out[f"{task}_0"] = task.replace("_", " ").strip()
return out
def main() -> int: def main() -> int:
parser = argparse.ArgumentParser(description=__doc__) parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--env", required=True, help="Environment family (libero, metaworld, ...)") parser.add_argument("--env", required=True, help="Environment family (libero, metaworld, ...)")
@@ -70,6 +85,8 @@ def main() -> int:
descriptions = _libero_descriptions(args.task) descriptions = _libero_descriptions(args.task)
elif args.env == "metaworld": elif args.env == "metaworld":
descriptions = _metaworld_descriptions(args.task) descriptions = _metaworld_descriptions(args.task)
elif args.env == "vlabench":
descriptions = _vlabench_descriptions(args.task)
else: else:
print( print(
f"[extract_task_descriptions] No description extractor for env '{args.env}'.", f"[extract_task_descriptions] No description extractor for env '{args.env}'.",
+27 -44
View File
@@ -26,9 +26,9 @@ with long-horizon reasoning, built on MuJoCo/dm_control.
from __future__ import annotations from __future__ import annotations
import contextlib import contextlib
import logging
from collections import defaultdict from collections import defaultdict
from collections.abc import Callable, Sequence from collections.abc import Callable, Sequence
from functools import partial
from typing import Any from typing import Any
import cv2 import cv2
@@ -40,6 +40,8 @@ from lerobot.types import RobotObservation
from .utils import _LazyAsyncVectorEnv from .utils import _LazyAsyncVectorEnv
logger = logging.getLogger(__name__)
ACTION_DIM = 7 # pos(3) + euler(3) + gripper(1) ACTION_DIM = 7 # pos(3) + euler(3) + gripper(1)
ACTION_LOW = -1.0 ACTION_LOW = -1.0
ACTION_HIGH = 1.0 ACTION_HIGH = 1.0
@@ -122,8 +124,6 @@ class VLABenchEnv(gym.Env):
robot: str = "franka", robot: str = "franka",
max_episode_steps: int = DEFAULT_MAX_EPISODE_STEPS, max_episode_steps: int = DEFAULT_MAX_EPISODE_STEPS,
action_mode: str = "eef", action_mode: str = "eef",
episode_index: int = 0,
n_envs: int = 1,
): ):
super().__init__() super().__init__()
self.task = task self.task = task
@@ -133,8 +133,6 @@ class VLABenchEnv(gym.Env):
self.robot = robot self.robot = robot
self._max_episode_steps = max_episode_steps self._max_episode_steps = max_episode_steps
self.action_mode = action_mode self.action_mode = action_mode
self.episode_index = episode_index
self.n_envs = n_envs
# Deferred — created on first reset() inside worker subprocess to avoid # Deferred — created on first reset() inside worker subprocess to avoid
# inheriting stale GPU/EGL contexts when AsyncVectorEnv spawns workers. # inheriting stale GPU/EGL contexts when AsyncVectorEnv spawns workers.
@@ -226,11 +224,12 @@ class VLABenchEnv(gym.Env):
break break
except PhysicsError as exc: except PhysicsError as exc:
last_exc = exc last_exc = exc
print( logger.warning(
f"[vlabench] PhysicsError on attempt {attempt}/" "PhysicsError on attempt %d/%d while building task '%s': %s. Retrying with fresh layout…",
f"{self._ENSURE_ENV_MAX_ATTEMPTS} while building task " attempt,
f"'{self.task}': {exc}. Retrying with fresh layout…", self._ENSURE_ENV_MAX_ATTEMPTS,
flush=True, self.task,
exc,
) )
np.random.seed(None) np.random.seed(None)
if self._env is None: if self._env is None:
@@ -453,7 +452,11 @@ class VLABenchEnv(gym.Env):
# Physics integrator diverged (e.g. mjWARN_BADQACC). Treat it as # Physics integrator diverged (e.g. mjWARN_BADQACC). Treat it as
# a graceful failed termination rather than a hard crash — the # a graceful failed termination rather than a hard crash — the
# rest of the multi-task eval should still run. # rest of the multi-task eval should still run.
print(f"[vlabench] PhysicsError during step on task '{self.task}': {exc}. Terminating episode.") logger.warning(
"PhysicsError during step on task '%s': %s. Terminating episode.",
self.task,
exc,
)
observation = self._get_obs() observation = self._get_obs()
info = {"task": self.task, "is_success": False, "physics_error": True} info = {"task": self.task, "is_success": False, "physics_error": True}
# Drop the stale env so the next reset() rebuilds it cleanly. # Drop the stale env so the next reset() rebuilds it cleanly.
@@ -495,31 +498,6 @@ class VLABenchEnv(gym.Env):
self._env = None self._env = None
# ---- Factory helpers ---------------------------------------------------------
def _make_env_fns(
*,
task: str,
n_envs: int,
gym_kwargs: dict[str, Any],
) -> list[Callable[[], VLABenchEnv]]:
"""Build n_envs factory callables for a single task."""
def _make_env(episode_index: int, **kwargs) -> VLABenchEnv:
return VLABenchEnv(
task=task,
episode_index=episode_index,
n_envs=n_envs,
**kwargs,
)
fns: list[Callable[[], VLABenchEnv]] = []
for episode_index in range(n_envs):
fns.append(partial(_make_env, episode_index, **gym_kwargs))
return fns
# ---- Main API ---------------------------------------------------------------- # ---- Main API ----------------------------------------------------------------
@@ -536,7 +514,7 @@ def create_vlabench_envs(
dict[suite_name][task_id] -> vec_env (env_cls([...]) with exactly n_envs factories) dict[suite_name][task_id] -> vec_env (env_cls([...]) with exactly n_envs factories)
Notes: Notes:
- n_envs is the number of rollouts *per task* (episode_index = 0..n_envs-1). - n_envs is the number of rollouts *per task*.
- `task` can be a suite name ("primitive", "composite"), a comma-separated list of - `task` can be a suite name ("primitive", "composite"), a comma-separated list of
suite names, or individual task names (e.g. "select_fruit,heat_food"). suite names, or individual task names (e.g. "select_fruit,heat_food").
""" """
@@ -550,7 +528,11 @@ def create_vlabench_envs(
if not task_groups: if not task_groups:
raise ValueError("`task` must contain at least one VLABench task or suite name.") raise ValueError("`task` must contain at least one VLABench task or suite name.")
print(f"Creating VLABench envs | task_groups={task_groups} | n_envs(per task)={n_envs}") logger.info(
"Creating VLABench envs | task_groups=%s | n_envs(per task)=%d",
task_groups,
n_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
@@ -563,14 +545,15 @@ def create_vlabench_envs(
tasks = SUITE_TASKS.get(group, [group]) tasks = SUITE_TASKS.get(group, [group])
for tid, task_name in enumerate(tasks): for tid, task_name in enumerate(tasks):
print(f"Building vec env | group={group} | task_id={tid} | task={task_name}") logger.info(
"Building vec env | group=%s | task_id=%d | task=%s",
fns = _make_env_fns( group,
task=task_name, tid,
n_envs=n_envs, task_name,
gym_kwargs=gym_kwargs,
) )
fns = [(lambda tn=task_name: VLABenchEnv(task=tn, **gym_kwargs)) for _ in range(n_envs)]
if is_async: if is_async:
lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata) lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata)
if cached_obs_space is None: if cached_obs_space is None: