feat(action): add RoboCasa mobile-manipulation action-policy support

Adds RoboCasa to the action-policy path: dataset loader, dataset factory,
Nano post-training config, embodiment registration, closed-loop evaluator,
and the LeRobot v2.1 -> v3.0 dataset converter.

RoboCasa is a mobile-manipulation benchmark, so unlike DROID and LIBERO the
base is not fixed. The loader exposes a 15-D raw action contract that passes
the native base command through unchanged:

    [base_motion(4), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)]

Observations are the agentview_left | eye_in_hand composite at full
resolution, with EEF proprioception prepended as a clean condition token.

  data/generator/action/datasets/     RoboCasaLeRobotDataset and its factory
  configs/.../action_policy_robocasa_nano.py
  simulation/robocasa/                closed-loop evaluator, next to LIBERO's
  scripts/convert_robocasa_to_lerobot_v30.sh

The action server gains support for the `use_state=True` contract the recipe
trains with: the current end-effector pose is prepended to the action sequence
as a clean conditioning frame, making it `chunk_length + 1` long. Without it a
`use_state` checkpoint was served an all-zero `chunk_length` action tensor and
the request's `"state"` field was ignored. Measured on target/atomic at iter
16000, 18 tasks x 20 rollouts, same weights either way: 226/360 without the
state token against 287/360 with it. Requests that omit "state" are unchanged.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: hyzhou-nv

cosmos_framework/configs/base/config.py

@@ -97,6 +97,7 @@ def make_config() -> Config:
     import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_droid_nano  # noqa: F401
     import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_libero_all_nano  # noqa: F401
     import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_libero_nano  # noqa: F401
+    import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_robocasa_nano  # noqa: F401
     import cosmos_framework.configs.base.experiment.action.posttrain_config.action_fd_droid_posttrain  # noqa: F401
     import cosmos_framework.configs.base.experiment.sft.vision_sft_nano  # noqa: F401
     import cosmos_framework.configs.base.experiment.sft.vision_sft_super  # noqa: F401

cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_robocasa_nano.py

@@ -0,0 +1,197 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""``action_policy_robocasa_nano`` — RoboCasa mobile-base action policy.
+
+Trains Cosmos3-Nano into a mobile-manipulation policy on the 18 RoboCasa ``target/atomic``
+tasks. The action contract is 15-dimensional and carries RoboCasa's native base command
+unchanged (``base_encoding="raw"``)::
+
+    [base_motion(4), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)]
+
+``base_motion`` is the command sent to the base controller, so the closed-loop client writes it
+straight back (``env[7:11] = action[0:4]``) and a recorded demonstration replays exactly. It is
+already normalised to [-1, 1], i.e. the same scale as the arm block, which matters because
+``action_normalization=None``: the state token is absolute, so delta statistics do not apply and
+the raw channel scales are what the loss sees.
+
+Observation is ``camera_set="left_wrist"`` — ``agentview_left`` and ``eye_in_hand`` concatenated
+horizontally at their native 256x256 each — plus end-effector proprioception prepended as a
+clean condition token (``use_state=True``). Action chunk length 32 at 20 fps, with
+``tokenizer.encode_exact_durations=[33]`` pinned to the matching 33-frame observation window
+(chunk + 1).
+
+Closed-loop evaluation must use the same contract::
+
+    ACTION_HORIZON=32 CAMERA_SET=left_wrist USE_STATE=1
+    USE_BASE_ACTION=1 BASE_ENCODING=raw RAW_ACTION_DIM=15
+"""
+
+import copy
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_framework.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG
+from cosmos_framework.data.generator.action.datasets.action_sft_dataset import get_action_robocasa_sft_dataset
+from cosmos_framework.data.generator.action.datasets.robocasa_lerobot_dataset import DEFAULT_ALL_ATOMIC_TASKS
+from cosmos_framework.data.generator.joint_dataloader import (
+    PackingDataLoader,
+    RankPartitionedDataLoader,
+)
+from cosmos_framework.utils.lazy_config import LazyCall as L
+from cosmos_framework.utils.lazy_config import LazyDict
+
+cs = ConfigStore.instance()
+
+
+def _lw_state_model_config() -> dict:
+    cfg = copy.deepcopy(NANO_MODEL_CONFIG)  # action_gen=True, max_action_dim=64
+    cfg["max_num_tokens_after_packing"] = 74000
+    cfg["activation_checkpointing"]["mode"] = "selective"
+    cfg["diffusion_expert_config"]["load_weights_from_pretrained"] = False
+    cfg["rectified_flow_training_config"]["loss_scale"] = 10.0
+    cfg["rectified_flow_training_config"]["image_loss_scale"] = None
+    # chunk_length=32 -> 33 observation frames; pin the VAE encode duration to match
+    # (mirrors action_policy_droid_nano, which also trains with a 32-step chunk).
+    cfg["tokenizer"]["encode_exact_durations"] = [33]
+    return cfg
+
+
+action_policy_robocasa_nano = LazyDict(
+    dict(
+        defaults=[
+            {"override /model": "mot_fsdp"},
+            {"override /data_train": None},
+            {"override /data_val": None},
+            {"override /optimizer": "fusedadamw"},
+            {"override /scheduler": "lambdalinear"},
+            {"override /checkpoint": "s3"},
+            {"override /callbacks": ["basic", "optimization", "job_monitor"]},
+            {"override /ema": "power"},
+            {"override /tokenizer": "wan2pt2_tokenizer"},
+            {"override /sound_tokenizer": None},
+            {"override /vlm_config": None},
+            {"override /ckpt_type": "dcp"},
+            "_self_",
+        ],
+        job=dict(
+            project="cosmos3",
+            group="action_sft",
+            name="action_policy_robocasa_nano",
+            wandb_mode="disabled",
+        ),
+        model=dict(config=_lw_state_model_config()),
+        optimizer=dict(
+            betas=[0.9, 0.99],
+            eps=1.0e-08,
+            fused=True,
+            keys_to_select=[
+                "moe_gen", "time_embedder", "vae2llm", "llm2vae",
+                "action2llm", "llm2action", "action_modality_embed",
+            ],
+            lr=5.0e-05,
+            lr_multipliers={"action2llm": 5.0, "llm2action": 5.0, "action_modality_embed": 5.0},
+            optimizer_type="FusedAdam",
+            weight_decay=0.05,
+        ),
+        scheduler=dict(
+            lr_scheduler_type="LambdaLinear",
+            cycle_lengths=[100],
+            f_max=[1.0], f_min=[0.0], f_start=[1.0e-06],
+            verbosity_interval=0, warm_up_steps=[0],
+        ),
+        trainer=dict(
+            distributed_parallelism="fsdp",
+            grad_accum_iter=1, logging_iter=1, max_iter=100, max_val_iter=None,
+            run_validation=False, run_validation_on_start=False,
+            save_zero_checkpoint=False, seed=42, timeout_period=999999999,
+            validation_iter=100,
+            compile_config=dict(recompile_limit=8, use_duck_shape=False),
+            cudnn=dict(benchmark=True, deterministic=False),
+            ddp=dict(broadcast_buffers=True, find_unused_parameters=False, static_graph=True),
+            grad_scaler_args=dict(enabled=False),
+            callbacks=dict(
+                dataloader_speed=dict(every_n=100, save_s3=False, step_size=1),
+                device_monitor=dict(every_n=200, log_memory_detail=True, save_s3=False, step_size=1, upload_every_n_mul=5),
+                grad_clip=dict(clip_norm=1.0, force_finite=True),
+                heart_beat=dict(every_n=200, save_s3=False, step_size=1, update_interval_in_minute=20),
+                iter_speed=dict(every_n=1, hit_thres=50, save_s3=False, save_s3_every_log_n=500),
+                low_precision=dict(update_iter=1),
+                manual_gc=dict(every_n=5, gc_level=1, warm_up=1),
+                param_count=dict(save_s3=False),
+                skip_nan_step=dict(max_consecutive_nan=100),
+                training_stats=dict(log_freq=100),
+            ),
+        ),
+        checkpoint=dict(
+            broadcast_via_filesystem=False,
+            dcp_async_mode_enabled=False,
+            enable_gcs_patch_in_boto3=True,
+            keys_not_to_resume=[],
+            keys_to_skip_loading=[
+                "net_ema.", "action2llm", "llm2action", "action_modality_embed", "action_pos_embed",
+            ],
+            load_ema_to_reg=False,
+            load_path="???",
+            load_training_state=False,
+            only_load_scheduler_state=False,
+            save_iter=100,
+            strict_resume=False,
+            verbose=True,
+            hf_export=dict(enabled=False, export_every_n=1, hf_repo_id=None,
+                           upload_to_object_store=dict(bucket="", credentials="", enabled=False)),
+            jit=dict(device="cuda", dtype="bfloat16", enabled=False, input_shape=None, strict=True),
+            load_from_object_store=dict(bucket="", credentials="", enabled=False),
+            save_to_object_store=dict(bucket="", credentials="", enabled=False),
+        ),
+        dataloader_train=L(PackingDataLoader)(
+            audio_sample_rate=48000,
+            dataset_name="action_robocasa",
+            max_samples_per_batch=128,
+            max_sequence_length=None,
+            patch_spatial=2,
+            sound_latent_fps=0,
+            tokenizer_spatial_compression_factor=16,
+            tokenizer_temporal_compression_factor=4,
+            dataloader=L(RankPartitionedDataLoader)(
+                batch_size=1, in_order=False, num_workers=4,
+                persistent_workers=True, pin_memory=True, prefetch_factor=4, sampler=None,
+                datasets=dict(
+                    robocasa=dict(
+                        ratio=1,
+                        dataset=L(get_action_robocasa_sft_dataset)(
+                            root="${oc.env:ROBOCASA_ROOT}",
+                            task_names=DEFAULT_ALL_ATOMIC_TASKS,   # all 18 atomic tasks (NavigateKitchen included)
+                            use_base_action=True,
+                            base_encoding="raw",   # 15D: native base_motion(4) + control_mode + arm
+                            fps=20,
+                            chunk_length=32,   # 32-step action chunk (DROID recipe uses 32); eval must use ACTION_HORIZON=32
+                            mode="wam",
+                            viewpoint="concat_view",
+                            camera_set="left_wrist",   # agentview_left + wrist, full-res 256x512
+                            use_state=True,   # EEF proprioception -> prepended clean condition token
+                            action_normalization=None,   # absolute state token: delta-stats do not apply
+                            split="train",
+                            split_val_ratio=0.01,
+                            iterable_shuffle=True,
+                            episode_shuffle_seed=42,
+                            resolution=None,   # keep 256x512 horizontal concat un-squished
+                            max_action_dim="${model.config.max_action_dim}",
+                            cfg_dropout_rate=0.1,
+                            format_prompt_as_json=True,
+                            tokenizer_config="${model.config.vlm_config.tokenizer}",
+                        ),
+                    ),
+                ),
+            ),
+        ),
+        dataloader_val=None,
+        upload_reproducible_setup=False,
+    ),
+    flags={"allow_objects": True},
+)
+
+
+for _item in [action_policy_robocasa_nano]:
+    _name = [k for k, v in globals().items() if v is _item][0]
+    cs.store(group="experiment", package="_global_", name=_name, node=_item)

cosmos_framework/data/generator/action/datasets/__init__.py

@@ -17,6 +17,7 @@
 from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset
 from cosmos_framework.data.generator.action.datasets.fractal_lerobot_dataset import FractalLeRobotDataset
 from cosmos_framework.data.generator.action.datasets.human_hand_pose_lerobot_dataset import HumanHandPoseLeRobotDataset
+from cosmos_framework.data.generator.action.datasets.robocasa_lerobot_dataset import RoboCasaLeRobotDataset
 from cosmos_framework.data.generator.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset
 from cosmos_framework.data.generator.action.datasets.robomind_franka_dataset import RoboMINDFrankaDataset
 from cosmos_framework.data.generator.action.datasets.robomind_ur_dataset import RoboMINDURDataset
@@ -30,6 +31,7 @@
     "DROIDMergedLeRobotDataset",
     "FractalLeRobotDataset",
     "HumanHandPoseLeRobotDataset",
+    "RoboCasaLeRobotDataset",
     "LIBEROLeRobotDataset",
     "RoboMINDFrankaDataset",
     "RoboMINDURDataset",

cosmos_framework/data/generator/action/datasets/action_sft_dataset.py

@@ -21,6 +21,10 @@
 
 from cosmos_framework.data.generator.action.datasets.droid_merged_lerobot_dataset import DROIDMergedLeRobotDataset
 from cosmos_framework.data.generator.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset
+from cosmos_framework.data.generator.action.datasets.robocasa_lerobot_dataset import (
+    DEFAULT_ALL_ATOMIC_TASKS,
+    RoboCasaLeRobotDataset,
+)
 from cosmos_framework.data.generator.action.datasets.libero_lerobot_dataset import LIBEROLeRobotDataset
 from cosmos_framework.data.generator.action.utils.transforms import ActionTransformPipeline
 
@@ -211,6 +215,89 @@ def get_action_droid_merged_lerobot_sft_dataset(
     return sft
 
 
+def get_action_robocasa_sft_dataset(
+    *,
+    root: str,
+    fps: float = 20.0,
+    chunk_length: int = 16,
+    mode: str = "wam",
+    viewpoint: str = "concat_view",
+    camera_set: str = "wrist_lr",
+    task_names: tuple[str, ...] | list[str] = DEFAULT_ALL_ATOMIC_TASKS,
+    use_state: bool = False,
+    use_base_action: bool = False,
+    base_encoding: str = "ego",
+    action_normalization: str | None = None,
+    split: str = "train",
+    split_val_ratio: float = 0.01,
+    split_seed: int = 42,
+    resolution: str | int = "256",
+    max_action_dim: int = 64,
+    tokenizer_config: dict | None = None,
+    cfg_dropout_rate: float = 0.1,
+    append_viewpoint_info: bool = True,
+    append_duration_fps_timestamps: bool = True,
+    append_resolution_info: bool = True,
+    append_idle_frames: bool = True,
+    format_prompt_as_json: bool = False,
+    iterable_shuffle: bool = False,
+    episode_shuffle_seed: int = 42,
+) -> Dataset:
+    """Build the RoboCasa fixed-base action-policy SFT dataset.
+
+    Feeds ``RoboCasaLeRobotDataset`` (10D ``[pos, rot6d, gripper]`` end-effector
+    deltas, concat_view wrist + two third-person cams) through
+    ``ActionTransformPipeline``. ``root`` is the RoboCasa atomic dir
+    (``.../target/atomic``); each task in ``task_names`` is discovered as
+    ``<root>/<task>/*/lerobot`` and registered as a separate LeRobot shard.
+    Defaults to ``action_normalization=None`` (like DROID); pass ``quantile_rot``
+    with a bundled stats file to normalize.
+
+    ``use_base_action=True`` widens the contract so the mobile base is representable
+    (required for the full 18-task ``DEFAULT_ALL_ATOMIC_TASKS`` set, which includes
+    NavigateKitchen); leave it False for the 10D fixed-base recipe. ``base_encoding``
+    then selects the base representation:
+
+    * ``"ego"`` (default, 20D) --
+      ``[base_pos(3), base_rot6d(6), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)]``
+    * ``"raw"`` (15D) --
+      ``[base_motion(4), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)]``
+
+    ``"raw"`` regresses RoboCasa's native normalised velocity command, so replay is an
+    identity round-trip and the base channels share the arm's scale; ``"ego"`` is retained
+    as the default so existing configs and checkpoints reproduce unchanged.
+    """
+    dataset: Dataset = RoboCasaLeRobotDataset(
+        root=root,
+        fps=fps,
+        chunk_length=chunk_length,
+        mode=mode,
+        viewpoint=viewpoint,
+        camera_set=camera_set,
+        task_names=task_names,
+        use_state=use_state,
+        use_base_action=use_base_action,
+        base_encoding=base_encoding,
+        action_normalization=action_normalization,
+        split=split,
+        split_val_ratio=split_val_ratio,
+        split_seed=split_seed,
+    )
+    transform = ActionTransformPipeline(
+        tokenizer_config=tokenizer_config,
+        cfg_dropout_rate=cfg_dropout_rate,
+        max_action_dim=max_action_dim,
+        append_viewpoint_info=append_viewpoint_info,
+        append_duration_fps_timestamps=append_duration_fps_timestamps,
+        append_resolution_info=append_resolution_info,
+        append_idle_frames=append_idle_frames,
+        format_prompt_as_json=format_prompt_as_json,
+    )
+    sft = ActionSFTDataset(dataset, transform, resolution)
+    if iterable_shuffle:
+        return ActionIterableShuffleDataset(sft, seed=episode_shuffle_seed)
+    return sft
+
 def get_action_libero_sft_dataset(
     *,
     root: str,