From 81ee2952a384c929632bc1f06f6a3fa3b61109dc Mon Sep 17 00:00:00 2001 From: Pepijn Date: Fri, 17 Apr 2026 12:06:30 +0100 Subject: [PATCH] fix(robocasa): restrict obj_registries to lightwheel by default MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CloseFridge (and most kitchen tasks) crashed at reset with `ValueError: Probabilities contain NaN` coming out of `sample_kitchen_object_helper`. RoboCasa's upstream default `obj_registries=("objaverse", "lightwheel")` normalizes per-registry candidate counts as probabilities; when a sampled category has zero mjcf paths in every configured registry (because the objaverse asset pack isn't on disk — ~30GB, skipped by our Docker build), the 0/0 divide yields NaNs and `rng.choice` raises. - Add `obj_registries: list[str] = ["lightwheel"]` to `RoboCasaEnv` config; thread it through `create_robocasa_envs`, `_make_env_fns`, and the gym.Env wrapper to the underlying `RoboCasaGymEnv` (which forwards to `create_env` → `robosuite.make` → kitchen env). - Default matches what `download_kitchen_assets --type objs_lw` actually ships, so the env works out of the box without a 30GB objaverse download. - Document the override (`--env.obj_registries='[objaverse,lightwheel]'`) for users who have downloaded the full asset set. Made-with: Cursor --- docs/source/robocasa.mdx | 14 +++++++++++++- src/lerobot/envs/configs.py | 7 +++++++ src/lerobot/envs/robocasa.py | 20 +++++++++++++++++++- 3 files changed, 39 insertions(+), 2 deletions(-) diff --git a/docs/source/robocasa.mdx b/docs/source/robocasa.mdx index e9e7b4e9f..8ded2c706 100644 --- a/docs/source/robocasa.mdx +++ b/docs/source/robocasa.mdx @@ -32,7 +32,13 @@ After following the LeRobot installation instructions: ```bash pip install -e ".[robocasa]" python -m robocasa.scripts.setup_macros -python -m robocasa.scripts.download_kitchen_assets +# Lightweight assets only (lightwheel registry, ~2GB). Enough for the +# default env out of the box. +python -m robocasa.scripts.download_kitchen_assets --type tex fixtures_lw objs_lw +# Optional: full objaverse/aigen registries (~30GB) for richer object +# variety. If you download these, pass them via `--env.obj_registries` +# at eval/train time (see below). +# python -m robocasa.scripts.download_kitchen_assets --type objaverse ``` @@ -44,6 +50,12 @@ export MUJOCO_GL=egl # for headless servers (HPC, cloud) +By default the env samples objects only from the `lightwheel` registry (the one shipped by `--type objs_lw`), which avoids a `Probabilities contain NaN` crash when the objaverse/aigen packs are absent. If you've downloaded the full asset set, enable it with: + +```bash +--env.obj_registries='[objaverse,lightwheel]' +``` + ## Policy inputs and outputs **Observations** (raw RoboCasa camera names are preserved verbatim): diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index 9bc7ff09c..68a1538c3 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -508,6 +508,12 @@ class RoboCasaEnv(EnvConfig): observation_height: int = 256 observation_width: int = 256 split: str | None = None + # Object-mesh registries to sample from. Upstream default is + # ("objaverse", "lightwheel"), but objaverse is ~30GB and the CI image + # only ships the lightwheel pack. Override to include objaverse once + # you've run `python -m robocasa.scripts.download_kitchen_assets + # --type objaverse` locally. + obj_registries: list[str] = field(default_factory=lambda: ["lightwheel"]) features: dict[str, PolicyFeature] = field( default_factory=lambda: {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(12,))} ) @@ -557,6 +563,7 @@ class RoboCasaEnv(EnvConfig): gym_kwargs=self.gym_kwargs, env_cls=env_cls, episode_length=self.episode_length, + obj_registries=tuple(self.obj_registries), ) diff --git a/src/lerobot/envs/robocasa.py b/src/lerobot/envs/robocasa.py index 4e31f4365..a39f05b9b 100644 --- a/src/lerobot/envs/robocasa.py +++ b/src/lerobot/envs/robocasa.py @@ -44,6 +44,16 @@ DEFAULT_CAMERAS = [ "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 users — including our CI image — 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` (0/0 divide in the probability +# normalization). Restricting to registries that are actually on disk +# avoids the NaN and matches what the asset download provides. +DEFAULT_OBJ_REGISTRIES: tuple[str, ...] = ("lightwheel",) + # Task-group shortcuts accepted as `--env.task`. When the user passes one of # these names, we expand it to the upstream RoboCasa task list and auto-set # the dataset split. Individual task names (optionally comma-separated) still @@ -132,6 +142,7 @@ class RoboCasaEnv(gym.Env): observation_height: int = 256, split: str | None = None, episode_length: int | None = None, + obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES, ): super().__init__() self.task = task @@ -140,6 +151,7 @@ class RoboCasaEnv(gym.Env): self.observation_width = observation_width self.observation_height = observation_height self.split = split + self.obj_registries = tuple(obj_registries) self.camera_name = _parse_camera_names(camera_name) @@ -197,12 +209,14 @@ class RoboCasaEnv(gym.Env): # RoboCasaGymEnv defaults split="test", which create_env rejects # (only None/"all"/"pretrain"/"target" are valid). Always pass a - # valid value so we don't hit that default. + # valid value so we don't hit that default. Extra kwargs are + # forwarded to the underlying kitchen env via create_env/robosuite.make. 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() @@ -286,6 +300,7 @@ def _make_env_fns( observation_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.""" @@ -299,6 +314,7 @@ def _make_env_fns( observation_height=observation_height, split=split, episode_length=episode_length, + obj_registries=obj_registries, **kwargs, ) @@ -312,6 +328,7 @@ def create_robocasa_envs( camera_name: str | Sequence[str] = ",".join(DEFAULT_CAMERAS), env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None, episode_length: int | None = None, + obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES, ) -> dict[str, dict[int, Any]]: """Create vectorized RoboCasa365 environments with a consistent return shape. @@ -362,6 +379,7 @@ def create_robocasa_envs( observation_height=observation_height, split=split, episode_length=episode_length, + obj_registries=obj_registries, ) if is_async: