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 \
'--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
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
+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
# time. Guard the call so missing dirs return [] instead of crashing (in case
# 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 "\
import pathlib; \
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}
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:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--env", required=True, help="Environment family (libero, metaworld, ...)")
@@ -70,6 +85,8 @@ def main() -> int:
descriptions = _libero_descriptions(args.task)
elif args.env == "metaworld":
descriptions = _metaworld_descriptions(args.task)
elif args.env == "vlabench":
descriptions = _vlabench_descriptions(args.task)
else:
print(
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
import contextlib
import logging
from collections import defaultdict
from collections.abc import Callable, Sequence
from functools import partial
from typing import Any
import cv2
@@ -40,6 +40,8 @@ from lerobot.types import RobotObservation
from .utils import _LazyAsyncVectorEnv
logger = logging.getLogger(__name__)
ACTION_DIM = 7 # pos(3) + euler(3) + gripper(1)
ACTION_LOW = -1.0
ACTION_HIGH = 1.0
@@ -122,8 +124,6 @@ class VLABenchEnv(gym.Env):
robot: str = "franka",
max_episode_steps: int = DEFAULT_MAX_EPISODE_STEPS,
action_mode: str = "eef",
episode_index: int = 0,
n_envs: int = 1,
):
super().__init__()
self.task = task
@@ -133,8 +133,6 @@ class VLABenchEnv(gym.Env):
self.robot = robot
self._max_episode_steps = max_episode_steps
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
# inheriting stale GPU/EGL contexts when AsyncVectorEnv spawns workers.
@@ -226,11 +224,12 @@ class VLABenchEnv(gym.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…",
flush=True,
logger.warning(
"PhysicsError on attempt %d/%d while building task '%s': %s. Retrying with fresh layout…",
attempt,
self._ENSURE_ENV_MAX_ATTEMPTS,
self.task,
exc,
)
np.random.seed(None)
if self._env is None:
@@ -453,7 +452,11 @@ class VLABenchEnv(gym.Env):
# 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.")
logger.warning(
"PhysicsError during step on task '%s': %s. Terminating episode.",
self.task,
exc,
)
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.
@@ -495,31 +498,6 @@ class VLABenchEnv(gym.Env):
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 ----------------------------------------------------------------
@@ -536,7 +514,7 @@ def create_vlabench_envs(
dict[suite_name][task_id] -> vec_env (env_cls([...]) with exactly n_envs factories)
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
suite names, or individual task names (e.g. "select_fruit,heat_food").
"""
@@ -550,7 +528,11 @@ def create_vlabench_envs(
if not task_groups:
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
cached_obs_space = None
@@ -563,14 +545,15 @@ def create_vlabench_envs(
tasks = SUITE_TASKS.get(group, [group])
for tid, task_name in enumerate(tasks):
print(f"Building vec env | group={group} | task_id={tid} | task={task_name}")
fns = _make_env_fns(
task=task_name,
n_envs=n_envs,
gym_kwargs=gym_kwargs,
logger.info(
"Building vec env | group=%s | task_id=%d | task=%s",
group,
tid,
task_name,
)
fns = [(lambda tn=task_name: VLABenchEnv(task=tn, **gym_kwargs)) for _ in range(n_envs)]
if is_async:
lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata)
if cached_obs_space is None: