mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-26 19:26:16 +00:00
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) <noreply@anthropic.com>
This commit is contained in:
@@ -310,3 +310,93 @@ jobs:
|
|||||||
name: metaworld-metrics
|
name: metaworld-metrics
|
||||||
path: /tmp/metaworld-artifacts/metrics.json
|
path: /tmp/metaworld-artifacts/metrics.json
|
||||||
if-no-files-found: warn
|
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
|
||||||
|
|||||||
@@ -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"]
|
||||||
@@ -81,6 +81,8 @@
|
|||||||
title: Meta-World
|
title: Meta-World
|
||||||
- local: envhub_isaaclab_arena
|
- local: envhub_isaaclab_arena
|
||||||
title: NVIDIA IsaacLab Arena Environments
|
title: NVIDIA IsaacLab Arena Environments
|
||||||
|
- local: vlabench
|
||||||
|
title: VLABench
|
||||||
title: "Benchmarks"
|
title: "Benchmarks"
|
||||||
- sections:
|
- sections:
|
||||||
- local: introduction_processors
|
- local: introduction_processors
|
||||||
|
|||||||
@@ -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
|
||||||
|
```
|
||||||
|
|
||||||
|
<Tip>
|
||||||
|
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)
|
||||||
|
```
|
||||||
|
|
||||||
|
</Tip>
|
||||||
|
|
||||||
|
## 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
|
||||||
|
```
|
||||||
@@ -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
|
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]"]
|
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]"]
|
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
|
||||||
all = [
|
all = [
|
||||||
@@ -247,6 +248,7 @@ all = [
|
|||||||
"lerobot[phone]",
|
"lerobot[phone]",
|
||||||
"lerobot[libero]; sys_platform == 'linux'",
|
"lerobot[libero]; sys_platform == 'linux'",
|
||||||
"lerobot[metaworld]",
|
"lerobot[metaworld]",
|
||||||
|
"lerobot[vlabench]; sys_platform == 'linux'",
|
||||||
"lerobot[sarm]",
|
"lerobot[sarm]",
|
||||||
"lerobot[peft]",
|
"lerobot[peft]",
|
||||||
# "lerobot[unitree_g1]", TODO: Unitree requires specific installation instructions for unitree_sdk2
|
# "lerobot[unitree_g1]", TODO: Unitree requires specific installation instructions for unitree_sdk2
|
||||||
|
|||||||
@@ -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")
|
@EnvConfig.register_subclass("isaaclab_arena")
|
||||||
@dataclass
|
@dataclass
|
||||||
class IsaaclabArenaEnv(HubEnvConfig):
|
class IsaaclabArenaEnv(HubEnvConfig):
|
||||||
|
|||||||
@@ -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()}
|
||||||
@@ -40,7 +40,7 @@ from .converters import (
|
|||||||
)
|
)
|
||||||
from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
|
from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
|
||||||
from .device_processor import DeviceProcessorStep
|
from .device_processor import DeviceProcessorStep
|
||||||
from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep
|
from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep, VLABenchProcessorStep
|
||||||
from .factory import (
|
from .factory import (
|
||||||
make_default_processors,
|
make_default_processors,
|
||||||
make_default_robot_action_processor,
|
make_default_robot_action_processor,
|
||||||
@@ -165,4 +165,5 @@ __all__ = [
|
|||||||
"to_relative_actions",
|
"to_relative_actions",
|
||||||
"UnnormalizerProcessorStep",
|
"UnnormalizerProcessorStep",
|
||||||
"VanillaObservationProcessorStep",
|
"VanillaObservationProcessorStep",
|
||||||
|
"VLABenchProcessorStep",
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -153,6 +153,35 @@ class LiberoProcessorStep(ObservationProcessorStep):
|
|||||||
return result
|
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
|
@dataclass
|
||||||
@ProcessorStepRegistry.register(name="isaaclab_arena_processor")
|
@ProcessorStepRegistry.register(name="isaaclab_arena_processor")
|
||||||
class IsaaclabArenaProcessorStep(ObservationProcessorStep):
|
class IsaaclabArenaProcessorStep(ObservationProcessorStep):
|
||||||
|
|||||||
Reference in New Issue
Block a user