mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-26 11:16:00 +00:00
feat(envs): add RoboCasa365 benchmark integration
Add RoboCasa365 (arXiv:2603.04356) as a new simulation benchmark with 365 everyday kitchen manipulation tasks across 2,500 diverse environments. New files: - src/lerobot/envs/robocasa.py: gym.Env wrapper with deferred env creation, flat 12D action / 16D state vectors, 3-camera support - docs/source/robocasa.mdx: user-facing documentation - docker/Dockerfile.benchmark.robocasa: CI benchmark image Modified files: - src/lerobot/envs/configs.py: RoboCasaEnv config (--env.type=robocasa) - pyproject.toml: robocasa optional dependency group - docs/source/_toctree.yml: sidebar entry - .github/workflows/benchmark_tests.yml: integration test job Refs: https://arxiv.org/abs/2603.04356, https://robocasa.ai Related: huggingface/lerobot#321 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
|
||||||
|
|
||||||
|
# ── ROBOCASA365 ──────────────────────────────────────────────────────────
|
||||||
|
# Isolated image: lerobot[robocasa] only (robocasa, robosuite, mujoco chain)
|
||||||
|
robocasa-integration-test:
|
||||||
|
name: RoboCasa365 — 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 RoboCasa365 benchmark image
|
||||||
|
uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses]
|
||||||
|
with:
|
||||||
|
context: .
|
||||||
|
file: docker/Dockerfile.benchmark.robocasa
|
||||||
|
push: false
|
||||||
|
load: true
|
||||||
|
tags: lerobot-benchmark-robocasa:ci
|
||||||
|
|
||||||
|
- name: Run RoboCasa365 smoke eval (1 episode)
|
||||||
|
if: env.HF_USER_TOKEN != ''
|
||||||
|
run: |
|
||||||
|
docker run --name robocasa-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-robocasa:ci \
|
||||||
|
bash -c "
|
||||||
|
hf auth login --token \"\$HF_USER_TOKEN\" --add-to-git-credential 2>/dev/null || true
|
||||||
|
lerobot-eval \
|
||||||
|
--policy.path=pepijn223/smolvla_robocasa \
|
||||||
|
--env.type=robocasa \
|
||||||
|
--env.task=CloseFridge \
|
||||||
|
--eval.batch_size=1 \
|
||||||
|
--eval.n_episodes=1 \
|
||||||
|
--eval.use_async_envs=false \
|
||||||
|
--policy.device=cuda \
|
||||||
|
--output_dir=/tmp/eval-artifacts
|
||||||
|
"
|
||||||
|
|
||||||
|
- name: Copy RoboCasa365 artifacts from container
|
||||||
|
if: always()
|
||||||
|
run: |
|
||||||
|
mkdir -p /tmp/robocasa-artifacts
|
||||||
|
docker cp robocasa-eval:/tmp/eval-artifacts/. /tmp/robocasa-artifacts/ 2>/dev/null || true
|
||||||
|
docker rm -f robocasa-eval || true
|
||||||
|
|
||||||
|
- name: Parse RoboCasa365 eval metrics
|
||||||
|
if: always()
|
||||||
|
run: |
|
||||||
|
python3 scripts/ci/parse_eval_metrics.py \
|
||||||
|
--artifacts-dir /tmp/robocasa-artifacts \
|
||||||
|
--env robocasa \
|
||||||
|
--task CloseFridge \
|
||||||
|
--policy pepijn223/smolvla_robocasa
|
||||||
|
|
||||||
|
- name: Upload RoboCasa365 rollout video
|
||||||
|
if: always()
|
||||||
|
uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
|
||||||
|
with:
|
||||||
|
name: robocasa-rollout-video
|
||||||
|
path: /tmp/robocasa-artifacts/videos/
|
||||||
|
if-no-files-found: warn
|
||||||
|
|
||||||
|
- name: Upload RoboCasa365 eval metrics
|
||||||
|
if: always()
|
||||||
|
uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
|
||||||
|
with:
|
||||||
|
name: robocasa-metrics
|
||||||
|
path: /tmp/robocasa-artifacts/metrics.json
|
||||||
|
if-no-files-found: warn
|
||||||
|
|||||||
@@ -0,0 +1,36 @@
|
|||||||
|
# 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 RoboCasa365 integration tests.
|
||||||
|
# Extends the nightly GPU image (which already has all extras installed)
|
||||||
|
# with the PR's source code and RoboCasa-specific asset setup.
|
||||||
|
#
|
||||||
|
# Build: docker build -f docker/Dockerfile.benchmark.robocasa -t lerobot-benchmark-robocasa .
|
||||||
|
# Run: docker run --gpus all --rm lerobot-benchmark-robocasa lerobot-eval ...
|
||||||
|
|
||||||
|
FROM huggingface/lerobot-gpu:latest
|
||||||
|
|
||||||
|
# Install robocasa and its dependencies
|
||||||
|
RUN pip install --no-cache-dir \
|
||||||
|
"robocasa @ git+https://github.com/robocasa/robocasa.git@v1.0.0" \
|
||||||
|
"robosuite @ git+https://github.com/ARISE-Initiative/robosuite.git"
|
||||||
|
|
||||||
|
# Set up robocasa macros and download kitchen assets
|
||||||
|
RUN python -m robocasa.scripts.setup_macros && \
|
||||||
|
python -m robocasa.scripts.download_kitchen_assets
|
||||||
|
|
||||||
|
# Overlay the PR's source code on top of the nightly image.
|
||||||
|
COPY --chown=user_lerobot:user_lerobot . .
|
||||||
|
|
||||||
|
CMD ["/bin/bash"]
|
||||||
@@ -79,6 +79,8 @@
|
|||||||
title: LIBERO
|
title: LIBERO
|
||||||
- local: metaworld
|
- local: metaworld
|
||||||
title: Meta-World
|
title: Meta-World
|
||||||
|
- local: robocasa
|
||||||
|
title: RoboCasa365
|
||||||
- local: envhub_isaaclab_arena
|
- local: envhub_isaaclab_arena
|
||||||
title: NVIDIA IsaacLab Arena Environments
|
title: NVIDIA IsaacLab Arena Environments
|
||||||
title: "Benchmarks"
|
title: "Benchmarks"
|
||||||
|
|||||||
@@ -0,0 +1,110 @@
|
|||||||
|
# RoboCasa365
|
||||||
|
|
||||||
|
RoboCasa365 is a large-scale simulation framework for training and benchmarking **generalist robots** in everyday kitchen tasks. It provides 365 diverse manipulation tasks across 2,500 kitchen environments, with over 3,200 object assets and 600+ hours of human demonstration data. The benchmark tests whether robots can handle the diversity and complexity of real-world household environments.
|
||||||
|
|
||||||
|
- Paper: [RoboCasa365: A Large-Scale Simulation Framework for Training and Benchmarking Generalist Robots](https://arxiv.org/abs/2603.04356)
|
||||||
|
- GitHub: [robocasa/robocasa](https://github.com/robocasa/robocasa)
|
||||||
|
- Project website: [robocasa.ai](https://robocasa.ai)
|
||||||
|
|
||||||
|
## Available tasks
|
||||||
|
|
||||||
|
RoboCasa365 includes **365 tasks** organized into atomic (single-skill) and composite (multi-step) categories:
|
||||||
|
|
||||||
|
| Category | Tasks | Description |
|
||||||
|
| --------- | ----- | ------------------------------------------------------------------------------- |
|
||||||
|
| Atomic | ~65 | Single-skill tasks: pick-and-place, door/drawer manipulation, appliance control |
|
||||||
|
| Composite | ~300 | Multi-step tasks across 60+ categories: cooking, cleaning, organizing, etc. |
|
||||||
|
|
||||||
|
**Atomic task examples:** `CloseFridge`, `OpenBlenderLid`, `PickPlaceCoffee`, `ManipulateStoveKnob`, `NavigateKitchen`, `TurnOnToaster`
|
||||||
|
|
||||||
|
**Composite task categories:** baking, boiling, brewing, chopping food, clearing table, defrosting food, loading dishwasher, making tea, microwaving food, washing dishes, and many more.
|
||||||
|
|
||||||
|
Pass individual task class names directly as `--env.task` (e.g., `CloseFridge`, `PickPlaceCoffee`). Multiple tasks can be comma-separated for multi-task evaluation.
|
||||||
|
|
||||||
|
## Installation
|
||||||
|
|
||||||
|
After following the LeRobot installation instructions:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
pip install -e ".[robocasa]"
|
||||||
|
python -m robocasa.scripts.setup_macros
|
||||||
|
python -m robocasa.scripts.download_kitchen_assets
|
||||||
|
```
|
||||||
|
|
||||||
|
<Tip>
|
||||||
|
RoboCasa365 requires MuJoCo for simulation. Set the rendering backend before training or evaluation:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export MUJOCO_GL=egl # for headless servers (HPC, cloud)
|
||||||
|
```
|
||||||
|
|
||||||
|
</Tip>
|
||||||
|
|
||||||
|
## Evaluation
|
||||||
|
|
||||||
|
### Single-task evaluation
|
||||||
|
|
||||||
|
Evaluate a policy on a single RoboCasa task:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
lerobot-eval \
|
||||||
|
--policy.path="your-policy-id" \
|
||||||
|
--env.type=robocasa \
|
||||||
|
--env.task=CloseFridge \
|
||||||
|
--eval.batch_size=1 \
|
||||||
|
--eval.n_episodes=20
|
||||||
|
```
|
||||||
|
|
||||||
|
### Multi-task evaluation
|
||||||
|
|
||||||
|
Evaluate across multiple tasks at once by passing a comma-separated list:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
lerobot-eval \
|
||||||
|
--policy.path="your-policy-id" \
|
||||||
|
--env.type=robocasa \
|
||||||
|
--env.task=CloseFridge,OpenBlenderLid,PickPlaceCoffee \
|
||||||
|
--eval.batch_size=1 \
|
||||||
|
--eval.n_episodes=20
|
||||||
|
```
|
||||||
|
|
||||||
|
- `--env.task` accepts individual task names (comma-separated).
|
||||||
|
- `--eval.batch_size` controls how many environments run in parallel.
|
||||||
|
- `--eval.n_episodes` sets how many episodes to run per task.
|
||||||
|
|
||||||
|
### Policy inputs and outputs
|
||||||
|
|
||||||
|
**Observations:**
|
||||||
|
|
||||||
|
- `observation.state` -- 16-dim proprioceptive state (base position, base quaternion, relative end-effector position, relative end-effector quaternion, gripper qpos)
|
||||||
|
- `observation.images.image` -- left agent view (`robot0_agentview_left`), 256x256 HWC uint8
|
||||||
|
- `observation.images.image2` -- wrist camera view (`robot0_eye_in_hand`), 256x256 HWC uint8
|
||||||
|
- `observation.images.image3` -- right agent view (`robot0_agentview_right`), 256x256 HWC uint8
|
||||||
|
|
||||||
|
**Actions:**
|
||||||
|
|
||||||
|
- Continuous control in `Box(-1, 1, shape=(12,))` -- base motion (4D) + control mode (1D) + end-effector position (3D) + end-effector rotation (3D) + gripper (1D)
|
||||||
|
|
||||||
|
### Recommended evaluation episodes
|
||||||
|
|
||||||
|
For reproducible benchmarking, use **20 episodes per task**. This matches the protocol used in published results.
|
||||||
|
|
||||||
|
## Training
|
||||||
|
|
||||||
|
### Example training command
|
||||||
|
|
||||||
|
```bash
|
||||||
|
lerobot-train \
|
||||||
|
--policy.type=smolvla \
|
||||||
|
--policy.repo_id=${HF_USER}/robocasa-test \
|
||||||
|
--policy.load_vlm_weights=true \
|
||||||
|
--dataset.repo_id=your-robocasa-dataset \
|
||||||
|
--env.type=robocasa \
|
||||||
|
--env.task=CloseFridge \
|
||||||
|
--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]"]
|
||||||
|
robocasa = ["lerobot[dataset]", "robocasa @ git+https://github.com/robocasa/robocasa.git@v1.0.0", "robosuite @ git+https://github.com/ARISE-Initiative/robosuite.git", "lerobot[scipy-dep]"]
|
||||||
|
|
||||||
# All
|
# All
|
||||||
all = [
|
all = [
|
||||||
|
|||||||
@@ -496,6 +496,88 @@ class MetaworldEnv(EnvConfig):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@EnvConfig.register_subclass("robocasa")
|
||||||
|
@dataclass
|
||||||
|
class RoboCasaEnv(EnvConfig):
|
||||||
|
task: str = "CloseFridge"
|
||||||
|
fps: int = 20
|
||||||
|
episode_length: int = 1000
|
||||||
|
obs_type: str = "pixels_agent_pos"
|
||||||
|
render_mode: str = "rgb_array"
|
||||||
|
camera_name: str = "robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right"
|
||||||
|
camera_name_mapping: dict[str, str] | None = None
|
||||||
|
observation_height: int = 256
|
||||||
|
observation_width: int = 256
|
||||||
|
split: str | None = None
|
||||||
|
features: dict[str, PolicyFeature] = field(
|
||||||
|
default_factory=lambda: {
|
||||||
|
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(12,)),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
features_map: dict[str, str] = field(
|
||||||
|
default_factory=lambda: {
|
||||||
|
ACTION: ACTION,
|
||||||
|
"agent_pos": OBS_STATE,
|
||||||
|
"pixels/image": f"{OBS_IMAGES}.image",
|
||||||
|
"pixels/image2": f"{OBS_IMAGES}.image2",
|
||||||
|
"pixels/image3": f"{OBS_IMAGES}.image3",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
def __post_init__(self):
|
||||||
|
if self.obs_type == "pixels":
|
||||||
|
for cam_key in ["pixels/image", "pixels/image2", "pixels/image3"]:
|
||||||
|
self.features[cam_key] = PolicyFeature(
|
||||||
|
type=FeatureType.VISUAL,
|
||||||
|
shape=(self.observation_height, self.observation_width, 3),
|
||||||
|
)
|
||||||
|
elif self.obs_type == "pixels_agent_pos":
|
||||||
|
for cam_key in ["pixels/image", "pixels/image2", "pixels/image3"]:
|
||||||
|
self.features[cam_key] = PolicyFeature(
|
||||||
|
type=FeatureType.VISUAL,
|
||||||
|
shape=(self.observation_height, self.observation_width, 3),
|
||||||
|
)
|
||||||
|
self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(16,))
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unsupported obs_type: {self.obs_type}")
|
||||||
|
|
||||||
|
if self.camera_name_mapping is not None:
|
||||||
|
# Update features_map to reflect custom camera name mapping
|
||||||
|
mapping = self.camera_name_mapping
|
||||||
|
cams = [c.strip() for c in self.camera_name.split(",") if c.strip()]
|
||||||
|
for cam in cams:
|
||||||
|
mapped = mapping.get(cam, cam)
|
||||||
|
self.features_map[f"pixels/{mapped}"] = f"{OBS_IMAGES}.{mapped}"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def gym_kwargs(self) -> dict:
|
||||||
|
kwargs: dict[str, Any] = {
|
||||||
|
"obs_type": self.obs_type,
|
||||||
|
"render_mode": self.render_mode,
|
||||||
|
"observation_height": self.observation_height,
|
||||||
|
"observation_width": self.observation_width,
|
||||||
|
}
|
||||||
|
if self.split is not None:
|
||||||
|
kwargs["split"] = self.split
|
||||||
|
return kwargs
|
||||||
|
|
||||||
|
def create_envs(self, n_envs: int, use_async_envs: bool = False):
|
||||||
|
from .robocasa import create_robocasa_envs
|
||||||
|
|
||||||
|
if self.task is None:
|
||||||
|
raise ValueError("RoboCasaEnv requires a task to be specified")
|
||||||
|
env_cls = _make_vec_env_cls(use_async_envs, n_envs)
|
||||||
|
return create_robocasa_envs(
|
||||||
|
task=self.task,
|
||||||
|
n_envs=n_envs,
|
||||||
|
camera_name=self.camera_name,
|
||||||
|
camera_name_mapping=self.camera_name_mapping,
|
||||||
|
gym_kwargs=self.gym_kwargs,
|
||||||
|
env_cls=env_cls,
|
||||||
|
episode_length=self.episode_length,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@EnvConfig.register_subclass("isaaclab_arena")
|
@EnvConfig.register_subclass("isaaclab_arena")
|
||||||
@dataclass
|
@dataclass
|
||||||
class IsaaclabArenaEnv(HubEnvConfig):
|
class IsaaclabArenaEnv(HubEnvConfig):
|
||||||
|
|||||||
@@ -0,0 +1,384 @@
|
|||||||
|
#!/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.
|
||||||
|
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
|
||||||
|
|
||||||
|
# Dimensions for the flat action/state vectors used by the LeRobot wrapper.
|
||||||
|
# These correspond to the PandaOmron robot in RoboCasa365.
|
||||||
|
OBS_STATE_DIM = 16 # base_pos(3) + base_quat(4) + ee_pos_rel(3) + ee_quat_rel(4) + gripper_qpos(2)
|
||||||
|
ACTION_DIM = 12 # base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
|
||||||
|
ACTION_LOW = -1.0
|
||||||
|
ACTION_HIGH = 1.0
|
||||||
|
|
||||||
|
# Default cameras for the PandaOmron robot.
|
||||||
|
DEFAULT_CAMERAS = [
|
||||||
|
"robot0_agentview_left",
|
||||||
|
"robot0_eye_in_hand",
|
||||||
|
"robot0_agentview_right",
|
||||||
|
]
|
||||||
|
|
||||||
|
# Map raw RoboCasa camera names to LeRobot convention names.
|
||||||
|
DEFAULT_CAMERA_NAME_MAPPING = {
|
||||||
|
"robot0_agentview_left": "image",
|
||||||
|
"robot0_eye_in_hand": "image2",
|
||||||
|
"robot0_agentview_right": "image3",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
|
||||||
|
"""Normalize camera_name into a non-empty list of strings."""
|
||||||
|
if isinstance(camera_name, str):
|
||||||
|
cams = [c.strip() for c in camera_name.split(",") if c.strip()]
|
||||||
|
elif isinstance(camera_name, (list | tuple)):
|
||||||
|
cams = [str(c).strip() for c in camera_name if str(c).strip()]
|
||||||
|
else:
|
||||||
|
raise TypeError(f"camera_name must be str or sequence[str], got {type(camera_name).__name__}")
|
||||||
|
if not cams:
|
||||||
|
raise ValueError("camera_name resolved to an empty list.")
|
||||||
|
return cams
|
||||||
|
|
||||||
|
|
||||||
|
def convert_state(raw_obs: dict[str, np.ndarray]) -> np.ndarray:
|
||||||
|
"""Concatenate RoboCasa robot state dict into a flat (16,) vector.
|
||||||
|
|
||||||
|
Layout: base_pos(3) + base_quat(4) + ee_pos_rel(3) + ee_quat_rel(4) + gripper_qpos(2)
|
||||||
|
"""
|
||||||
|
return np.concatenate(
|
||||||
|
[
|
||||||
|
raw_obs["robot0_base_pos"], # (3,)
|
||||||
|
raw_obs["robot0_base_quat"], # (4,)
|
||||||
|
raw_obs["robot0_base_to_eef_pos"], # (3,)
|
||||||
|
raw_obs["robot0_base_to_eef_quat"], # (4,)
|
||||||
|
raw_obs["robot0_gripper_qpos"], # (2,)
|
||||||
|
],
|
||||||
|
axis=-1,
|
||||||
|
).astype(np.float32)
|
||||||
|
|
||||||
|
|
||||||
|
def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
|
||||||
|
"""Split a flat (12,) action vector into a RoboCasa action dict.
|
||||||
|
|
||||||
|
Layout: base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
|
||||||
|
"""
|
||||||
|
return {
|
||||||
|
"action.base_motion": flat_action[0:4],
|
||||||
|
"action.control_mode": flat_action[4:5],
|
||||||
|
"action.end_effector_position": flat_action[5:8],
|
||||||
|
"action.end_effector_rotation": flat_action[8:11],
|
||||||
|
"action.gripper_close": flat_action[11:12],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class RoboCasaEnv(gym.Env):
|
||||||
|
"""LeRobot gym.Env wrapper for RoboCasa365 kitchen environments.
|
||||||
|
|
||||||
|
Wraps the RoboCasaGymEnv from the robocasa package and converts its
|
||||||
|
dict-based observations and actions into flat arrays compatible with
|
||||||
|
the LeRobot evaluation pipeline.
|
||||||
|
"""
|
||||||
|
|
||||||
|
metadata = {"render_modes": ["rgb_array"], "render_fps": 20}
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
task: str,
|
||||||
|
camera_name: str | Sequence[str] = ",".join(DEFAULT_CAMERAS),
|
||||||
|
camera_name_mapping: dict[str, str] | None = None,
|
||||||
|
obs_type: str = "pixels_agent_pos",
|
||||||
|
render_mode: str = "rgb_array",
|
||||||
|
observation_width: int = 256,
|
||||||
|
observation_height: int = 256,
|
||||||
|
split: str | None = None,
|
||||||
|
episode_length: int | None = None,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.task = task
|
||||||
|
self.obs_type = obs_type
|
||||||
|
self.render_mode = render_mode
|
||||||
|
self.observation_width = observation_width
|
||||||
|
self.observation_height = observation_height
|
||||||
|
self.split = split
|
||||||
|
|
||||||
|
self.camera_name = _parse_camera_names(camera_name)
|
||||||
|
if camera_name_mapping is None:
|
||||||
|
camera_name_mapping = dict(DEFAULT_CAMERA_NAME_MAPPING)
|
||||||
|
self.camera_name_mapping = camera_name_mapping
|
||||||
|
|
||||||
|
self._max_episode_steps = episode_length if episode_length is not None else 1000
|
||||||
|
|
||||||
|
# Deferred — created on first reset() inside the worker subprocess
|
||||||
|
# to avoid inheriting stale GPU/EGL contexts across fork().
|
||||||
|
self._env = None
|
||||||
|
self.task_description = ""
|
||||||
|
|
||||||
|
# Build observation space
|
||||||
|
images = {}
|
||||||
|
for cam in self.camera_name:
|
||||||
|
mapped = self.camera_name_mapping.get(cam, cam)
|
||||||
|
images[mapped] = spaces.Box(
|
||||||
|
low=0,
|
||||||
|
high=255,
|
||||||
|
shape=(self.observation_height, self.observation_width, 3),
|
||||||
|
dtype=np.uint8,
|
||||||
|
)
|
||||||
|
|
||||||
|
if self.obs_type == "pixels":
|
||||||
|
self.observation_space = spaces.Dict({"pixels": spaces.Dict(images)})
|
||||||
|
elif self.obs_type == "pixels_agent_pos":
|
||||||
|
self.observation_space = spaces.Dict(
|
||||||
|
{
|
||||||
|
"pixels": spaces.Dict(images),
|
||||||
|
"agent_pos": spaces.Box(
|
||||||
|
low=-np.inf,
|
||||||
|
high=np.inf,
|
||||||
|
shape=(OBS_STATE_DIM,),
|
||||||
|
dtype=np.float32,
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unsupported obs_type '{self.obs_type}'. Use 'pixels' or 'pixels_agent_pos'.")
|
||||||
|
|
||||||
|
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 RoboCasaGymEnv 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 robocasa.wrappers.gym_wrapper import RoboCasaGymEnv
|
||||||
|
|
||||||
|
kwargs: dict[str, Any] = {
|
||||||
|
"env_name": self.task,
|
||||||
|
"camera_widths": self.observation_width,
|
||||||
|
"camera_heights": self.observation_height,
|
||||||
|
}
|
||||||
|
if self.split is not None:
|
||||||
|
kwargs["split"] = self.split
|
||||||
|
|
||||||
|
self._env = RoboCasaGymEnv(**kwargs)
|
||||||
|
|
||||||
|
# Extract task description from environment metadata
|
||||||
|
assert self._env is not None
|
||||||
|
ep_meta = self._env.env.get_ep_meta()
|
||||||
|
self.task_description = ep_meta.get("lang", self.task)
|
||||||
|
|
||||||
|
def _format_raw_obs(self, raw_obs: dict) -> RobotObservation:
|
||||||
|
"""Convert RoboCasaGymEnv observation dict to LeRobot format."""
|
||||||
|
# Extract camera images (RoboCasaGymEnv provides "video.<cam>" keys)
|
||||||
|
images = {}
|
||||||
|
for cam in self.camera_name:
|
||||||
|
video_key = f"video.{cam}"
|
||||||
|
if video_key in raw_obs:
|
||||||
|
mapped = self.camera_name_mapping.get(cam, cam)
|
||||||
|
images[mapped] = raw_obs[video_key]
|
||||||
|
|
||||||
|
if self.obs_type == "pixels":
|
||||||
|
return {"pixels": images}
|
||||||
|
|
||||||
|
# Extract state from raw_obs (state.* keys from PandaOmronKeyConverter)
|
||||||
|
agent_pos = np.concatenate(
|
||||||
|
[
|
||||||
|
raw_obs.get("state.base_position", np.zeros(3)),
|
||||||
|
raw_obs.get("state.base_rotation", np.zeros(4)),
|
||||||
|
raw_obs.get("state.end_effector_position_relative", np.zeros(3)),
|
||||||
|
raw_obs.get("state.end_effector_rotation_relative", np.zeros(4)),
|
||||||
|
raw_obs.get("state.gripper_qpos", np.zeros(2)),
|
||||||
|
],
|
||||||
|
axis=-1,
|
||||||
|
).astype(np.float32)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"pixels": images,
|
||||||
|
"agent_pos": agent_pos,
|
||||||
|
}
|
||||||
|
|
||||||
|
def render(self) -> np.ndarray:
|
||||||
|
self._ensure_env()
|
||||||
|
assert self._env is not None
|
||||||
|
return self._env.render()
|
||||||
|
|
||||||
|
def reset(self, seed=None, **kwargs):
|
||||||
|
self._ensure_env()
|
||||||
|
assert self._env is not None
|
||||||
|
super().reset(seed=seed)
|
||||||
|
raw_obs, info = self._env.reset(seed=seed)
|
||||||
|
|
||||||
|
# Update task description on each reset (may change per episode)
|
||||||
|
ep_meta = self._env.env.get_ep_meta()
|
||||||
|
self.task_description = ep_meta.get("lang", self.task)
|
||||||
|
|
||||||
|
observation = self._format_raw_obs(raw_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}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Convert flat action to RoboCasa dict format
|
||||||
|
action_dict = convert_action(action)
|
||||||
|
raw_obs, reward, done, truncated, info = self._env.step(action_dict)
|
||||||
|
|
||||||
|
is_success = bool(info.get("success", False))
|
||||||
|
terminated = done or is_success
|
||||||
|
info.update(
|
||||||
|
{
|
||||||
|
"task": self.task,
|
||||||
|
"done": done,
|
||||||
|
"is_success": is_success,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
observation = self._format_raw_obs(raw_obs)
|
||||||
|
if terminated:
|
||||||
|
self.reset()
|
||||||
|
|
||||||
|
return observation, reward, terminated, truncated, info
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
if self._env is not None:
|
||||||
|
self._env.close()
|
||||||
|
|
||||||
|
|
||||||
|
# ---- Main API ----------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
def _make_env_fns(
|
||||||
|
*,
|
||||||
|
task: str,
|
||||||
|
n_envs: int,
|
||||||
|
camera_names: list[str],
|
||||||
|
camera_name_mapping: dict[str, str] | None,
|
||||||
|
obs_type: str,
|
||||||
|
render_mode: str,
|
||||||
|
observation_width: int,
|
||||||
|
observation_height: int,
|
||||||
|
split: str | None,
|
||||||
|
episode_length: int | None,
|
||||||
|
) -> list[Callable[[], RoboCasaEnv]]:
|
||||||
|
"""Build n_envs factory callables for a single task."""
|
||||||
|
|
||||||
|
def _make_env(**kwargs) -> RoboCasaEnv:
|
||||||
|
return RoboCasaEnv(
|
||||||
|
task=task,
|
||||||
|
camera_name=camera_names,
|
||||||
|
camera_name_mapping=camera_name_mapping,
|
||||||
|
obs_type=obs_type,
|
||||||
|
render_mode=render_mode,
|
||||||
|
observation_width=observation_width,
|
||||||
|
observation_height=observation_height,
|
||||||
|
split=split,
|
||||||
|
episode_length=episode_length,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
return [partial(_make_env) for _ in range(n_envs)]
|
||||||
|
|
||||||
|
|
||||||
|
def create_robocasa_envs(
|
||||||
|
task: str,
|
||||||
|
n_envs: int,
|
||||||
|
gym_kwargs: dict[str, Any] | None = None,
|
||||||
|
camera_name: str | Sequence[str] = ",".join(DEFAULT_CAMERAS),
|
||||||
|
camera_name_mapping: dict[str, str] | None = None,
|
||||||
|
env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
|
||||||
|
episode_length: int | None = None,
|
||||||
|
) -> dict[str, dict[int, Any]]:
|
||||||
|
"""Create vectorized RoboCasa365 environments with a consistent return shape.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict[task_name][task_id] -> vec_env (env_cls([...]) with exactly n_envs factories)
|
||||||
|
Notes:
|
||||||
|
- n_envs is the number of rollouts *per task* (parallel environments).
|
||||||
|
- `task` can be a single task or a comma-separated list of tasks.
|
||||||
|
"""
|
||||||
|
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 {})
|
||||||
|
obs_type = gym_kwargs.pop("obs_type", "pixels_agent_pos")
|
||||||
|
render_mode = gym_kwargs.pop("render_mode", "rgb_array")
|
||||||
|
observation_width = gym_kwargs.pop("observation_width", 256)
|
||||||
|
observation_height = gym_kwargs.pop("observation_height", 256)
|
||||||
|
split = gym_kwargs.pop("split", None)
|
||||||
|
|
||||||
|
camera_names = _parse_camera_names(camera_name)
|
||||||
|
task_names = [t.strip() for t in str(task).split(",") if t.strip()]
|
||||||
|
if not task_names:
|
||||||
|
raise ValueError("`task` must contain at least one RoboCasa task name.")
|
||||||
|
|
||||||
|
print(f"Creating RoboCasa envs | tasks={task_names} | n_envs(per task)={n_envs}")
|
||||||
|
|
||||||
|
is_async = env_cls is gym.vector.AsyncVectorEnv
|
||||||
|
|
||||||
|
cached_obs_space: spaces.Space | None = None
|
||||||
|
cached_act_space: spaces.Space | None = None
|
||||||
|
out: dict[str, dict[int, Any]] = defaultdict(dict)
|
||||||
|
|
||||||
|
for _tid, task_name in enumerate(task_names):
|
||||||
|
fns = _make_env_fns(
|
||||||
|
task=task_name,
|
||||||
|
n_envs=n_envs,
|
||||||
|
camera_names=camera_names,
|
||||||
|
camera_name_mapping=camera_name_mapping,
|
||||||
|
obs_type=obs_type,
|
||||||
|
render_mode=render_mode,
|
||||||
|
observation_width=observation_width,
|
||||||
|
observation_height=observation_height,
|
||||||
|
split=split,
|
||||||
|
episode_length=episode_length,
|
||||||
|
)
|
||||||
|
|
||||||
|
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[task_name][0] = lazy
|
||||||
|
else:
|
||||||
|
out[task_name][0] = env_cls(fns)
|
||||||
|
print(f"Built vec env | task={task_name} | n_envs={n_envs}")
|
||||||
|
|
||||||
|
return {name: dict(task_map) for name, task_map in out.items()}
|
||||||
Reference in New Issue
Block a user