diff --git a/CLAUDE.md b/CLAUDE.md index c912c7d2..9d38f7ce 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -43,13 +43,15 @@ Dependency management is **`uv` (>= 0.8.4) only** — `pyproject.toml`/`uv.lock` ```bash uv sync --extra dev --extra libero # standard dev setup (matches CI) -uv sync --all-extras # everything (libero + urdf now co-install on numpy 2.x) +uv sync --all-extras # everything (libero + robocasa + urdf co-install on the shared robosuite-1.5 master / numpy 2.x stack); Linux-only (trt) source .venv/bin/activate ``` Re-run `uv sync` whenever `pyproject.toml`/`uv.lock` change. Add deps with `uv add `; lock with `uv lock`. -Installable extras: `dev` (pre-commit, sphinx, pytest), `libero` (sim env — pulls a forked LIBERO from `shuheng-liu/LIBERO`, runs on numpy 2.x + gymnasium), `urdf` (rerun ≥0.28, numpy 2.x), `trt` (TensorRT, Linux/Win x86_64 only). +Installable extras: `dev` (pre-commit, sphinx, pytest), `libero` (sim env — pulls a forked LIBERO from `shuheng-liu/LIBERO`, on robosuite 1.5 master + numpy 2.x + gymnasium), `robocasa` (RoboCasa365 kitchen sim — co-installs with `libero` on the shared robosuite stack; see the RoboCasa365 note below), `urdf` (rerun ≥0.28, numpy 2.x), `trt` (TensorRT, Linux/Win x86_64 only). + +**RoboCasa365** (`envs/robocasa.py`) is a first-class extra that co-installs with `libero` on a shared robosuite stack: `uv sync --extra robocasa`. Two non-obvious things make it resolve: (1) robocasa needs `MujocoEnv(load_model_on_init=...)`, added on robosuite **master** *after* the 1.5.2 PyPI release, so `[tool.uv.sources]` repins `robosuite` to a master commit — which still self-reports "1.5.2" (matching the extras' pins) and is validated to also run LIBERO; (2) robocasa is pulled from the `shuheng-liu/robocasa` packaging fork (mirroring the `shuheng-liu/LIBERO` / `egl_probe` forks) that drops upstream's `lerobot==0.3.3` / `tianshou` / `opencv-python` / `hidapi` deps and loosens its `numpy`/`numba`/`scipy`/`mujoco` pins + import-time version asserts, since uv can't `--no-deps` a single package in a lock. Kitchen assets (~5-10GB) are a separate runtime step — `python -m robocasa.scripts.download_kitchen_assets` — then run headless with `MUJOCO_GL=egl`. NOTE: a full `uv lock` / `uv sync --all-extras` must run on Linux (the `trt` extra's TensorRT sdist can't build on macOS arm64); `uv sync --extra robocasa` from the committed lock works anywhere. ## Common commands @@ -117,7 +119,7 @@ Key invariant on `TrainPipelineConfig`: `batch_size == dataloader_batch_size * g - `configs/` — dataclass configs (train, eval, policies, envs, optim, deployment, libero, ros2lerobot) - `datasets/` — LeRobot-compatible datasets, `WeightedDatasetMixture` (heterogeneous co-training), VQA datasets, v1→v2 / v2→v2.1 converters under `v2/`, `v21/` -- `envs/` — gym/gymnasium envs (currently LIBERO); `factory.make_envs()` +- `envs/` — gym/gymnasium envs (LIBERO in `libero.py`, RoboCasa365 in `robocasa.py`); `factory.make_envs()` dispatches per `env.type`. Both return `dict[group][task_id] -> VectorEnv` so the env-agnostic eval pipeline (`scripts/eval.py`) gives per-task success rates + `grid_summary` wandb videos for free. - `optim/` — optimizer + LR-scheduler dataclass-configured factories - `planner/` — high-level planner using `prompts.yaml` - `policies/` — `pi0`, `pi05`, `pi05_mem`, `pi06`, `pi07/{high_level_planner,low_level}` (current π0.7 impl: Gemma 3 backbone + SpaceTime SigLIP video encoder; note `low_level/` — not `low_level_planner/`, since the low-level policy is a controller, not a planner), `pi07_paligemma/{high_level_planner,low_level}` (legacy PaliGemma variant of π0.7 — kept for older checkpoints; a fix targeting π0.7 usually needs to land in `pi07/`, not here), `value`. Each subdir has a `configuration_*.py` and `modeling_*.py`. Vision backbone wrappers: `paligemma_with_expert.py` (pi0/pi05/pi05_mem/pi07_paligemma) and `gemma3_with_expert.py` (pi06/pi07). diff --git a/configs/examples/pi05_robocasa_eval_config.json b/configs/examples/pi05_robocasa_eval_config.json new file mode 100644 index 00000000..c0be6bf4 --- /dev/null +++ b/configs/examples/pi05_robocasa_eval_config.json @@ -0,0 +1,174 @@ +{ + "dataset_mixture": { + "datasets": [ + { + "repo_id": "lerobot/droid_100" + } + ], + "weights": [ + 1.0 + ], + "action_freq": 20.0, + "image_resample_strategy": "nearest", + "vector_resample_strategy": "nearest" + }, + "policy": { + "type": "pi05", + "pretrained_path": "TensorAuto/tPi0.5-libero", + "n_obs_steps": 1, + "input_features": { + "camera0": { + "shape": [ + 3, + 224, + 224 + ], + "type": "VISUAL" + }, + "camera1": { + "shape": [ + 3, + 224, + 224 + ], + "type": "VISUAL" + }, + "camera2": { + "shape": [ + 3, + 224, + 224 + ], + "type": "VISUAL" + }, + "state": { + "shape": [ + 32 + ], + "type": "STATE" + } + }, + "output_features": { + "actions": { + "shape": [ + 32 + ], + "type": "ACTION" + } + }, + "normalization_mapping": { + "VISUAL": "IDENTITY", + "STATE": "MIN_MAX", + "ACTION": "MEAN_STD" + }, + "chunk_size": 10, + "n_action_steps": 10, + "max_state_dim": 32, + "max_action_dim": 32, + "proj_width": 1024, + "num_steps": 10, + "attention_implementation": "eager", + "freeze_vision_encoder": false, + "train_expert_only": false, + "prompt_max_length": 256, + "discrete_action_max_length": 60, + "optimizer_lr": 2.5e-05, + "optimizer_betas": [ + 0.9, + 0.95 + ], + "optimizer_eps": 1e-08, + "optimizer_weight_decay": 1e-10, + "scheduler_warmup_steps": 1000, + "scheduler_decay_steps": 30000, + "scheduler_decay_lr": 2.5e-06 + }, + "output_dir": "outputs/pi05_robocasa_eval", + "resume": false, + "seed": 1000, + "resolution": [ + 224, + 224 + ], + "num_cams": 3, + "max_state_dim": 32, + "max_action_dim": 32, + "action_chunk": 10, + "loss_weighting": { + "MSE": 1.0, + "CE": 1.0 + }, + "num_workers": 4, + "batch_size": 2, + "gradient_accumulation_steps": 1, + "dataloader_batch_size": 2, + "prefetch_factor": 8, + "steps": 100, + "log_freq": 1, + "save_checkpoint": true, + "save_freq": 100, + "use_policy_training_preset": true, + "trace_nans": false, + "optimizer": { + "type": "adamw", + "lr": 2.5e-05, + "weight_decay": 1e-10, + "grad_clip_norm": 10.0, + "betas": [ + 0.9, + 0.95 + ], + "eps": 1e-08 + }, + "env": { + "type": "robocasa", + "task": "CloseFridge", + "fps": 20, + "max_parallel_tasks": 1, + "episode_length": 1000, + "obs_type": "pixels_agent_pos", + "render_mode": "rgb_array", + "camera_name": "robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right", + "observation_height": 256, + "observation_width": 256, + "visualization_height": 512, + "visualization_width": 512, + "split": null, + "obj_registries": [ + "lightwheel" + ], + "metadata": { + "robot_type": "PandaOmron", + "control_mode": "ee" + } + }, + "eval": { + "n_episodes": 2, + "batch_size": 2, + "use_async_envs": true, + "max_episodes_rendered": 2, + "grid_size": null, + "control_mode": "ee" + }, + "scheduler": { + "type": "cosine_decay_with_warmup", + "num_warmup_steps": 1000, + "num_decay_steps": 30000, + "peak_lr": 2.5e-05, + "decay_lr": 2.5e-06 + }, + "wandb": { + "enable": true, + "entity": "wyautox-autox", + "project": "pi05", + "run_id": null, + "name": null, + "notes": "RoboCasa eval plumbing smoke (CloseFridge). Swap policy.pretrained_path for a RoboCasa-trained checkpoint for meaningful success rates.", + "tags": [], + "group": null, + "job_type": null, + "mode": null, + "allow_resume": true, + "disable_artifact": false + } +} diff --git a/pyproject.toml b/pyproject.toml index 2885f80a..e8fd0ecf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -127,9 +127,11 @@ libero = [ "egl-probe", "robomimic==0.2.0", # robosuite 1.5.2 (composite-controller framework) so LIBERO shares a venv with - # RoboCasa, which needs >=1.5. The LIBERO fork is ported to the 1.5 controller API - # (see load_arm_controller_config); robomimic 0.2.0 does not pin robosuite, so it - # co-installs unchanged. + # RoboCasa. Sourced from a pinned robosuite *master* commit (see [tool.uv.sources]) + # because robocasa needs MujocoEnv(load_model_on_init=...), added on master after + # the 1.5.2 PyPI release; master still self-reports "1.5.2" so this pin matches. + # The LIBERO fork (1.5 controller API, see load_arm_controller_config) is validated + # on this commit; robomimic 0.2.0 does not pin robosuite, so it co-installs unchanged. "robosuite==1.5.2", "thop==0.1.1.post2209072238", "mujoco>=3.3.5", @@ -143,6 +145,20 @@ libero = [ "mujoco>=3.1.6 ; sys_platform == 'linux'", "pyopengl==3.1.10 ; sys_platform == 'linux'", ] +robocasa = [ + # RoboCasa365 kitchen sim for simulated eval, co-installed on the shared + # robosuite-1.5 (master) stack from the `libero` extra. Uses the + # shuheng-liu/robocasa packaging fork (see [tool.uv.sources]) that drops + # upstream's lerobot==0.3.3 / tianshou / opencv-python / hidapi deps and loosens + # its numpy/numba/scipy/mujoco pins so it resolves in the shared lock (upstream's + # setup.py can't be `--no-deps`'d inside a uv lock). robosuite is not in robocasa's + # deps (installed separately), so pin it here matching `libero`. Kitchen assets + # (~5-10GB) are a separate runtime download: + # `python -m robocasa.scripts.download_kitchen_assets`. + "robocasa", + "robosuite==1.5.2", + "mujoco>=3.3.5", +] urdf = [ "rerun-sdk>=0.28.2", ] @@ -156,6 +172,15 @@ libero = { git = "https://github.com/shuheng-liu/LIBERO" , branch = "master" } # which CMake >= 4 rejects. This fork raises the floor to 3.5 so it builds on CMake 4 # without pinning the build's cmake; the tag keeps the source reproducible. egl-probe = { git = "https://github.com/shuheng-liu/egl_probe", tag = "v1.0.1-cmake4" } +# robocasa main needs robosuite *master* (MujocoEnv(load_model_on_init=...), added after +# the 1.5.2 PyPI release); pin the exact commit for reproducibility. Master self-reports +# version "1.5.2", which satisfies the `robosuite==1.5.2` pins in the extras above, and +# LIBERO is validated to construct+step on this commit (so the shared stack still works). +robosuite = { git = "https://github.com/ARISE-Initiative/robosuite", rev = "85abee228d1c43ab1939bce33028099945d453b4" } +# OpenTau packaging fork of robocasa: drops upstream's lerobot==0.3.3 / tianshou / +# opencv-python / hidapi deps and loosens its numpy/numba/scipy/mujoco pins + import-time +# version asserts so it co-installs on the shared robosuite-1.5 / numpy-2 stack. +robocasa = { git = "https://github.com/shuheng-liu/robocasa", rev = "f7db21c11f25408d3a59a5e878bb5c7ca9030c4d" } [tool.uv] # `extra-build-dependencies` injects a cmake wheel into egl-probe's PEP 517 build diff --git a/src/opentau/envs/configs.py b/src/opentau/envs/configs.py index 8c30cc2a..6faa68c3 100644 --- a/src/opentau/envs/configs.py +++ b/src/opentau/envs/configs.py @@ -307,3 +307,109 @@ def gym_kwargs(self) -> dict: "task_ids": task_ids, "control_freq": self.fps, } + + +@EnvConfig.register_subclass("robocasa") +@dataclass +class RoboCasaEnv(EnvConfig): + r"""Configuration for the RoboCasa365 kitchen environment. + + RoboCasa runs on robosuite 1.5 (shared with LIBERO since the libero extra was + bumped to robosuite 1.5.2), so it co-installs in the same venv. The default + robot is the PandaOmron mobile manipulator — hence the 12-D action and 16-D + state, distinct from LIBERO's 7-D/8-D. Set ``metadata.robot_type`` / + ``eval.control_mode`` to select the matching per-(robot_type, control_mode) + projection head when evaluating a co-trained policy. + + Args: + task: A RoboCasa task name (e.g. ``"CloseFridge"``), a comma-separated + list of task names, or a benchmark-group shortcut + (``atomic_seen``/``composite_seen``/``composite_unseen``/ + ``pretrain50``/``pretrain100``/``pretrain200``/``pretrain300``), which + auto-expands to the upstream task list and auto-sets ``split``. + fps: RoboCasa control frequency (Hz); also the ``render_fps`` for videos. + episode_length: Maximum steps per episode (``_max_episode_steps``). + obs_type: ``"pixels"`` or ``"pixels_agent_pos"``. + render_mode: Rendering mode for the environment. + camera_name: Comma-separated raw RoboCasa camera names to render. The + wrapper remaps them to ``camera0``/``camera1``/... so the policy input + structure matches LIBERO regardless of the raw names; when the policy + was trained with a larger ``cfg.num_cams``, ``preprocess_observation`` + zero-fills the remaining slots. + observation_height: Height of observation images. + observation_width: Width of observation images. + visualization_height: Height of visualization frames. + visualization_width: Width of visualization frames. + split: RoboCasa dataset split (``None``/``"all"``/``"pretrain"``/ + ``"target"``). Left ``None`` unless a task-group shortcut sets it. + obj_registries: Object-mesh registries to sample assets from. Defaults to + ``["lightwheel"]`` (the pack the asset downloader ships by default); + add ``"objaverse"`` only after downloading that ~30GB pack. + features: Mapping from logical feature names to ``PolicyFeature`` definitions. + features_map: Mapping from environment keys to standardized OpenTau keys. + """ + + 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" + observation_height: int = 256 + observation_width: int = 256 + visualization_height: int = 512 + visualization_width: int = 512 + split: str | None = None + obj_registries: list[str] = field(default_factory=lambda: ["lightwheel"]) + 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, + } + ) + + def __post_init__(self): + if self.fps <= 0: + raise ValueError(f"RoboCasa env.fps (control frequency in Hz) must be positive, got {self.fps}") + if self.obs_type not in ("pixels", "pixels_agent_pos"): + raise ValueError(f"Unsupported obs_type: {self.obs_type}") + + # The wrapper remaps the i-th raw camera to ``camera{i}``; mirror that in + # the feature map using OpenTau's ``image`` / ``image2`` / ... convention + # (camera0 -> image, camera1 -> image2, ...), matching LIBERO. + cams = [c.strip() for c in self.camera_name.split(",") if c.strip()] + for i, cam in enumerate(cams): + self.features[f"pixels/{cam}"] = PolicyFeature( + type=FeatureType.VISUAL, + shape=(self.observation_height, self.observation_width, 3), + ) + mapped = "image" if i == 0 else f"image{i + 1}" + self.features_map[f"pixels/{cam}"] = f"{OBS_IMAGES}.{mapped}" + + if self.obs_type == "pixels_agent_pos": + self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(16,)) + + @property + def gym_kwargs(self) -> dict: + r"""Return the keyword arguments used to construct the RoboCasa environment. + + Task resolution and per-rank sharding live in ``create_robocasa_envs`` (they + need the ``robocasa`` package for group expansion), so this stays sim-free + and only carries the obs/render parameters plus an optional ``split``. + """ + kwargs: dict = { + "obs_type": self.obs_type, + "render_mode": self.render_mode, + "observation_height": self.observation_height, + "observation_width": self.observation_width, + "visualization_height": self.visualization_height, + "visualization_width": self.visualization_width, + } + if self.split is not None: + kwargs["split"] = self.split + return kwargs diff --git a/src/opentau/envs/factory.py b/src/opentau/envs/factory.py index 3fc32c5c..2b98f2fb 100644 --- a/src/opentau/envs/factory.py +++ b/src/opentau/envs/factory.py @@ -22,15 +22,17 @@ import gymnasium as gym from opentau.configs.train import TrainPipelineConfig -from opentau.envs.configs import EnvConfig, LiberoEnv +from opentau.envs.configs import EnvConfig, LiberoEnv, RoboCasaEnv def make_env_config(env_type: str, **kwargs) -> EnvConfig: r"""Factory method to create an environment config based on the env_type. - Right now, only 'libero' is supported. + Supports 'libero' and 'robocasa'. """ if env_type == "libero": return LiberoEnv(**kwargs) + elif env_type == "robocasa": + return RoboCasaEnv(**kwargs) else: raise ValueError(f"Env type '{env_type}' is not available.") @@ -79,6 +81,22 @@ def make_envs( env_cls=env_cls, ) + # RoboCasa, like LIBERO, is multi-task: build one vec env per task so eval + # reports per-task success and per-task grid videos, and so tasks shard + # disjointly across accelerator ranks (handled inside create_robocasa_envs). + if isinstance(cfg, RoboCasaEnv): + from opentau.envs.robocasa import create_robocasa_envs + + return create_robocasa_envs( + task=cfg.task, + n_envs=n_envs, + camera_name=cfg.camera_name, + gym_kwargs=cfg.gym_kwargs, + env_cls=env_cls, + episode_length=cfg.episode_length, + obj_registries=tuple(cfg.obj_registries), + ) + try: importlib.import_module(cfg.import_name) except ModuleNotFoundError as e: diff --git a/src/opentau/envs/robocasa.py b/src/opentau/envs/robocasa.py new file mode 100644 index 00000000..d2499d9d --- /dev/null +++ b/src/opentau/envs/robocasa.py @@ -0,0 +1,535 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# Copyright 2026 Tensor Auto Inc. 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. + +r"""Environment wrapper for RoboCasa365 kitchen tasks. + +Ported from upstream LeRobot's ``lerobot/envs/robocasa.py`` and reshaped to +OpenTau's LIBERO conventions: cameras are remapped to ``camera0``/``camera1``/... +so :func:`opentau.envs.utils.preprocess_observation` and the policy's +``num_cams`` zero-fill path consume them exactly like LIBERO, and the vec-env +builder shards tasks across accelerator ranks so distributed eval and the +``_rank{N}``-strip uniqueness assumption in +:func:`opentau.scripts.eval.collect_grid_summary_videos` both hold. + +The underlying simulator (``robocasa`` / ``robosuite`` 1.5) is imported lazily +inside :meth:`RoboCasaEnv._ensure_env` and :func:`_resolve_tasks`, so importing +this module (e.g. in the CPU test suite) never requires the sim to be installed. +""" + +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 opentau.utils.accelerate_utils import acc_print, get_proc_accelerator + +# Flat action/state vector dimensions for the PandaOmron mobile manipulator +# (RoboCasa365's default robot). +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 PandaOmron cameras (raw RoboCasa names). The wrapper remaps these to +# ``camera0``/``camera1``/``camera2`` so the OpenTau policy input structure +# matches LIBERO's (``camera{i}`` keys + ``img_is_pad``), independent of how +# many real cameras a given config renders. +DEFAULT_CAMERAS = [ + "robot0_agentview_left", + "robot0_eye_in_hand", + "robot0_agentview_right", +] + +# Object-mesh registries to sample from. RoboCasa's upstream default is +# ("objaverse", "lightwheel"), but the objaverse pack is huge (~30GB) and most +# setups only download the lightwheel pack (`--type objs_lw` in +# `download_kitchen_assets`). When a sampled object category has zero candidates +# in every registry, robocasa crashes with `ValueError: Probabilities contain +# NaN`. Restricting to registries that are actually on disk avoids that. +DEFAULT_OBJ_REGISTRIES: tuple[str, ...] = ("lightwheel",) + +# Task-group shortcuts accepted as ``env.task``. A group name expands to the +# upstream RoboCasa task list and auto-sets the dataset split; individual task +# names (optionally comma-separated) take precedence and only match exactly. +_TASK_GROUP_SPLITS = { + "atomic_seen": "target", + "composite_seen": "target", + "composite_unseen": "target", + "pretrain50": "pretrain", + "pretrain100": "pretrain", + "pretrain200": "pretrain", + "pretrain300": "pretrain", +} + + +def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]: + """Normalize camera_name into a non-empty list of strings. + + Local copy of the LIBERO helper so importing this module never pulls in + ``opentau.envs.libero`` (which imports the ``libero`` package at top level). + """ + 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 _default_camera_name_mapping(camera_names: Sequence[str]) -> dict[str, list[str]]: + """Map each raw RoboCasa camera to a positional ``camera{i}`` key. + + Mirrors ``LiberoEnv``'s ``camera_name_mapping``: the first rendered camera + becomes ``camera0``, the second ``camera1``, and so on, so the policy + receives a consistent ``camera{i}`` structure regardless of the raw names. + """ + return {cam: [f"camera{i}"] for i, cam in enumerate(camera_names)} + + +def _import_robocasa_with_version_shim() -> None: + """Import the top-level ``robocasa`` package despite its over-strict pins. + + robocasa 1.0.1 hard-asserts ``mujoco.__version__ == "3.3.1"`` and + ``numpy.__version__ in ["2.2.5"]`` at import time. OpenTau shares LIBERO's + stack (robosuite 1.5.2 + ``mujoco>=3.3.5`` + ``numpy>=2.2.6``) — and robosuite + 1.5.2 itself only needs ``mujoco>=3.3.0`` / ``numpy>=1.13.3`` — so those + equality pins are too strict to satisfy here. Spoof the two version strings + only for the one-time package import, then restore them. A robocasa packaging + fork that drops the asserts makes this shim unnecessary; until then it lets + the integration run on a stock ``pip install robocasa``. No-op once imported. + """ + import sys + + if "robocasa" in sys.modules: + return + import mujoco + import numpy + + saved = (mujoco.__version__, numpy.__version__) + try: + mujoco.__version__ = "3.3.1" + numpy.__version__ = "2.2.5" + import robocasa # noqa: F401 (runs robocasa's version asserts once) + finally: + mujoco.__version__, numpy.__version__ = saved + + +def _resolve_tasks(task: str) -> tuple[list[str], str | None]: + """Resolve an ``env.task`` value to ``(task_names, split_override)``. + + If ``task`` is a known task-group name (e.g. ``atomic_seen``, + ``pretrain100``), expand it via + ``robocasa.utils.dataset_registry.{TARGET,PRETRAINING}_TASKS`` and return the + matching split. Otherwise treat ``task`` as a single task or comma-separated + list and leave the split untouched (``None``). The ``robocasa`` import is + deferred to here so only group expansion requires the sim package. + """ + key = task.strip() + if key in _TASK_GROUP_SPLITS: + _import_robocasa_with_version_shim() + from robocasa.utils.dataset_registry import PRETRAINING_TASKS, TARGET_TASKS + + combined = {**TARGET_TASKS, **PRETRAINING_TASKS} + if key not in combined: + raise ValueError( + f"Task group '{key}' is not available in this version of robocasa. " + f"Known groups: {sorted(combined.keys())}." + ) + return list(combined[key]), _TASK_GROUP_SPLITS[key] + + names = [t.strip() for t in task.split(",") if t.strip()] + if not names: + raise ValueError("`task` must contain at least one RoboCasa task name.") + return names, None + + +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): + r"""Gym wrapper for RoboCasa365 kitchen environments. + + Wraps ``RoboCasaGymEnv`` from the ``robocasa`` package and converts its + dict-based observations/actions into the flat arrays OpenTau expects. Raw + camera frames are remapped to ``camera{i}`` keys (see + ``camera_name_mapping``) so the policy input structure matches LIBERO. + """ + + metadata = {"render_modes": ["rgb_array"], "render_fps": 20} + + def __init__( + self, + task: str, + camera_name: str | Sequence[str] = ",".join(DEFAULT_CAMERAS), + obs_type: str = "pixels_agent_pos", + render_mode: str = "rgb_array", + observation_width: int = 256, + observation_height: int = 256, + visualization_width: int = 512, + visualization_height: int = 512, + split: str | None = None, + episode_length: int | None = None, + obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES, + episode_index: int = 0, + camera_name_mapping: dict[str, list[str]] | None = None, + ): + r"""Initialize the RoboCasaEnv. + + Args: + task: RoboCasa task name (e.g. ``"CloseFridge"``). + camera_name: Raw RoboCasa camera name(s); comma-separated string or + sequence. Both count and order are driven by this value. + obs_type: ``"pixels"`` or ``"pixels_agent_pos"``. + render_mode: Rendering mode for the environment. + observation_width: Width of observation images. + observation_height: Height of observation images. + visualization_width: Width of visualization frames. + visualization_height: Height of visualization frames. + split: RoboCasa dataset split (``None``/``"all"``/``"pretrain"``/``"target"``). + episode_length: Max steps per episode (``_max_episode_steps``); defaults to 1000. + obj_registries: Object-mesh registries to sample assets from. + episode_index: Per-worker index (``0..n_envs-1``) used to spread the + ``reset`` seed so each sub-env explores a distinct layout. + camera_name_mapping: Optional mapping from raw camera names to + positional ``camera{i}`` keys; defaults to first→``camera0``, etc. + """ + 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.visualization_width = visualization_width + self.visualization_height = visualization_height + self.split = split + self.obj_registries = tuple(obj_registries) + self.episode_index = int(episode_index) + + self.camera_name = _parse_camera_names(camera_name) + if camera_name_mapping is None: + camera_name_mapping = _default_camera_name_mapping(self.camera_name) + self.camera_name_mapping = camera_name_mapping + for cam in self.camera_name_mapping: + assert not isinstance(self.camera_name_mapping[cam], str), ( + "camera_name_mapping values must be lists of strings; " + f"got string {self.camera_name_mapping[cam]} for {cam} instead" + ) + + self._max_episode_steps = episode_length if episode_length is not None else 1000 + + # Deferred — created on first reset()/render()/step() inside the worker + # subprocess to avoid inheriting a stale GPU/EGL context across fork(). + self._env: Any = None + self.task_description = "" + + images = {} + for cam in self.camera_name: + for mapped_cam in self.camera_name_mapping[cam]: + images[mapped_cam] = 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: + r"""Create the underlying ``RoboCasaGymEnv`` on first use. + + Called inside the worker subprocess (after fork/spawn) so each worker + gets its own clean rendering context rather than inheriting a stale one + from the parent (which crashes with ``AsyncVectorEnv``). + """ + if self._env is not None: + return + _import_robocasa_with_version_shim() + from robocasa.wrappers.gym_wrapper import RoboCasaGymEnv + + # RoboCasaGymEnv defaults split="test", which create_env rejects (only + # None/"all"/"pretrain"/"target" are valid). Always pass a valid value. + self._env = RoboCasaGymEnv( + env_name=self.task, + camera_widths=self.observation_width, + camera_heights=self.observation_height, + split=self.split if self.split is not None else "all", + obj_registries=self.obj_registries, + ) + + 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[str, Any]) -> dict[str, Any]: + r"""Convert a ``RoboCasaGymEnv`` observation dict to OpenTau format.""" + # RoboCasaGymEnv emits camera frames under "video.". + images: dict[str, np.ndarray] = {} + for cam in self.camera_name: + key = f"video.{cam}" + if key not in raw_obs: + continue + frame = raw_obs[key] + for mapped_cam in self.camera_name_mapping[cam]: + images[mapped_cam] = frame + + if self.obs_type == "pixels": + return {"pixels": images} + + # `state.*` keys come from PandaOmronKeyConverter inside the wrapper. + 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: + r"""Render the environment and return an RGB array for video recording.""" + self._ensure_env() + assert self._env is not None + return self._env.render() + + def reset(self, seed=None, **kwargs) -> tuple[dict[str, Any], dict[str, Any]]: + r"""Reset the environment, deriving a per-worker seed from ``episode_index``.""" + self._ensure_env() + assert self._env is not None + super().reset(seed=seed) + # Spread the seed across workers so n_envs factories don't all roll the + # same scene: shift an explicit seed by episode_index; with no seed fall + # back to episode_index so each worker is still distinct. + worker_seed = seed + self.episode_index if seed is not None else self.episode_index + raw_obs, _ = self._env.reset(seed=worker_seed) + + 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[dict[str, Any], float, bool, bool, dict[str, Any]]: + r"""Take a step; remaps RoboCasa's ``info["success"]`` to ``is_success``.""" + 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}" + ) + # Policies may emit a padded action wider than ACTION_DIM; keep the + # leading RoboCasa dims (mirrors LiberoEnv). + if len(action) > ACTION_DIM: + action = action[:ACTION_DIM] + + action_dict = convert_action(action) + raw_obs, reward, done, truncated, info = self._env.step(action_dict) + + # RoboCasa reports success under "success"; OpenTau's rollout reads + # "is_success". Bridge the two here. + 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: + # Auto-reset on terminal so the next rollout step starts a fresh + # episode (the rollout masks post-done data). Mirrors LiberoEnv.step; + # the task / done / is_success keys set above are what eval consumes. + self.reset() + + return observation, reward, terminated, truncated, info + + def close(self): + r"""Close the environment and release any resources.""" + if self._env is not None: + self._env.close() + + +def _make_env_fns( + *, + task: str, + n_envs: int, + camera_names: list[str], + obs_type: str, + render_mode: str, + observation_width: int, + observation_height: int, + visualization_width: int, + visualization_height: int, + split: str | None, + episode_length: int | None, + obj_registries: Sequence[str], +) -> list[Callable[[], RoboCasaEnv]]: + """Build ``n_envs`` factory callables for a single task. + + Each factory carries a distinct ``episode_index`` (``0..n_envs-1``) so + ``RoboCasaEnv.reset()`` derives a per-worker seed from the rollout seed. + """ + + def _make_env(episode_index: int) -> RoboCasaEnv: + return RoboCasaEnv( + task=task, + camera_name=camera_names, + obs_type=obs_type, + render_mode=render_mode, + observation_width=observation_width, + observation_height=observation_height, + visualization_width=visualization_width, + visualization_height=visualization_height, + split=split, + episode_length=episode_length, + obj_registries=obj_registries, + episode_index=episode_index, + ) + + return [partial(_make_env, i) for i in range(n_envs)] + + +# main API entry point +def create_robocasa_envs( + task: str, + n_envs: int, + gym_kwargs: dict[str, Any] | None = None, + camera_name: str | Sequence[str] = ",".join(DEFAULT_CAMERAS), + env_cls: type[gym.vector.SyncVectorEnv] | type[gym.vector.AsyncVectorEnv] | None = None, + episode_length: int | None = None, + obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES, +) -> dict[str, dict[int, gym.vector.VectorEnv]]: + r"""Create vectorized RoboCasa365 environments with a consistent return shape. + + Returns: + ``dict[task_name][0] -> vec_env`` (``env_cls([...])`` with ``n_envs`` + factories). Each distinct task is its own group, so eval reports a + per-task ``Success/{task}`` and a per-task ``Eval Videos/{task}_0`` grid. + + ``task`` can be a single task name (``CloseFridge``), a comma-separated list + (``CloseFridge,PickPlaceCoffee``), or a benchmark-group shortcut + (``atomic_seen``/``composite_seen``/``composite_unseen``/``pretrain50``…), + which auto-expands and auto-sets the dataset ``split``. + + When run under an accelerator with multiple processes, tasks are sharded + round-robin (``idx % num_processes == process_index``) so each rank evaluates + a disjoint subset — matching LIBERO and keeping the ``_rank{N}``-strip + uniqueness assumption in ``collect_grid_summary_videos`` valid. + """ + 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) + visualization_width = gym_kwargs.pop("visualization_width", 512) + visualization_height = gym_kwargs.pop("visualization_height", 512) + split = gym_kwargs.pop("split", None) + + camera_names = _parse_camera_names(camera_name) + task_names, group_split = _resolve_tasks(str(task)) + if group_split is not None and split is None: + split = group_split + + # Shard tasks across accelerator ranks (round-robin), so distributed eval + # spreads tasks disjointly and per-task video keys stay unique after the + # `_rank{N}` strip in `collect_grid_summary_videos`. + accelerator = get_proc_accelerator() + if accelerator is not None: + task_names = [ + t + for idx, t in enumerate(task_names) + if idx % accelerator.num_processes == accelerator.process_index + ] + acc_print( + f"Creating RoboCasa envs | tasks={task_names} | split={split} | " + f"n_envs(per task)={n_envs} | rank={accelerator.process_index}" + ) + else: + acc_print(f"Creating RoboCasa envs | tasks={task_names} | split={split} | n_envs(per task)={n_envs}") + + # No tasks on this rank (more ranks than tasks) → empty dict, like LIBERO. + if not task_names: + acc_print("No RoboCasa tasks assigned to this rank, returning empty dict.") + return {} + + out: dict[str, dict[int, Any]] = defaultdict(dict) + for task_name in task_names: + fns = _make_env_fns( + task=task_name, + n_envs=n_envs, + camera_names=camera_names, + obs_type=obs_type, + render_mode=render_mode, + observation_width=observation_width, + observation_height=observation_height, + visualization_width=visualization_width, + visualization_height=visualization_height, + split=split, + episode_length=episode_length, + obj_registries=obj_registries, + ) + out[task_name][0] = env_cls(fns) + acc_print(f"Built vec env | task={task_name} | n_envs={n_envs}") + + # return plain dicts for predictability + return {name: dict(task_map) for name, task_map in out.items()} diff --git a/src/opentau/scripts/eval.py b/src/opentau/scripts/eval.py index 4cdf0855..5db4e547 100644 --- a/src/opentau/scripts/eval.py +++ b/src/opentau/scripts/eval.py @@ -823,6 +823,17 @@ def eval_policy_all( """ start_t = time.time() + # `recording_root` records rollouts through the LIBERO-specific dataset recorder + # (``libero_dataset_recorder``, with a hardcoded ``LIBERO_TASKS`` list), so it only + # makes sense for a LIBERO env. Fail fast for any other env rather than silently + # mislabeling the recorded dataset — until an env-aware rollout recorder exists. + if cfg.eval.recording_root is not None and not isinstance(cfg.env, LiberoEnv): + raise NotImplementedError( + f"eval.recording_root is only supported for the LIBERO env (it uses the LIBERO " + f"dataset recorder), but env.type={cfg.env.type!r}. Unset recording_root, or add " + f"an env-aware rollout recorder." + ) + # Flatten envs into list of (task_group, task_id, env) tasks = [(tg, tid, vec) for tg, group in envs.items() for tid, vec in group.items()] diff --git a/tests/envs/test_factory.py b/tests/envs/test_factory.py index aa90775f..1f08ed25 100644 --- a/tests/envs/test_factory.py +++ b/tests/envs/test_factory.py @@ -20,7 +20,7 @@ import pytest from opentau.configs.train import TrainPipelineConfig -from opentau.envs.configs import LiberoEnv +from opentau.envs.configs import LiberoEnv, RoboCasaEnv from opentau.envs.factory import make_env_config, make_envs @@ -38,6 +38,11 @@ def test_make_env_config_libero(self): assert config.task == "libero_10" assert config.task_ids is None + def test_make_env_config_robocasa(self): + config = make_env_config("robocasa") + assert isinstance(config, RoboCasaEnv) + assert config.task == "CloseFridge" + class TestMakeEnv: """Test cases for make_env function""" diff --git a/tests/envs/test_robocasa.py b/tests/envs/test_robocasa.py new file mode 100644 index 00000000..d9d7ae5c --- /dev/null +++ b/tests/envs/test_robocasa.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python + +# Copyright 2026 Tensor Auto Inc. 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. + +r"""CPU-only tests for the RoboCasa env integration. + +These never import the ``robocasa`` / ``robosuite`` sim packages: the wrapper +defers those imports into ``_ensure_env`` / ``_resolve_tasks``, so config +registration, the factory dispatch, per-rank task sharding, and the pure +action/camera helpers are all exercisable without a GPU or the sim installed. +Real sim rollouts are validated separately on a CUDA box. +""" + +from unittest.mock import Mock, patch + +import numpy as np +import pytest + +from opentau.configs.train import TrainPipelineConfig +from opentau.configs.types import FeatureType +from opentau.constants import ACTION, OBS_IMAGES, OBS_STATE +from opentau.envs.configs import EnvConfig, RoboCasaEnv +from opentau.envs.factory import make_env_config, make_envs +from opentau.envs.robocasa import ( + ACTION_DIM, + OBS_STATE_DIM, + _default_camera_name_mapping, + _import_robocasa_with_version_shim, + _parse_camera_names, + _resolve_tasks, + convert_action, + create_robocasa_envs, +) + + +class TestRoboCasaConfig: + """Config registration, defaults, and feature wiring.""" + + def test_registered_in_choice_registry(self): + assert "robocasa" in EnvConfig._choice_registry + assert EnvConfig._choice_registry["robocasa"] is RoboCasaEnv + assert RoboCasaEnv().type == "robocasa" + + def test_default_values(self): + cfg = RoboCasaEnv() + assert cfg.task == "CloseFridge" + assert cfg.fps == 20 + assert cfg.episode_length == 1000 + assert cfg.obs_type == "pixels_agent_pos" + assert cfg.observation_height == 256 + assert cfg.observation_width == 256 + assert cfg.split is None + assert cfg.obj_registries == ["lightwheel"] + assert len(_parse_camera_names(cfg.camera_name)) == 3 + + def test_action_and_state_features(self): + cfg = RoboCasaEnv() + assert cfg.features["action"].type == FeatureType.ACTION + assert cfg.features["action"].shape == (12,) + # pixels_agent_pos adds a 16-D proprio state. + assert cfg.features["agent_pos"].type == FeatureType.STATE + assert cfg.features["agent_pos"].shape == (16,) + assert cfg.features_map["action"] == ACTION + assert cfg.features_map["agent_pos"] == OBS_STATE + + def test_camera_features_map_uses_image_convention(self): + """The 3 raw cameras map to OpenTau's image/image2/image3 keys.""" + cfg = RoboCasaEnv() + cams = _parse_camera_names(cfg.camera_name) + assert cfg.features_map[f"pixels/{cams[0]}"] == f"{OBS_IMAGES}.image" + assert cfg.features_map[f"pixels/{cams[1]}"] == f"{OBS_IMAGES}.image2" + assert cfg.features_map[f"pixels/{cams[2]}"] == f"{OBS_IMAGES}.image3" + for cam in cams: + assert cfg.features[f"pixels/{cam}"].type == FeatureType.VISUAL + assert cfg.features[f"pixels/{cam}"].shape == (256, 256, 3) + + def test_pixels_only_omits_agent_pos(self): + cfg = RoboCasaEnv(obs_type="pixels") + assert "agent_pos" not in cfg.features + + @pytest.mark.parametrize("bad_fps", [0, -5]) + def test_rejects_nonpositive_fps(self, bad_fps): + with pytest.raises(ValueError, match="must be positive"): + RoboCasaEnv(fps=bad_fps) + + def test_rejects_unsupported_obs_type(self): + with pytest.raises(ValueError, match="Unsupported obs_type"): + RoboCasaEnv(obs_type="state") + + def test_gym_kwargs_carries_obs_params_and_split(self): + cfg = RoboCasaEnv(split="pretrain") + kwargs = cfg.gym_kwargs + assert kwargs["obs_type"] == "pixels_agent_pos" + assert kwargs["observation_height"] == 256 + assert kwargs["observation_width"] == 256 + assert kwargs["visualization_height"] == 512 + assert kwargs["split"] == "pretrain" + + def test_gym_kwargs_omits_split_when_none(self): + assert "split" not in RoboCasaEnv().gym_kwargs + + +class TestMakeEnvConfigRoboCasa: + """``make_env_config`` dispatch.""" + + def test_make_env_config_robocasa(self): + cfg = make_env_config("robocasa") + assert isinstance(cfg, RoboCasaEnv) + assert cfg.task == "CloseFridge" + + def test_make_env_config_robocasa_with_overrides(self): + cfg = make_env_config("robocasa", task="OpenDrawer", split="target") + assert cfg.task == "OpenDrawer" + assert cfg.split == "target" + + +class TestPureHelpers: + """Helpers that never touch the sim.""" + + def test_convert_action_layout(self): + flat = np.arange(ACTION_DIM, dtype=np.float32) + out = convert_action(flat) + np.testing.assert_array_equal(out["action.base_motion"], flat[0:4]) + np.testing.assert_array_equal(out["action.control_mode"], flat[4:5]) + np.testing.assert_array_equal(out["action.end_effector_position"], flat[5:8]) + np.testing.assert_array_equal(out["action.end_effector_rotation"], flat[8:11]) + np.testing.assert_array_equal(out["action.gripper_close"], flat[11:12]) + + def test_parse_camera_names(self): + assert _parse_camera_names("a, b ,c") == ["a", "b", "c"] + assert _parse_camera_names(["x", "y"]) == ["x", "y"] + with pytest.raises(ValueError): + _parse_camera_names(" , ") + + def test_default_camera_name_mapping_is_positional(self): + mapping = _default_camera_name_mapping(["cam_a", "cam_b", "cam_c"]) + assert mapping == {"cam_a": ["camera0"], "cam_b": ["camera1"], "cam_c": ["camera2"]} + + def test_resolve_single_and_comma_tasks_no_split(self): + # Concrete task names never import robocasa and leave split untouched. + assert _resolve_tasks("CloseFridge") == (["CloseFridge"], None) + names, split = _resolve_tasks("CloseFridge, PickPlaceCoffee") + assert names == ["CloseFridge", "PickPlaceCoffee"] + assert split is None + + def test_resolve_empty_task_raises(self): + with pytest.raises(ValueError, match="at least one RoboCasa task"): + _resolve_tasks(" , ") + + def test_constants(self): + assert ACTION_DIM == 12 + assert OBS_STATE_DIM == 16 + + def test_version_shim_is_noop_when_robocasa_already_imported(self): + # When robocasa is already in sys.modules the shim must return early + # without importing mujoco/numpy, so CPU-only runs (no sim) never touch + # those modules. Inject a placeholder and assert it returns cleanly. + import sys + + had = "robocasa" in sys.modules + saved = sys.modules.get("robocasa") + sys.modules["robocasa"] = object() + try: + _import_robocasa_with_version_shim() # must not raise / import anything + finally: + if had: + sys.modules["robocasa"] = saved + else: + sys.modules.pop("robocasa", None) + + +def _mock_accelerator(num_processes: int, process_index: int) -> Mock: + acc = Mock() + acc.num_processes = num_processes + acc.process_index = process_index + return acc + + +class TestCreateRoboCasaEnvs: + """``create_robocasa_envs`` return shape and per-rank task sharding. + + ``env_cls`` is mocked so the env factories are never invoked — no + ``RoboCasaEnv`` is constructed and the sim is never imported. + """ + + def test_returns_one_vec_env_per_task(self): + sentinel = Mock(name="vec_env") + env_cls = Mock(return_value=sentinel) + with patch("opentau.envs.robocasa.get_proc_accelerator", return_value=None): + out = create_robocasa_envs(task="A,B", n_envs=3, env_cls=env_cls) + assert set(out.keys()) == {"A", "B"} + assert out["A"][0] is sentinel and out["B"][0] is sentinel + # Each task built one vec env from exactly n_envs factory callables. + assert env_cls.call_count == 2 + fns = env_cls.call_args[0][0] + assert len(fns) == 3 + assert all(callable(f) for f in fns) + + @pytest.mark.parametrize( + ("process_index", "expected"), + [(0, {"A", "C"}), (1, {"B", "D"})], + ) + def test_round_robin_task_sharding(self, process_index, expected): + env_cls = Mock(return_value=Mock()) + acc = _mock_accelerator(num_processes=2, process_index=process_index) + with patch("opentau.envs.robocasa.get_proc_accelerator", return_value=acc): + out = create_robocasa_envs(task="A,B,C,D", n_envs=1, env_cls=env_cls) + assert set(out.keys()) == expected + + def test_more_ranks_than_tasks_returns_empty(self): + env_cls = Mock(return_value=Mock()) + acc = _mock_accelerator(num_processes=4, process_index=2) + with patch("opentau.envs.robocasa.get_proc_accelerator", return_value=acc): + out = create_robocasa_envs(task="A", n_envs=1, env_cls=env_cls) + assert out == {} + env_cls.assert_not_called() + + def test_rejects_bad_n_envs(self): + # n_envs is validated before any accelerator call, so no patch needed. + with pytest.raises(ValueError, match="positive int"): + create_robocasa_envs(task="A", n_envs=0, env_cls=Mock()) + + def test_rejects_non_callable_env_cls(self): + with pytest.raises(ValueError, match="env_cls must be a callable"): + create_robocasa_envs(task="A", n_envs=1, env_cls=None) + + +class TestMakeEnvsDispatch: + """``make_envs`` routes RoboCasa configs to ``create_robocasa_envs``.""" + + @pytest.fixture + def mock_train_cfg(self): + return Mock(spec=TrainPipelineConfig) + + def test_make_envs_dispatches_to_create_robocasa_envs(self, mock_train_cfg): + expected = {"CloseFridge": {0: Mock()}} + with patch("opentau.envs.robocasa.create_robocasa_envs", return_value=expected) as mock_create: + result = make_envs(RoboCasaEnv(), mock_train_cfg, n_envs=2, use_async_envs=False) + + assert result is expected + mock_create.assert_called_once() + kwargs = mock_create.call_args.kwargs + assert kwargs["task"] == "CloseFridge" + assert kwargs["n_envs"] == 2 + assert kwargs["episode_length"] == 1000 + assert kwargs["obj_registries"] == ("lightwheel",) + # SyncVectorEnv path: env_cls is the class itself. + import gymnasium as gym + + assert kwargs["env_cls"] is gym.vector.SyncVectorEnv diff --git a/tests/scripts/test_eval.py b/tests/scripts/test_eval.py index 9ae62e9b..1e6f40e3 100644 --- a/tests/scripts/test_eval.py +++ b/tests/scripts/test_eval.py @@ -13,8 +13,12 @@ # limitations under the License. from pathlib import Path +from unittest.mock import Mock -from opentau.scripts.eval import collect_grid_summary_videos +import pytest + +from opentau.envs.configs import RoboCasaEnv +from opentau.scripts.eval import collect_grid_summary_videos, eval_policy_all def _touch(p: Path): @@ -35,3 +39,15 @@ def test_collect_grid_summary_videos_strips_rank_and_excludes_clips(tmp_path): def test_collect_grid_summary_videos_missing_dir(tmp_path): assert collect_grid_summary_videos(tmp_path / "does_not_exist") == [] + + +def test_eval_policy_all_rejects_recording_root_for_non_libero_env(): + """eval.recording_root drives the LIBERO-only dataset recorder, so it must fail + fast (before any rollout) for a non-LIBERO env rather than silently mislabel the + recorded dataset.""" + cfg = Mock() + cfg.eval.recording_root = "/tmp/robocasa-rollout" + cfg.env = RoboCasaEnv() # not a LIBERO env + + with pytest.raises(NotImplementedError, match="recording_root"): + eval_policy_all({}, policy=Mock(), n_episodes=1, cfg=cfg) diff --git a/uv.lock b/uv.lock index 0b3115db..23530208 100644 --- a/uv.lock +++ b/uv.lock @@ -485,21 +485,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, ] -[[package]] -name = "daqp" -version = "0.8.7" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/91/2d/fd713bbd3660d5c5306599dd138a341ff70cff246c7ccffc7a8d586e7004/daqp-0.8.7.tar.gz", hash = "sha256:62cfd3208a9841ffb6a87d145bce930a0e87ecea802a5ffed5d0146c4b133a54", size = 37292, upload-time = "2026-05-19T17:21:44.084Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3c/00/97043ee6f2b13a63f0b59f776055dc199621ca8773f9d627fae3d2bfabe6/daqp-0.8.7-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:67e99835efe68e3ba1bc721074e726081e65bfc51f5a81b6ac9407979ad4d33e", size = 166267, upload-time = "2026-05-19T17:34:17.97Z" }, - { url = "https://files.pythonhosted.org/packages/e2/fe/4947398babec3a517f43e10555d2d37f8dea19dbc9a09bb398080088e3ac/daqp-0.8.7-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:de5540765884e56d0fb8bf945a8d500aa9ca268b506535c4dac37ff48c8bf9d5", size = 156307, upload-time = "2026-05-19T17:34:19.415Z" }, - { url = "https://files.pythonhosted.org/packages/bd/af/a8e544f8e9dd143dbc0cb2afdbd0221f2a272dd662323ab1243857269d5e/daqp-0.8.7-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2307d8e96417c6965e4bd6e7857e80d1ca0c52bec2cba13173dec736252d52c4", size = 838020, upload-time = "2026-05-19T17:21:32.219Z" }, - { url = "https://files.pythonhosted.org/packages/80/a5/c80745b09d54eb83ddef70dacc74499360395f3d9f7cbba8bfc95165ed40/daqp-0.8.7-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4be0d4256100a4fd35d01907ae1d4ad164e8a152707e7ddfebe386b082652242", size = 776438, upload-time = "2026-05-19T18:15:19.014Z" }, - { url = "https://files.pythonhosted.org/packages/0c/06/d11d5bb0ad9d14ee2b0db09eeba796ea56168c7ce20494e22a5f25482c5d/daqp-0.8.7-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:62dd2c6d14c8baefbc489ae6ba178debf36388eebd12bfa6fefed0599675fc69", size = 802887, upload-time = "2026-05-19T18:15:20.662Z" }, - { url = "https://files.pythonhosted.org/packages/a0/dd/2a12b44e054c4e416ac27b0e115c518cec79c7efe62f5306b12113815657/daqp-0.8.7-cp310-cp310-win32.whl", hash = "sha256:97ba8d791ae9fa27233d2e35458753e51e6868f8bb69a969db2f5a58845b26a6", size = 112174, upload-time = "2026-05-19T17:28:59.071Z" }, - { url = "https://files.pythonhosted.org/packages/a3/97/c790956cfda45fa350a7edcc04e52d91807af6872a87d8887db15140a29b/daqp-0.8.7-cp310-cp310-win_amd64.whl", hash = "sha256:dfcf324bb1b1092c2732d6ec1080cf526a52a915a52aa2a0161e874dc6c2d57e", size = 135871, upload-time = "2026-05-19T17:29:00.191Z" }, -] - [[package]] name = "datasets" version = "4.5.0" @@ -1439,6 +1424,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b2/bc/4689e1ba0c073c196b594471eb21be0aa51d9e64b911728aa13cd85ef0ae/llvmlite-0.46.0-cp310-cp310-win_amd64.whl", hash = "sha256:e7a34d4aa6f9a97ee006b504be6d2b8cb7f755b80ab2f344dda1ef992f828559", size = 38138651, upload-time = "2025-12-08T18:14:45.845Z" }, ] +[[package]] +name = "lxml" +version = "6.1.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/3b/aab6728cae887456f409b4d75e8a01856e4f04bd510de38052a47768b680/lxml-6.1.1.tar.gz", hash = "sha256:ba96ae44888e0185281e937633a743ea90d5a196c6000f82565ebb0580012d40", size = 4197430, upload-time = "2026-05-18T19:19:06.424Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/da/dbe4dfc01ac226fb0504fad035f4d69f3202f3502e20e68537631daddd96/lxml-6.1.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:09dd5b7075dc2f7709654a46543ba1ea3c2e217b2ed8fbd413a8a945a0f40f60", size = 8541124, upload-time = "2026-05-18T19:17:11.589Z" }, + { url = "https://files.pythonhosted.org/packages/78/20/f7095ed9fc2c025f9cfe71cc6ec9f1feb05624edc1812423b5f1aecf3d4b/lxml-6.1.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:f6ac4ef4d82dff54670227a69c67782ae0b811b5cf6b17954f1e8f7502fc0d1d", size = 4602783, upload-time = "2026-05-18T19:17:20.888Z" }, + { url = "https://files.pythonhosted.org/packages/4a/a4/65c63ca98bd129f6cff7b8c2fa48953ab058cc6005b541354e7dd54d8000/lxml-6.1.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:556e94a63c9b04716f8e4de2abb65775061f846e89331b6c5be79183a24f98ea", size = 5002687, upload-time = "2026-05-18T19:17:01.738Z" }, + { url = "https://files.pythonhosted.org/packages/96/1d/ab7a5c4b5a394d98a94e2d0fc67bab8297597426770dd4978370fbdaf531/lxml-6.1.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5c6bf403fbb3b3e348a561a5f4f0b9961835657981c802a1df03653eef8a9074", size = 5155099, upload-time = "2026-05-18T19:17:05.159Z" }, + { url = "https://files.pythonhosted.org/packages/d0/b1/07603bfeeb891a2596d5c2a68f7d2f70f7d11c841ebe391412c69c2857b0/lxml-6.1.1-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1dde6131244bba38a17c745836ba190bc753fd73c9291666287fd0a3fa3dcf30", size = 5057225, upload-time = "2026-05-18T19:17:08.117Z" }, + { url = "https://files.pythonhosted.org/packages/7a/16/cb391ee4b90186fa16d9ebcbe3ea96c71b8da3b0686386c8dcbcc3c67d44/lxml-6.1.1-cp310-cp310-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:98fc784c2c1440667aeedf8465bdfe10208acf0ead656a2c68627299f546b315", size = 5287643, upload-time = "2026-05-18T19:17:11.507Z" }, + { url = "https://files.pythonhosted.org/packages/eb/d6/b619717f918fd76747448fdbaee0e769edbc70e659b5b5d0112b7020b7a3/lxml-6.1.1-cp310-cp310-manylinux_2_28_i686.whl", hash = "sha256:add8cf6ddf9a65116119a28ece0f7886e30af27ba724a7594305f1d1b58a92a1", size = 5412445, upload-time = "2026-05-18T19:17:22.182Z" }, + { url = "https://files.pythonhosted.org/packages/c6/80/12bc5390ac0a3edeb579d9535e5049a5dda663438728e179d52fb319c33a/lxml-6.1.1-cp310-cp310-manylinux_2_31_armv7l.whl", hash = "sha256:cf9d57306d848218f3601fee7601fab1a327c942d56e2e97610583cb4dd74206", size = 4770864, upload-time = "2026-05-18T19:17:26.851Z" }, + { url = "https://files.pythonhosted.org/packages/0b/59/6500c09da3137f54f020e908d81cfc5ee3e8888e908fd380207afad7c2e6/lxml-6.1.1-cp310-cp310-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:88136950da4d13c318bde414ce10219931937851327f44328f2df4d2c4614067", size = 5359594, upload-time = "2026-05-18T19:17:32.527Z" }, + { url = "https://files.pythonhosted.org/packages/f2/9b/f64b4cc6b7ebcf75d95af3cde934d254b5f2f10d4163928d838d86b6eb48/lxml-6.1.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:cecdd5dfdc87b1fd87dbf81d4b037a544f47f4c744200a67013771682d67686a", size = 5107713, upload-time = "2026-05-18T19:17:04.402Z" }, + { url = "https://files.pythonhosted.org/packages/16/19/c7388ad5d3a72315d2832dc1458cbf4f2af7f2b990b606ff4876efd04511/lxml-6.1.1-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:cd312b9692e831d2ffcad61eab31d91d4b4655a962e61de8fb410472cbcd37aa", size = 4803973, upload-time = "2026-05-18T19:17:06.545Z" }, + { url = "https://files.pythonhosted.org/packages/3f/22/76197f0bbf165f0b9e75be59be4997e5259cde973f12f098c1b54c7f5d60/lxml-6.1.1-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:5b7328b46d49fc9477d91ae8f6d55340347d827b7734ba3ea33faae0efef1383", size = 5349925, upload-time = "2026-05-18T19:17:09.743Z" }, + { url = "https://files.pythonhosted.org/packages/24/52/d2a0cfeccb9bcdc47c7ee05cdae5d69b48c9acf20997790a6338bb0d0b3b/lxml-6.1.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:37a58976370f36d9329d118ad0b953c5aeb9119ac9c6a4e258942a225d0573a1", size = 5309825, upload-time = "2026-05-18T19:17:13.831Z" }, + { url = "https://files.pythonhosted.org/packages/19/4a/b30944266776c2f49749ef2445aa7e78898194134b80ad776386f61b56ae/lxml-6.1.1-cp310-cp310-win32.whl", hash = "sha256:cea3f4c1af79af13cdb2da0c028111d8f8522d4f22a000c82385535f24e5cf3a", size = 3598402, upload-time = "2026-05-18T19:17:08.21Z" }, + { url = "https://files.pythonhosted.org/packages/9e/97/33691c66a4d7ec1a5a98e7c909a5b83ee45c7f7ba4cf92b1c4cf26e98079/lxml-6.1.1-cp310-cp310-win_amd64.whl", hash = "sha256:3abf332af33a74288675d936fe861fd4344da0dd6622193fbc4f2bfbb35536b5", size = 4021295, upload-time = "2026-05-18T19:17:28.638Z" }, + { url = "https://files.pythonhosted.org/packages/d0/5f/26a4dd0e12b9456ff7b12a21af5b491eb6629680d1edd73f4140fd386bcf/lxml-6.1.1-cp310-cp310-win_arm64.whl", hash = "sha256:8dadbe5b217ff35b6a8d16610dd710219b59b76d13f0e3f0d9f36786206e4485", size = 3667717, upload-time = "2026-05-19T19:22:44.474Z" }, +] + [[package]] name = "lz4" version = "4.4.5" @@ -1571,21 +1580,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2c/19/04f9b178c2d8a15b076c8b5140708fa6ffc5601fb6f1e975537072df5b2a/mergedeep-1.3.4-py3-none-any.whl", hash = "sha256:70775750742b25c0d8f36c55aed03d24c3384d17c951b3175d898bd778ef0307", size = 6354, upload-time = "2021-02-05T18:55:29.583Z" }, ] -[[package]] -name = "mink" -version = "0.0.13" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "mujoco" }, - { name = "numpy" }, - { name = "qpsolvers", extra = ["daqp"] }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/9d/a3/dcb9ba6099894bd5ae50a3edce4b5b5bdac42a648b919cd06a177f58b0d8/mink-0.0.13.tar.gz", hash = "sha256:481f6187dc3fd320e2c0b6e8393d79c52bc70272835d641ecdf6e5fdf0e6b835", size = 693562, upload-time = "2025-09-12T21:29:46.821Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/66/3e/3340943f4c614c584efd56ea2ac42cf5358f469a81b424e4acf4fccaa33e/mink-0.0.13-py3-none-any.whl", hash = "sha256:07b9e780f79937b082de01a0fea64b33aa43160ad088217998a212fc0bfb1173", size = 914736, upload-time = "2025-09-12T21:29:42.051Z" }, -] - [[package]] name = "ml-dtypes" version = "0.5.4" @@ -2251,6 +2245,11 @@ libero = [ { name = "robosuite" }, { name = "thop" }, ] +robocasa = [ + { name = "mujoco" }, + { name = "robocasa" }, + { name = "robosuite" }, +] trt = [ { name = "tensorrt", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'x86_64' and sys_platform == 'win32')" }, ] @@ -2290,6 +2289,7 @@ requires-dist = [ { name = "mediapipe", specifier = ">=0.10.0" }, { name = "mujoco", marker = "sys_platform == 'linux' and extra == 'libero'", specifier = ">=3.1.6" }, { name = "mujoco", marker = "extra == 'libero'", specifier = ">=3.3.5" }, + { name = "mujoco", marker = "extra == 'robocasa'", specifier = ">=3.3.5" }, { name = "myst-parser", marker = "extra == 'dev'", specifier = ">=4.0.0" }, { name = "ninja", marker = "extra == 'libero'" }, { name = "numba", specifier = ">=0.62.0" }, @@ -2323,8 +2323,10 @@ requires-dist = [ { name = "pyzmq", specifier = ">=26.2.1" }, { name = "rerun-sdk", specifier = ">=0.21.0" }, { name = "rerun-sdk", marker = "extra == 'urdf'", specifier = ">=0.28.2" }, + { name = "robocasa", marker = "extra == 'robocasa'", git = "https://github.com/shuheng-liu/robocasa?rev=f7db21c11f25408d3a59a5e878bb5c7ca9030c4d" }, { name = "robomimic", marker = "extra == 'libero'", specifier = "==0.2.0" }, - { name = "robosuite", marker = "extra == 'libero'", specifier = "==1.5.2" }, + { name = "robosuite", marker = "extra == 'libero'", git = "https://github.com/ARISE-Initiative/robosuite?rev=85abee228d1c43ab1939bce33028099945d453b4" }, + { name = "robosuite", marker = "extra == 'robocasa'", git = "https://github.com/ARISE-Initiative/robosuite?rev=85abee228d1c43ab1939bce33028099945d453b4" }, { name = "rosbags", specifier = ">=0.10.4" }, { name = "scikit-image", specifier = ">=0.23.2" }, { name = "scikit-learn", specifier = ">=1.7.1" }, @@ -2343,7 +2345,7 @@ requires-dist = [ { name = "wandb", specifier = ">=0.27.0" }, { name = "zarr", specifier = ">=2.17.0" }, ] -provides-extras = ["dev", "libero", "urdf", "trt"] +provides-extras = ["dev", "libero", "robocasa", "urdf", "trt"] [[package]] name = "orderly-set" @@ -2648,6 +2650,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2f/02/8559b1f26ee0d502c74f9cca5c0d2fd97e967e083e006bbbb4e97f3a043a/pydantic_core-2.41.5-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:d3a978c4f57a597908b7e697229d996d77a6d3c94901e9edee593adada95ce1a", size = 2147009, upload-time = "2025-11-04T13:43:23.286Z" }, ] +[[package]] +name = "pygame" +version = "2.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/49/cc/08bba60f00541f62aaa252ce0cfbd60aebd04616c0b9574f755b583e45ae/pygame-2.6.1.tar.gz", hash = "sha256:56fb02ead529cee00d415c3e007f75e0780c655909aaa8e8bf616ee09c9feb1f", size = 14808125, upload-time = "2024-09-29T13:41:34.698Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/0b/334c7c50a2979e15f2a027a41d1ca78ee730d5b1c7f7f4b26d7cb899839d/pygame-2.6.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:9beeb647e555afb5657111fa83acb74b99ad88761108eaea66472e8b8547b55b", size = 13109297, upload-time = "2024-09-29T14:25:34.709Z" }, + { url = "https://files.pythonhosted.org/packages/dc/48/f8b1069788d1bd42e63a960d74d3355242480b750173a42b2749687578ca/pygame-2.6.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:10e3d2a55f001f6c0a6eb44aa79ea7607091c9352b946692acedb2ac1482f1c9", size = 12375837, upload-time = "2024-09-29T14:25:50.538Z" }, + { url = "https://files.pythonhosted.org/packages/bc/33/a1310386b8913ce1bdb90c33fa536970e299ad57eb35785f1d71ea1e2ad3/pygame-2.6.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:816e85000c5d8b02a42b9834f761a5925ef3377d2924e3a7c4c143d2990ce5b8", size = 13607860, upload-time = "2024-09-29T11:10:44.173Z" }, + { url = "https://files.pythonhosted.org/packages/88/0f/4e37b115056e43714e7550054dd3cd7f4d552da54d7fc58a2fb1407acda5/pygame-2.6.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:8a78fd030d98faab4a8e27878536fdff7518d3e062a72761c552f624ebba5a5f", size = 14304696, upload-time = "2024-09-29T11:39:46.724Z" }, + { url = "https://files.pythonhosted.org/packages/11/b3/de6ed93ae483cf3bac8f950a955e83f7ffe59651fd804d100fff65d66d6c/pygame-2.6.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:da3ad64d685f84a34ebe5daacb39fff14f1251acb34c098d760d63fee768f50c", size = 13977684, upload-time = "2024-09-29T11:39:49.921Z" }, + { url = "https://files.pythonhosted.org/packages/d3/05/d86440aa879708c41844bafc6b3eb42c6d8cf54082482499b53139133e2a/pygame-2.6.1-cp310-cp310-win32.whl", hash = "sha256:9dd5c054d4bd875a8caf978b82672f02bec332f52a833a76899220c460bb4b58", size = 10251775, upload-time = "2024-09-29T11:40:34.952Z" }, + { url = "https://files.pythonhosted.org/packages/38/88/8de61324775cf2c844a51d8db14a8a6d2a9092312f27678f6eaa3a460376/pygame-2.6.1-cp310-cp310-win_amd64.whl", hash = "sha256:00827aba089355925902d533f9c41e79a799641f03746c50a374dc5c3362e43d", size = 10618801, upload-time = "2024-09-29T12:13:25.284Z" }, +] + [[package]] name = "pygments" version = "2.19.2" @@ -2975,9 +2992,6 @@ wheels = [ ] [package.optional-dependencies] -daqp = [ - { name = "daqp" }, -] quadprog = [ { name = "quadprog" }, ] @@ -3074,6 +3088,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d8/13/e65e3bcabc69f93e3b768d78e7c82070edb1d82147dba81ca799a5de4e9b/rerun_sdk-0.29.0-cp310-abi3-win_amd64.whl", hash = "sha256:bd3fc3a9f526935674ed6a021d20f802f40cdfa904e6247a779713cde038385b", size = 108256716, upload-time = "2026-01-30T10:05:00.921Z" }, ] +[[package]] +name = "robocasa" +version = "1.0.1" +source = { git = "https://github.com/shuheng-liu/robocasa?rev=f7db21c11f25408d3a59a5e878bb5c7ca9030c4d#f7db21c11f25408d3a59a5e878bb5c7ca9030c4d" } +dependencies = [ + { name = "gymnasium" }, + { name = "h5py" }, + { name = "imageio" }, + { name = "lxml" }, + { name = "mujoco" }, + { name = "numba" }, + { name = "numpy" }, + { name = "pillow" }, + { name = "pygame" }, + { name = "pynput" }, + { name = "pyyaml" }, + { name = "scipy" }, + { name = "termcolor" }, + { name = "tqdm" }, +] + [[package]] name = "robomimic" version = "0.2.0" @@ -3097,9 +3132,8 @@ sdist = { url = "https://files.pythonhosted.org/packages/3d/c3/44b1d1ea4bcb4bbed [[package]] name = "robosuite" version = "1.5.2" -source = { registry = "https://pypi.org/simple" } +source = { git = "https://github.com/ARISE-Initiative/robosuite?rev=85abee228d1c43ab1939bce33028099945d453b4#85abee228d1c43ab1939bce33028099945d453b4" } dependencies = [ - { name = "mink" }, { name = "mujoco" }, { name = "numba" }, { name = "numpy" }, @@ -3112,10 +3146,6 @@ dependencies = [ { name = "termcolor" }, { name = "tqdm" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/27/96/13c4c63860a3809740efa7406d01ff8ea3ebbdbd5d35eec0e8ead486f654/robosuite-1.5.2.tar.gz", hash = "sha256:9f128d57e4f090a9d78d93b246bd2a862b7fa6fc2ccb837aac06cf0ab19fde78", size = 155629116, upload-time = "2025-12-24T22:31:47.394Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c8/54/3c0940c92d381eb5837cd29bf2c9fbf0a9dff313a2b2779e3902665276b3/robosuite-1.5.2-py3-none-any.whl", hash = "sha256:b9e02df55de8949739de6ba1d196d2ba59cf33c51a76efcfbe5c2f73cdcf99ce", size = 156827641, upload-time = "2025-12-24T22:31:40.166Z" }, -] [[package]] name = "rosbags"