feat(robocasa): raw camera names + benchmark-group task shortcuts

Align the LeRobot env with RoboCasa's native conventions so policies
trained on the upstream datasets don't need a --rename_map at eval
time, and expose the standard task groups as first-class --env.task
values.

- Preserve raw RoboCasa camera names (e.g. robot0_agentview_left)
  as observation.images.<name> end-to-end. Drops camera_name_mapping
  and DEFAULT_CAMERA_NAME_MAPPING; features/features_map are now
  built dynamically from the parsed camera list.
- Accept benchmark-group names as --env.task: atomic_seen,
  composite_seen, composite_unseen, pretrain50/100/200/300. Expanded
  lazily via robocasa.utils.dataset_registry and auto-sets the
  split ("target" | "pretrain").
- Update CI smoke-eval rename_map to map raw cam names to the
  camera1/2/3 keys expected by lerobot/smolvla_robocasa.

Made-with: Cursor
This commit is contained in:
Pepijn
2026-04-17 10:28:10 +01:00
parent 14c62b34d5
commit 77b1c399cb
3 changed files with 90 additions and 117 deletions
+1 -1
View File
@@ -366,7 +366,7 @@ jobs:
--eval.n_episodes=1 \ --eval.n_episodes=1 \
--eval.use_async_envs=false \ --eval.use_async_envs=false \
--policy.device=cuda \ --policy.device=cuda \
'--rename_map={\"observation.images.image\": \"observation.images.camera1\", \"observation.images.image2\": \"observation.images.camera2\", \"observation.images.image3\": \"observation.images.camera3\"}' \ '--rename_map={\"observation.images.robot0_agentview_left\": \"observation.images.camera1\", \"observation.images.robot0_eye_in_hand\": \"observation.images.camera2\", \"observation.images.robot0_agentview_right\": \"observation.images.camera3\"}' \
--output_dir=/tmp/eval-artifacts --output_dir=/tmp/eval-artifacts
" "
+17 -35
View File
@@ -505,49 +505,32 @@ class RoboCasaEnv(EnvConfig):
obs_type: str = "pixels_agent_pos" obs_type: str = "pixels_agent_pos"
render_mode: str = "rgb_array" render_mode: str = "rgb_array"
camera_name: str = "robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right" 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_height: int = 256
observation_width: int = 256 observation_width: int = 256
split: str | None = None split: str | None = None
features: dict[str, PolicyFeature] = field( features: dict[str, PolicyFeature] = field(
default_factory=lambda: { default_factory=lambda: {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(12,))}
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",
}
) )
features_map: dict[str, str] = field(default_factory=lambda: {ACTION: ACTION, "agent_pos": OBS_STATE})
def __post_init__(self): def __post_init__(self):
if self.obs_type == "pixels": if self.obs_type not in ("pixels", "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),
)
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}") raise ValueError(f"Unsupported obs_type: {self.obs_type}")
if self.camera_name_mapping is not None: # Preserve raw RoboCasa camera names end-to-end (e.g.
# Update features_map to reflect custom camera name mapping # `observation.images.robot0_agentview_left`). This matches the
mapping = self.camera_name_mapping # naming convention used by the RoboCasa datasets on the Hub, so
cams = [c.strip() for c in self.camera_name.split(",") if c.strip()] # trained policies don't need a `--rename_map` at eval time.
for cam in cams: cams = [c.strip() for c in self.camera_name.split(",") if c.strip()]
mapped = mapping.get(cam, cam) for cam in cams:
self.features_map[f"pixels/{mapped}"] = f"{OBS_IMAGES}.{mapped}" self.features[f"pixels/{cam}"] = PolicyFeature(
type=FeatureType.VISUAL,
shape=(self.observation_height, self.observation_width, 3),
)
self.features_map[f"pixels/{cam}"] = f"{OBS_IMAGES}.{cam}"
if self.obs_type == "pixels_agent_pos":
self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(16,))
@property @property
def gym_kwargs(self) -> dict: def gym_kwargs(self) -> dict:
@@ -571,7 +554,6 @@ class RoboCasaEnv(EnvConfig):
task=self.task, task=self.task,
n_envs=n_envs, n_envs=n_envs,
camera_name=self.camera_name, camera_name=self.camera_name,
camera_name_mapping=self.camera_name_mapping,
gym_kwargs=self.gym_kwargs, gym_kwargs=self.gym_kwargs,
env_cls=env_cls, env_cls=env_cls,
episode_length=self.episode_length, episode_length=self.episode_length,
+72 -81
View File
@@ -35,18 +35,27 @@ ACTION_DIM = 12 # base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gr
ACTION_LOW = -1.0 ACTION_LOW = -1.0
ACTION_HIGH = 1.0 ACTION_HIGH = 1.0
# Default cameras for the PandaOmron robot. # Default PandaOmron cameras. We surface these raw names directly as
# `observation.images.<name>` so the LeRobot dataset/policy keys match
# RoboCasa's native convention (no implicit renaming).
DEFAULT_CAMERAS = [ DEFAULT_CAMERAS = [
"robot0_agentview_left", "robot0_agentview_left",
"robot0_eye_in_hand", "robot0_eye_in_hand",
"robot0_agentview_right", "robot0_agentview_right",
] ]
# Map raw RoboCasa camera names to LeRobot convention names. # Task-group shortcuts accepted as `--env.task`. When the user passes one of
DEFAULT_CAMERA_NAME_MAPPING = { # these names, we expand it to the upstream RoboCasa task list and auto-set
"robot0_agentview_left": "image", # the dataset split. Individual task names (optionally comma-separated) still
"robot0_eye_in_hand": "image2", # take precedence; this only triggers on an exact group-name match.
"robot0_agentview_right": "image3", _TASK_GROUP_SPLITS = {
"atomic_seen": "target",
"composite_seen": "target",
"composite_unseen": "target",
"pretrain50": "pretrain",
"pretrain100": "pretrain",
"pretrain200": "pretrain",
"pretrain300": "pretrain",
} }
@@ -63,21 +72,30 @@ def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
return cams return cams
def convert_state(raw_obs: dict[str, np.ndarray]) -> np.ndarray: def _resolve_tasks(task: str) -> tuple[list[str], str | None]:
"""Concatenate RoboCasa robot state dict into a flat (16,) vector. """Resolve a `--env.task` value to (task_names, split_override).
Layout: base_pos(3) + base_quat(4) + ee_pos_rel(3) + ee_quat_rel(4) + gripper_qpos(2) 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).
""" """
return np.concatenate( key = task.strip()
[ if key in _TASK_GROUP_SPLITS:
raw_obs["robot0_base_pos"], # (3,) from robocasa.utils.dataset_registry import PRETRAINING_TASKS, TARGET_TASKS
raw_obs["robot0_base_quat"], # (4,)
raw_obs["robot0_base_to_eef_pos"], # (3,) combined = {**TARGET_TASKS, **PRETRAINING_TASKS}
raw_obs["robot0_base_to_eef_quat"], # (4,) if key not in combined:
raw_obs["robot0_gripper_qpos"], # (2,) raise ValueError(
], f"Task group '{key}' is not available in this version of robocasa. "
axis=-1, f"Known groups: {sorted(combined.keys())}."
).astype(np.float32) )
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]: def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
@@ -97,9 +115,9 @@ def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
class RoboCasaEnv(gym.Env): class RoboCasaEnv(gym.Env):
"""LeRobot gym.Env wrapper for RoboCasa365 kitchen environments. """LeRobot gym.Env wrapper for RoboCasa365 kitchen environments.
Wraps the RoboCasaGymEnv from the robocasa package and converts its Wraps RoboCasaGymEnv from the robocasa package and converts its
dict-based observations and actions into flat arrays compatible with dict-based observations and actions into the flat arrays LeRobot expects.
the LeRobot evaluation pipeline. Raw RoboCasa camera names are preserved verbatim under `pixels/<cam>`.
""" """
metadata = {"render_modes": ["rgb_array"], "render_fps": 20} metadata = {"render_modes": ["rgb_array"], "render_fps": 20}
@@ -108,7 +126,6 @@ class RoboCasaEnv(gym.Env):
self, self,
task: str, task: str,
camera_name: str | Sequence[str] = ",".join(DEFAULT_CAMERAS), camera_name: str | Sequence[str] = ",".join(DEFAULT_CAMERAS),
camera_name_mapping: dict[str, str] | None = None,
obs_type: str = "pixels_agent_pos", obs_type: str = "pixels_agent_pos",
render_mode: str = "rgb_array", render_mode: str = "rgb_array",
observation_width: int = 256, observation_width: int = 256,
@@ -125,27 +142,23 @@ class RoboCasaEnv(gym.Env):
self.split = split self.split = split
self.camera_name = _parse_camera_names(camera_name) 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 self._max_episode_steps = episode_length if episode_length is not None else 1000
# Deferred — created on first reset() inside the worker subprocess # Deferred — created on first reset() inside the worker subprocess
# to avoid inheriting stale GPU/EGL contexts across fork(). # to avoid inheriting stale GPU/EGL contexts across fork().
self._env = None self._env: Any = None
self.task_description = "" self.task_description = ""
# Build observation space images = {
images = {} cam: spaces.Box(
for cam in self.camera_name:
mapped = self.camera_name_mapping.get(cam, cam)
images[mapped] = spaces.Box(
low=0, low=0,
high=255, high=255,
shape=(self.observation_height, self.observation_width, 3), shape=(self.observation_height, self.observation_width, 3),
dtype=np.uint8, dtype=np.uint8,
) )
for cam in self.camera_name
}
if self.obs_type == "pixels": if self.obs_type == "pixels":
self.observation_space = spaces.Dict({"pixels": spaces.Dict(images)}) self.observation_space = spaces.Dict({"pixels": spaces.Dict(images)})
@@ -182,37 +195,28 @@ class RoboCasaEnv(gym.Env):
return return
from robocasa.wrappers.gym_wrapper import RoboCasaGymEnv from robocasa.wrappers.gym_wrapper import RoboCasaGymEnv
# RoboCasaGymEnv has a broken default split="test" (invalid for create_env # RoboCasaGymEnv defaults split="test", which create_env rejects
# which only accepts None/"all"/"pretrain"/"target"). Always pass a valid # (only None/"all"/"pretrain"/"target" are valid). Always pass a
# value so we don't hit that default. # valid value so we don't hit that default.
kwargs: dict[str, Any] = { self._env = RoboCasaGymEnv(
"env_name": self.task, env_name=self.task,
"camera_widths": self.observation_width, camera_widths=self.observation_width,
"camera_heights": self.observation_height, camera_heights=self.observation_height,
"split": self.split if self.split is not None else "all", split=self.split if self.split is not None else "all",
} )
self._env = RoboCasaGymEnv(**kwargs)
# Extract task description from environment metadata
assert self._env is not None
ep_meta = self._env.env.get_ep_meta() ep_meta = self._env.env.get_ep_meta()
self.task_description = ep_meta.get("lang", self.task) self.task_description = ep_meta.get("lang", self.task)
def _format_raw_obs(self, raw_obs: dict) -> RobotObservation: def _format_raw_obs(self, raw_obs: dict) -> RobotObservation:
"""Convert RoboCasaGymEnv observation dict to LeRobot format.""" """Convert RoboCasaGymEnv observation dict to LeRobot format."""
# Extract camera images (RoboCasaGymEnv provides "video.<cam>" keys) # RoboCasaGymEnv emits camera frames under "video.<cam>".
images = {} images = {cam: raw_obs[f"video.{cam}"] for cam in self.camera_name if f"video.{cam}" in raw_obs}
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": if self.obs_type == "pixels":
return {"pixels": images} return {"pixels": images}
# Extract state from raw_obs (state.* keys from PandaOmronKeyConverter) # `state.*` keys come from PandaOmronKeyConverter inside the wrapper.
agent_pos = np.concatenate( agent_pos = np.concatenate(
[ [
raw_obs.get("state.base_position", np.zeros(3)), raw_obs.get("state.base_position", np.zeros(3)),
@@ -224,10 +228,7 @@ class RoboCasaEnv(gym.Env):
axis=-1, axis=-1,
).astype(np.float32) ).astype(np.float32)
return { return {"pixels": images, "agent_pos": agent_pos}
"pixels": images,
"agent_pos": agent_pos,
}
def render(self) -> np.ndarray: def render(self) -> np.ndarray:
self._ensure_env() self._ensure_env()
@@ -240,7 +241,6 @@ class RoboCasaEnv(gym.Env):
super().reset(seed=seed) super().reset(seed=seed)
raw_obs, info = self._env.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() ep_meta = self._env.env.get_ep_meta()
self.task_description = ep_meta.get("lang", self.task) self.task_description = ep_meta.get("lang", self.task)
@@ -257,19 +257,12 @@ class RoboCasaEnv(gym.Env):
f"but got shape {action.shape} with ndim={action.ndim}" f"but got shape {action.shape} with ndim={action.ndim}"
) )
# Convert flat action to RoboCasa dict format
action_dict = convert_action(action) action_dict = convert_action(action)
raw_obs, reward, done, truncated, info = self._env.step(action_dict) raw_obs, reward, done, truncated, info = self._env.step(action_dict)
is_success = bool(info.get("success", False)) is_success = bool(info.get("success", False))
terminated = done or is_success terminated = done or is_success
info.update( info.update({"task": self.task, "done": done, "is_success": is_success})
{
"task": self.task,
"done": done,
"is_success": is_success,
}
)
observation = self._format_raw_obs(raw_obs) observation = self._format_raw_obs(raw_obs)
if terminated: if terminated:
@@ -282,15 +275,11 @@ class RoboCasaEnv(gym.Env):
self._env.close() self._env.close()
# ---- Main API ----------------------------------------------------------------
def _make_env_fns( def _make_env_fns(
*, *,
task: str, task: str,
n_envs: int, n_envs: int,
camera_names: list[str], camera_names: list[str],
camera_name_mapping: dict[str, str] | None,
obs_type: str, obs_type: str,
render_mode: str, render_mode: str,
observation_width: int, observation_width: int,
@@ -304,7 +293,6 @@ def _make_env_fns(
return RoboCasaEnv( return RoboCasaEnv(
task=task, task=task,
camera_name=camera_names, camera_name=camera_names,
camera_name_mapping=camera_name_mapping,
obs_type=obs_type, obs_type=obs_type,
render_mode=render_mode, render_mode=render_mode,
observation_width=observation_width, observation_width=observation_width,
@@ -322,7 +310,6 @@ def create_robocasa_envs(
n_envs: int, n_envs: int,
gym_kwargs: dict[str, Any] | None = None, gym_kwargs: dict[str, Any] | None = None,
camera_name: str | Sequence[str] = ",".join(DEFAULT_CAMERAS), 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, env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
episode_length: int | None = None, episode_length: int | None = None,
) -> dict[str, dict[int, Any]]: ) -> dict[str, dict[int, Any]]:
@@ -330,9 +317,14 @@ def create_robocasa_envs(
Returns: Returns:
dict[task_name][task_id] -> vec_env (env_cls([...]) with exactly n_envs factories) 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:
- `task` can be a single task or a comma-separated list of tasks. - a single task name (e.g. `CloseFridge`)
- a comma-separated list of task names (e.g. `CloseFridge,PickPlaceCoffee`)
- 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 the dataset `split` ("target" or "pretrain").
""" """
if env_cls is None or not callable(env_cls): 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.") raise ValueError("env_cls must be a callable that wraps a list of environment factory callables.")
@@ -347,11 +339,11 @@ def create_robocasa_envs(
split = gym_kwargs.pop("split", None) split = gym_kwargs.pop("split", None)
camera_names = _parse_camera_names(camera_name) camera_names = _parse_camera_names(camera_name)
task_names = [t.strip() for t in str(task).split(",") if t.strip()] task_names, group_split = _resolve_tasks(str(task))
if not task_names: if group_split is not None and split is None:
raise ValueError("`task` must contain at least one RoboCasa task name.") split = group_split
print(f"Creating RoboCasa envs | tasks={task_names} | n_envs(per task)={n_envs}") print(f"Creating RoboCasa envs | tasks={task_names} | split={split} | n_envs(per task)={n_envs}")
is_async = env_cls is gym.vector.AsyncVectorEnv is_async = env_cls is gym.vector.AsyncVectorEnv
@@ -359,12 +351,11 @@ def create_robocasa_envs(
cached_act_space: spaces.Space | None = None cached_act_space: spaces.Space | None = None
out: dict[str, dict[int, Any]] = defaultdict(dict) out: dict[str, dict[int, Any]] = defaultdict(dict)
for _tid, task_name in enumerate(task_names): for task_name in task_names:
fns = _make_env_fns( fns = _make_env_fns(
task=task_name, task=task_name,
n_envs=n_envs, n_envs=n_envs,
camera_names=camera_names, camera_names=camera_names,
camera_name_mapping=camera_name_mapping,
obs_type=obs_type, obs_type=obs_type,
render_mode=render_mode, render_mode=render_mode,
observation_width=observation_width, observation_width=observation_width,