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:
Pepijn
2026-04-14 11:25:25 +02:00
parent a656a982af
commit 69ed70edeb
7 changed files with 705 additions and 0 deletions
+90
View File
@@ -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
+36
View File
@@ -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"]
+2
View File
@@ -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"
+110
View File
@@ -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
```
+1
View File
@@ -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 = [
+82
View File
@@ -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):
+384
View File
@@ -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()}