diff --git a/.github/workflows/benchmark_tests.yml b/.github/workflows/benchmark_tests.yml index 79d5614b2..f4cbe859c 100644 --- a/.github/workflows/benchmark_tests.yml +++ b/.github/workflows/benchmark_tests.yml @@ -310,3 +310,93 @@ jobs: name: metaworld-metrics path: /tmp/metaworld-artifacts/metrics.json if-no-files-found: warn + + # ── 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 diff --git a/docker/Dockerfile.benchmark.robocasa b/docker/Dockerfile.benchmark.robocasa new file mode 100644 index 000000000..81f036438 --- /dev/null +++ b/docker/Dockerfile.benchmark.robocasa @@ -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"] diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 3dcba5993..b48bbf757 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -79,6 +79,8 @@ title: LIBERO - local: metaworld title: Meta-World + - local: robocasa + title: RoboCasa365 - local: envhub_isaaclab_arena title: NVIDIA IsaacLab Arena Environments title: "Benchmarks" diff --git a/docs/source/robocasa.mdx b/docs/source/robocasa.mdx new file mode 100644 index 000000000..0f0309352 --- /dev/null +++ b/docs/source/robocasa.mdx @@ -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 +``` + + +RoboCasa365 requires MuJoCo for simulation. Set the rendering backend before training or evaluation: + +```bash +export MUJOCO_GL=egl # for headless servers (HPC, cloud) +``` + + + +## 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 +``` diff --git a/pyproject.toml b/pyproject.toml index 2f12840e7..83277b503 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -206,6 +206,7 @@ aloha = ["lerobot[dataset]", "gym-aloha>=0.1.2,<0.2.0", "lerobot[scipy-dep]"] pusht = ["lerobot[dataset]", "gym-pusht>=0.1.5,<0.2.0", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead libero = ["lerobot[dataset]", "lerobot[transformers-dep]", "hf-libero>=0.1.3,<0.2.0; sys_platform == 'linux'", "lerobot[scipy-dep]"] metaworld = ["lerobot[dataset]", "metaworld==3.0.0", "lerobot[scipy-dep]"] +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 = [ diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index 2a7c52d45..eb140f879 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -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") @dataclass class IsaaclabArenaEnv(HubEnvConfig): diff --git a/src/lerobot/envs/robocasa.py b/src/lerobot/envs/robocasa.py new file mode 100644 index 000000000..945b7e646 --- /dev/null +++ b/src/lerobot/envs/robocasa.py @@ -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." 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()}