diff --git a/cosmos_framework/configs/base/config.py b/cosmos_framework/configs/base/config.py index b3bd2259..9e950946 100644 --- a/cosmos_framework/configs/base/config.py +++ b/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 diff --git a/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_robocasa_nano.py b/cosmos_framework/configs/base/experiment/action/posttrain_config/action_policy_robocasa_nano.py new file mode 100644 index 00000000..9a43cd70 --- /dev/null +++ b/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) diff --git a/cosmos_framework/data/generator/action/datasets/__init__.py b/cosmos_framework/data/generator/action/datasets/__init__.py index 140904ca..6797cbaf 100644 --- a/cosmos_framework/data/generator/action/datasets/__init__.py +++ b/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", diff --git a/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py b/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py index c952fb8c..ae916c42 100644 --- a/cosmos_framework/data/generator/action/datasets/action_sft_dataset.py +++ b/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 + ``//*/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, diff --git a/cosmos_framework/data/generator/action/datasets/robocasa_lerobot_dataset.py b/cosmos_framework/data/generator/action/datasets/robocasa_lerobot_dataset.py new file mode 100644 index 00000000..43946952 --- /dev/null +++ b/cosmos_framework/data/generator/action/datasets/robocasa_lerobot_dataset.py @@ -0,0 +1,564 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""RoboCasa (PandaOmron) LeRobot action-policy dataset. + +Reads the RoboCasa ``PhysicalAI-Robotics-Manipulation-Kitchen-Demos`` LeRobot +export (codebase_version v2.1) through the official ``lerobot`` reader (via +:class:`BaseActionLeRobotDataset`), so the v2.1 per-episode layout is handled by +the library rather than a hand-rolled parquet reader. + +RoboCasa's native 12D action is +``[base_motion(4), control_mode(1), eef_pos(3), eef_rot_axisangle(3), gripper(1)]`` +(see ``meta/modality.json``). + +With ``use_base_action=False`` the base channels are dropped and only the arm delta is kept, +re-encoding rotation axis-angle -> rot6d. That 10D contract suits task sets whose base never +moves; it reduces the embodiment to a stationary Franka arm, i.e. the LIBERO/DROID regime: + + 12D -> slice [5:8]+[8:11]+[11] -> [pos(3), rot_axisangle(3), gripper(1)] (7D) + -> axisangle->rot6d -> [pos(3), rot6d(6), gripper(1)] (10D) + +The stored ``eef_pos``/``eef_rot`` are already per-frame OSC deltas, so this is a +``frame_wise_relative`` (``backward_framewise``) action — identical semantics to +``LIBEROLeRobotDataset._build_frame_wise_action``, only the slice indices differ. + +Cameras: RoboCasa exports 3 views (``robot0_eye_in_hand`` wrist + ``agentview_left`` +/ ``agentview_right`` third-person). ``concat_view`` tiles them wrist-on-top, +left/right-on-bottom (mirrors ``DROIDLeRobotDataset._compose_multi_view``). +""" + +from __future__ import annotations + +import glob +import os +import random +from typing import Any + +import torch +import torch.nn.functional as F + +from cosmos_framework.data.generator.action.datasets.cosmos3_action_lerobot import ( + ActionNormalization, + ActionSpec, + BaseActionLeRobotDataset, + Gripper, + Pos, + Rot, + build_action_spec, +) +from cosmos_framework.data.generator.action.utils.pose_utils import PoseConvention, convert_rotation +from cosmos_framework.data.generator.action.utils.viewpoint_utils import Viewpoint +from cosmos_framework.utils import log + +# LeRobot column names (constant across all RoboCasa atomic tasks; see meta/info.json). +_ACTION_FEATURE = "action" +_STATE_FEATURE = "observation.state" +_IMAGE_FEATURES = { + "wrist": "observation.images.robot0_eye_in_hand", + "left": "observation.images.robot0_agentview_left", + "right": "observation.images.robot0_agentview_right", +} + +# Native 12D action layout (meta/modality.json): the arm-only contract keeps only +# the arm end-effector delta + gripper. +_EEF_POS = slice(5, 8) # end_effector_position delta +_EEF_ROT = slice(8, 11) # end_effector_rotation delta (axis-angle) +_GRIPPER = slice(11, 12) # gripper_close + +# observation.state 16D layout (meta/modality.json): base_position[0:3], +# base_rotation[3:7] (quat), end_effector_position_relative[7:10] (base frame, meters), +# end_effector_rotation_relative[10:14] (quat, xyzw — verified: base_rot=[0,0,.707,.707] +# is a 90 deg yaw about z, w last), gripper_qpos[14:16] (two finger joint positions). +# For the EEF proprioception token we keep only the arm eef pose + a 1D gripper opening. +_STATE_EEF_POS = slice(7, 10) # absolute eef position (base frame) +_STATE_EEF_ROT = slice(10, 14) # absolute eef rotation (quaternion, xyzw) +_STATE_GRIPPER = slice(14, 16) # gripper qpos (two fingers) + + +# All 18 ``target/atomic`` tasks. Base-motion prevalence measured over 40 episodes each: +# NavigateKitchen is the only genuinely mobile task (100% of episodes, 2.6 m mean net +# displacement); PickPlaceDrawerToCounter moves in 48% (<=0.71 m); five more move in +# 2-10% of episodes (<=0.4 m); the remaining ten never move the base (<6 mm noise). +# Training on this set requires ``use_base_action=True`` so the base channels exist. +DEFAULT_ALL_ATOMIC_TASKS: tuple[str, ...] = ( + "CloseBlenderLid", + "CloseFridge", + "CloseToasterOvenDoor", + "CoffeeSetupMug", + "NavigateKitchen", + "OpenCabinet", + "OpenDrawer", + "OpenStandMixerHead", + "PickPlaceCounterToCabinet", + "PickPlaceCounterToStove", + "PickPlaceDrawerToCounter", + "PickPlaceSinkToCounter", + "PickPlaceToasterToCounter", + "SlideDishwasherRack", + "TurnOffStove", + "TurnOnElectricKettle", + "TurnOnMicrowave", + "TurnOnSinkFaucet", +) + + +# The 65 ``pretrain/atomic`` tasks (every atomic task RoboCasa365 ships a pretrain split for; +# the 18 ``target/atomic`` names above are a subset). Human demos only -- the registry also +# lists machine-generated ``mg`` paths for 60 of these, but they are not part of the local +# download. Same 12D action / 16D state / 3-camera contract as target/atomic, so this loader +# handles them unchanged; only ``root`` and ``task_names`` differ. +DEFAULT_PRETRAIN_ATOMIC_TASKS: tuple[str, ...] = ( + "AdjustToasterOvenTemperature", "AdjustWaterTemperature", "CheesyBread", "CloseBlenderLid", + "CloseCabinet", "CloseDishwasher", "CloseDrawer", "CloseElectricKettleLid", "CloseFridge", + "CloseFridgeDrawer", "CloseMicrowave", "CloseOven", "CloseStandMixerHead", + "CloseToasterOvenDoor", "CoffeeServeMug", "CoffeeSetupMug", "LowerHeat", "MakeIcedCoffee", + "NavigateKitchen", "OpenBlenderLid", "OpenCabinet", "OpenDishwasher", "OpenDrawer", + "OpenElectricKettleLid", "OpenFridge", "OpenFridgeDrawer", "OpenMicrowave", "OpenOven", + "OpenStandMixerHead", "OpenToasterOvenDoor", "PackDessert", "PickPlaceCabinetToCounter", + "PickPlaceCounterToBlender", "PickPlaceCounterToCabinet", "PickPlaceCounterToDrawer", + "PickPlaceCounterToMicrowave", "PickPlaceCounterToOven", "PickPlaceCounterToSink", + "PickPlaceCounterToStandMixer", "PickPlaceCounterToStove", "PickPlaceCounterToToasterOven", + "PickPlaceDrawerToCounter", "PickPlaceFridgeDrawerToShelf", "PickPlaceFridgeShelfToDrawer", + "PickPlaceMicrowaveToCounter", "PickPlaceSinkToCounter", "PickPlaceStoveToCounter", + "PickPlaceToasterOvenToCounter", "PickPlaceToasterToCounter", "PreheatOven", + "SlideDishwasherRack", "SlideOvenRack", "SlideToasterOvenRack", "StartCoffeeMachine", + "TurnOffMicrowave", "TurnOffSinkFaucet", "TurnOffStove", "TurnOnBlender", + "TurnOnElectricKettle", "TurnOnMicrowave", "TurnOnSinkFaucet", "TurnOnStove", + "TurnOnToaster", "TurnOnToasterOven", "TurnSinkSpout", +) + +# The 235 ``pretrain/composite`` tasks: multi-stage, long-horizon kitchen activities. +# Episodes average ~1565 frames (78 s) against ~142 for atomic, so by frame count composite +# outweighs pretrain/atomic roughly 18:1 (27.6M vs 1.5M). Mixing the two unweighted therefore +# yields a ~95% composite corpus -- intentional here, but worth remembering when reading +# per-task numbers. Composite demos drive the mobile base far more than atomic ones, which is +# why they require ``use_base_action=True``. +DEFAULT_PRETRAIN_COMPOSITE_TASKS: tuple[str, ...] = ( + "AddIceCubes", "AddLemonToFish", "AddMarshmallow", "AddSugarCubes", "AddSweetener", + "AdjustHeat", "AfterwashSorting", "AirDryFruit", "AlcoholServingPrep", "AlignSilverware", + "ArrangeBreadBowl", "ArrangeBuffetDessert", "ArrangeDrinkware", "ArrangeTeaAccompaniments", + "ArrangeUtensilsByType", "ArrangeVegetables", "AssembleCookingArray", "BalancedMealPrep", + "BeverageOrganization", "BeverageSorting", "BlendIngredients", "BlendMarinade", + "BowlAndCup", "BreadAndCheese", "BreadSetupSlicing", "BuildAppetizerPlate", "ButterOnPan", + "CandleCleanup", "CerealAndBowl", "ChooseMeasuringCup", "ChooseRipeFruit", "CleanBoard", + "CleanMicrowave", "ClearClutter", "ClearCuttingBoard", "ClearFoodWaste", "ClearFreezer", + "ClearReceptaclesForCleaning", "ClearSink", "ClearSinkArea", "ClearSinkSpace", + "CollectWashingSupplies", "ColorfulSalsa", "CondimentCollection", "CookieDoughPrep", + "CoolBakedCake", "CoolKettle", "CreateChildFriendlyFridge", "CupcakeCleanup", + "CutBuffetPizza", "DateNight", "DefrostByCategory", "DeliverBrewedCoffee", "DeliverStraw", + "DessertAssembly", "DessertUpgrade", "DisplayMeatVariety", "DistributeChicken", + "DivideBasins", "DivideBuffetTrays", "DrainVeggies", "DrinkwareConsolidation", "DryDishes", + "DryDrinkware", "DumpLeftovers", "FillBlenderJug", "FillKettle", "FilterMicrowavableItem", + "FoodCleanup", "FreezeBottledWaters", "FreezeCookedFood", "FreezeIceTray", + "FreshProduceOrganization", "FryingPanAdjustment", "GarnishCupcake", "GatherCuttingTools", + "GatherMarinadeIngredients", "GatherVegetables", "GetToastedBread", "HeatMug", + "HeatMultipleWater", "HotDogSetup", "JuiceFruitReamer", "KettleBoiling", + "LemonSeasoningFish", "LineUpCondiments", "LoadCondimentsInFridge", "LoadDishwasher", + "LoadFridgeByType", "LoadFridgeFifo", "LoadPreparedFood", "MakeFruitBowl", + "MakeLoadedPotato", "MatchCupAndDrink", "MaximizeFreezerSpace", "MealPrepStaging", + "MeatSkewerAssembly", "MeatTransfer", "MicrowaveCorrectMeal", "MicrowaveDefrostMeat", + "MicrowaveThawing", "MicrowaveThawingFridge", "MixCakeFrosting", "MixedFruitPlatter", + "MoveFreezerToFridge", "MoveFridgeToFreezer", "MoveToCounter", "MoveToFreezerDrawer", + "MultistepSteaming", "OrganizeBakingIngredients", "OrganizeCleaningSupplies", + "OrganizeCoffeeCondiments", "OrganizeCondiments", "OrganizeMetallicUtensils", + "OrganizeMugsByHandle", "OrganizeVegetables", "OvenBroilFish", "PackFoodByTemp", + "PackFruitContainer", "PackIdenticalLunches", "PastryDisplay", "PlaceBeveragesTogether", + "PlaceDishesBySink", "PlaceEqualIceCubes", "PlaceFoodInBowls", "PlaceIceInCup", + "PlaceMeatInMarinade", "PlaceMicrowaveSafeItem", "PlaceStraw", "PlaceVegetablesEvenly", + "PlaceVeggiesInDrawer", "PlateSteakMeal", "PlateStoreDinner", "PortionFruitBowl", + "PortionInTupperware", "PortionOnSize", "PortionYogurt", "PreRinseStation", "PreSoakPan", + "PreheatPot", "PrepForSanitizing", "PrepForTenderizing", "PrepFridgeForCleaning", + "PrepMarinatingMeat", "PrepSinkForCleaning", "PrepareBroilingStation", + "PrepareCheeseStation", "PrepareCocktailStation", "PrepareCoffee", "PrepareDishwasher", + "PrepareDrinkStation", "PrepareSausageCheese", "PrepareSmoothie", "PrepareSoupServing", + "PrepareStoringLeftovers", "PrepareToast", "PressChicken", "PrewashFoodAssembly", + "PrewashFoodSorting", "QuickThaw", "RearrangeFridgeItems", "RecycleBottlesBySize", + "RecycleSodaCans", "RecycleStackedYogurt", "RefillCondimentStation", "ReheatMeal", + "RemoveBroiledFish", "RemoveCuttingBoardItems", "ReorganizeFrozenVegetables", + "ResetCabinetDoors", "RestockBowls", "RestockCannedFood", "RestockPantry", + "RestockSinkSupplies", "RetrieveIceTray", "RetrieveMeat", "ReturnHeatedFood", + "ReturnWashingSupplies", "RinseBowls", "RinseCuttingBoard", "RinseSinkBasin", "RotatePan", + "SanitizePrepCuttingBoard", "ScalePortioning", "ScrubBowl", "ScrubCuttingBoard", + "SearingMeat", "SeasoningSpiceSetup", "SeasoningSteak", "ServeSteak", "ServeTea", + "ServeWarmCroissant", "SetBowlsForSoup", "SetUpCuttingStation", "SetUpSpiceStation", + "SetupBowls", "SetupButterPlate", "SetupFruitBowl", "SetupFrying", "SetupSodaBowl", + "SetupWineGlasses", "ShakePan", "SimmeringSauce", "SizeSorting", "SnackSorting", + "SoakSponge", "SortingCleanup", "SpicyMarinade", "StackBowlsCabinet", "StackBowlsInSink", + "StackCans", "StartElectricKettle", "SteamInMicrowave", "StirVegetables", + "StockingBreakfastFoods", "StoreDumplings", "StoreLeftoversByType", "StoreLeftoversInBowl", + "StrainerSetup", "SweetSavoryToastSetup", "SweetenCoffee", "SweetenHotChocolate", + "ThawInSink", "TiltPan", "ToastBagel", "ToastBaguette", "ToastOnCorrectRack", + "TongBuffetSetup", "TransportCookware", "TurnOffSimmeredSauceHeat", "VeggieDipPrep", + "WarmCroissant", "WashFish", "WashLettuce", "YogurtDelightPrep", +) + +# Convenience union for the 300-task pretrain soup (atomic + composite). Both live under one +# converted root, so the loader still takes a single ``root``. +DEFAULT_PRETRAIN_ALL_TASKS: tuple[str, ...] = ( + DEFAULT_PRETRAIN_ATOMIC_TASKS + DEFAULT_PRETRAIN_COMPOSITE_TASKS +) + +# Mobile-base action encoding (``use_base_action=True``). The Omron base is a +# JOINT_VELOCITY part with joints [forward, side, yaw, torso_height]; torso_height is +# never actuated in target/atomic (std 0 across all 18 tasks), so the base is a planar +# 3-DoF (x, y, yaw) system. Rather than regress the raw velocity command we derive a +# frame-wise EGO-FRAME base pose delta from ``observation.state`` (base_position + +# base_rotation) and encode it like the EEF delta (pos + rot6d). Because the deltas are +# expressed in the base frame at each step, no world->body yaw rotation is needed at +# inference; the closed-loop client converts back with +# ``base_motion[0:3] = (ego_delta / dt) / BASE_MAX_VELOCITY``. +# +# ``base_motion`` is a NORMALISED command in [-1, 1] (the JOINT_VELOCITY controller maps +# it onto the joint velocity limits), not a physical velocity. Calibrated on +# NavigateKitchen (least-squares through the origin of physical ego velocity against the +# recorded command, restricted to |command| > 0.2): full deflection corresponds to +# ~0.60 m/s forward, ~0.64 m/s lateral, ~1.25 rad/s yaw. Fit quality: yaw corr 0.997 +# (6.9% residual); the translation channels are corr 0.92-0.96 with ~26% residual, which +# is controller lag — the commanded velocity is not reached instantaneously. Closed-loop +# feedback absorbs that per step. +BASE_MAX_VELOCITY: tuple[float, float, float] = (0.601, 0.640, 1.251) +_STATE_BASE_POS = slice(0, 3) # absolute base position (world frame) +_STATE_BASE_ROT = slice(3, 7) # absolute base rotation (quaternion, xyzw) +_CONTROL_MODE = slice(4, 5) # +1 = base mode (arm tracks the moving base), -1 = arm mode + +# ``base_encoding="raw"``: regress the native command instead of inverting the controller. +# The ego encoding above derives the base delta from the ACHIEVED pose in +# ``observation.state``, so recovering the command at inference requires inverting a +# controller with real inertia -- structurally lossy (yaw residual 6.9 %, translation +# 26 %), and a recorded demo can never be replayed exactly. Passing ``base_motion`` +# through is an identity round-trip: the client writes ``env[7:11] = action[0:4]``. +# +# It also fixes a loss-weighting defect. Measured RMS per step: the ego translation delta +# is 0.018 (NavigateKitchen) to 0.002 (PickPlaceDrawerToCounter) in METRES, while the arm's +# ``eef_pos`` is a NORMALISED command with RMS 0.21-0.52 -- a 12x to 260x scale gap that, +# under ``action_normalization=None``, makes the base translation channels contribute +# ~1/144 to ~1/68000 of the arm's gradient. ``base_motion`` is normalised to [-1, 1] +# (RMS 0.49-0.61 on NavigateKitchen), i.e. the same scale as the arm block. +_BASE_MOTION = slice(0, 4) # [forward, side, yaw, torso_height]; torso is never actuated + + +class RoboCasaLeRobotDataset(BaseActionLeRobotDataset): + """RoboCasa manipulation dataset. + + Actions are ``[pos_delta(3), rot6d_delta(6), gripper(1)]`` (10D, arm only) or, with + ``use_base_action=True``, widened to include the mobile base — 15D for + ``base_encoding="raw"`` and 20D for ``"ego"``. Observation is a ``concat_view`` + composite selected by ``camera_set``. Reads v2.1 exports via the official lerobot + backend; each of ``task_names`` is discovered under ``root`` as + ``//*/lerobot`` and registered as a separate shard. + """ + + def __init__( + self, + root: str, + fps: float = 20.0, + chunk_length: int = 16, + split_seed: int = 42, + split_val_ratio: float = 0.01, + split: str = "train", + mode: str = "wam", + pose_convention: PoseConvention = "backward_framewise", + rotation_format: str = "rot6d", + action_normalization: ActionNormalization | None = None, + tolerance_s: float = 1e-4, + viewpoint: Viewpoint = "concat_view", + task_names: tuple[str, ...] | list[str] = DEFAULT_ALL_ATOMIC_TASKS, + use_state: bool = False, + enable_fast_init: bool = False, + camera_set: str = "wrist_lr", + use_base_action: bool = False, + base_encoding: str = "ego", + ) -> None: + if rotation_format != "rot6d": + raise NotImplementedError(f"RoboCasa loader only supports rotation_format='rot6d', got {rotation_format!r}.") + super().__init__( + fps=fps, + chunk_length=chunk_length, + split_seed=split_seed, + split_val_ratio=split_val_ratio, + split=split, + mode=mode, + embodiment_type="robocasa", + viewpoint=viewpoint, + pose_convention=pose_convention, + rotation_format=rotation_format, + action_normalization=action_normalization, + tolerance_s=tolerance_s, + enable_fast_init=enable_fast_init, + ) + self._use_state = use_state + # use_base_action: widen the arm-only 10D contract so mobile tasks + # (NavigateKitchen et al.) are representable. ``base_encoding`` selects HOW: + # + # "ego" (default, 20D) -- state-derived ego-frame base pose delta: + # [base_pos(3), base_rot6d(6), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)] + # "raw" (15D) -- the native base_motion velocity command, passed through: + # [base_motion(4), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)] + # + # "ego" is kept as the default so existing checkpoints/configs reproduce bit-for-bit. + # See the ``BASE_MAX_VELOCITY`` block above for why "raw" is the better contract. + self._use_base_action = use_base_action + if base_encoding not in ("ego", "raw"): + raise ValueError(f"Unsupported base_encoding={base_encoding!r}. Use 'ego' or 'raw'.") + self._base_encoding = base_encoding + self._image_features = _IMAGE_FEATURES + # camera_set: + # "wrist_lr" (default): concat_view = wrist (top) + L/R (bottom, squished to half). + # "left_wrist" : LIBERO-style = agentview_left + wrist, full-res, horizontally + # concatenated ([T,C,H,2W]); right cam DROPPED, no downscaling. + if camera_set not in ("wrist_lr", "left_wrist"): + raise ValueError(f"Unsupported camera_set={camera_set!r}. Use 'wrist_lr' or 'left_wrist'.") + self._camera_set = camera_set + + # Discover one LeRobot shard root per requested task under ``root``. + self._all_shard_roots = self._discover_shard_roots(root, task_names) + if not self._all_shard_roots: + raise FileNotFoundError( + f"No RoboCasa lerobot shards found under {root!r} for tasks {list(task_names)}. " + f"Expected //*/lerobot directories." + ) + log.info(f"RoboCasaLeRobotDataset: {len(self._all_shard_roots)} task shard(s): {self._all_shard_roots}") + + # delta_timestamps: chunk_length per-frame action deltas; chunk_length+1 + # observation frames (one more image/state than transitions). + observation_ts = [i * self._dt for i in range(0, self._chunk_length + 1)] + action_ts = [i * self._dt for i in range(0, self._chunk_length)] + self._delta_timestamps: dict[str, list[float]] = {_ACTION_FEATURE: action_ts} + if self._use_state: + self._delta_timestamps[_STATE_FEATURE] = observation_ts + if self._camera_set == "left_wrist": + # LIBERO-style: only agentview_left + wrist (right dropped). + self._delta_timestamps[self._image_features["wrist"]] = observation_ts + self._delta_timestamps[self._image_features["left"]] = observation_ts + else: + if self._viewpoint in ("wrist_view", "concat_view"): + self._delta_timestamps[self._image_features["wrist"]] = observation_ts + if self._viewpoint in ("third_person_view", "concat_view"): + self._delta_timestamps[self._image_features["left"]] = observation_ts + self._delta_timestamps[self._image_features["right"]] = observation_ts + + self._register_sources() + + @staticmethod + def _discover_shard_roots(root: str, task_names: tuple[str, ...] | list[str]) -> list[str]: + """Resolve ``//*/lerobot`` for each requested task. + + ``root`` may point at the atomic dir (``.../target/atomic``) or directly + at a single ``lerobot`` dir. Missing tasks are skipped with a warning. + """ + root = root.rstrip("/") + if os.path.basename(root) == "lerobot" and os.path.isdir(os.path.join(root, "meta")): + return [root] + shard_roots: list[str] = [] + for task in task_names: + matches = sorted(glob.glob(os.path.join(root, task, "*", "lerobot"))) + matches = [m for m in matches if os.path.isdir(os.path.join(m, "meta"))] + if not matches: + log.warning(f"RoboCasaLeRobotDataset: no lerobot shard for task {task!r} under {root!r}; skipping.") + continue + shard_roots.append(matches[0]) + return shard_roots + + # ---- action / spec ----------------------------------------------------- + + @property + def action_dim(self) -> int: + if self._use_base_action: + if self._base_encoding == "raw": + # [base_motion(4), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)] + return 15 + # [base_pos(3), base_rot6d(6), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)] + return 20 + return 10 # [pos(3), rot6d(6), gripper(1)] + + def _build_action_spec(self) -> ActionSpec: + if self._use_base_action: + if self._base_encoding == "raw": + # base_motion is a NORMALISED velocity command in [-1, 1], the same units + # as the arm's eef_pos block -- so it is declared as ``Pos`` (magnitude-based + # idle detection, ``|v| < eps_t``) and joins the arm translation in one L2. + # NOT ``Joint``: that branch is frame-DIFF based, so a constant cruise + # command (diff == 0) would be misread as idle. The 4th channel + # (torso height) is never actuated and contributes 0 to the norm. + return build_action_spec(Pos(dim=4), Gripper(), Pos(), Rot("rot6d"), Gripper()) + # Base pose delta first (Pos+Rot), then the mode flag, then the arm block. + # ``control_mode`` is a +/-1 channel regressed like the gripper (continuous + # rectified flow, no discrete head); the client thresholds it at 0. + return build_action_spec(Pos(), Rot("rot6d"), Gripper(), Pos(), Rot("rot6d"), Gripper()) + return build_action_spec(Pos(), Rot("rot6d"), Gripper()) + + def _build_frame_wise_action(self, raw_action: torch.Tensor) -> torch.Tensor: + """RoboCasa 12D per-frame action -> 10D ``[pos(3), rot6d(6), gripper(1)]``. + + Drops base_motion[0:4] + control_mode[4]; re-encodes the eef axis-angle + delta to rot6d. Mirrors ``LIBEROLeRobotDataset._build_frame_wise_action`` + with RoboCasa slice indices. + """ + raw = raw_action.float() # [chunk, 12] + translation = raw[:, _EEF_POS] # [chunk, 3] + rotation_matrix = convert_rotation(raw[:, _EEF_ROT], input_format="axisangle", output_format="matrix") + rotation = convert_rotation(rotation_matrix, input_format="matrix", output_format="rot6d") # [chunk, 6] + gripper = raw[:, _GRIPPER] # [chunk, 1] + return torch.cat([translation, rotation, gripper], dim=-1) # [chunk, 10] + + def _build_initial_state(self, state_seq: torch.Tensor) -> torch.Tensor: + """Current EEF proprioception -> 10D ``[pos(3), rot6d(6), gripper(1)]`` token. + + Mirrors DROID's ``use_state`` path (``midtrain`` branch): the frame BEFORE the + action chunk (index 0 of the ``chunk_length+1`` observation window) is prepended + as a CLEAN conditioning action token — same width as the action so it flows + through the shared ``action2llm`` embedding, gets ``sigma=0`` (never noised) and + is excluded from the flow-matching loss (``condition_frame_indexes_action=[0]``). + + RoboCasa ``observation.state`` eef fields are in the (fixed) base frame; the + rotation is a ``quat_xyzw`` re-encoded to rot6d to match the 10D action layout. + Gripper = signed finger opening (``qpos[0]-qpos[1]``), a 1D summary of the + two-finger state. NOTE: this token is ABSOLUTE pose while the predicted actions + are per-frame deltas — identical to DROID ``midtrain`` — so pair with + ``action_normalization=None`` (the quantile_rot delta-stats do not apply to an + absolute pose). + """ + s0 = state_seq[-self._chunk_length - 1].float() # [16]; earliest = current pre-chunk state + pos = s0[_STATE_EEF_POS] # [3] + quat = s0[_STATE_EEF_ROT].unsqueeze(0) # [1,4] xyzw + matrix = convert_rotation(quat, input_format="quat_xyzw", output_format="matrix") + rot6d = convert_rotation(matrix, input_format="matrix", output_format="rot6d").reshape(6) # [6] + grip = s0[14:15] - s0[15:16] # [1] signed finger opening + return torch.cat([pos, rot6d, grip], dim=-1) # [10] + + def _build_base_delta(self, state_seq: torch.Tensor) -> torch.Tensor: + """Frame-wise EGO-FRAME base pose delta -> ``[chunk, 9]`` ``[pos(3), rot6d(6)]``. + + ``observation.state`` carries the base pose in the WORLD frame + (``base_position``, ``base_rotation`` quat_xyzw). The delta between consecutive + frames is re-expressed in the base frame at the earlier step, i.e. + + T_delta = T_base[t]^-1 @ T_base[t+1] + + so the representation is ego-centric and matches the ``backward_framewise`` + convention already used for the EEF. This removes any dependence on world + heading: the closed-loop client can map straight back to the base velocity + command with ``(ego_delta / dt) * BASE_VELOCITY_SCALE`` without a yaw rotation. + + The state window is ``chunk_length+1`` frames (one more than the transitions), + so the ``chunk_length`` deltas line up one-to-one with the action chunk. + """ + s = state_seq[-self._chunk_length - 1 :].float() # [chunk+1, 16] + pos = s[:, _STATE_BASE_POS] # [chunk+1, 3] world + quat = s[:, _STATE_BASE_ROT] # [chunk+1, 4] xyzw, world + rot = convert_rotation(quat, input_format="quat_xyzw", output_format="matrix") # [chunk+1,3,3] + + r_prev = rot[:-1] # [chunk,3,3] + r_next = rot[1:] # [chunk,3,3] + # Translation delta rotated into the previous base frame. + d_world = (pos[1:] - pos[:-1]).unsqueeze(-1) # [chunk,3,1] + d_ego = torch.matmul(r_prev.transpose(-1, -2), d_world).squeeze(-1) # [chunk,3] + # Relative rotation expressed in the previous base frame. + r_rel = torch.matmul(r_prev.transpose(-1, -2), r_next) # [chunk,3,3] + rot6d = convert_rotation(r_rel, input_format="matrix", output_format="rot6d") # [chunk,6] + return torch.cat([d_ego, rot6d], dim=-1) # [chunk, 9] + + # ---- video ------------------------------------------------------------- + + def _compose_multi_view(self, sample: dict[str, Any]) -> torch.Tensor: + """Tile wrist (top) + left/right third-person (bottom) into one frame. + + Layout per frame (all source views are square, same size): + ┌──────────────┐ + │ wrist │ (H, W) + ├───────┬──────┤ + │ left │ right │ (H/2, W/2) each + └───────┴──────┘ + Output height is 3H/2 (mirrors ``DROIDLeRobotDataset._compose_multi_view``). + """ + wrist = sample[self._image_features["wrist"]] # [T,C,H,W] + left = sample[self._image_features["left"]] # [T,C,H,W] + right = sample[self._image_features["right"]] # [T,C,H,W] + + _, _, h_w, w_w = wrist.shape + half_h, half_w = h_w // 2, w_w // 2 + left = F.interpolate(left, size=(half_h, half_w), mode="bilinear", align_corners=False) + right = F.interpolate(right, size=(half_h, half_w), mode="bilinear", align_corners=False) + bottom = torch.cat([left, right], dim=-1) # [T,C,H/2,W] + return torch.cat([wrist, bottom], dim=-2) # [T,C,3H/2,W] + + def _compose_left_wrist(self, sample: dict[str, Any]) -> torch.Tensor: + """LIBERO-style: agentview_left (left) + wrist (right), full-res, horizontal. + + Both source views are native 256x256; concatenated along width -> [T,C,H,2W] + (256x512). No downscaling, right camera dropped. Mirrors + ``LIBEROLeRobotDataset`` concat_view ordering (third-person | wrist). + """ + left = sample[self._image_features["left"]] # [T,C,H,W] + wrist = sample[self._image_features["wrist"]] # [T,C,H,W] + return torch.cat([left, wrist], dim=-1) # [T,C,H,2W] + + + # ---- sample build ------------------------------------------------------ + + def __getitem__(self, idx: int) -> dict[str, Any]: + mode, _, _, sample = self._fetch_sample(idx) + + action = self._build_frame_wise_action(sample[_ACTION_FEATURE]) # [chunk, 10] + + if self._use_base_action: + raw = sample[_ACTION_FEATURE].float() + control_mode = raw[:, _CONTROL_MODE] # [chunk, 1]; +/-1, regressed like the gripper + if self._base_encoding == "raw": + # Native normalised velocity command, passed through unchanged -> [chunk, 15]. + base_block = raw[:, _BASE_MOTION] # [chunk, 4] + else: + # State-derived ego-frame pose delta -> [chunk, 20]. + base_block = self._build_base_delta(sample[_STATE_FEATURE]) # [chunk, 9] + action = torch.cat([base_block, control_mode, action], dim=-1) + + # EEF proprioception: prepend the current eef pose as a clean conditioning + # token (DROID use_state parity). Compute idle_frames on the delta-only chunk + # BEFORE prepending so the absolute-pose token does not skew idle detection. + state_extras: dict[str, Any] = {} + if self._use_state: + idle = self._compute_idle_frames(action) + if idle is not None: + state_extras["idle_frames"] = idle + initial_state = self._build_initial_state(sample[_STATE_FEATURE]) # [10] + if self._use_base_action: + # Widen the conditioning token to the action contract (20D ego / 15D raw). + # The base block is zero-filled on purpose: under "ego" the absolute base + # pose is world-frame and the kitchen layout is re-randomised per episode, + # and under "raw" ``observation.state`` carries no base VELOCITY at all — + # only pose. Either way there is no transferable signal to put here, unlike + # the EEF pose, which is base-relative. + pad = self.action_dim - 10 + initial_state = torch.cat( + [torch.zeros(pad, dtype=initial_state.dtype), initial_state], dim=-1 + ) # [action_dim] + action = torch.cat([initial_state.unsqueeze(0), action], dim=0) # [chunk+1, action_dim] + + if self._skip_video_loading: + video = None + elif self._camera_set == "left_wrist": + video = self._compose_left_wrist(sample) + elif self._viewpoint == "concat_view": + video = self._compose_multi_view(sample) + elif self._viewpoint == "wrist_view": + video = sample[self._image_features["wrist"]] + else: # third_person_view + video = sample[self._image_features["left"]] + + ai_caption = sample["task"] + + extras: dict[str, Any] = {} + if self._camera_set == "left_wrist": + extras["additional_view_description"] = ( + "The left half is a third-person view of the scene. " + "The right half is from the wrist-mounted camera." + ) + elif self._viewpoint == "concat_view": + extras["additional_view_description"] = ( + "The top row is from the wrist-mounted camera. " + "The bottom row contains two horizontally concatenated third-person views of the scene." + ) + return self._build_result( + mode=mode, video=video, action=action, ai_caption=ai_caption, **state_extras, **extras + ) diff --git a/cosmos_framework/data/generator/action/utils/domain_utils.py b/cosmos_framework/data/generator/action/utils/domain_utils.py index b559596b..11db233c 100644 --- a/cosmos_framework/data/generator/action/utils/domain_utils.py +++ b/cosmos_framework/data/generator/action/utils/domain_utils.py @@ -27,6 +27,7 @@ "drawanything": 21, "behavior1k_lerobot": 22, # BEHAVIOR-1K R1Pro mobile bimanual (23D joint action) "maniparena": 23, # ManipArena x2robot/ex001_6r dual-arm; own 20D EE-direct action projection + "robocasa": 24, } @@ -56,6 +57,7 @@ # because their raw width is set per-dataset at construction time. Inference # in inverse_dynamics/WAM modes is not supported for these domains until # canonical widths are added here. + "robocasa": 10, } diff --git a/cosmos_framework/scripts/action_policy_server_robocasa.py b/cosmos_framework/scripts/action_policy_server_robocasa.py new file mode 100644 index 00000000..c9d4187e --- /dev/null +++ b/cosmos_framework/scripts/action_policy_server_robocasa.py @@ -0,0 +1,1410 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +"""HTTP inference server for RoboCasa action policies using OmniMoTModel. + +Mirrors ``action_policy_server_libero`` (same endpoints, same request format) with the +additions the RoboCasa recipe needs: an optional ``"state"`` field carrying the current +end-effector pose, which is prepended to the action sequence as the clean conditioning +frame at index 0 for checkpoints trained with ``use_state=True``. + +The server exposes two endpoints: + +- POST /predict: run policy inference. + - Input: {"image": "", "prompt": "", "domain_name": "", "image_size": } + - Output: {"action": [[a0, a1, ...], ...], "video": ["", ...]} +- GET /info: model / runtime info (run_name, checkpoint, sampling params, ...). + +The server can load either a training-time DCP checkpoint or a consolidated +Hugging Face/safetensors checkpoint directory. To match the standard OSS +inference flow, export DCP checkpoints as a separate step first: + + PYTHONPATH=. python -m cosmos_framework.scripts.export_model \ + --checkpoint-path /path/to/job/checkpoints/iter_000020000 \ + --config-file /path/to/train-output/config.yaml \ + -o /path/to/train-output/model + +Example: + + PYTHONPATH=. python -m cosmos_framework.scripts.action_policy_server_robocasa \ + --checkpoint-path /path/to/train-output/model \ + --port 8000 + +Direct DCP loading also works when given the matching training config: + + PYTHONPATH=. python -m cosmos_framework.scripts.action_policy_server_robocasa \ + --checkpoint-path /path/to/job/checkpoints/iter_000020000 \ + --config-file /path/to/train-output/config.yaml \ + --port 8000 +""" + +from cosmos_framework.inference.common.init import ( # noqa: F401 (is_rank0 may be useful for logging) + init_script, + is_rank0, +) + +init_script() + +import base64 +import binascii +import datetime +import io +import json +import threading +import time +import traceback +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import pydantic +import torch +import tyro +from omegaconf import DictConfig +from PIL import Image + +# Action-specific helpers live in the in-tree project tree. Imports stay as +# `projects.cosmos3.vfm.*` and are auto-rewritten to `cosmos3._src.vfm.*` by the +# cosmos-framework release script. +from cosmos_framework.data.generator.action.utils.action_processing import ( + ActionProcessingRecord, + make_batched_action_processing_fields, +) +from cosmos_framework.data.generator.action.utils.domain_utils import get_domain_id +from cosmos_framework.data.generator.action.utils.json_formatter import ActionPromptJsonFormatter +from cosmos_framework.data.generator.action.utils.transforms import ( + build_sequence_plan_from_mode, + find_closest_target_size, + reflection_pad_to_target, + remove_reflection_padding, +) +from cosmos_framework.inference.args import OmniSetupArgs, OmniSetupOverrides +from cosmos_framework.inference.common.args import CheckpointOverrides, ConfigFileType, tyro_cli +from cosmos_framework.inference.common.config import deserialize_config_dict +from cosmos_framework.inference.common.init import init_output_dir +from cosmos_framework.inference.inference import OmniInference +from cosmos_framework.scripts.action_policy_server_utils import ( + DEFAULT_FALLBACK_OUTPUT_DIR, + disable_runtime_ema_for_frozen_config, + get_local_ip, + maybe_init_distributed, +) +from cosmos_framework.utils import log +from cosmos_framework.utils.lazy_config import instantiate +from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution + +_DEFAULT_ACTION_CHUNK_SIZE = 32 # RoboCasa recipe chunk_length; only a fallback when the config omits it +ActionNormalization = Literal["auto", "meanstd", "minmax", "quantile", "quantile_rot"] +ResolvedActionNormalization = Literal["meanstd", "minmax", "quantile", "quantile_rot"] + +_DURATION_FPS_TEMPLATE = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS." +_RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution." + +# Viewpoint tag for the concat_view (third-person + wrist) composite the client sends; +# matches the dataset's _VIEWPOINT_BY_CAMERA["concat_view"]. Used only when the experiment +# trains with JSON-structured prompts (format_prompt_as_json=True). +_JSON_VIEWPOINT = "concat_view" + + +# --------------------------------------------------------------------------- +# Pre/post processing helpers (copied verbatim from the previous server, with +# minor adjustments so config-introspection helpers also accept plain dicts). +# --------------------------------------------------------------------------- + + +def _augment_prompt_with_metadata( + prompt: str, + *, + t_frames: int, + fps: int, + height: int, + width: int, + append_duration_fps: bool = True, + append_resolution_info: bool = True, +) -> str: + """Append duration/FPS and resolution metadata to match training-time augmentation. + + Mirrors ``DurationFPSTextTimeStamps`` and ``ResolutionTextInfo`` augmentors + from the Action training transform pipeline. Only appends each piece when the + corresponding flag is ``True`` (matching the training config). + """ + if append_duration_fps: + duration = t_frames / fps + sep = " " if prompt.rstrip().endswith(".") else ". " + prompt = prompt + sep + _DURATION_FPS_TEMPLATE.format(duration=duration, fps=fps) + if append_resolution_info: + sep = " " if prompt.rstrip().endswith(".") else ". " + prompt = prompt + sep + _RESOLUTION_TEMPLATE.format(height=height, width=width) + return prompt + + +def _extract_bool_from_config(config: Any, key: str, default: bool) -> bool: + """Recursively search dataloader_train config for a boolean flag.""" + + def _search(obj: Any) -> bool | None: + if isinstance(obj, (DictConfig, dict)): + if key in obj: + val = obj[key] + if isinstance(val, bool): + return val + iterable = obj.values() + for v in iterable: + result = _search(v) + if result is not None: + return result + return None + + try: + if isinstance(config, dict): + dl_train = config.get("dataloader_train") + else: + dl_train = getattr(config, "dataloader_train", None) + if dl_train is not None: + result = _search(dl_train) + if result is not None: + return result + except Exception: + pass + return default + + +def _extract_str_from_config(config: Any, key: str) -> str | None: + """Recursively search dataloader_train config for a string field.""" + + def _search(obj: Any) -> str | None: + if isinstance(obj, (DictConfig, dict)): + if key in obj: + val = obj[key] + if isinstance(val, str): + return val + iterable = obj.values() + for v in iterable: + result = _search(v) + if result is not None: + return result + return None + + try: + if isinstance(config, dict): + dl_train = config.get("dataloader_train") + else: + dl_train = getattr(config, "dataloader_train", None) + if dl_train is not None: + return _search(dl_train) + except Exception: + pass + return None + + +def _extract_chunk_length_from_config(config: Any) -> int | None: + """Try to extract chunk_length from the experiment's dataloader config. + + Recursively searches ``config.dataloader_train`` for ``chunk_length`` or + ``num_action_per_chunk`` to determine the action chunk size the model was + trained with. Returns ``None`` when neither key is found. + """ + + def _search(obj: Any, keys: tuple[str, ...] = ("chunk_length", "num_action_per_chunk")) -> int | None: + if isinstance(obj, (DictConfig, dict)): + for key in keys: + if key in obj: + val = obj[key] + if isinstance(val, int): + return val + for v in obj.values(): + result = _search(v, keys) + if result is not None: + return result + return None + + try: + if isinstance(config, dict): + dl_train = config.get("dataloader_train") + else: + dl_train = getattr(config, "dataloader_train", None) + if dl_train is not None: + return _search(dl_train) + except Exception: + pass + return None + + +def _strip_data_url_prefix(b64: str) -> str: + # Accept "data:image/png;base64,...." as well as raw base64. + if "," in b64 and b64[:64].lower().startswith("data:"): + return b64.split(",", 1)[1].strip() + return b64.strip() + + +def _b64decode_loose(b64: str) -> bytes: + """ + Decode base64 permissively. + + The simulator/client may include whitespace/newlines, omit padding, or use urlsafe base64. + """ + s = _strip_data_url_prefix(b64) + s = "".join(s.split()) # remove whitespace/newlines + pad = (-len(s)) % 4 + if pad: + s = s + ("=" * pad) + try: + return base64.b64decode(s, validate=False) + except binascii.Error: + return base64.urlsafe_b64decode(s) + + +def _decode_base64_png_to_rgb_uint8(image_b64: str) -> torch.Tensor: + """ + Returns a tensor with shape (3, H, W), dtype uint8, RGB. + """ + try: + raw = _b64decode_loose(image_b64) + except (binascii.Error, ValueError) as e: + raise ValueError(f"Invalid base64 image: {e}") from e + + with Image.open(io.BytesIO(raw)) as img: + img = img.convert("RGB") + # Pillow can expose a read-only view; make it writable to avoid PyTorch warnings/UB. + arr = np.asarray(img, dtype=np.uint8).copy() + if arr.ndim != 3 or arr.shape[2] != 3: + raise ValueError(f"Expected RGB image, got shape {arr.shape}") + chw = torch.from_numpy(arr).permute(2, 0, 1).contiguous() # [3,H,W] + return chw + + +def _video_tensor_to_pil_images(video_c_t_h_w: torch.Tensor) -> list[Image.Image]: + """ + Convert (C, T, H, W) float tensor in [-1,1] or [0,1] to a list of PIL RGB frames. + """ + if video_c_t_h_w.dim() != 4: + raise ValueError(f"Expected (C,T,H,W), got {tuple(video_c_t_h_w.shape)}") + if int(video_c_t_h_w.shape[0]) != 3: + raise ValueError(f"Expected C=3 RGB, got C={int(video_c_t_h_w.shape[0])}") + + images: list[Image.Image] = [] + t = int(video_c_t_h_w.shape[1]) + for ti in range(t): + frame = video_c_t_h_w[:, ti].detach().cpu().float() + if frame.min().item() < 0.0: + frame = (frame + 1.0) / 2.0 + frame = frame.clamp(0.0, 1.0) + frame_uint8 = (frame * 255.0).round().to(torch.uint8) # [3,H,W] + hwc = frame_uint8.permute(1, 2, 0).numpy() # [H,W,3] + images.append(Image.fromarray(hwc)) + return images + + +def _save_gif(frames: list[Image.Image], path: Path, fps: int) -> None: + if not frames: + return + path.parent.mkdir(parents=True, exist_ok=True) + duration_ms = max(1, int(round(1000.0 / float(fps)))) + frames[0].save( + path, + save_all=True, + append_images=frames[1:], + duration=duration_ms, + loop=0, + optimize=False, + ) + + print(f"Saved gif to {path}") + + +def _save_policy_request_dump( + *, + dump_root: Path, + request_id: int, + request_json: dict[str, Any], + obs_chw_uint8: torch.Tensor, + pred_action: list[list[float]], + pred_video_c_t_h_w: torch.Tensor | None, + fps: int, +) -> None: + """ + Dump input observation, predicted actions, and rollout video for offline debugging. + Creates: + - request.json + - observation.png + - action_output.json + - rollout.gif (if pred_video provided) + - rollout_frames/frame_XXX.png (if pred_video provided) + """ + ts = datetime.datetime.now(tz=datetime.timezone.utc).strftime("%Y%m%d_%H%M%S_%fZ") + out_dir = dump_root / f"{ts}_req{request_id:06d}" + out_dir.mkdir(parents=True, exist_ok=False) + + # Save request JSON (without the base64 image to save space, image is saved separately) + request_json_copy = {k: v for k, v in request_json.items() if k != "image"} + request_json_copy["image"] = "" + (out_dir / "request.json").write_text(json.dumps(request_json_copy, indent=2), encoding="utf-8") + + # Save observation image + obs_hwc = obs_chw_uint8.permute(1, 2, 0).cpu().numpy() # [H,W,3] + obs_img = Image.fromarray(obs_hwc) + obs_img.save(out_dir / "observation.png") + + # Save predicted actions + action_output = {"action": pred_action} + (out_dir / "action_output.json").write_text(json.dumps(action_output, indent=2), encoding="utf-8") + + if pred_video_c_t_h_w is None: + return + + # Save rollout video + frames = _video_tensor_to_pil_images(pred_video_c_t_h_w) + _save_gif(frames, out_dir / "rollout.gif", fps=fps) + + # Save individual frames + frames_dir = out_dir / "rollout_frames" + frames_dir.mkdir(parents=True, exist_ok=True) + for i, frame in enumerate(frames): + frame.save(frames_dir / f"frame_{i:03d}.png") + + +def _save_failed_request_dump( + *, + dump_root: Path, + request_id: int, + request_json: dict[str, Any], + error: str, +) -> None: + """ + Dump request + error even if inference fails. + Best-effort: will try to decode and save observation image if present. + """ + ts = datetime.datetime.now(tz=datetime.timezone.utc).strftime("%Y%m%d_%H%M%S_%fZ") + out_dir = dump_root / f"{ts}_req{request_id:06d}_ERROR" + out_dir.mkdir(parents=True, exist_ok=False) + + # Save request JSON (without base64 image) + request_json_copy = {k: v for k, v in request_json.items() if k != "image"} + if "image" in request_json: + request_json_copy["image"] = "" + (out_dir / "request.json").write_text(json.dumps(request_json_copy, indent=2), encoding="utf-8") + (out_dir / "error.txt").write_text(error, encoding="utf-8") + + try: + image_b64 = request_json.get("image") + if isinstance(image_b64, str): + img_chw_uint8 = _decode_base64_png_to_rgb_uint8(image_b64) + obs_hwc = img_chw_uint8.permute(1, 2, 0).cpu().numpy() # [H,W,3] + Image.fromarray(obs_hwc).save(out_dir / "observation.png") + except Exception: + # Ignore any dump failures. + return + + +def _ts() -> str: + return datetime.datetime.now(tz=datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ") + + +def _load_introspection_config_dict(setup_args: OmniSetupArgs) -> dict: + """Load the full experiment config as a plain dict for prompt-augmentation + introspection. + + For ``MODULE`` (``.py``) configs, ``OmniInference.create`` saves the + structured config to ``setup_args.output_dir / 'config.yaml'`` via + ``cosmos_framework.inference.common.config.save_config`` immediately after model load. For + ``YAML`` / ``JSON`` configs we just deserialize the source file directly. + """ + if setup_args.config_file_type == ConfigFileType.MODULE: + saved = Path(setup_args.output_dir) / "config.yaml" + if saved.exists(): + return deserialize_config_dict(saved) + # Fallback: re-parse the .py module without instantiating anything. + import importlib + + from cosmos_framework.inference.common.config import unstructure_config + from cosmos_framework.utils import config_helper + + config_module = importlib.import_module(config_helper.get_config_module(setup_args.config_file)) + config = config_module.make_config() + config = config_helper.override( + config, ["--", f"experiment={setup_args.experiment}", *setup_args.experiment_overrides] + ) + return unstructure_config(config, invalid="ignore") + + return deserialize_config_dict(Path(setup_args.config_file)) + + +# --------------------------------------------------------------------------- +# CLI args +# --------------------------------------------------------------------------- + + +class ActionServerArgs(pydantic.BaseModel): + model_config = pydantic.ConfigDict(extra="forbid", use_attribute_docstrings=True) + + # We deliberately do NOT expose the full ``OmniSetupOverrides`` (i.e. no + # ``setup: SetupOverrides`` field). That would surface every + # ``setup.sample_overrides.*`` knob (``--guidance``, ``--seed``, + # ``--num-steps``, ...) which collides with our own server-level defaults. + # Instead we reuse the shared checkpoint/config args and build + # ``OmniSetupOverrides`` programmatically in ``build_setup_overrides``. + + checkpoint: tyro.conf.OmitArgPrefixes[CheckpointOverrides] = CheckpointOverrides.model_construct() + """Checkpoint and config loading configuration. ``use_ema_weights`` lives here and + defaults True at inference (suppressed from CLI) -> evals load net_ema by default.""" + + output_dir: Path | None = None + """Output directory for ``OmniInference`` (saved config.yaml, benchmarks). + Defaults to ``--dump-dir`` if set, else ``/tmp/cosmos3_action_server``.""" + + # ----- single-rank parallelism / sampler ---------------------------------- + sampler: Literal["unipc", "edm"] = "unipc" + """Diffusion sampler used by ``OmniInference``.""" + + # ----- sampling defaults (per-request, used when client doesn't override) -- + seed: int = 0 + """Random seed for ``model.generate_samples_from_batch``.""" + guidance: float = 1.0 + """Guidance scale for denoising.""" + num_steps: int = 30 + """Number of denoising steps.""" + fps: int = 20 + """Frames per second used for both prompt augmentation and rollout encoding.""" + + # ----- action policy parameters ------------------------------------------- + action_chunk_size: int | None = None + """Number of action steps to predict. Defaults to ``chunk_length`` / + ``num_action_per_chunk`` from the experiment config (or 16).""" + max_action_dim: int | None = None + """Maximum action dimension. Defaults to ``model.config.max_action_dim`` + from the experiment config (or 64).""" + raw_action_dim: int | None = None + """Unpadded action dimension used for action-channel masking. Inferred + from action stats when omitted.""" + + # ----- action denormalization --------------------------------------------- + action_stats_path: Path | None = None + """Path to action stats JSON for denormalizing predicted actions.""" + action_normalization: ActionNormalization = "auto" + """Action normalization to invert. ``auto`` reads ``action_normalization`` + from the experiment config (default ``minmax`` if unspecified).""" + + # ----- prompt format ------------------------------------------------------ + format_prompt_as_json: bool | None = None + """Serve prompts as structured JSON (matching training ``format_prompt_as_json``). + ``None`` reads the flag from the experiment config; set explicitly to override when + the eval experiment differs from the checkpoint's training prompt format.""" + + # ----- debug dumps -------------------------------------------------------- + dump_dir: Path | None = None + """If set, dump observations, predicted actions, and rollout videos under + this directory for offline debugging.""" + dump_every: int = 1 + """Dump every N-th request (only used when ``--dump-dir`` is set).""" + + # ----- HTTP server -------------------------------------------------------- + host: str = "0.0.0.0" + """HTTP host to bind.""" + port: int = 8000 + """HTTP port to bind.""" + http_400_on_error: bool = False + """If set, return HTTP 400 on inference errors. Default is HTTP 200 with an + empty action list, matching the legacy simulator client expectations.""" + + # ----- developer utilities ------------------------------------------------ + run_validation: bool = False + """If set, run a one-shot validation/training batch through the model on + startup (developer debugging only).""" + + def build_setup_overrides(self) -> OmniSetupOverrides: + """Build an ``OmniSetupOverrides`` from checkpoint and server fields. + + Required fields (``checkpoint_path``) must be present; optional fields + keep their ``OmniSetupOverrides`` defaults when not specified by the + user. + """ + if not getattr(self.checkpoint, "checkpoint_path", ""): + raise ValueError("--checkpoint-path is required") + + output_dir = self.output_dir or self.dump_dir or DEFAULT_FALLBACK_OUTPUT_DIR + + base = OmniSetupOverrides.model_validate(self.checkpoint.model_dump()) + base.output_dir = output_dir + base.sampler = self.sampler + return base + + +# --------------------------------------------------------------------------- +# Service implementation (predict path is a verbatim port from the previous +# evaluation/action/http_inference_server.py). +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ActionServerConfig: + """Internal snapshot of the resolved CLI args, threaded through the service. + + Kept as a plain dataclass (vs. carrying ``ActionServerArgs`` directly) so + request-time code can ``replace(...)`` individual fields after model load. + """ + + seed: int + guidance: float + num_steps: int + fps: int + action_chunk_size: int + max_action_dim: int + raw_action_dim: int | None + dump_dir: Path | None + dump_every: int + http_400_on_error: bool + action_stats_path: Path | None + action_normalization: ActionNormalization + experiment_name: str + checkpoint_dir: str + + +class ActionModelService: + def __init__(self, args: ActionServerArgs) -> None: + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for OmniMoTModel inference in this repo.") + + # OmniInference internally calls into FSDP / DTensor parallelize utilities + # that expect a process group; create a single-rank PG when not under + # torchrun (init_script() only inits the PG when WORLD_SIZE>1). + maybe_init_distributed() + + setup_overrides = args.build_setup_overrides() + setup_args = setup_overrides.build_setup() + init_output_dir(setup_args.output_dir) + setup_args = disable_runtime_ema_for_frozen_config(setup_args) + + # Surface the resolved max_action_dim into the experiment config when + # the user explicitly overrode it on the CLI; matches the previous + # ``experiment_opts=[f"model.config.max_action_dim={...}"]`` plumbing. + if args.max_action_dim is not None: + setup_args.experiment_overrides = [ + *setup_args.experiment_overrides, + f"model.config.max_action_dim={int(args.max_action_dim)}", + ] + + log.info( + f"[action-server] loading model: config_file='{setup_args.config_file}' " + f"({setup_args.config_file_type}) experiment='{setup_args.experiment}' " + f"checkpoint_path='{setup_args.checkpoint_path}'" + ) + + # OmniInference dispatches between MODULE (.py) and YAML/JSON loaders. + pipe = OmniInference.create(setup_args) + self.pipe: OmniInference = pipe + self.model = pipe.model + self.model.eval() + # OmniInference always uses OmniSetupArgs at runtime, but the base class + # types the attribute as the more general SetupArgs. + assert isinstance(pipe.setup_args, OmniSetupArgs) + self.setup_args: OmniSetupArgs = pipe.setup_args + self.experiment_config: dict = _load_introspection_config_dict(self.setup_args) + + # Resolve action_chunk_size: CLI arg > experiment config > default. + if args.action_chunk_size is not None: + resolved_chunk_size = int(args.action_chunk_size) + else: + config_chunk = _extract_chunk_length_from_config(self.experiment_config) + if config_chunk is not None: + resolved_chunk_size = config_chunk + log.info( + f"[action-server] --action-chunk-size not specified, " + f"using chunk_length={resolved_chunk_size} from experiment config" + ) + else: + resolved_chunk_size = _DEFAULT_ACTION_CHUNK_SIZE + log.info( + f"[action-server] --action-chunk-size not specified and not found in experiment config, " + f"using default={resolved_chunk_size}" + ) + + # Resolve max_action_dim: CLI > model config > default (64). + if args.max_action_dim is not None: + resolved_max_action_dim = int(args.max_action_dim) + else: + model_max_action_dim = getattr(self.model.config, "max_action_dim", None) + if isinstance(model_max_action_dim, int): + resolved_max_action_dim = model_max_action_dim + log.info( + f"[action-server] --max-action-dim not specified, " + f"using max_action_dim={resolved_max_action_dim} from model config" + ) + else: + resolved_max_action_dim = 64 + log.info( + f"[action-server] --max-action-dim not specified and not found in model config, " + f"using default={resolved_max_action_dim}" + ) + + self.cfg = ActionServerConfig( + seed=int(args.seed), + guidance=float(args.guidance), + num_steps=int(args.num_steps), + fps=int(args.fps), + action_chunk_size=resolved_chunk_size, + max_action_dim=resolved_max_action_dim, + raw_action_dim=int(args.raw_action_dim) if args.raw_action_dim is not None else None, + dump_dir=args.dump_dir, + dump_every=int(args.dump_every), + http_400_on_error=bool(args.http_400_on_error), + action_stats_path=args.action_stats_path, + action_normalization=args.action_normalization, + experiment_name=setup_args.experiment or "", + checkpoint_dir=setup_args.checkpoint_path, + ) + + self._lock = threading.Lock() + self._req_id_lock = threading.Lock() + self._req_id = 0 + + # Whether the checkpoint was trained with proprioception prepended as a clean + # conditioning action frame. A request must then carry "state": serving such a model + # without it does not raise anywhere downstream, it just drops a conditioning signal + # the policy depends on, so the mismatch is rejected explicitly in _prep_policy_item. + self.requires_state = _extract_bool_from_config(self.experiment_config, "use_state", default=False) + if self.requires_state: + log.info("[action-server] checkpoint trained with use_state=True; requests must carry 'state'") + + self.append_duration_fps = _extract_bool_from_config( + self.experiment_config, "append_duration_fps", default=True + ) + self.append_resolution_info = _extract_bool_from_config( + self.experiment_config, "append_resolution_info", default=True + ) + # When the experiment trains with format_prompt_as_json=True, the caption is a + # structured JSON dict (ActionPromptJsonFormatter) and the legacy string appenders + # are skipped. Mirror that at serve time so the prompt format matches training. The + # CLI flag overrides the config when the eval experiment differs from the checkpoint. + if args.format_prompt_as_json is not None: + self.format_prompt_as_json = bool(args.format_prompt_as_json) + else: + self.format_prompt_as_json = _extract_bool_from_config( + self.experiment_config, "format_prompt_as_json", default=False + ) + self._prompt_json_formatter = ( + ActionPromptJsonFormatter(caption_key="ai_caption") if self.format_prompt_as_json else None + ) + log.info( + f"[action-server] prompt augmentation: " + f"append_duration_fps={self.append_duration_fps}, append_resolution_info={self.append_resolution_info}, " + f"format_prompt_as_json={self.format_prompt_as_json}" + ) + + # Action denormalization stats. + self.action_min: torch.Tensor | None = None + self.action_range: torch.Tensor | None = None + self.action_mean: torch.Tensor | None = None + self.action_std: torch.Tensor | None = None + self.action_normalization: ResolvedActionNormalization = "minmax" + self.raw_action_dim: int | None = self.cfg.raw_action_dim + self._load_action_normalization_stats() + if self.raw_action_dim is None: + self.raw_action_dim = 7 + + if args.run_validation: + self._run_developer_validation() + + # ------------------------------------------------------------------ + # Action denormalization + # ------------------------------------------------------------------ + + def _load_action_normalization_stats(self) -> None: + """Load action denormalization tensors from ``cfg.action_stats_path``. + + Populates ``self.action_normalization``, the relevant tensor attributes + (``action_mean`` / ``action_std`` for meanstd; ``action_min`` / + ``action_range`` for minmax / quantile / quantile_rot), and infers + ``self.raw_action_dim`` when not explicitly set. + """ + if self.cfg.action_stats_path is None: + return + + self.action_normalization = self._resolve_action_normalization(self.cfg.action_normalization) + stats_path = Path(self.cfg.action_stats_path) + if not stats_path.is_absolute(): + stats_path = Path.cwd() / stats_path + with open(stats_path) as f: + raw_stats = json.load(f) + if not isinstance(raw_stats, dict): + raise ValueError(f"Action stats file must contain a dict: {stats_path}") + stats_key = "global_raw" if self.action_normalization == "quantile_rot" else "global" + stats = raw_stats.get(stats_key, raw_stats) + if not isinstance(stats, dict): + raise ValueError(f"Action stats file must contain a dict or {stats_key} stats dict: {stats_path}") + if self.action_normalization == "meanstd": + if "mean" not in stats or "std" not in stats: + raise ValueError(f"Mean/std action normalization requires 'mean' and 'std' in {stats_path}") + self.action_mean = torch.tensor(stats["mean"], dtype=torch.float32) # [D] + action_std = torch.tensor(stats["std"], dtype=torch.float32) # [D] + self.action_std = torch.clamp(action_std, min=1e-8) # [D] + stats_dim = int(self.action_mean.shape[0]) + stats_summary = f"mean={self.action_mean.tolist()}, std={self.action_std.tolist()}" + elif self.action_normalization in ("quantile", "quantile_rot"): + if "q01" not in stats or "q99" not in stats: + raise ValueError(f"Quantile action normalization requires 'q01' and 'q99' in {stats_path}") + self.action_min = torch.tensor(stats["q01"], dtype=torch.float32) # [D] + action_max = torch.tensor(stats["q99"], dtype=torch.float32) # [D] + action_range = action_max - self.action_min # [D] + self.action_range = torch.clamp(action_range, min=1e-6) # [D] + stats_dim = int(self.action_min.shape[0]) + stats_summary = f"q01={self.action_min.tolist()}, q99={action_max.tolist()}" + else: + if "min" not in stats or "max" not in stats: + raise ValueError(f"Min/max action normalization requires 'min' and 'max' in {stats_path}") + self.action_min = torch.tensor(stats["min"], dtype=torch.float32) # [D] + action_max = torch.tensor(stats["max"], dtype=torch.float32) # [D] + action_range = action_max - self.action_min # [D] + self.action_range = torch.clamp(action_range, min=1e-6) # [D] + stats_dim = int(self.action_min.shape[0]) + stats_summary = f"min={self.action_min.tolist()}, max={action_max.tolist()}" + if self.raw_action_dim is None: + self.raw_action_dim = stats_dim + if stats_dim != self.raw_action_dim: + raise ValueError(f"Action stats dimension {stats_dim} does not match raw_action_dim={self.raw_action_dim}") + log.info( + f"[action-server] Loaded action stats for denormalization from {stats_path}: " + f"normalization={self.action_normalization}, {stats_summary}" + ) + + def _resolve_action_normalization( + self, requested_normalization: ActionNormalization + ) -> ResolvedActionNormalization: + """Resolve auto action normalization from the loaded experiment config.""" + if requested_normalization != "auto": + return requested_normalization + + configured_normalization = _extract_str_from_config(self.experiment_config, "action_normalization") + if configured_normalization is None: + return "minmax" + if configured_normalization in ("meanstd", "minmax", "quantile", "quantile_rot"): + return configured_normalization # type: ignore[return-value] + raise ValueError( + "action_policy_server_robocasa.py can denormalize action_normalization='minmax', 'meanstd', " + "'quantile', or 'quantile_rot'; " + f"loaded experiment config requested {configured_normalization!r}. " + "Pass --action-normalization explicitly if this checkpoint should use a supported method." + ) + + def _denormalize_action(self, action: torch.Tensor) -> torch.Tensor: + """Invert the configured action normalization.""" + if self.action_normalization == "meanstd": + if self.action_mean is None or self.action_std is None: + return action + action_dim = self.action_mean.shape[0] + normalized = action[..., :action_dim] # [...,D] + action_mean = self.action_mean.to(action.device) # [D] + action_std = self.action_std.to(action.device) # [D] + return normalized * action_std + action_mean # [...,D] + + if self.action_min is None or self.action_range is None: + return action + action_dim = self.action_min.shape[0] + normalized = action[..., :action_dim] # [...,D] + action_min = self.action_min.to(action.device) # [D] + action_range = self.action_range.to(action.device) # [D] + return (normalized + 1.0) / 2.0 * action_range + action_min # [...,D] + + # ------------------------------------------------------------------ + # HTTP plumbing + # ------------------------------------------------------------------ + + def _should_dump(self, request_id: int) -> bool: + """Decide whether to dump this request. + + We always dump the first request (request_id == 1) as a quick sanity + check that dumping is wired up, then dump every N-th request controlled + by ``dump_every``. + """ + if self.cfg.dump_dir is None: + return False + n = int(self.cfg.dump_every) + if n <= 0: + return False + return request_id == 1 or (request_id % n == 0) + + def get_info(self) -> dict[str, Any]: + """Return model / server info for the /info endpoint. + + Includes all runtime-relevant config so clients can record reproducible + params.json without needing to know CLI flags. + """ + return { + "run_name": self.cfg.experiment_name, + "checkpoint": self.cfg.checkpoint_dir, + "config_file": str(self.setup_args.config_file), + "config_file_type": str(self.setup_args.config_file_type), + "guidance": self.cfg.guidance, + "num_steps": self.cfg.num_steps, + "fps": self.cfg.fps, + "seed": self.cfg.seed, + "action_chunk_size": self.cfg.action_chunk_size, + "max_action_dim": self.cfg.max_action_dim, + "raw_action_dim": self.cfg.raw_action_dim, + "action_stats_path": str(self.cfg.action_stats_path) if self.cfg.action_stats_path else None, + } + + # ------------------------------------------------------------------ + # Predict + # ------------------------------------------------------------------ + + def _input_video_key(self) -> str: + input_video_key = getattr(self.model, "input_video_key", None) + if input_video_key is None: + input_video_key = getattr(self.model, "config", None).input_video_key # type: ignore[union-attr] + return input_video_key + + def _build_json_prompt(self, prompt: str, *, video: torch.Tensor, image_size: torch.Tensor) -> str: + """Reproduce the training-time JSON prompt for format_prompt_as_json=True runs. + + Runs the same ``ActionPromptJsonFormatter`` the training pipeline uses (after + spatial resize/pad), then ``json.dumps`` the dict exactly as + ``TextTokenizerTransform`` does before tokenization. ``idle_frames=0`` matches the + modal active-manipulation chunk (the policy should keep moving); ``viewpoint`` and + the zero ``action`` (total-frame count) mirror the RoboCasa concat_view dataset.""" + data_dict: dict[str, Any] = { + "ai_caption": prompt, + "viewpoint": _JSON_VIEWPOINT, + "video": video, # post-pad [C,T,H,W]; formatter reads T for duration + "image_size": image_size, # post-pad [H,W]; formatter reads resolution + "conditioning_fps": torch.tensor(self.cfg.fps, dtype=torch.long), + "mode": "wam", + # Zero action chunk: only its frame count (chunk length) is read, for " out of ". + "action": torch.zeros((self.cfg.action_chunk_size, self.cfg.max_action_dim), dtype=torch.float32), + "idle_frames": torch.tensor(0, dtype=torch.long), + } + formatted = self._prompt_json_formatter(data_dict)["ai_caption"] + return json.dumps(formatted) if isinstance(formatted, dict) else str(formatted) + + def _prep_policy_item(self, req: dict[str, Any]) -> dict[str, Any]: + """Validate one request and build the per-sample model inputs (video pad, + prompt augmentation, sequence_plan). Shared by predict_policy (batch=1) and + predict_policy_batch (batch=N) so the two paths stay byte-identical per item.""" + image_b64 = req.get("image") + if not isinstance(image_b64, str): + raise ValueError("'image' must be a base64 string") + prompt = req.get("prompt") + if not isinstance(prompt, str): + raise ValueError("'prompt' must be a string") + domain_name = req.get("domain_name") + if not isinstance(domain_name, str): + raise ValueError("'domain_name' must be a string") + image_size = req.get("image_size") + if not isinstance(image_size, int) or image_size <= 0: + raise ValueError("'image_size' must be a positive integer") + + # Optional proprioception: a raw ``raw_action_dim``-wide state token (the current eef + # pose) prepended as the clean conditioning action frame at index 0, matching the + # ``use_state`` dataset path. + state_token: torch.Tensor | None = None + raw_state = req.get("state") + if raw_state is None and self.requires_state: + raise ValueError( + "this checkpoint was trained with use_state=True, so 'state' is required: send the " + f"current eef pose as a {self.raw_action_dim}-wide list. Omitting it does not fail " + "downstream, it just drops the conditioning frame the policy was trained on." + ) + if raw_state is not None: + state_token = torch.as_tensor(raw_state, dtype=torch.float32).reshape(-1) + if self.raw_action_dim is not None and state_token.shape[0] != self.raw_action_dim: + raise ValueError( + f"'state' has width {state_token.shape[0]} but raw_action_dim={self.raw_action_dim}" + ) + + img_chw_uint8 = _decode_base64_png_to_rgb_uint8(image_b64) + img_h, img_w = img_chw_uint8.shape[-2:] + # Multi-view (non-square) images: scale proportionally, matching height to image_size. + if img_h != image_size: + scale = image_size / img_h + new_w = int(round(img_w * scale)) + hwc = img_chw_uint8.permute(1, 2, 0).cpu().numpy() + resized = Image.fromarray(hwc).resize((new_w, image_size), resample=Image.Resampling.BILINEAR) + arr = np.asarray(resized, dtype=np.uint8).copy() + img_chw_uint8 = torch.from_numpy(arr).permute(2, 0, 1).contiguous() + + t_frames = self.cfg.action_chunk_size + 1 + _, final_h, final_w = img_chw_uint8.shape + video_c_t_h_w_uint8 = img_chw_uint8.unsqueeze(1).repeat(1, t_frames, 1, 1) # [3,T,H,W] + resolution = get_vision_data_resolution((final_h, final_w)) + target_w, target_h = find_closest_target_size(final_h, final_w, resolution) + pad_dict: dict[str, Any] = {"video": video_c_t_h_w_uint8} + reflection_pad_to_target(pad_dict, ["video"], True, target_w, target_h) + action_length = self.cfg.action_chunk_size + (1 if state_token is not None else 0) + sequence_plan = build_sequence_plan_from_mode( + mode="wam", + video_length=self.cfg.action_chunk_size + 1, + action_length=action_length, + has_text=True, + ) + if self._prompt_json_formatter is not None: + augmented_prompt = self._build_json_prompt( + prompt, video=pad_dict["video"], image_size=pad_dict["image_size"] + ) + else: + augmented_prompt = _augment_prompt_with_metadata( + prompt, + t_frames=t_frames, + fps=self.cfg.fps, + height=final_h, + width=final_w, + append_duration_fps=self.append_duration_fps, + append_resolution_info=self.append_resolution_info, + ) + return { + "img_chw_uint8": img_chw_uint8, + "video_padded": pad_dict["video"], + "padded_image_size": pad_dict["image_size"], + "augmented_prompt": augmented_prompt, + "sequence_plan": sequence_plan, + "domain_name": domain_name, + "image_size": image_size, + "state_token": state_token, + } + + def _build_action_input(self, state_token: "torch.Tensor | None") -> "torch.Tensor": + """Zeros placeholder [T, max_action_dim] for the diffusion start. A state token becomes + row 0, the clean conditioning frame, so T = chunk+1. It is written raw (unnormalized) to + match ``action_normalization=None``; a normalized-state run would need the forward + normalizer applied here too.""" + t = self.cfg.action_chunk_size + (1 if state_token is not None else 0) + action_t_d = torch.zeros((t, self.cfg.max_action_dim), dtype=torch.float32) + if state_token is not None: + action_t_d[0, : state_token.shape[0]] = state_token + return action_t_d + + def predict_policy_batch(self, reqs: list[dict[str, Any]]) -> dict[str, Any]: + """Batched policy inference: N requests -> ONE diffusion forward (batch_size=N) + -> N denormalized action chunks. Skips vision decode (the vectorized eval client + only needs actions), so it is ~N x faster than N serial /predict calls.""" + t0 = time.monotonic() + if not isinstance(reqs, list) or not reqs: + raise ValueError("'items' must be a non-empty list of policy requests") + preps = [self._prep_policy_item(r) for r in reqs] + n = len(preps) + input_video_key = self._input_video_key() + batch: dict[str, Any] = { + input_video_key: [[p["video_padded"]] for p in preps], + **make_batched_action_processing_fields( + ActionProcessingRecord(raw_action_dim=self.raw_action_dim, action_normalizer=None), + batch_size=n, + ), + "action": [[self._build_action_input(p["state_token"])] for p in preps], + "mode": ["wam"] * n, + "ai_caption": [p["augmented_prompt"] for p in preps], + "prompt": [p["augmented_prompt"] for p in preps], + "conditioning_fps": [torch.tensor(self.cfg.fps, dtype=torch.long) for _ in preps], + "image_size": torch.stack([p["padded_image_size"] for p in preps]).to(device="cuda"), + "domain_id": [torch.tensor(get_domain_id(p["domain_name"]), dtype=torch.long) for p in preps], + "sequence_plan": [p["sequence_plan"] for p in preps], + } + t_inf0 = time.monotonic() + with self._lock: + with torch.inference_mode(): + samples = self.model.generate_samples_from_batch( + batch, + guidance=self.cfg.guidance, + seed=[self.cfg.seed] * n, + num_steps=self.cfg.num_steps, + has_negative_prompt=False, + ) + t_inf1 = time.monotonic() + actions: list[list[list[float]]] = [] + for i in range(n): + pred = samples["action"][i].float().squeeze(0) # [T,D] + pred = self._denormalize_action(pred) + if preps[i]["state_token"] is not None: + pred = pred[1:] # drop the prepended clean state-conditioning frame + actions.append(pred.detach().cpu().numpy().tolist()) + log.info( + f"[action-server] predict_batch n={n} steps={self.cfg.num_steps} " + f"ms_total={(time.monotonic() - t0) * 1000.0:.1f} ms_infer={(t_inf1 - t_inf0) * 1000.0:.1f}" + ) + return {"actions": actions} + + def predict_policy(self, req: dict[str, Any]) -> dict[str, Any]: + """ + Run policy inference: given an observation image and prompt, predict actions. + + Input request format: + { + "image": "", + "prompt": "", + "domain_name": "", + "image_size": + } + + Output format: + { + "action": [[a0_0, a0_1, ...], ..., [aN_0, aN_1, ...]], + "video": ["", ...] # List of T base64-encoded PNG frames + } + + All action dimensions are returned. Video is the decoded predicted rollout as base64 PNGs. + """ + t0 = time.monotonic() + + # Get or assign request ID + injected_id = req.get("request_id", None) + if isinstance(injected_id, int) and injected_id > 0: + request_id = int(injected_id) + else: + with self._req_id_lock: + self._req_id += 1 + request_id = int(self._req_id) + + # Per-item preprocessing (validation, decode/resize/pad, prompt, sequence_plan). + t_decode0 = time.monotonic() + prep = self._prep_policy_item(req) + t_decode1 = time.monotonic() + img_chw_uint8 = prep["img_chw_uint8"] + video_padded = prep["video_padded"] + padded_image_size = prep["padded_image_size"] + augmented_prompt = prep["augmented_prompt"] + sequence_plan = prep["sequence_plan"] + domain_name = prep["domain_name"] + image_size = prep["image_size"] + + # Action: zeros noise start; row 0 holds the clean state-conditioning token when + # proprioception was supplied. + action_t_d = self._build_action_input(prep["state_token"]) # [T,action_dim] + + input_video_key = self._input_video_key() + + batch: dict[str, Any] = { + input_video_key: [[video_padded]], + # Provide BOTH raw_action_dim and the action_processing_record the model + # needs to externalize (invert) the generated action; building the batch + # by hand previously omitted the record -> "cannot be externalized". + **make_batched_action_processing_fields( + ActionProcessingRecord(raw_action_dim=self.raw_action_dim, action_normalizer=None), + batch_size=1, + ), + "action": [[action_t_d]], + "mode": ["wam"], + "ai_caption": [augmented_prompt], + "prompt": [augmented_prompt], + "conditioning_fps": [torch.tensor(self.cfg.fps, dtype=torch.long)], + "image_size": padded_image_size.unsqueeze(0).to(device="cuda"), + "domain_id": [torch.tensor(get_domain_id(domain_name), dtype=torch.long)], + "sequence_plan": [sequence_plan], + } + + if getattr(self.model, "training", False): + log.warning(f"[action-server] request_id={request_id} WARNING: model.training=True") + + log.info( + f"[action-server] request_id={request_id} mode=policy " + f"prompt={augmented_prompt!r} domain_name={domain_name!r} image_size={image_size} " + f"img={tuple(img_chw_uint8.shape)} steps={self.cfg.num_steps} guidance={self.cfg.guidance}" + ) + + # Run inference + t_inf0 = time.monotonic() + with self._lock: + with torch.inference_mode(): + samples = self.model.generate_samples_from_batch( + batch, + guidance=self.cfg.guidance, + seed=[self.cfg.seed], + num_steps=self.cfg.num_steps, + has_negative_prompt=False, + ) + pred_action = samples["action"][0] # [T,D] or [1,T,D] + + # Decode vision for rollout video (samples["vision"] is a list; take first sample) + pred_video_c_t_h_w = self.model.decode(samples["vision"][0]).squeeze(0) # [C,T,H,W] + + # Remove reflection padding so the reported video matches the original resolution + pred_video_c_t_h_w = remove_reflection_padding(pred_video_c_t_h_w, padded_image_size) + t_inf1 = time.monotonic() + + # Extract actions: return all dimensions — (T, D) or (1, T, D) + pred_action = pred_action.float().squeeze(0) # [T,D] + pred_action = self._denormalize_action(pred_action) + if prep["state_token"] is not None: + pred_action = pred_action[1:] # drop the prepended clean state-conditioning frame + pred_action_np = pred_action.detach().cpu().numpy() # [T,D] + pred_action_list = pred_action_np.tolist() # List of [a0, a1, ..., aD] + + # Convert video to base64-encoded PNG frames + pred_video_frames = _video_tensor_to_pil_images(pred_video_c_t_h_w) + pred_video_b64: list[str] = [] + for frame in pred_video_frames: + buf = io.BytesIO() + frame.save(buf, format="PNG") + pred_video_b64.append(base64.b64encode(buf.getvalue()).decode("ascii")) + + # Optional offline debug dump + if self._should_dump(request_id): + dump_dir = self.cfg.dump_dir + assert dump_dir is not None + dump_root = Path(dump_dir) + dump_root.mkdir(parents=True, exist_ok=True) + try: + log.info(f"[action-server] request_id={request_id} dumping to {str(dump_root)}") + _save_policy_request_dump( + dump_root=dump_root, + request_id=request_id, + request_json=req, + obs_chw_uint8=img_chw_uint8, + pred_action=pred_action_list, + pred_video_c_t_h_w=pred_video_c_t_h_w, + fps=int(self.cfg.fps), + ) + except Exception as e: + # Never fail serving a request due to dump failures + log.error(f"[action-server] dump failed for request_id={request_id}: {e}") + + dt_total_ms = (time.monotonic() - t0) * 1000.0 + dt_decode_ms = (t_decode1 - t_decode0) * 1000.0 + dt_inf_ms = (t_inf1 - t_inf0) * 1000.0 + log.info( + f"[action-server] request_id={request_id} done action_steps={len(pred_action_list)} " + f"video_frames={len(pred_video_b64)} " + f"ms_total={dt_total_ms:.1f} ms_decode={dt_decode_ms:.1f} ms_infer={dt_inf_ms:.1f}" + ) + return {"action": pred_action_list, "video": pred_video_b64} + + # ------------------------------------------------------------------ + # Developer validation (optional, --run-validation) + # ------------------------------------------------------------------ + + def _run_developer_validation(self) -> None: + """Run a single validation/training batch through the model on startup.""" + # Re-instantiate config so we can spin up dataloaders without the + # inference-time freezes applied to the model config. + if self.setup_args.config_file_type != ConfigFileType.MODULE: + log.warning( + "[action-server] --run-validation requires a .py config-file (got " + f"{self.setup_args.config_file_type}); skipping." + ) + return + + try: + config = self.setup_args.load_config() + except Exception as e: + log.warning(f"[action-server] --run-validation could not load config: {e}; skipping.") + return + + try: + val_dataset = instantiate(config.dataloader_val) + train_dataset = instantiate(config.dataloader_train) + except Exception as e: + log.warning(f"[action-server] --run-validation could not instantiate datasets: {e}; skipping.") + return + + val_batch = next(iter(val_dataset)) # pyrefly: ignore[no-matching-overload] + train_batch = next(iter(train_dataset)) # pyrefly: ignore[no-matching-overload] + + with torch.inference_mode(): + self.model.training_step(val_batch, 0) + sample_num = 1 + val_result = self.model.generate_samples_from_batch( + val_batch, + guidance=1.0, + seed=[0 for _ in range(sample_num)], + num_steps=8, + n_sample=sample_num, + ) + video_mse_list = [] + action_mse_list = [] + for i in range(sample_num): + val_video = self.model.decode(val_result["vision"][i]).detach().cpu() + val_video_mse = torch.nn.functional.mse_loss(val_video, val_batch["video"][i].cpu()) + val_action = val_result["action"][i].detach().cpu() + val_action_mse = torch.nn.functional.mse_loss(val_action[:, :6], val_batch["action"][i][0][:, :6].cpu()) + video_mse_list.append(val_video_mse.item()) + action_mse_list.append(val_action_mse.item()) + log.info(f"Val video MSE: {np.mean(video_mse_list)}") + log.info(f"Val action MSE: {np.mean(action_mse_list)}") + + self.model.training_step(train_batch, 0) + train_result = self.model.generate_samples_from_batch( + train_batch, + guidance=1.0, + seed=list(range(20)), + num_steps=8, + n_sample=20, + ) + train_video = self.model.decode(train_result["vision"][0]) + train_action = train_result["action"][0] + train_action_mse = torch.nn.functional.mse_loss(train_action[:, :6], train_batch["action"][0][:, :6]) + train_video_mse = torch.nn.functional.mse_loss(train_video, train_batch["video"][0]) + log.info(f"Train action MSE: {train_action_mse}; Train video MSE: {train_video_mse}") + + +# --------------------------------------------------------------------------- +# HTTP handler +# --------------------------------------------------------------------------- + + +class _ActionHandler(BaseHTTPRequestHandler): + """ + ThreadingHTTPServer handler. + + The service instance is injected via ``server.service``. + """ + + server: ThreadingHTTPServer # type: ignore[assignment] + + def _send_json(self, status_code: int, payload: dict[str, Any]) -> None: + body = json.dumps(payload).encode("utf-8") + self.send_response(status_code) + self.send_header("Content-Type", "application/json") + # Avoid caches/proxies returning stale results. + self.send_header("Cache-Control", "no-store") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + try: + self.wfile.write(body) + except (BrokenPipeError, ConnectionResetError): + # Client closed the connection (often due to request timeout). + # Avoid noisy tracebacks; nothing to do server-side. + return + + def do_GET(self) -> None: # noqa: N802 + if self.path == "/info": + service: ActionModelService = getattr(self.server, "service") # type: ignore[attr-defined] + self._send_json(200, service.get_info()) + elif self.path == "/": + self._send_json(200, {"status": "ok"}) + else: + self._send_json(404, {"error": "Not found"}) + + def do_POST(self) -> None: # noqa: N802 + if self.path not in ("/", "/predict", "/predict_batch"): + self._send_json(404, {"error": "Not found"}) + return + + content_type = (self.headers.get("Content-Type") or "").split(";", 1)[0].strip().lower() + if content_type != "application/json": + self._send_json(415, {"error": "Content-Type must be application/json"}) + return + + try: + length = int(self.headers.get("Content-Length") or "0") + except ValueError: + self._send_json(400, {"error": "Invalid Content-Length"}) + return + + body = self.rfile.read(max(0, length)) + try: + req = json.loads(body.decode("utf-8")) + except Exception as e: + self._send_json(400, {"error": f"Invalid JSON: {e}"}) + return + + if not isinstance(req, dict): + self._send_json(400, {"error": "JSON body must be an object"}) + return + + service: ActionModelService = getattr(self.server, "service") # type: ignore[attr-defined] + + # Generate a per-request id at the HTTP layer to correlate logs and dumps. + req_id_lock: threading.Lock | None = getattr(self.server, "_req_id_lock", None) # type: ignore[attr-defined] + if req_id_lock is None: + req_id_lock = threading.Lock() + setattr(self.server, "_req_id_lock", req_id_lock) # type: ignore[attr-defined] + setattr(self.server, "_req_id", 0) # type: ignore[attr-defined] + with req_id_lock: + next_id = int(getattr(self.server, "_req_id")) + 1 # type: ignore[attr-defined] + setattr(self.server, "_req_id", next_id) # type: ignore[attr-defined] + req["request_id"] = next_id + + log.info( + f"[action-server] HTTP request_id={next_id} from={self.client_address[0]}:{self.client_address[1]} " + f"path={self.path} bytes={length}" + ) + + is_batch = self.path == "/predict_batch" + try: + if is_batch: + out = service.predict_policy_batch(req.get("items", [])) + else: + out = service.predict_policy(req) + except Exception as e: + err = str(e) + traceback.print_exc() + + payload = ( + {"actions": [], "error": err} + if is_batch + else {"action": [], "error": err, "request_id": req.get("request_id")} + ) + log.error(f"[action-server] request_id={req.get('request_id')} ERROR: {err}") + + # Dump failed request for offline debugging if enabled. + if service.cfg.dump_dir is not None: + try: + dump_root = Path(service.cfg.dump_dir) + dump_root.mkdir(parents=True, exist_ok=True) + _save_failed_request_dump( + dump_root=dump_root, + request_id=int(req.get("request_id") or 0), + request_json=req, + error=err, + ) + except Exception: + pass + + status = 400 if service.cfg.http_400_on_error else 200 + self._send_json(status, payload) + return + + self._send_json(200, out) + + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 + # Silence default request logging (the simulator can be chatty). + return + + +# --------------------------------------------------------------------------- +# Entrypoint +# --------------------------------------------------------------------------- + + +def serve(args: ActionServerArgs) -> None: + if args.dump_dir is not None: + # Create dump dir up front so it's obvious where outputs will go. + dump_root = Path(args.dump_dir).resolve() + dump_root.mkdir(parents=True, exist_ok=True) + log.info(f"[action-server] dump_root={str(dump_root)} dump_every={args.dump_every}") + + service = ActionModelService(args) + + local_ip = get_local_ip() + log.info( + f"[action-server] starting host={args.host} port={int(args.port)} " + f"experiment_name={service.cfg.experiment_name!r} " + f"steps={service.cfg.num_steps} guidance={service.cfg.guidance} fps={service.cfg.fps} " + f"action_chunk_size={service.cfg.action_chunk_size} max_action_dim={service.cfg.max_action_dim} " + f"raw_action_dim={service.cfg.raw_action_dim} " + f"dump_dir={service.cfg.dump_dir} dump_every={service.cfg.dump_every} " + f"http_400_on_error={service.cfg.http_400_on_error}" + ) + log.info(f"[action-server] Server accessible at: http://{local_ip}:{int(args.port)}/") + log.info("[action-server] Endpoints:") + log.info(" - GET / : Health check") + log.info(" - GET /info : Model info (run_name, checkpoint, sampling params)") + log.info(" - POST /predict: Policy inference (image + prompt + domain_name + image_size -> action)") + + httpd: ThreadingHTTPServer = ThreadingHTTPServer((args.host, int(args.port)), _ActionHandler) + setattr(httpd, "service", service) + httpd.serve_forever() + + +def main() -> None: + args = tyro_cli( + ActionServerArgs, + description=__doc__, + config=( + tyro.conf.OmitArgPrefixes, + tyro.conf.CascadeSubcommandArgs, + tyro.conf.OmitSubcommandPrefixes, + ), + ) + serve(args) + + +if __name__ == "__main__": + main() diff --git a/cosmos_framework/scripts/convert_robocasa_to_lerobot_v30.sh b/cosmos_framework/scripts/convert_robocasa_to_lerobot_v30.sh new file mode 100755 index 00000000..c49089a5 --- /dev/null +++ b/cosmos_framework/scripts/convert_robocasa_to_lerobot_v30.sh @@ -0,0 +1,89 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + +# One-time data prep for the RoboCasa recipes: convert a released RoboCasa export from LeRobot +# v2.1 to v3.0. Every task found under SRC_ROOT is converted, so this works for any split -- +# target/atomic, pretrain/atomic, pretrain/composite, ... +# +# SRC_ROOT=/path/to/robocasa/datasets/v1.0/target/atomic \ +# bash $(python -c "import cosmos_framework, pathlib; print(pathlib.Path(cosmos_framework.__file__).parent)")/scripts/convert_robocasa_to_lerobot_v30.sh +# +# Why this is needed: RoboCasa publishes LeRobot **v2.1**, and the `lerobot` pinned by +# cosmos-framework is v3.0-only -- it rejects the v2.1 layout outright. Conversion is done on +# COPIES; the released dataset is never modified. +# +# CPU-only and idempotent: a task whose destination already has meta/tasks.parquet (a v3.0-only +# marker) is skipped, so re-running after an interruption resumes rather than redoing the work. +# +# Output layout, which is what ROBOCASA_ROOT must point at: +# ///lerobot/ + +set -euo pipefail + +# Deliberately does NOT cd to its own directory: the script lives inside the installed package, +# while the converted dataset belongs wherever the caller is working (typically a cookbook +# recipe folder). Both roots resolve against the caller's cwd. +# +# No default for SRC_ROOT: the RoboCasa release is site-specific, so it must be given. +SRC_ROOT="${SRC_ROOT:?set SRC_ROOT=/path/to/robocasa/datasets/v1.0/target/atomic (see https://robocasa.ai)}" +V30_ROOT="${V30_ROOT:-$PWD/data/robocasa_v30}" + +# Discover every task under SRC_ROOT rather than hardcoding a list: a task is any directory +# holding a //lerobot export. Narrow it with TASKS_OVERRIDE="A B C" if wanted. +# `find -L` so a SRC_ROOT assembled from symlinked task directories still resolves; without it +# such a layout yields no tasks and the run stops on the check below. +if [[ -n "${TASKS_OVERRIDE:-}" ]]; then + read -ra TASKS <<< "${TASKS_OVERRIDE}" +else + mapfile -t TASKS < <( + find -L "${SRC_ROOT}" -mindepth 3 -maxdepth 3 -type d -name lerobot -printf '%h\n' \ + | xargs -r -n1 dirname | xargs -r -n1 basename | sort -u + ) +fi +(( ${#TASKS[@]} )) || { + echo "ERROR: no //lerobot exports found under ${SRC_ROOT}" >&2 + echo " Is SRC_ROOT pointing at a RoboCasa split directory?" >&2 + exit 1 +} +echo "discovered ${#TASKS[@]} task(s)" + +echo "== v2.1 -> v3.0 conversion ==" +echo " from: ${SRC_ROOT}" +echo " to : ${V30_ROOT}" +for t in "${TASKS[@]}"; do + src="$(ls -d "${SRC_ROOT}/${t}/"*/lerobot 2>/dev/null | head -1 || true)" + if [[ -z "${src}" ]]; then echo " ${t}: no source shard, skipping"; continue; fi + date="$(basename "$(dirname "${src}")")" + dst_parent="${V30_ROOT}/${t}/${date}" + dst="${dst_parent}/lerobot" + if [[ -f "${dst}/meta/tasks.parquet" ]]; then echo " ${t}: already v3.0, skip"; continue; fi + echo " ${t}: copy + convert (${src})" + mkdir -p "${dst_parent}" + rm -rf "${dst}" "${dst_parent}/lerobot_old" "${dst_parent}/lerobot_v30" + cp -r "${src}" "${dst}" + python -m lerobot.datasets.v30.convert_dataset_v21_to_v30 \ + --repo-id=lerobot --root="${dst_parent}" --push-to-hub false + rm -rf "${dst_parent}/lerobot_old" + # Carry over RoboCasa-specific meta the converter ignores. The loader does not need these, + # but keeping them makes the converted copy self-describing. + for f in modality.json embodiment.json; do + [[ -f "${src}/meta/${f}" ]] && cp "${src}/meta/${f}" "${dst}/meta/${f}" || true + done +done + +echo "== verify ==" +missing=0 +for t in "${TASKS[@]}"; do + if ls "${V30_ROOT}/${t}/"*/lerobot/meta/tasks.parquet >/dev/null 2>&1; then + echo " OK ${t}" + else + echo " MISS ${t}"; missing=$((missing + 1)) + fi +done +if (( missing )); then + echo "ERROR: ${missing} task(s) failed to convert; see the log above." >&2 + exit 1 +fi +echo "== done: ${#TASKS[@]} tasks under ${V30_ROOT} ==" +echo " train with: ROBOCASA_ROOT=${V30_ROOT} bash launch_sft_action_policy_robocasa_nano.sh" diff --git a/cosmos_framework/simulation/robocasa/__init__.py b/cosmos_framework/simulation/robocasa/__init__.py new file mode 100644 index 00000000..503ec1b1 --- /dev/null +++ b/cosmos_framework/simulation/robocasa/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 + diff --git a/cosmos_framework/simulation/robocasa/closed_loop_eval.py b/cosmos_framework/simulation/robocasa/closed_loop_eval.py new file mode 100644 index 00000000..f46381a2 --- /dev/null +++ b/cosmos_framework/simulation/robocasa/closed_loop_eval.py @@ -0,0 +1,377 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""Closed-loop evaluation for RoboCasa using the Action HTTP inference server. + +# Example (mobile-base recipe: 15-D raw base action, agentview_left | eye_in_hand composite): +MUJOCO_GL=egl PYTHONPATH=. python cosmos_framework/simulation/robocasa/closed_loop_eval.py \ + --server-url http://localhost:8900 \ + --dataset-dir /path/to/robocasa/datasets/v1.0/target/atomic/CloseFridge//lerobot \ + --output-dir results/robocasa_closed_loop/CloseFridge \ + --num-test-episodes 50 \ + --action-horizon 32 \ + --camera-set left_wrist \ + --use-state \ + --use-base-action --base-encoding raw \ + --success-latch 1 --seed 0 --image-size 256 --cam-size 256 + +TWO PYTHON ENVIRONMENTS ARE REQUIRED. robosuite/robocasa and cosmos-framework cannot share one +venv (conflicting numpy/mujoco pins), so this script -- which only drives the simulator -- runs +in the robosuite/robocasa venv, while the policy is served by +``cosmos_framework.scripts.action_policy_server_robocasa`` from the cosmos-framework venv. Only +empty namespace packages are imported from ``cosmos_framework`` here, so ``PYTHONPATH=.`` is +enough and cosmos-framework need not be installed in the simulator venv. + +The evaluation contract -- ``--action-horizon``, ``--camera-set``, ``--use-state``, +``--base-encoding`` -- MUST match the training recipe. A mismatch does not raise; it silently +degrades the policy. + +Protocol: env args are rebuilt from the dataset's ``extras/dataset_meta.json`` so they match how +the demonstrations were recorded, held-out scenes come from the official ``target`` object split +and official layout/style combinations, success uses the environment's own ``_check_success`` +with first-success latching, and the rollout horizon is RoboCasa's official per-task value. + +Writes ``results.json`` (one record per rollout) into ``--output-dir``. +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import imageio.v2 as imageio +from PIL import Image + +import robosuite +import robocasa # noqa: F401 (registers RoboCasa envs like CloseToasterOvenDoor; does NOT pull lerobot) + +from cosmos_framework.simulation.robocasa.eval_utils import ( + decode_10d_to_env12, + decode_15d_to_env12, + decode_20d_to_env12, + decode_pred_video, + predict, +) + +CAMS = ["robot0_agentview_left", "robot0_agentview_right", "robot0_eye_in_hand"] + + + + +def get_env_metadata(dsdir: str) -> dict: + with open(Path(dsdir) / "extras" / "dataset_meta.json") as f: + return json.load(f)["env_args"] + + + + + + + + + + + + +CAMERA_SET = "wrist_lr" # set from --camera-set in main(); switches compose layout +USE_BASE_ACTION = False # set from --use-base-action; widens the state token to the mobile contract +BASE_ENCODING = "ego" # set from --base-encoding; "ego" = 20D pose delta, "raw" = 15D base_motion +_ACTION_DIM_BY_ENCODING = {"ego": 20, "raw": 15} + + +def _upright(img: np.ndarray) -> np.ndarray: + """Flip a raw robosuite camera observation to match the recorded dataset. + + MuJoCo renders with the OpenGL convention (origin bottom-left), so + ``obs["_image"]`` comes back vertically mirrored, while the RoboCasa + LeRobot videos the policy was trained on are stored upright. Feeding the + unflipped frame does not raise -- the policy simply acts on an upside-down + world and success collapses to ~0, so this must not be made optional. + """ + return img[::-1] + + +def compose(obs) -> np.ndarray: + if CAMERA_SET == "left_wrist": + # LIBERO-style: [agentview_left | wrist], full-res, horizontal -> 256x512. + # Mirrors RoboCasaLeRobotDataset._compose_left_wrist (left | wrist, no downscale). + l = _upright(obs[f"{CAMS[0]}_image"]) + w = _upright(obs[f"{CAMS[2]}_image"]) + return np.concatenate([l, w], axis=1) # [H, 2W, 3] + if CAMERA_SET == "lrw": + # All three cameras side by side at full res -> 256x768. + # Mirrors RoboCasaLeRobotDataset._compose_lrw (left | right | wrist). + l = _upright(obs[f"{CAMS[0]}_image"]) + r = _upright(obs[f"{CAMS[1]}_image"]) + w = _upright(obs[f"{CAMS[2]}_image"]) + return np.concatenate([l, r, w], axis=1) # [H, 3W, 3] + w = _upright(obs[f"{CAMS[2]}_image"]) + l = _upright(obs[f"{CAMS[0]}_image"]) + r = _upright(obs[f"{CAMS[1]}_image"]) + h, wd = w.shape[:2] + hh, hw = h // 2, wd // 2 + lh = np.asarray(Image.fromarray(l).resize((hw, hh), Image.BILINEAR)) + rh = np.asarray(Image.fromarray(r).resize((hw, hh), Image.BILINEAR)) + return np.concatenate([w, np.concatenate([lh, rh], axis=1)], axis=0) + + +def build_state_token(obs) -> list[float]: + """Current EEF proprioception -> 10D ``[pos(3), rot6d(6), gripper(1)]`` token, matching + ``RoboCasaLeRobotDataset._build_initial_state``. The stored observation.state eef fields + map to robosuite obs keys (robocasa lerobot_utils): end_effector_position_relative = + robot0_base_to_eef_pos, end_effector_rotation_relative = robot0_base_to_eef_quat (xyzw), + gripper_qpos = robot0_gripper_qpos. So these obs are already in the (base) frame the model + was trained on — no manual transform needed.""" + import robosuite.utils.transform_utils as T + pos = np.asarray(obs["robot0_base_to_eef_pos"], dtype=np.float32).reshape(3) + quat = np.asarray(obs["robot0_base_to_eef_quat"], dtype=np.float32).reshape(4) # xyzw + m = T.quat2mat(quat) # [3,3]; robosuite quat is xyzw + rot6d = np.concatenate([m[:, 0], m[:, 1]]).astype(np.float32) # [col0, col1] (matches convert_rotation) + qpos = np.asarray(obs["robot0_gripper_qpos"], dtype=np.float32).reshape(-1) + grip = np.array([qpos[0] - qpos[1]], dtype=np.float32) # signed finger opening + token = np.concatenate([pos, rot6d, grip]).astype(np.float32) # [10] + if USE_BASE_ACTION: + # The conditioning token must match the action width (20D ego / 15D raw). The base + # block is zero-filled exactly as in training (RoboCasaLeRobotDataset.__getitem__): + # under "ego" the absolute base pose is world-frame and carries no transferable + # signal, and under "raw" the observation has no base velocity to put there at all. + pad = _ACTION_DIM_BY_ENCODING[BASE_ENCODING] - 10 + token = np.concatenate([np.zeros(pad, dtype=np.float32), token]) + return token.tolist() + + + +# RoboCasa365 "Atomic-Seen" split (leaderboard protocol). Authoritative source is the +# installed robocasa's ``robocasa/utils/env_utils.py::create_env`` with ``split="target"``: +# +# obj_instance_split = "target" +# layout_and_style_ids = list(zip(range(1, 11), range(1, 11))) # 10 target kitchens +# robots = "PandaOmron" (composite controller default -> HYBRID_MOBILE_BASE) +# +# and ``docs/benchmarking/benchmarking_overview.md``: "For all experiments, we randomly +# sample 50 scenarios to run evaluation rollouts on." Horizons come from +# ``dataset_registry.py`` (v1.0.1 bumped every horizon by 1.5x). +# +# NOTE: ``robocasa/utils/eval_utils.py::create_eval_env`` still defaults to PandaMobile + +# OSC_POSE + obj_instance_split="B" + the 5 off-diagonal combos above. That is the older +# v0.x generalization protocol, NOT Atomic-Seen -- do not use it for leaderboard numbers. +ATOMIC_SEEN_LAYOUT_STYLE = tuple(zip(range(1, 11), range(1, 11))) +ATOMIC_SEEN_OBJ_SPLIT = "target" +ATOMIC_SEEN_NUM_ROLLOUTS = 50 + +# The three RoboCasa365 target splits are defined authoritatively by +# ``dataset_registry.TARGET_TASKS`` -- 18 atomic_seen + 16 composite_seen + 16 +# composite_unseen = the "50 target tasks" the leaderboard reports. Do NOT infer the +# composite split from which tasks have a pretrain entry: that gives 17/16, not 16/16. +# FIXED composite subset -- 12 seen + 6 unseen = 18 tasks, matching atomic_seen's task count +# so the three splits cost the same and the numbers sit on the same scale. +# + + +def official_horizon(task: str): + """Official per-task rollout horizon from robocasa's registry (source of truth), with a + hardcoded fallback for a few tasks if the import is unavailable.""" + try: + from robocasa.utils.dataset_registry_utils import get_task_horizon + return int(get_task_horizon(task)) + except Exception: + return _OFFICIAL_HORIZON_FALLBACK.get(task) + + +def make_env(dataset_dir: str, cam_size: int, *, obj_split=None, layout_style=None, seed=None): + env_meta = get_env_metadata(dataset_dir) + ek = dict(env_meta["env_kwargs"]) + ek["env_name"] = env_meta["env_name"] + if seed is not None: + # Fixed seed -> reproducible test-scene sequence (self.rng.choice over layout/style + + # object/placement sampling). Lets different checkpoints see the SAME 20 scenes/task. + ek["seed"] = seed + ek["has_renderer"] = False + ek["has_offscreen_renderer"] = True + ek["use_camera_obs"] = True + ek["camera_names"] = CAMS + ek["camera_widths"] = cam_size + ek["camera_heights"] = cam_size + ek["control_freq"] = 20 + ek["ignore_done"] = True + # Held-out test protocol: override object split + layout/style set so env.reset() (test + # from the demo model.xml, so exact reproduction is unaffected by these overrides. + if obj_split is not None: + ek["obj_instance_split"] = obj_split + if layout_style is not None: + ek["layout_and_style_ids"] = [list(ls) for ls in layout_style] + ek["layout_ids"] = None + ek["style_ids"] = None + ek.pop("renderer", None) + return robosuite.make(**ek) + + + +def check_success(env) -> bool: + try: + return bool(env._check_success()) + except Exception: + return False + + + + +def run_policy(env, *, server_url, image_size, action_horizon, max_steps, + latch, timeout, use_state=False, save_png=None, save_video=None, + gen_video_path=None, video_fps=20) -> tuple[bool, int, str]: + obs = env.reset() # a freshly sampled held-out scene + # Read the language annotation AFTER reset: get_ep_meta() describes the scene currently + # loaded, so reading it before reset returns the previous episode's metadata (and nothing + # at all on the first rollout, leaving the policy unconditioned). + try: + prompt = env.get_ep_meta().get("lang", "") or "" + except Exception: + prompt = "" + init_comp = compose(obs).astype(np.uint8) + if save_png is not None: + Image.fromarray(init_comp).save(save_png) + frames = [init_comp] # composite (what the policy sees + third-person robot view) + gen_frames = [] # model's GENERATED video (flow-matching vision branch), accumulated over ALL chunks + streak, queue, success, done_steps = 0, [], False, max_steps + for step in range(max_steps): + if not queue: + comp = compose(obs) + state_token = build_state_token(obs) if use_state else None + result = predict(server_url, comp, prompt, image_size, timeout, state=state_token) + # Accumulate the generated video from EVERY inference chunk (one clip per re-plan), + # so the saved gen video covers the whole rollout, not just the first chunk. + if gen_video_path is not None and result.get("video"): + gen_frames.extend(decode_pred_video(result["video"])) + acts = result["action"] + queue = acts[:action_horizon] if action_horizon > 0 else list(acts) + a = np.asarray(queue.pop(0)) + # Decoder is chosen by --base-encoding (explicit), not by guessing from the width: + # raw -> 15D [base_motion(4), control_mode(1), arm(10)] (identity base round-trip) + # ego -> 20D [base_pos(3), base_rot6d(6), control_mode(1), arm(10)] + # Without --use-base-action it is the 10D fixed-base contract. The width the server + # returned is asserted against that choice so a mismatched flag fails loudly instead + # of silently decoding garbage. + if USE_BASE_ACTION: + want = _ACTION_DIM_BY_ENCODING[BASE_ENCODING] + if a.shape[-1] != want: + raise ValueError( + f"--base-encoding={BASE_ENCODING} expects {want}D actions but the server " + f"returned {a.shape[-1]}D. Check that the flag matches the checkpoint." + ) + env_action = (decode_15d_to_env12 if BASE_ENCODING == "raw" else decode_20d_to_env12)(a, False) + else: + if a.shape[-1] != 10: + raise ValueError( + f"expected the 10D fixed-base contract but the server returned " + f"{a.shape[-1]}D; pass --use-base-action (and --base-encoding)." + ) + env_action = decode_10d_to_env12(a, False) + obs, _, _, _ = env.step(env_action) + frames.append(compose(obs).astype(np.uint8)) + if check_success(env): + streak += 1 + if streak >= latch: + success, done_steps = True, step + 1 + break + else: + streak = 0 + if save_video is not None: + imageio.mimwrite(save_video, frames, fps=video_fps, macro_block_size=None) + if gen_video_path is not None and gen_frames: + imageio.mimwrite(gen_video_path, gen_frames, fps=video_fps, macro_block_size=None) + return success, done_steps, prompt + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--server-url", default="http://127.0.0.1:8912") + ap.add_argument("--dataset-dir", required=True, help="v2.1 lerobot dir WITH extras/") + ap.add_argument("--num-test-episodes", type=int, default=0, help="test-scene episodes (random reset, same env_args)") + ap.add_argument("--max-steps", type=int, default=400) + ap.add_argument("--save-gen-video", action="store_true", + help="save the model's GENERATED video (flow-matching vision branch, from the " + "first inference) per rollout, for comparing video-generation quality") + ap.add_argument("--action-horizon", type=int, default=16) + ap.add_argument("--image-size", type=int, default=256) + ap.add_argument("--cam-size", type=int, default=256) + ap.add_argument("--success-latch", type=int, default=1, + help="consecutive _check_success() steps to declare success; 1 = official " + "run_random_rollouts (first success). Higher only to reject transient flukes.") + ap.add_argument("--seed", type=int, default=None, + help="fixed env seed for reproducible test scenes (same scenes across checkpoints); " + "omit for official nondeterministic sampling") + ap.add_argument("--camera-set", default="wrist_lr", choices=["wrist_lr", "left_wrist", "lrw"], + help="wrist_lr = 3-cam squished (384x256); left_wrist = LIBERO-style [left|wrist] 256x512; " + "lrw = 3-cam full-res side by side [left|right|wrist] 256x768") + ap.add_argument("--base-encoding", choices=("ego", "raw"), default="ego", + help="mobile-base action contract: 'ego' (20D state-derived pose delta, " + "the original) or 'raw' (15D native base_motion command)") + ap.add_argument("--use-base-action", action="store_true", + help="20D mobile-base checkpoint: widen the state token and decode base_motion") + ap.add_argument("--use-state", action="store_true", + help="send EEF proprioception (robot0_base_to_eef_pos/quat + gripper_qpos) as the " + "clean conditioning token; MUST match a use_state=True trained checkpoint") + ap.add_argument("--timeout", type=float, default=600) + ap.add_argument("--output-dir", default=".") + args = ap.parse_args() + + global CAMERA_SET, USE_BASE_ACTION, BASE_ENCODING + CAMERA_SET = args.camera_set + USE_BASE_ACTION = args.use_base_action + BASE_ENCODING = args.base_encoding + + out = Path(args.output_dir) + out.mkdir(parents=True, exist_ok=True) + + env = make_env( + args.dataset_dir, args.cam_size, + obj_split=ATOMIC_SEEN_OBJ_SPLIT, + layout_style=ATOMIC_SEEN_LAYOUT_STYLE, + seed=args.seed, + ) + print(f"[eval] official protocol: rollouts={args.num_test_episodes} seed={args.seed}", flush=True) + + # Effective rollout horizon: official per-task (robocasa registry) or the fixed --max-steps. + task_name = get_env_metadata(args.dataset_dir)["env_name"] + h = official_horizon(task_name) + eff_max_steps = h if h else args.max_steps + print(f"[eval] {task_name}: official horizon {eff_max_steps}", flush=True) + + # wait for server + import requests + t0 = time.time() + while time.time() - t0 < args.timeout: + try: + if requests.get(f"{args.server_url}/info", timeout=5).ok: + print(f"[train-scene] server ready at {args.server_url}", flush=True) + break + except requests.RequestException: + time.sleep(3) + else: + raise RuntimeError("server not ready") + + results = [] + + # Official protocol: rollouts in freshly sampled held-out scenes. + for t in range(args.num_test_episodes): + pol_ok, pol_steps, prompt = run_policy( + env, server_url=args.server_url, + image_size=args.image_size, action_horizon=args.action_horizon, + max_steps=eff_max_steps, latch=args.success_latch, + timeout=args.timeout, use_state=args.use_state, + save_png=str(out / f"rollout{t:02d}_init.png"), + save_video=str(out / f"rollout{t:02d}.mp4"), + gen_video_path=str(out / f"rollout{t:02d}_generated.mp4") if args.save_gen_video else None) + print(f"[eval] rollout {t:02d} success={pol_ok} steps={pol_steps} prompt={prompt!r}", flush=True) + results.append({"ep": t, "policy": pol_ok, "steps": pol_steps, "prompt": prompt}) + + env.close() + n_ok = sum(1 for r in results if r["policy"]) + print(f"\n{task_name}: {n_ok}/{len(results)} successful rollouts", flush=True) + (out / "results.json").write_text(json.dumps(results, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/cosmos_framework/simulation/robocasa/eval_utils.py b/cosmos_framework/simulation/robocasa/eval_utils.py new file mode 100644 index 00000000..4df59745 --- /dev/null +++ b/cosmos_framework/simulation/robocasa/eval_utils.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""Shared helpers for RoboCasa closed-loop evaluation. + +Action decoders that turn the policy's output into RoboCasa's native 12-D env action, the HTTP +call to the policy server, and the base64 -> frames decoder for the model's generated video. +Imported by ``closed_loop_eval.py``; not a script. +""" + +import base64 +import io +import json + +import numpy as np +import requests +from PIL import Image +from scipy.spatial.transform import Rotation as R + + + +# ----------------------------- action decode ----------------------------- + +def rot6d_to_matrix(r6: np.ndarray) -> np.ndarray: + """6D rotation (first two columns, Zhou 2019) -> 3x3 via Gram-Schmidt.""" + a1, a2 = r6[:3], r6[3:6] + b1 = a1 / (np.linalg.norm(a1) + 1e-8) + a2 = a2 - np.dot(b1, a2) * b1 + b2 = a2 / (np.linalg.norm(a2) + 1e-8) + b3 = np.cross(b1, b2) + return np.stack([b1, b2, b3], axis=1) + + +def decode_10d_to_env12(a10: np.ndarray, gripper_flip: bool) -> np.ndarray: + """[pos(3), rot6d(6), gripper(1)] -> env 12D [Δpos, Δrotvec, grip, base4=0, mode=-1].""" + a10 = np.asarray(a10, dtype=np.float64) + pos = a10[:3] + rotvec = R.from_matrix(rot6d_to_matrix(a10[3:9])).as_rotvec() + grip = -a10[9] if gripper_flip else a10[9] + env = np.zeros(12, dtype=np.float32) + env[0:3] = pos + env[3:6] = rotvec + env[6] = float(np.clip(grip, -1.0, 1.0)) + env[7:11] = 0.0 # base_motion (fixed base) + env[11] = -1.0 # control_mode = arm + return env + + +# Normalised base command in [-1, 1] maps onto these full-deflection velocities +# (calibrated on NavigateKitchen; see BASE_MAX_VELOCITY in robocasa_lerobot_dataset.py). +BASE_MAX_VELOCITY = (0.601, 0.640, 1.251) # forward m/s, side m/s, yaw rad/s +CONTROL_DT = 1.0 / 20.0 + + +def decode_20d_to_env12(a20: np.ndarray, gripper_flip: bool) -> np.ndarray: + """Mobile-base 20D policy output -> env 12D action. + + Policy layout (dataset order, base first):: + + [base_pos(3), base_rot6d(6), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)] + + Env layout (arm first — NOT the same order):: + + [eef_pos(3), eef_rotvec(3), gripper(1), base_motion(4), control_mode(1)] + + The base block is an ego-frame pose delta, so it converts straight to the normalised + velocity command without any world->body rotation: ``cmd = (delta / dt) / v_max``. + ``base_motion[3]`` (torso height) stays 0 — it is never actuated in target/atomic. + When the mode channel says arm-only, the base command is forced to zero so a small + spurious base delta cannot creep in. + """ + a20 = np.asarray(a20, dtype=np.float64) + env = decode_10d_to_env12(a20[10:20], gripper_flip) # arm block reuses the 10D path + + base_mode = float(a20[9]) > 0.0 + env[11] = 1.0 if base_mode else -1.0 + if base_mode: + d_pos = a20[0:3] # ego-frame translation delta (forward, side, up) + yaw = float(R.from_matrix(rot6d_to_matrix(a20[3:9])).as_rotvec()[2]) + cmd = np.array( + [ + d_pos[0] / CONTROL_DT / BASE_MAX_VELOCITY[0], + d_pos[1] / CONTROL_DT / BASE_MAX_VELOCITY[1], + yaw / CONTROL_DT / BASE_MAX_VELOCITY[2], + ] + ) + env[7:10] = np.clip(cmd, -1.0, 1.0) + env[10] = 0.0 # torso height: never actuated + return env + + +def decode_15d_to_env12(a15: np.ndarray, gripper_flip: bool) -> np.ndarray: + """Raw-base 15D policy output -> env 12D action. + + Policy layout (``base_encoding='raw'``):: + + [base_motion(4), control_mode(1), eef_pos(3), eef_rot6d(6), gripper(1)] + + The base block IS the env's native ``base_motion`` command, so this is a straight + copy — no dt, no velocity calibration, no controller inversion, and replaying a + recorded demo through this path is exact. Contrast ``decode_20d_to_env12``, which + has to invert a lagging velocity controller and therefore cannot round-trip. + + ``base_motion`` is already normalised to [-1, 1] in the data; the clip only guards + against the policy overshooting the valid command range. When the mode channel says + arm-only the base command is zeroed, matching the 20D path. + """ + a15 = np.asarray(a15, dtype=np.float64) + env = decode_10d_to_env12(a15[5:15], gripper_flip) # arm block reuses the 10D path + + base_mode = float(a15[4]) > 0.0 + env[11] = 1.0 if base_mode else -1.0 + env[7:11] = np.clip(a15[0:4], -1.0, 1.0) if base_mode else 0.0 + return env + + +# ----------------------------- observation ----------------------------- + + + +def b64_png(img: np.ndarray) -> str: + buf = io.BytesIO() + Image.fromarray(img.astype(np.uint8)).save(buf, format="PNG") + return base64.b64encode(buf.getvalue()).decode("ascii") + + +def predict(server_url: str, composite: np.ndarray, prompt: str, image_size: int, timeout: float, + state: list[float] | None = None) -> dict: + payload = {"image": b64_png(composite), "prompt": prompt, + "domain_name": "robocasa", "image_size": image_size} + if state is not None: + payload["state"] = state # 10D eef proprioception -> clean conditioning token + resp = requests.post(f"{server_url}/predict", json=payload, + headers={"Content-Type": "application/json"}, timeout=timeout) + resp.raise_for_status() + result = resp.json() + if result.get("error"): + raise RuntimeError(f"server error: {result['error']}") + return result + + +def decode_pred_video(video_b64_list) -> list[np.ndarray]: + frames = [] + for b in video_b64_list or []: + raw = base64.b64decode(b.split(",", 1)[-1]) + frames.append(np.asarray(Image.open(io.BytesIO(raw)).convert("RGB"))) + return frames + + +# ----------------------------- rollout ----------------------------- + + + + +