From 0a3195749ed612ad4ff3b3b91b910882f590c837 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Tue, 14 Apr 2026 11:28:17 +0200 Subject: [PATCH] feat(sim): VLABench benchmark integration Add VLABench (language-conditioned manipulation with long-horizon reasoning) as a new simulation benchmark, following the established LIBERO/MetaWorld patterns. Co-Authored-By: Claude Opus 4.6 (1M context) --- .github/workflows/benchmark_tests.yml | 90 ++++++ docker/Dockerfile.benchmark.vlabench | 30 ++ docs/source/_toctree.yml | 2 + docs/source/vlabench.mdx | 162 ++++++++++ pyproject.toml | 2 + src/lerobot/envs/configs.py | 73 +++++ src/lerobot/envs/vlabench.py | 405 +++++++++++++++++++++++++ src/lerobot/processor/__init__.py | 3 +- src/lerobot/processor/env_processor.py | 29 ++ 9 files changed, 795 insertions(+), 1 deletion(-) create mode 100644 docker/Dockerfile.benchmark.vlabench create mode 100644 docs/source/vlabench.mdx create mode 100644 src/lerobot/envs/vlabench.py diff --git a/.github/workflows/benchmark_tests.yml b/.github/workflows/benchmark_tests.yml index 79d5614b2..aaeb014f9 100644 --- a/.github/workflows/benchmark_tests.yml +++ b/.github/workflows/benchmark_tests.yml @@ -310,3 +310,93 @@ jobs: name: metaworld-metrics path: /tmp/metaworld-artifacts/metrics.json if-no-files-found: warn + + # ── VLABENCH ───────────────────────────────────────────────────────────── + # Isolated image: lerobot[vlabench] only (VLABench, mujoco==3.2.2, dm-control chain) + vlabench-integration-test: + name: VLABench — build image + 1-episode eval + runs-on: + group: aws-g6-4xlarge-plus + env: + HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }} + + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + lfs: true + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 # zizmor: ignore[unpinned-uses] + with: + cache-binary: false + + - name: Login to Docker Hub + uses: docker/login-action@v3 # zizmor: ignore[unpinned-uses] + with: + username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }} + password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }} + + - name: Build VLABench benchmark image + uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses] + with: + context: . + file: docker/Dockerfile.benchmark.vlabench + push: false + load: true + tags: lerobot-benchmark-vlabench:ci + + - name: Run VLABench smoke eval (1 episode) + if: env.HF_USER_TOKEN != '' + run: | + docker run --name vlabench-eval --gpus all \ + --shm-size=4g \ + -e HF_HOME=/tmp/hf \ + -e HF_USER_TOKEN="${HF_USER_TOKEN}" \ + -e HF_HUB_DOWNLOAD_TIMEOUT=300 \ + -e MUJOCO_GL=egl \ + lerobot-benchmark-vlabench:ci \ + bash -c " + hf auth login --token \"\$HF_USER_TOKEN\" --add-to-git-credential 2>/dev/null || true + lerobot-eval \ + --policy.path=your-vlabench-policy \ + --env.type=vlabench \ + --env.task=select_fruit \ + --eval.batch_size=1 \ + --eval.n_episodes=1 \ + --eval.use_async_envs=false \ + --policy.device=cuda \ + --output_dir=/tmp/eval-artifacts + " + + - name: Copy VLABench artifacts from container + if: always() + run: | + mkdir -p /tmp/vlabench-artifacts + docker cp vlabench-eval:/tmp/eval-artifacts/. /tmp/vlabench-artifacts/ 2>/dev/null || true + docker rm -f vlabench-eval || true + + - name: Parse VLABench eval metrics + if: always() + run: | + python3 scripts/ci/parse_eval_metrics.py \ + --artifacts-dir /tmp/vlabench-artifacts \ + --env vlabench \ + --task select_fruit \ + --policy your-vlabench-policy + + - name: Upload VLABench rollout video + if: always() + uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses] + with: + name: vlabench-rollout-video + path: /tmp/vlabench-artifacts/videos/ + if-no-files-found: warn + + - name: Upload VLABench eval metrics + if: always() + uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses] + with: + name: vlabench-metrics + path: /tmp/vlabench-artifacts/metrics.json + if-no-files-found: warn diff --git a/docker/Dockerfile.benchmark.vlabench b/docker/Dockerfile.benchmark.vlabench new file mode 100644 index 000000000..e93f6a956 --- /dev/null +++ b/docker/Dockerfile.benchmark.vlabench @@ -0,0 +1,30 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Benchmark image for VLABench integration tests. +# Extends the nightly GPU image with the PR's source code and VLABench setup. +# +# Build: docker build -f docker/Dockerfile.benchmark.vlabench -t lerobot-benchmark-vlabench . +# Run: docker run --gpus all --rm lerobot-benchmark-vlabench lerobot-eval ... + +FROM huggingface/lerobot-gpu:latest + +# Install VLABench and download simulation assets. +RUN pip install --no-cache-dir vlabench mujoco==3.2.2 dm-control==1.0.22 && \ + python -c "from VLABench.utils import download_assets; download_assets()" 2>/dev/null || true + +# Overlay the PR's source code on top of the nightly image. +COPY --chown=user_lerobot:user_lerobot . . + +CMD ["/bin/bash"] diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 3dcba5993..87fbe5876 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -81,6 +81,8 @@ title: Meta-World - local: envhub_isaaclab_arena title: NVIDIA IsaacLab Arena Environments + - local: vlabench + title: VLABench title: "Benchmarks" - sections: - local: introduction_processors diff --git a/docs/source/vlabench.mdx b/docs/source/vlabench.mdx new file mode 100644 index 000000000..20107276f --- /dev/null +++ b/docs/source/vlabench.mdx @@ -0,0 +1,162 @@ +# VLABench + +VLABench is a large-scale benchmark for **language-conditioned robotic manipulation with long-horizon reasoning**. It provides 100 task categories across 2000+ objects, evaluating six dimensions of robot intelligence: mesh & texture understanding, spatial reasoning, world knowledge transfer, semantic instruction comprehension, physical law understanding, and long-horizon reasoning. VLABench is built on MuJoCo/dm_control and uses a Franka Panda 7-DOF arm. + +- Paper: [VLABench: A Large-Scale Benchmark for Language-Conditioned Robotics Manipulation with Long-Horizon Reasoning](https://arxiv.org/abs/2412.18194) +- GitHub: [OpenMOSS/VLABench](https://github.com/OpenMOSS/VLABench) +- Project website: [vlabench.github.io](https://vlabench.github.io) + +## Available tasks + +VLABench includes **two task suites** covering **100 task categories**: + +| Suite | CLI name | Tasks | Description | +| --------- | ----------- | ----- | ---------------------------------------------------------------- | +| Primitive | `primitive` | 21 | Single/few skill combinations (select, insert, physics QA) | +| Composite | `composite` | 22 | Multi-step reasoning and long-horizon planning (cook, rearrange) | + +### Primitive tasks + +Includes `select_fruit`, `select_toy`, `select_chemistry_tube`, `add_condiment`, `select_book`, `select_painting`, `select_drink`, `insert_flower`, `select_billiards`, `select_ingredient`, `select_mahjong`, `select_poker`, and physical reasoning tasks (`density_qa`, `friction_qa`, `magnetism_qa`, `reflection_qa`, `simple_cuestick_usage`, `simple_seesaw_usage`, `sound_speed_qa`, `thermal_expansion_qa`, `weight_qa`). + +### Composite tasks + +Includes `cluster_billiards`, `cluster_book`, `cluster_drink`, `cluster_toy`, `cook_dishes`, `cool_drink`, `find_unseen_object`, `get_coffee`, `hammer_nail`, `heat_food`, `make_juice`, `play_mahjong`, `play_math_game`, `play_poker`, `play_snooker`, `rearrange_book`, `rearrange_chemistry_tube`, `set_dining_table`, `set_study_table`, `store_food`, `take_chemistry_experiment`, `use_seesaw_complex`. + +### Evaluation tracks + +VLABench defines five standard evaluation tracks: + +| Track | Focus | +| ----- | ----------------------------- | +| 1 | In-distribution task learning | +| 2 | Cross-category generalization | +| 3 | Commonsense reasoning | +| 4 | Semantic instruction | +| 6 | Unseen texture robustness | + +## Installation + +After following the LeRobot installation instructions: + +```bash +pip install -e ".[vlabench]" +``` + +VLABench also requires downloading simulation assets: + +```bash +# Clone the VLABench repo and download assets +git clone https://github.com/OpenMOSS/VLABench.git +cd VLABench +python scripts/download_assets.py +``` + + +VLABench requires Linux (`sys_platform == 'linux'`) and Python 3.10+. Set the MuJoCo rendering backend before training or evaluation: + +```bash +export MUJOCO_GL=egl # for headless servers (HPC, cloud) +``` + + + +## Evaluation + +### Default evaluation (recommended) + +Evaluate on a single task (10 episodes): + +```bash +lerobot-eval \ + --policy.path="your-policy-id" \ + --env.type=vlabench \ + --env.task=select_fruit \ + --eval.batch_size=1 \ + --eval.n_episodes=10 +``` + +### Suite-wide evaluation + +Evaluate across all primitive tasks: + +```bash +lerobot-eval \ + --policy.path="your-policy-id" \ + --env.type=vlabench \ + --env.task=primitive \ + --eval.batch_size=1 \ + --eval.n_episodes=10 \ + --env.max_parallel_tasks=1 +``` + +### Multi-suite evaluation + +Evaluate across both primitive and composite suites: + +```bash +lerobot-eval \ + --policy.path="your-policy-id" \ + --env.type=vlabench \ + --env.task=primitive,composite \ + --eval.batch_size=1 \ + --eval.n_episodes=10 \ + --env.max_parallel_tasks=1 +``` + +### Individual task evaluation + +Evaluate on specific tasks by name: + +```bash +lerobot-eval \ + --policy.path="your-policy-id" \ + --env.type=vlabench \ + --env.task=select_fruit,heat_food \ + --eval.batch_size=1 \ + --eval.n_episodes=10 +``` + +## Policy inputs and outputs + +**Observations:** + +- `observation.state` — 7-dim end-effector state (position xyz, euler xyz, gripper) +- `observation.images.image` — front camera view, 480x480 HWC uint8 +- `observation.images.second_image` — second camera view, 480x480 HWC uint8 +- `observation.images.wrist_image` — wrist camera view, 480x480 HWC uint8 + +**Actions:** + +- Continuous control in `Box(-1, 1, shape=(7,))` — 3D position + 3D euler orientation + 1D gripper + +### Recommended evaluation episodes + +For reproducible benchmarking, use **10 episodes per task**. For the full primitive suite this gives 210 episodes; for the full composite suite, 220 episodes. + +## Training + +### Datasets + +Pre-collected VLABench datasets in LeRobot format are available on the Hugging Face Hub: + +- Primitive tasks: [VLABench/vlabench_primitive_ft_lerobot_video](https://huggingface.co/datasets/VLABench/vlabench_primitive_ft_lerobot_video) (5,000 episodes, 128 tasks, 480x480 images) +- Composite tasks: [VLABench/vlabench_composite_ft_lerobot_video](https://huggingface.co/datasets/VLABench/vlabench_composite_ft_lerobot_video) (5,977 episodes, 167 tasks, 224x224 images) + +### Example training command + +```bash +lerobot-train \ + --policy.type=smolvla \ + --policy.repo_id=${HF_USER}/vlabench-test \ + --policy.load_vlm_weights=true \ + --dataset.repo_id=VLABench/vlabench_primitive_ft_lerobot_video \ + --env.type=vlabench \ + --env.task=select_fruit \ + --output_dir=./outputs/ \ + --steps=100000 \ + --batch_size=4 \ + --eval.batch_size=1 \ + --eval.n_episodes=1 \ + --eval_freq=1000 +``` diff --git a/pyproject.toml b/pyproject.toml index 2f12840e7..10206f55a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -206,6 +206,7 @@ aloha = ["lerobot[dataset]", "gym-aloha>=0.1.2,<0.2.0", "lerobot[scipy-dep]"] pusht = ["lerobot[dataset]", "gym-pusht>=0.1.5,<0.2.0", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead libero = ["lerobot[dataset]", "lerobot[transformers-dep]", "hf-libero>=0.1.3,<0.2.0; sys_platform == 'linux'", "lerobot[scipy-dep]"] metaworld = ["lerobot[dataset]", "metaworld==3.0.0", "lerobot[scipy-dep]"] +vlabench = ["lerobot[dataset]", "vlabench>=0.1.0; sys_platform == 'linux'", "mujoco==3.2.2", "dm-control==1.0.22", "lerobot[scipy-dep]"] # All all = [ @@ -247,6 +248,7 @@ all = [ "lerobot[phone]", "lerobot[libero]; sys_platform == 'linux'", "lerobot[metaworld]", + "lerobot[vlabench]; sys_platform == 'linux'", "lerobot[sarm]", "lerobot[peft]", # "lerobot[unitree_g1]", TODO: Unitree requires specific installation instructions for unitree_sdk2 diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index 2a7c52d45..044d94bb9 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -496,6 +496,79 @@ class MetaworldEnv(EnvConfig): ) +@EnvConfig.register_subclass("vlabench") +@dataclass +class VLABenchEnv(EnvConfig): + task: str = "select_fruit" + fps: int = 10 + episode_length: int = 500 + obs_type: str = "pixels_agent_pos" + render_mode: str = "rgb_array" + render_resolution: tuple[int, int] = (480, 480) + robot: str = "franka" + action_mode: str = "eef" + features: dict[str, PolicyFeature] = field( + default_factory=lambda: { + ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,)), + } + ) + features_map: dict[str, str] = field( + default_factory=lambda: { + ACTION: ACTION, + "agent_pos": OBS_STATE, + "pixels/image": f"{OBS_IMAGES}.image", + "pixels/second_image": f"{OBS_IMAGES}.second_image", + "pixels/wrist_image": f"{OBS_IMAGES}.wrist_image", + } + ) + + def __post_init__(self): + h, w = self.render_resolution + if self.obs_type == "pixels": + self.features["pixels/image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(h, w, 3)) + self.features["pixels/second_image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(h, w, 3)) + self.features["pixels/wrist_image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(h, w, 3)) + elif self.obs_type == "pixels_agent_pos": + self.features["pixels/image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(h, w, 3)) + self.features["pixels/second_image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(h, w, 3)) + self.features["pixels/wrist_image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(h, w, 3)) + self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(7,)) + else: + raise ValueError(f"Unsupported obs_type: {self.obs_type}") + + @property + def gym_kwargs(self) -> dict: + return { + "obs_type": self.obs_type, + "render_mode": self.render_mode, + "render_resolution": self.render_resolution, + "robot": self.robot, + "max_episode_steps": self.episode_length, + "action_mode": self.action_mode, + } + + def create_envs(self, n_envs: int, use_async_envs: bool = False): + from .vlabench import create_vlabench_envs + + if self.task is None: + raise ValueError("VLABenchEnv requires a task to be specified") + env_cls = _make_vec_env_cls(use_async_envs, n_envs) + return create_vlabench_envs( + task=self.task, + n_envs=n_envs, + gym_kwargs=self.gym_kwargs, + env_cls=env_cls, + ) + + def get_env_processors(self): + from lerobot.processor.env_processor import VLABenchProcessorStep + + return ( + PolicyProcessorPipeline(steps=[VLABenchProcessorStep()]), + PolicyProcessorPipeline(steps=[]), + ) + + @EnvConfig.register_subclass("isaaclab_arena") @dataclass class IsaaclabArenaEnv(HubEnvConfig): diff --git a/src/lerobot/envs/vlabench.py b/src/lerobot/envs/vlabench.py new file mode 100644 index 000000000..edd93f76d --- /dev/null +++ b/src/lerobot/envs/vlabench.py @@ -0,0 +1,405 @@ +#!/usr/bin/env python + +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""VLABench environment wrapper for LeRobot. + +VLABench is a large-scale benchmark for language-conditioned robotic manipulation +with long-horizon reasoning, built on MuJoCo/dm_control. + +- Paper: https://arxiv.org/abs/2412.18194 +- GitHub: https://github.com/OpenMOSS/VLABench +- Website: https://vlabench.github.io +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Callable, Sequence +from functools import partial +from typing import Any + +import gymnasium as gym +import numpy as np +from gymnasium import spaces + +from lerobot.types import RobotObservation + +from .utils import _LazyAsyncVectorEnv + +ACTION_DIM = 7 # pos(3) + euler(3) + gripper(1) +ACTION_LOW = -1.0 +ACTION_HIGH = 1.0 + +# Default max episode steps per task type +DEFAULT_MAX_EPISODE_STEPS = 500 + +# VLABench task suites +PRIMITIVE_TASKS = [ + "select_fruit", + "select_toy", + "select_chemistry_tube", + "add_condiment", + "select_book", + "select_painting", + "select_drink", + "insert_flower", + "select_billiards", + "select_ingredient", + "select_mahjong", + "select_poker", + # Physical series + "density_qa", + "friction_qa", + "magnetism_qa", + "reflection_qa", + "simple_cuestick_usage", + "simple_seesaw_usage", + "sound_speed_qa", + "thermal_expansion_qa", + "weight_qa", +] + +COMPOSITE_TASKS = [ + "cluster_billiards", + "cluster_book", + "cluster_drink", + "cluster_toy", + "cook_dishes", + "cool_drink", + "find_unseen_object", + "get_coffee", + "hammer_nail", + "heat_food", + "make_juice", + "play_mahjong", + "play_math_game", + "play_poker", + "play_snooker", + "rearrange_book", + "rearrange_chemistry_tube", + "set_dining_table", + "set_study_table", + "store_food", + "take_chemistry_experiment", + "use_seesaw_complex", +] + +SUITE_TASKS: dict[str, list[str]] = { + "primitive": PRIMITIVE_TASKS, + "composite": COMPOSITE_TASKS, +} + + +class VLABenchEnv(gym.Env): + """Gymnasium wrapper for VLABench environments. + + Wraps the dm_control-based VLABench simulator behind a standard gym.Env interface. + Supports multiple cameras (front, second, wrist) and end-effector control. + """ + + metadata = {"render_modes": ["rgb_array"], "render_fps": 10} + + def __init__( + self, + task: str = "select_fruit", + obs_type: str = "pixels_agent_pos", + render_mode: str = "rgb_array", + render_resolution: tuple[int, int] = (480, 480), + 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 + self.obs_type = obs_type + self.render_mode = render_mode + self.render_resolution = render_resolution + 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. + self._env = None + self._physics = None + self.task_description = "" # populated on first reset + + h, w = self.render_resolution + + if self.obs_type == "state": + raise NotImplementedError( + "The 'state' observation type is not supported in VLABenchEnv. " + "Please use 'pixels' or 'pixels_agent_pos'." + ) + elif self.obs_type == "pixels": + self.observation_space = spaces.Dict( + { + "pixels": spaces.Dict( + { + "image": spaces.Box(low=0, high=255, shape=(h, w, 3), dtype=np.uint8), + "second_image": spaces.Box(low=0, high=255, shape=(h, w, 3), dtype=np.uint8), + "wrist_image": spaces.Box(low=0, high=255, shape=(h, w, 3), dtype=np.uint8), + } + ), + } + ) + elif self.obs_type == "pixels_agent_pos": + self.observation_space = spaces.Dict( + { + "pixels": spaces.Dict( + { + "image": spaces.Box(low=0, high=255, shape=(h, w, 3), dtype=np.uint8), + "second_image": spaces.Box(low=0, high=255, shape=(h, w, 3), dtype=np.uint8), + "wrist_image": spaces.Box(low=0, high=255, shape=(h, w, 3), dtype=np.uint8), + } + ), + "agent_pos": spaces.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float64), + } + ) + else: + raise ValueError(f"Unsupported obs_type: {self.obs_type}") + + self.action_space = spaces.Box( + low=ACTION_LOW, high=ACTION_HIGH, shape=(ACTION_DIM,), dtype=np.float32 + ) + + 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). + """ + if self._env is not None: + return + + 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 + self._physics = env.physics + + # Extract task description from the dm_control task + task_obj = env.task + if hasattr(task_obj, "task_description"): + self.task_description = task_obj.task_description + elif hasattr(task_obj, "language_instruction"): + self.task_description = task_obj.language_instruction + else: + self.task_description = self.task + + def _get_obs(self) -> dict: + """Get current observation from the environment.""" + assert self._env is not None + + obs = self._env.get_observation() + + # Extract camera images — VLABench returns (n_cameras, C, H, W) or individual arrays + images = {} + if "rgb" in obs: + rgb = obs["rgb"] + if isinstance(rgb, np.ndarray): + if rgb.ndim == 4: + # (n_cameras, C, H, W) → transpose each to (H, W, C) + if rgb.shape[0] >= 1: + images["image"] = np.transpose(rgb[0], (1, 2, 0)).astype(np.uint8) + if rgb.shape[0] >= 2: + images["second_image"] = np.transpose(rgb[1], (1, 2, 0)).astype(np.uint8) + if rgb.shape[0] >= 3: + images["wrist_image"] = np.transpose(rgb[2], (1, 2, 0)).astype(np.uint8) + elif rgb.ndim == 3: + # Single camera (C, H, W) or (H, W, C) + if rgb.shape[0] == 3: # CHW + images["image"] = np.transpose(rgb, (1, 2, 0)).astype(np.uint8) + else: # HWC + images["image"] = rgb.astype(np.uint8) + + # Fill missing cameras with zeros + h, w = self.render_resolution + for key in ["image", "second_image", "wrist_image"]: + if key not in images: + images[key] = np.zeros((h, w, 3), dtype=np.uint8) + + # Extract end-effector state + ee_state = obs.get("ee_state", np.zeros(7, dtype=np.float64)) + + if self.obs_type == "pixels": + return {"pixels": images} + elif self.obs_type == "pixels_agent_pos": + return { + "pixels": images, + "agent_pos": ee_state.astype(np.float64), + } + else: + raise ValueError(f"Unknown obs_type: {self.obs_type}") + + def reset(self, seed=None, **kwargs) -> tuple[RobotObservation, dict[str, Any]]: + self._ensure_env() + assert self._env is not None + super().reset(seed=seed) + + self._env.reset() + + observation = self._get_obs() + info = {"is_success": False} + return observation, info + + def step(self, action: np.ndarray) -> tuple[RobotObservation, float, bool, bool, dict[str, Any]]: + self._ensure_env() + assert self._env is not None + + if action.ndim != 1: + raise ValueError( + f"Expected action to be 1-D (shape (action_dim,)), " + f"but got shape {action.shape} with ndim={action.ndim}" + ) + + # VLABench uses dm_control TimeStep — convert action format + if self.action_mode == "eef": + # action: pos(3) + euler(3) + gripper(1) → absolute EEF + timestep = self._env.step(action) + elif self.action_mode == "joint" or self.action_mode == "delta_eef": + timestep = self._env.step(action) + else: + raise ValueError(f"Unknown action_mode: {self.action_mode}") + + # Extract reward from dm_control timestep + reward = float(timestep.reward) if timestep.reward is not None else 0.0 + + # Check success via the task's termination condition + is_success = False + if hasattr(self._env, "task") and hasattr(self._env.task, "should_terminate_episode"): + is_success = bool(self._env.task.should_terminate_episode(self._physics)) + + terminated = is_success + truncated = False + info = { + "task": self.task, + "is_success": is_success, + } + + observation = self._get_obs() + + if terminated: + self.reset() + + return observation, reward, terminated, truncated, info + + def render(self) -> np.ndarray: + self._ensure_env() + obs = self._get_obs() + return obs["pixels"]["image"] + + def close(self): + if self._env is not None: + self._env.close() + self._env = None + self._physics = 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 ---------------------------------------------------------------- + + +def create_vlabench_envs( + task: str, + n_envs: int, + gym_kwargs: dict[str, Any] | None = None, + env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None, +) -> dict[str, dict[int, Any]]: + """ + Create vectorized VLABench environments with a consistent return shape. + + Returns: + 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). + - `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"). + """ + if env_cls is None or not callable(env_cls): + raise ValueError("env_cls must be a callable that wraps a list of environment factory callables.") + if not isinstance(n_envs, int) or n_envs <= 0: + raise ValueError(f"n_envs must be a positive int; got {n_envs}.") + + gym_kwargs = dict(gym_kwargs or {}) + task_groups = [t.strip() for t in task.split(",") if t.strip()] + 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}") + + is_async = env_cls is gym.vector.AsyncVectorEnv + cached_obs_space = None + cached_act_space = None + out: dict[str, dict[int, Any]] = defaultdict(dict) + + for group in task_groups: + # Check if it's a suite name, otherwise treat as individual task + 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, + ) + + if is_async: + lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space) + if cached_obs_space is None: + cached_obs_space = lazy.observation_space + cached_act_space = lazy.action_space + out[group][tid] = lazy + else: + out[group][tid] = env_cls(fns) + + return {group: dict(task_map) for group, task_map in out.items()} diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 3688a4b8c..21acd8db2 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -40,7 +40,7 @@ from .converters import ( ) from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep from .device_processor import DeviceProcessorStep -from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep +from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep, VLABenchProcessorStep from .factory import ( make_default_processors, make_default_robot_action_processor, @@ -165,4 +165,5 @@ __all__ = [ "to_relative_actions", "UnnormalizerProcessorStep", "VanillaObservationProcessorStep", + "VLABenchProcessorStep", ] diff --git a/src/lerobot/processor/env_processor.py b/src/lerobot/processor/env_processor.py index 75cbb79de..3c7645863 100644 --- a/src/lerobot/processor/env_processor.py +++ b/src/lerobot/processor/env_processor.py @@ -153,6 +153,35 @@ class LiberoProcessorStep(ObservationProcessorStep): return result +@dataclass +@ProcessorStepRegistry.register(name="vlabench_processor") +class VLABenchProcessorStep(ObservationProcessorStep): + """ + Processes VLABench observations into the LeRobot format. + + VLABench returns end-effector state as a 7-dim vector: + [pos_x, pos_y, pos_z, euler_x, euler_y, euler_z, gripper]. + This is already in the format expected by the datasets, so we pass it + through as observation.state with no additional conversion. + + Images are passed through without transformation (VLABench cameras + are already oriented correctly). + """ + + def _process_observation(self, observation): + """Processes observations from VLABench — minimal transform needed.""" + return observation.copy() + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + """VLABench state is already in the expected format.""" + return features + + def observation(self, observation): + return self._process_observation(observation) + + @dataclass @ProcessorStepRegistry.register(name="isaaclab_arena_processor") class IsaaclabArenaProcessorStep(ObservationProcessorStep):