From 422c666cedcea2801e8c238b6cdfbac398b952e1 Mon Sep 17 00:00:00 2001 From: Lukas Zbinden Date: Sun, 26 Jul 2026 21:27:51 +0200 Subject: [PATCH 1/2] Add Cosmos-H-Dreams tabletop distillation recipe --- .../action/configs/action_conditioned/data.py | 89 ++ ...B_action_conditioned_rectify_flow_gr00t.py | 132 ++ .../datasets/gr00t_dreams/data/dataset.py | 44 +- .../datasets/gr00t_dreams/groot_configs.py | 79 +- .../inference/inference_jhu_dvrk_warmup.py | 412 ++++++ .../predict2/distill/utils/config_helper.py | 94 +- .../_src/predict2/interactive/configs/data.py | 18 + .../experiment/exp_action_self_forcing.py | 59 +- .../configs/experiment/exp_action_warmup.py | 45 +- ...orial_teacher_training_and_self_forcing.md | 1156 ++++++++++++++++- scripts/compute_openh_action_stats.py | 881 +++++++++++++ scripts/extract_jhu_inference_manifest.py | 315 +++++ .../tabletop/01_train_short_teacher_h13.sh | 55 + .../02_fine_anneal_short_teacher_h13.sh | 52 + .../tabletop/03_train_long_teacher_h73.sh | 53 + .../tabletop/04_phase0_teacher_cache_h73.sh | 71 + .../tabletop/05_warmup_student_h73.sh | 58 + train_scripts/tabletop/06_self_forcing_h73.sh | 54 + train_scripts/tabletop/README.md | 43 + 19 files changed, 3681 insertions(+), 29 deletions(-) create mode 100644 cosmos_predict2/_src/predict2/action/inference/inference_jhu_dvrk_warmup.py create mode 100644 scripts/compute_openh_action_stats.py create mode 100644 scripts/extract_jhu_inference_manifest.py create mode 100755 train_scripts/tabletop/01_train_short_teacher_h13.sh create mode 100755 train_scripts/tabletop/02_fine_anneal_short_teacher_h13.sh create mode 100755 train_scripts/tabletop/03_train_long_teacher_h73.sh create mode 100755 train_scripts/tabletop/04_phase0_teacher_cache_h73.sh create mode 100755 train_scripts/tabletop/05_warmup_student_h73.sh create mode 100755 train_scripts/tabletop/06_self_forcing_h73.sh create mode 100644 train_scripts/tabletop/README.md diff --git a/cosmos_predict2/_src/predict2/action/configs/action_conditioned/data.py b/cosmos_predict2/_src/predict2/action/configs/action_conditioned/data.py index f7f8809..8b73808 100644 --- a/cosmos_predict2/_src/predict2/action/configs/action_conditioned/data.py +++ b/cosmos_predict2/_src/predict2/action/configs/action_conditioned/data.py @@ -276,6 +276,8 @@ def build_webdataset(webdataset_instance, **kwargs): # ============================================================================ from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.dataset import MixedLeRobotDataset from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.groot_configs import ( + JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS, + JHU_DVRK_MONO_FINETUNE_VAL_DATASET_SPECS, MAX_ACTION_DIM, OPEN_H_DATASET_SPECS, ) @@ -310,6 +312,67 @@ def build_webdataset(webdataset_instance, **kwargs): drop_last=True, ) +# ============================================================================ +# JHU dVRK monocular reference tabletop mixture +# ============================================================================ +jhu_dvrk_mono_finetune_train_dataset = L(MixedLeRobotDataset)( + dataset_specs=JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS, + num_frames=13, + data_split="train", + max_action_dim=MAX_ACTION_DIM, + downscaled_res=False, + test_split_ratio=0.02, +) +jhu_dvrk_mono_finetune_val_dataset = L(MixedLeRobotDataset)( + dataset_specs=JHU_DVRK_MONO_FINETUNE_VAL_DATASET_SPECS, + num_frames=13, + data_split="test", + max_action_dim=MAX_ACTION_DIM, + downscaled_res=False, + test_split_ratio=0.02, +) +jhu_dvrk_mono_finetune_train_dataloader = L(DataLoader)( + dataset=jhu_dvrk_mono_finetune_train_dataset, + sampler=L(get_sampler)(dataset=jhu_dvrk_mono_finetune_train_dataset), + batch_size=1, + drop_last=True, +) +jhu_dvrk_mono_finetune_val_dataloader = L(DataLoader)( + dataset=jhu_dvrk_mono_finetune_val_dataset, + sampler=L(get_sampler)(dataset=jhu_dvrk_mono_finetune_val_dataset), + batch_size=1, + drop_last=True, +) + +jhu_dvrk_mono_finetune_h73_train_dataset = L(MixedLeRobotDataset)( + dataset_specs=JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS, + num_frames=73, + data_split="train", + max_action_dim=MAX_ACTION_DIM, + downscaled_res=False, + test_split_ratio=0.02, +) +jhu_dvrk_mono_finetune_h73_val_dataset = L(MixedLeRobotDataset)( + dataset_specs=JHU_DVRK_MONO_FINETUNE_VAL_DATASET_SPECS, + num_frames=73, + data_split="test", + max_action_dim=MAX_ACTION_DIM, + downscaled_res=False, + test_split_ratio=0.02, +) +jhu_dvrk_mono_finetune_h73_train_dataloader = L(DataLoader)( + dataset=jhu_dvrk_mono_finetune_h73_train_dataset, + sampler=L(get_sampler)(dataset=jhu_dvrk_mono_finetune_h73_train_dataset), + batch_size=1, + drop_last=True, +) +jhu_dvrk_mono_finetune_h73_val_dataloader = L(DataLoader)( + dataset=jhu_dvrk_mono_finetune_h73_val_dataset, + sampler=L(get_sampler)(dataset=jhu_dvrk_mono_finetune_h73_val_dataset), + batch_size=1, + drop_last=True, +) + # ============================================================================ # SutureBot Dataset Configuration @@ -453,6 +516,32 @@ def register_training_and_val_data(): node=open_h_multi_val_dataloader, ) + # JHU dVRK monocular tabletop reference recipe (short and long horizons). + cs.store( + group="data_train", + package="dataloader_train", + name="jhu_dvrk_mono_finetune_train", + node=jhu_dvrk_mono_finetune_train_dataloader, + ) + cs.store( + group="data_val", + package="dataloader_val", + name="jhu_dvrk_mono_finetune_val", + node=jhu_dvrk_mono_finetune_val_dataloader, + ) + cs.store( + group="data_train", + package="dataloader_train", + name="jhu_dvrk_mono_finetune_h73_train", + node=jhu_dvrk_mono_finetune_h73_train_dataloader, + ) + cs.store( + group="data_val", + package="dataloader_val", + name="jhu_dvrk_mono_finetune_h73_val", + node=jhu_dvrk_mono_finetune_h73_val_dataloader, + ) + # ============================================================================ # SutureBot dataset (20D actions zero-padded to 44D) # ============================================================================ diff --git a/cosmos_predict2/_src/predict2/action/configs/action_conditioned/experiment/exp_2B_action_conditioned_rectify_flow_gr00t.py b/cosmos_predict2/_src/predict2/action/configs/action_conditioned/experiment/exp_2B_action_conditioned_rectify_flow_gr00t.py index 1d6edee..8cfb0ba 100644 --- a/cosmos_predict2/_src/predict2/action/configs/action_conditioned/experiment/exp_2B_action_conditioned_rectify_flow_gr00t.py +++ b/cosmos_predict2/_src/predict2/action/configs/action_conditioned/experiment/exp_2B_action_conditioned_rectify_flow_gr00t.py @@ -16,9 +16,11 @@ # Configs for resuming from stage3 training import functools +import os from hydra.core.config_store import ConfigStore +from cosmos_predict2._src.imaginaire.functional.lr_scheduler import LambdaWarmUpCosineScheduler from cosmos_predict2._src.imaginaire.lazy_config import LazyCall as L from cosmos_predict2._src.imaginaire.lazy_config import LazyDict from cosmos_predict2._src.imaginaire.utils.checkpoint_db import get_checkpoint_path @@ -37,6 +39,19 @@ DEFAULT_CHECKPOINT = MODEL_CHECKPOINTS[ModelKey()] # This uses post_trained=True by default +_TABLETOP_OUTPUT_ROOT = os.environ.get("IMAGINAIRE_OUTPUT_ROOT", "imaginaire/output") +_TABLETOP_CHSS_CHECKPOINT = os.environ.get( + "CHSS_CHECKPOINT_DIR", + "checkpoints/cosmos-h-surgical-simulator", +) + + +def _tabletop_teacher_checkpoint(run_name: str, iteration: int) -> str: + return ( + f"{_TABLETOP_OUTPUT_ROOT}/cosmos_predict2_action_conditioned/" + f"official_runs_vid2vid/{run_name}/checkpoints/iter_{iteration:09d}" + ) + _TRAINER_DEBUG_CONFIG = dict( max_iter=1000, logging_iter=50, @@ -951,6 +966,111 @@ def build_debug_runs(job): flags={"allow_objects": True}, ) +# ============================================================================= +# JHU dVRK monocular reference tabletop recipe +# ============================================================================= +AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_FINETUNE_13FRAME_8NODES_OSS = LazyDict( + dict( + defaults=[ + "/experiment/2b_bridge_action_conditioned_oss", + {"override /net": "cosmos_v1_2B_action_chunk_conditioned"}, + {"override /data_train": "jhu_dvrk_mono_finetune_train"}, + {"override /data_val": "jhu_dvrk_mono_finetune_val"}, + "_self_", + ], + job=dict( + group="official_runs_vid2vid", + name="cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss", + project="cosmos_predict2_action_conditioned", + ), + checkpoint=dict( + load_path=_TABLETOP_CHSS_CHECKPOINT, + load_training_state=False, + strict_resume=False, + ), + model=dict( + config=dict( + state_t=1 + 12 // 4, + net=dict(action_dim=44), + ), + ), + dataloader_train=dict(batch_size=16), + optimizer=dict(lr=1.6e-4, weight_decay=0.1), + trainer=dict(max_iter=16000), + ), + flags={"allow_objects": True}, +) + +_JHU_H13_TEACHER_RUN = ( + "cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_" + "finetune_13frame_8nodes_release_oss" +) +AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_FINETUNE_13FRAME_8NODES_OSS_FINE_ANNEAL_4K = LazyDict( + dict( + defaults=[ + f"/experiment/{_JHU_H13_TEACHER_RUN}", + "_self_", + ], + job=dict( + group="official_runs_vid2vid", + name=f"{_JHU_H13_TEACHER_RUN}_fine_anneal_4k", + project="cosmos_predict2_action_conditioned", + ), + checkpoint=dict( + load_path=_tabletop_teacher_checkpoint(_JHU_H13_TEACHER_RUN, 16000), + load_training_state=False, + strict_resume=False, + ), + scheduler=L(LambdaWarmUpCosineScheduler)( + warm_up_steps=[100], + f_start=[0.10], + f_max=[1.00], + f_min=[0.05], + cycle_lengths=[4000], + ), + trainer=dict(max_iter=4000), + ), + flags={"allow_objects": True}, +) + +AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_TABLETOP_H73_8NODES_OSS = LazyDict( + dict( + defaults=[ + f"/experiment/{_JHU_H13_TEACHER_RUN}", + {"override /data_train": "jhu_dvrk_mono_finetune_h73_train"}, + {"override /data_val": "jhu_dvrk_mono_finetune_h73_val"}, + "_self_", + ], + job=dict( + group="official_runs_vid2vid", + name=f"{_JHU_H13_TEACHER_RUN}_h73_tabletop", + project="cosmos_predict2_action_conditioned", + ), + checkpoint=dict( + load_path=_tabletop_teacher_checkpoint(f"{_JHU_H13_TEACHER_RUN}_fine_anneal_4k", 4000), + load_training_state=False, + strict_resume=False, + ), + model=dict( + config=dict( + state_t=1 + 72 // 4, + net=dict(action_dim=44), + ), + ), + dataloader_train=dict(batch_size=4), + optimizer=dict(lr=4e-5, weight_decay=0.1), + scheduler=L(LambdaWarmUpCosineScheduler)( + warm_up_steps=[1000], + f_start=[0.10], + f_max=[1.00], + f_min=[0.05], + cycle_lengths=[5000], + ), + trainer=dict(max_iter=5000), + ), + flags={"allow_objects": True}, +) + cs = ConfigStore.instance() @@ -1005,6 +1125,18 @@ def build_debug_runs(job): AC_CHUNK_SINGLE_VIEW_2B_SUTUREBOT_13FRAME_NODES_OSS, *build_debug_runs(AC_CHUNK_SINGLE_VIEW_2B_SUTUREBOT_13FRAME_NODES_OSS), ], + [ + AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_FINETUNE_13FRAME_8NODES_OSS, + *build_debug_runs(AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_FINETUNE_13FRAME_8NODES_OSS), + ], + [ + AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_FINETUNE_13FRAME_8NODES_OSS_FINE_ANNEAL_4K, + *build_debug_runs(AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_FINETUNE_13FRAME_8NODES_OSS_FINE_ANNEAL_4K), + ], + [ + AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_TABLETOP_H73_8NODES_OSS, + *build_debug_runs(AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_TABLETOP_H73_8NODES_OSS), + ], ]: cs.store(group="experiment", package="_global_", name=f"{_item['job']['name']}", node=_item) if _item_wo_resume is not None: diff --git a/cosmos_predict2/_src/predict2/action/datasets/gr00t_dreams/data/dataset.py b/cosmos_predict2/_src/predict2/action/datasets/gr00t_dreams/data/dataset.py index 5973f6c..f3088e7 100644 --- a/cosmos_predict2/_src/predict2/action/datasets/gr00t_dreams/data/dataset.py +++ b/cosmos_predict2/_src/predict2/action/datasets/gr00t_dreams/data/dataset.py @@ -1457,13 +1457,22 @@ def __init__( self, *args, data_split="full", + test_split_ratio: float = 0.05, modality_filename: str | None = None, exclude_splits: list[str] | None = None, **kwargs, ): + """Wrap ``LeRobotSingleDataset`` with a deterministic train/test split. + + The split is taken from the trailing end of ``_all_steps``. + ``data_split="full"`` skips partitioning entirely. + """ + if not 0.0 < test_split_ratio < 1.0: + raise ValueError(f"test_split_ratio must be in (0, 1), got {test_split_ratio}") # Store data_split BEFORE calling super().__init__() because # _get_all_steps_cmr_filtered needs it for cache path generation self.data_split = data_split + self.test_split_ratio = test_split_ratio super().__init__( *args, modality_filename=modality_filename, @@ -1474,11 +1483,16 @@ def __init__( if data_split == "full": pass elif data_split == "train": - self._all_steps = self._all_steps[: -len(self) // 20] + n_test = max(1, int(len(self) * test_split_ratio)) + self._all_steps = self._all_steps[:-n_test] elif data_split == "test": - self._all_steps = self._all_steps[-len(self) // 20 :] + n_test = max(1, int(len(self) * test_split_ratio)) + self._all_steps = self._all_steps[-n_test:] - print(f"Dataset is split into {data_split} data, with {len(self._all_steps)} steps.") + print( + f"Dataset is split into {data_split} data (test_split_ratio={test_split_ratio:.4f}), " + f"with {len(self._all_steps)} steps." + ) def _get_trajectories(self) -> tuple[np.ndarray, np.ndarray]: """Get the trajectories in the dataset.""" @@ -1730,11 +1744,19 @@ class MixedLeRobotDataset(torch.utils.data.Dataset): - ``embodiment`` (str): Embodiment tag string (must be in EMBODIMENT_REGISTRY or be one of the built-in embodiments). - ``mix_ratio`` (float, optional): Relative sampling weight. Default 1.0. + - ``data_split_override`` (str, optional): Per-spec override of the + global ``data_split``. + - ``test_split_ratio_override`` (float, optional): Per-spec override + of the global ``test_split_ratio``. + - ``exclude_splits`` (list[str], optional): Episode-level split names + from ``meta/info.json`` to exclude. num_frames: Number of video frames per sample (e.g. 13 = 1 context + 12 pred). - data_split: One of ``"train"``, ``"test"``, ``"full"``. + data_split: One of ``"train"``, ``"test"``, ``"full"``. Applied to every + spec unless overridden by ``data_split_override``. max_action_dim: All action tensors are zero-padded to this dimension. Default 44 (CMR Versius conditioning dimension). downscaled_res: If True, use 256x256 resolution for all videos. + test_split_ratio: Default held-out fraction for each sub-dataset. Example:: @@ -1752,6 +1774,7 @@ def __init__( data_split: str = "train", max_action_dim: int = 44, downscaled_res: bool = False, + test_split_ratio: float = 0.05, ): from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.groot_configs import ( construct_modality_config_and_transforms, @@ -1775,7 +1798,13 @@ def __init__( embodiment = raw_embodiment.value if isinstance(raw_embodiment, EmbodimentTag) else raw_embodiment mix_ratio = spec.get("mix_ratio", 1.0) - print(f"\n[{i}] Loading: embodiment={embodiment}, mix_ratio={mix_ratio}") + spec_data_split = spec.get("data_split_override", data_split) + spec_test_split_ratio = spec.get("test_split_ratio_override", test_split_ratio) + + print( + f"\n[{i}] Loading: embodiment={embodiment}, mix_ratio={mix_ratio}, " + f"data_split={spec_data_split}, test_split_ratio={spec_test_split_ratio}" + ) print(f" path={path}") config, train_transform, test_transform = construct_modality_config_and_transforms( @@ -1789,7 +1818,7 @@ def __init__( if isinstance(config, dict) and "modality_filename" in config: modality_filename = config.pop("modality_filename") - transform = train_transform if data_split in ("train", "full") else test_transform + transform = train_transform if spec_data_split in ("train", "full") else test_transform # Per-dataset episode filtering (e.g., exclude "fail", "bad_frames" splits) exclude_splits = spec.get("exclude_splits", None) @@ -1799,7 +1828,8 @@ def __init__( modality_configs=config, transforms=transform, embodiment_tag=embodiment, - data_split=data_split, + data_split=spec_data_split, + test_split_ratio=spec_test_split_ratio, modality_filename=modality_filename, exclude_splits=exclude_splits, ) diff --git a/cosmos_predict2/_src/predict2/action/datasets/gr00t_dreams/groot_configs.py b/cosmos_predict2/_src/predict2/action/datasets/gr00t_dreams/groot_configs.py index aa5aaec..50eebf4 100644 --- a/cosmos_predict2/_src/predict2/action/datasets/gr00t_dreams/groot_configs.py +++ b/cosmos_predict2/_src/predict2/action/datasets/gr00t_dreams/groot_configs.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os + from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.dataset import ModalityConfig from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.embodiment_tags import EmbodimentTag from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.transform.base import ComposedModalityTransform @@ -603,12 +605,12 @@ def _dual_arm_eef_configs( # | **Total** | | **8.000** | **100%** | # # Frame counts sourced from gr00t-H exp62/exp79/exp86 configs and info.json. -# Dataset paths below are cluster defaults; override in experiment configs. +# Dataset paths below are portable defaults; override them with environment variables. # ============================================================================= -# Base paths (override per cluster / experiment) -_OPEN_H_BASE = "/lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_holoscan/datasets/Open-H" -_JHU_BASE = "/lustre/fsw/portfolios/healthcareeng/users/lzbinden/cache/huggingface/lerobot/jhu" +# Base paths (override per cluster / experiment). +_OPEN_H_BASE = os.environ.get("OPEN_H_DATA_ROOT", "datasets/open_h") +_JHU_BASE = os.environ.get("JHU_DATA_ROOT", "datasets/jhu") _LSCR_BASE = f"{_OPEN_H_BASE}/Surgical/JHU/LSCR" _STANFORD_BASE = ( f"{_OPEN_H_BASE}/Surgical/Stanford/Collaborative Haptics and Robotics in Medicine Lab/Real Robot (dVRK)" @@ -798,6 +800,75 @@ def _dual_arm_eef_configs( }, # 12,020 frames ] +# ============================================================================= +# JHU dVRK monocular reference tabletop fine-tuning mixture +# ============================================================================= +# Point this environment variable at a directory containing the nine LeRobot +# datasets listed below. Frame-count mix ratios keep sampling approximately +# frame-proportional instead of heavily oversampling the smallest subsets. +_JHU_TABLETOP_DATA_ROOT = os.environ.get("JHU_TABLETOP_DATA_ROOT", "datasets/jhu_tabletop") + +_JHU_DVRK_MONO_FINETUNE_NON_OOD_SPECS: list[dict] = [ + { + "path": f"{_JHU_TABLETOP_DATA_ROOT}/hf_suturebot", + "embodiment": EmbodimentTag.JHU_DVRK_MONO, + "mix_ratio": 516334.0, + "test_split_ratio_override": 0.01, + }, + { + "path": f"{_JHU_TABLETOP_DATA_ROOT}/knot_tying", + "embodiment": EmbodimentTag.JHU_DVRK_MONO, + "mix_ratio": 209253.0, + "test_split_ratio_override": 0.01, + }, + { + "path": f"{_JHU_TABLETOP_DATA_ROOT}/suture_bot_success", + "embodiment": EmbodimentTag.JHU_DVRK_MONO, + "mix_ratio": 1557.0, + }, + { + "path": f"{_JHU_TABLETOP_DATA_ROOT}/suture_bot_failure", + "embodiment": EmbodimentTag.JHU_DVRK_MONO, + "mix_ratio": 8793.0, + }, + { + "path": f"{_JHU_TABLETOP_DATA_ROOT}/cosmos_fail_filtered", + "embodiment": EmbodimentTag.JHU_DVRK_MONO, + "mix_ratio": 12948.0, + }, + { + "path": f"{_JHU_TABLETOP_DATA_ROOT}/cosmos_throw_fail_demo", + "embodiment": EmbodimentTag.JHU_DVRK_MONO, + "mix_ratio": 54581.0, + "test_split_ratio_override": 0.01, + }, + { + "path": f"{_JHU_TABLETOP_DATA_ROOT}/cosmos_knot_fail_demo", + "embodiment": EmbodimentTag.JHU_DVRK_MONO, + "mix_ratio": 30502.0, + }, + { + "path": f"{_JHU_TABLETOP_DATA_ROOT}/suturebot_act_throw_eval", + "embodiment": EmbodimentTag.JHU_DVRK_MONO, + "mix_ratio": 11548.0, + }, +] + +_JHU_DVRK_MONO_FINETUNE_OOD_SPEC: dict = { + "path": f"{_JHU_TABLETOP_DATA_ROOT}/ood", + "embodiment": EmbodimentTag.JHU_DVRK_MONO, + "mix_ratio": 227990.0, + "data_split_override": "full", +} + +JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS: list[dict] = [ + *_JHU_DVRK_MONO_FINETUNE_NON_OOD_SPECS, + _JHU_DVRK_MONO_FINETUNE_OOD_SPEC, +] + +# OOD trajectories are intentionally train-only in the reference recipe. +JHU_DVRK_MONO_FINETUNE_VAL_DATASET_SPECS: list[dict] = list(_JHU_DVRK_MONO_FINETUNE_NON_OOD_SPECS) + # Derived: the set of all Open-H embodiment tag strings. # Used by dataset.py to enforce stats_cosmos.json requirement. # Includes both EMBODIMENT_REGISTRY keys (non-CMR) and all tags from the specs. diff --git a/cosmos_predict2/_src/predict2/action/inference/inference_jhu_dvrk_warmup.py b/cosmos_predict2/_src/predict2/action/inference/inference_jhu_dvrk_warmup.py new file mode 100644 index 0000000..2895c0d --- /dev/null +++ b/cosmos_predict2/_src/predict2/action/inference/inference_jhu_dvrk_warmup.py @@ -0,0 +1,412 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +JHU dVRK Mono variant of the GR00T warmup inference script (Phase 0 of the +self-forcing pipeline). Generates teacher trajectory caches over the 9-dataset +JHU dVRK mono training mixture (frame-proportional weighted via +``JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS``) instead of a single dataset. + +Output layout matches the upstream warmup cache convention so it is consumed +unchanged by ``ActionDatasetSFWarmup`` in Phase 1: + + / + latents/.pt # teacher denoising-trajectory latents at query_steps + images/.png # first conditioning frame + actions/.json # the (num_frames-1)-step, zero-padded 44D action chunk + videos/.mp4 # the num_frames-frame ground-truth video clip + indices.json # the global ordered index list used by this run + # (seed + strategy + len + first/last few indices) + +Where ```` is the *virtual* MixedLeRobotDataset index (NOT the slot index +in this rank's shard). This way file names are still globally unique, the +identity check ``cache[vidx] == dataset[vidx]`` continues to work, and the +``ActionDatasetSFWarmup`` glob loader doesn't care that virtual indices are +sparse / non-contiguous after random sampling. + +Sampling strategies (--sample_strategy): + + - ``random`` (default, recommended): draw ``total_samples`` virtual indices + via ``np.random.permutation`` over the FULL virtual range. Mirrors what + the warmup trainer's ``DistributedSampler(shuffle=True)`` would draw at + training time, so the cache is frame-proportional across the mixture. + - ``uniform``: ``np.linspace(0, len(dataset) - 1, total_samples)``. Evenly + spaced; deterministic; gives every subset its frame-proportional share + within +/-1 sample. + - ``sequential``: legacy upstream behaviour. Walks indices [0, total_samples) + contiguously. With MixedLeRobotDataset this samples ONLY the first subset + in the specs list and leaves the other 8 subsets untouched. Kept for + backward-compat / debugging; do NOT use this for production caches on + multi-subset mixtures. + +The global index list is built ONCE (seeded by --indices_seed for ``random``) +and then sharded across SLURM ranks via --start/--end, where these now slice +into the index LIST not the dataset directly. So rank N processes the indices +``all_indices[start:end]`` regardless of where those land in the virtual range. + +Already-existing samples in ```` are skipped, so the script can be +re-run / requeued safely (each rank only generates the ones it owns and that +are not yet on disk). + +Example: + + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. python \\ + cosmos_predict2/_src/predict2/action/inference/inference_jhu_dvrk_warmup.py \\ + --experiment cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss \\ + --ckpt_path /path/to/iter_000005000/model_ema_bf16.pt \\ + --save_root datasets/jhu_dvrk_mono_warmup_4step_h73_tabletop \\ + --resolution 288,512 --guidance 0 --num_frames 73 --chunk_size 72 \\ + --sample_strategy random --total_samples 10000 --indices_seed 0 \\ + --start 0 --end 10000 \\ + --query_steps 0,9,18,27,34 +""" + +import argparse +import json +import os + +import mediapy +import numpy as np +import torch +import tqdm +from loguru import logger + +from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.dataset import MixedLeRobotDataset +from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.groot_configs import ( + JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS, + MAX_ACTION_DIM, +) +from cosmos_predict2._src.predict2.action.inference.inference_pipeline import ( + ActionVideo2WorldInference, +) + + +def parse_arguments() -> argparse.Namespace: + """Parses command-line arguments for the JHU dVRK warmup cache generator.""" + parser = argparse.ArgumentParser(description="JHU dVRK Mono Phase 0 warmup-cache generator") + parser.add_argument("--experiment", type=str, required=True, help="Experiment config name") + parser.add_argument("--chunk_size", type=int, default=12, help="Action chunk size (must match teacher training)") + parser.add_argument( + "--num_frames", + type=int, + default=13, + help="Clip horizon to cache (1 context + num_frames-1 prediction frames). " + "Defaults to 13 (state_t=4, the baseline SF cascade). Set to 73 for the " + "reference tabletop h73 student (state_t=19). Constraint: (num_frames-1) must be " + "divisible by 4 (valid horizons 13, 25, 49, 73, ...). Keep --chunk_size = " + "num_frames-1 and select the matching long-horizon teacher checkpoint/experiment.", + ) + parser.add_argument("--guidance", type=float, default=0.0, help="Classifier-free guidance scale (0 = no CFG)") + parser.add_argument("--seed", type=int, default=0, help="Random seed") + parser.add_argument( + "--ckpt_path", + type=str, + default="", + help="Path to the teacher checkpoint (.pt file or DCP dir). If empty, falls back to the experiment's load_path.", + ) + parser.add_argument("--s3_cred", type=str, default="credentials/s3_checkpoint.secret") + parser.add_argument( + "--resolution", + type=str, + default="288,512", + help="Resolution of the rendered video, format H,W. Default 288,512 matches the JHU dVRK teacher training.", + ) + parser.add_argument( + "--save_root", + type=str, + default="datasets/jhu_dvrk_mono_warmup_4step", + help="Output directory (relative to repo root). Sub-dirs latents/ images/ actions/ videos/ are created.", + ) + parser.add_argument( + "--start", + type=int, + default=0, + help=( + "Start position (inclusive) for this rank's shard. Slices into the " + "GLOBAL ORDERED INDEX LIST built per --sample_strategy, NOT into the " + "dataset directly. So rank N processes all_indices[start:end]." + ), + ) + parser.add_argument( + "--end", + type=int, + default=10000, + help="End position (exclusive) for this rank's shard. See --start.", + ) + parser.add_argument( + "--total_samples", + type=int, + default=10000, + help=( + "Size of the GLOBAL index list (the cache target size across all ranks). " + "Each rank slices [start:end] of this list. Default 10000." + ), + ) + parser.add_argument( + "--sample_strategy", + type=str, + default="random", + choices=["sequential", "uniform", "random"], + help=( + "How to draw the global index list. 'random' (default): permutation over " + "the full virtual range, seeded by --indices_seed. Mirrors the warmup " + "trainer's DistributedSampler(shuffle=True). 'uniform': np.linspace. " + "'sequential': legacy upstream pattern; on MixedLeRobotDataset this " + "draws ONLY the first subset and is generally not what you want." + ), + ) + parser.add_argument( + "--indices_seed", + type=int, + default=0, + help=( + "Seed for the --sample_strategy=random permutation. All ranks must use " + "the same seed so they slice the same global ordered list. Default 0." + ), + ) + parser.add_argument( + "--num_latent_conditional_frames", + type=int, + default=1, + help="Number of latent conditional frames (warmup uses 1 = single image conditioning).", + ) + parser.add_argument( + "--query_steps", + type=lambda x: [int(i) for i in x.split(",")], + default=[0, 9, 18, 27, 34], + help="Denoising-step indices at which to snapshot the teacher's latent trajectory.", + ) + parser.add_argument( + "--context_parallel_size", + type=int, + default=1, + help="Context parallel size (default 1 = no CP). Set to 8 if launching on 8 GPUs jointly.", + ) + parser.add_argument( + "--skip_existing", + action="store_true", + default=True, + help="Skip indices whose cache files already exist on disk. Default True (safe re-run / requeue).", + ) + parser.add_argument( + "--no_skip_existing", + action="store_false", + dest="skip_existing", + help="Force regenerate even if cache files already exist.", + ) + return parser.parse_args() + + +def _is_already_cached(save_root: str, idx: int) -> bool: + """Return True iff all 4 cache artefacts for ``idx`` already exist on disk.""" + return ( + os.path.exists(os.path.join(save_root, "latents", f"{idx}.pt")) + and os.path.exists(os.path.join(save_root, "images", f"{idx}.png")) + and os.path.exists(os.path.join(save_root, "actions", f"{idx}.json")) + and os.path.exists(os.path.join(save_root, "videos", f"{idx}.mp4")) + ) + + +def build_global_index_list( + dataset_len: int, + total_samples: int, + strategy: str, + seed: int, +) -> np.ndarray: + """Construct the global ordered list of virtual indices to cache. + + All ranks call this with identical arguments and slice [start:end] of the + returned array. So the deterministic-by-seed property of ``random`` (and + the inherent determinism of ``uniform`` / ``sequential``) gives every + rank the same global view, just a different shard of it. + """ + n = min(total_samples, dataset_len) + if strategy == "sequential": + return np.arange(n, dtype=np.int64) + if strategy == "uniform": + # endpoint=True: linspace covers [0, dataset_len - 1] inclusive. + return np.linspace(0, dataset_len - 1, n, dtype=np.int64) + if strategy == "random": + rng = np.random.default_rng(seed) + # Note: np.random.default_rng(seed) is a fresh stream per call but + # deterministic by seed. So all ranks get the same permutation. + perm = rng.permutation(dataset_len) + return perm[:n].astype(np.int64) + raise ValueError(f"Unknown sample_strategy={strategy!r}") + + +def _save_indices_manifest( + save_root: str, + all_indices: np.ndarray, + args: argparse.Namespace, + dataset_len: int, +) -> None: + """Persist the global ordered index list (and how it was built) to disk. + + Writes ``/indices.json``. The full list is written each rank-0 + write but with a small head/tail preview for human inspection of the + first few writers; rank > 0 will overwrite with identical content (since + the list is deterministic-by-seed).""" + manifest = { + "sample_strategy": args.sample_strategy, + "indices_seed": args.indices_seed, + "total_samples": int(args.total_samples), + "dataset_len": int(dataset_len), + "n_indices": int(len(all_indices)), + "first_10": [int(x) for x in all_indices[:10]], + "last_10": [int(x) for x in all_indices[-10:]], + "min": int(all_indices.min()) if len(all_indices) else None, + "max": int(all_indices.max()) if len(all_indices) else None, + "all_indices": [int(x) for x in all_indices], + } + out_path = os.path.join(save_root, "indices.json") + try: + with open(out_path, "w") as f: + json.dump(manifest, f, indent=2) + except Exception as exc: + # Non-fatal: caching itself is what matters; log and continue. + logger.warning(f"Failed to write {out_path}: {exc}") + + +def main(): + torch.enable_grad(False) + args = parse_arguments() + + # Reproducibility for the dataset's internal sample shuffling. + np.random.seed(args.seed) + torch.manual_seed(args.seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(args.seed) + + # Build the JHU dVRK Mono training mixture (9 datasets, frame-proportional + # weighting via mix_ratio = total_frames per subset). Using data_split="train" + # ensures we cache from the same partition the teacher saw during fine-tuning, + # honoring per-spec test_split_ratio_override (0%-2%) and the ood spec's + # data_split_override="full". + # The VAE temporal_compression_ratio=4 collapses 4 pixel frames into 1 + # latent frame, so (num_frames-1) must be divisible by 4. The teacher's + # state_t MUST match this clip horizon (state_t = 1 + (num_frames-1)//4), + # e.g. the reference tabletop h73 teacher uses state_t=19. + if (args.num_frames - 1) % 4 != 0: + raise ValueError( + f"--num_frames must satisfy (num_frames-1) % 4 == 0 " + f"(got {args.num_frames}); valid horizons are 13, 25, 49, 73, ..." + ) + dataset = MixedLeRobotDataset( + dataset_specs=JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS, + num_frames=args.num_frames, + data_split="train", + max_action_dim=MAX_ACTION_DIM, + downscaled_res=False, + test_split_ratio=0.02, + ) + logger.info( + f"JHU dVRK Mono cache horizon: num_frames={args.num_frames} " + f"(state_t={1 + (args.num_frames - 1) // 4}), chunk_size={args.chunk_size}" + ) + logger.info( + f"MixedLeRobotDataset built: virtual size = {len(dataset)}, " + f"max_action_dim = {MAX_ACTION_DIM}, " + f"requesting indices [{args.start}, {args.end})" + ) + + # Initialize the inference handler with context parallel support + video2world_cli = ActionVideo2WorldInference( + args.experiment, + args.ckpt_path, + args.s3_cred, + context_parallel_size=args.context_parallel_size, + ) + + mem_bytes = torch.cuda.memory_allocated(device=torch.device("cuda" if torch.cuda.is_available() else "cpu")) + logger.info(f"GPU memory after model dcp.load: {mem_bytes / (1024**3):.2f} GB") + + save_root = args.save_root + os.makedirs(os.path.join(save_root, "latents"), exist_ok=True) + os.makedirs(os.path.join(save_root, "images"), exist_ok=True) + os.makedirs(os.path.join(save_root, "actions"), exist_ok=True) + os.makedirs(os.path.join(save_root, "videos"), exist_ok=True) + + # Build the GLOBAL ordered index list once per rank (all ranks build the + # same list because all inputs are deterministic / seed-controlled). Then + # this rank slices [start:end] of that list. Cache files are named by the + # actual VIRTUAL index (not the slot in the shard) so file names remain + # globally unique across ranks and the identity check still works. + all_indices = build_global_index_list( + dataset_len=len(dataset), + total_samples=args.total_samples, + strategy=args.sample_strategy, + seed=args.indices_seed, + ) + logger.info( + f"Global index list built: strategy={args.sample_strategy}, " + f"seed={args.indices_seed}, n={len(all_indices)}, " + f"min={int(all_indices.min())}, max={int(all_indices.max())}, " + f"head={all_indices[:5].tolist()}, tail={all_indices[-5:].tolist()}" + ) + _save_indices_manifest(save_root, all_indices, args, len(dataset)) + + end = min(args.end, len(all_indices)) + start = max(0, min(args.start, end)) + rank_indices = all_indices[start:end].tolist() + logger.info( + f"This rank: processing slot range [{start}, {end}) " + f"({len(rank_indices)} virtual indices, " + f"first={rank_indices[0] if rank_indices else 'n/a'}, " + f"last={rank_indices[-1] if rank_indices else 'n/a'})" + ) + + n_skipped = 0 + n_generated = 0 + for vidx in tqdm.tqdm(rank_indices, desc=f"Phase 0 cache slots [{start},{end})"): + if args.skip_existing and _is_already_cached(save_root, vidx): + n_skipped += 1 + continue + + data = dataset[vidx] + # data["video"] is shape (C, T, H, W) in uint8. Match the upstream warmup + # convention: feed the FIRST frame as the conditioning image; the action + # chunk and the num_frames-frame ground-truth video are saved alongside. + img_np_array = data["video"][:, 0, :, :].permute(1, 2, 0).cpu().numpy() + video_np_array = data["video"].permute(1, 2, 3, 0).cpu().numpy() + action = data["action"].cpu().numpy() + + next_img_array, video_clamped, latents_to_save = video2world_cli.step_inference_with_latents( + img_array=img_np_array, + action=action, + guidance=args.guidance, + seed=args.seed, + num_latent_conditional_frames=args.num_latent_conditional_frames, + query_steps=args.query_steps, + ) + + for k in latents_to_save: + latents_to_save[k] = latents_to_save[k].squeeze(0).cpu() + + torch.save(latents_to_save, os.path.join(save_root, "latents", f"{vidx}.pt")) + mediapy.write_image(os.path.join(save_root, "images", f"{vidx}.png"), img_np_array) + mediapy.write_video(os.path.join(save_root, "videos", f"{vidx}.mp4"), video_np_array) + with open(os.path.join(save_root, "actions", f"{vidx}.json"), "w") as f: + json.dump(action.tolist(), f, indent=4) + n_generated += 1 + + logger.info( + f"Phase 0 cache shard slots [{start},{end}) complete: " + f"generated {n_generated}, skipped (already cached) {n_skipped}" + ) + + +if __name__ == "__main__": + main() diff --git a/cosmos_predict2/_src/predict2/distill/utils/config_helper.py b/cosmos_predict2/_src/predict2/distill/utils/config_helper.py index c326f0e..cfd202d 100644 --- a/cosmos_predict2/_src/predict2/distill/utils/config_helper.py +++ b/cosmos_predict2/_src/predict2/distill/utils/config_helper.py @@ -15,10 +15,20 @@ from cosmos_predict2._src.imaginaire.utils.checkpoint_db import get_checkpoint_path -def build_no_s3_run(job: dict, local_path: bool = False) -> dict: +def build_no_s3_run( + job: dict, + local_path: bool = False, + resumable: bool = False, + load_training_state: bool | None = None, + wandb_mode: str = "offline", +) -> dict: """ Make a copy of the input config that doesn't require S3 for checkpointing and I/O in the callbacks. + + ``resumable=True`` uses a stable output name across scheduler requeues. + ``load_training_state`` controls only the initial load; subsequent local + resumes still recover the complete training state. """ # If local_path is True, use the local path as the load path if local_path: @@ -27,16 +37,24 @@ def build_no_s3_run(job: dict, local_path: bool = False) -> dict: model_url = f"s3://bucket/{job['checkpoint']['load_path']}/model" load_path = get_checkpoint_path(model_url) defaults = job.get("defaults", []) + job_name = ( + f"{job['job']['name']}_no_s3_resumable" + if resumable + else f"{job['job']['name']}_no_s3" + "_${now:%Y-%m-%d}_${now:%H-%M-%S}" + ) + if load_training_state is None: + load_training_state = resumable no_s3_run = dict( defaults=defaults + ["_self_"] if "_self_" not in defaults else defaults, job=dict( - name=f"{job['job']['name']}_no_s3" + "_${now:%Y-%m-%d}_${now:%H-%M-%S}", - wandb_mode="offline", + name=job_name, + wandb_mode=wandb_mode, ), checkpoint=dict( save_to_object_store=dict(enabled=False, credentials=""), load_from_object_store=dict(enabled=False), load_path=load_path, + load_training_state=load_training_state, ), trainer=dict( straggler_detection=dict(enabled=False), @@ -55,6 +73,76 @@ def build_no_s3_run(job: dict, local_path: bool = False) -> dict: return no_s3_run +def build_no_s3_run_v2( + job: dict, + local_path: bool = False, + resumable: bool = False, + load_training_state: bool | None = None, + wandb_mode: str = "offline", +) -> dict: + """Return a no-object-store run while preserving all source overrides. + + Unlike the legacy helper, this starts from a deep copy of the complete + experiment. This is required by warmup and Self Forcing recipes whose + action dimensions, checkpoint paths, and optimizer settings are supplied + as nested experiment overrides. + """ + from copy import deepcopy + + from omegaconf import OmegaConf + + try: + job_dict = OmegaConf.to_container(job, resolve=False) + except Exception: + job_dict = dict(job) + no_s3_run = deepcopy(job_dict) + + if local_path: + load_path = job_dict["checkpoint"]["load_path"] + else: + model_url = f"s3://bucket/{job_dict['checkpoint']['load_path']}/model" + load_path = get_checkpoint_path(model_url) + + job_name = ( + f"{job_dict['job']['name']}_no_s3_resumable" + if resumable + else f"{job_dict['job']['name']}_no_s3" + "_${now:%Y-%m-%d}_${now:%H-%M-%S}" + ) + if load_training_state is None: + load_training_state = resumable + + deep_update_config_dict( + no_s3_run, + dict( + job=dict(name=job_name, wandb_mode=wandb_mode), + checkpoint=dict( + save_to_object_store=dict(enabled=False, credentials=""), + load_from_object_store=dict(enabled=False), + load_path=load_path, + load_training_state=load_training_state, + ), + trainer=dict( + straggler_detection=dict(enabled=False), + callbacks=dict( + heart_beat=dict(save_s3=False), + iter_speed=dict(save_s3=False), + device_monitor=dict(save_s3=False), + every_n_sample_reg=dict(save_s3=False), + every_n_sample_ema=dict(save_s3=False), + wandb=dict(save_s3=False), + wandb_10x=dict(save_s3=False), + dataloader_speed=dict(save_s3=False), + ), + ), + ), + ) + + defaults = no_s3_run.get("defaults", []) + if "_self_" not in defaults: + no_s3_run["defaults"] = defaults + ["_self_"] + return no_s3_run + + def deep_update_config_dict(dst: dict, src: dict) -> dict: """ Updates nested dictionaries in the config dictionary (dst) with the values in src dictionary. diff --git a/cosmos_predict2/_src/predict2/interactive/configs/data.py b/cosmos_predict2/_src/predict2/interactive/configs/data.py index 80446a5..33caa84 100644 --- a/cosmos_predict2/_src/predict2/interactive/configs/data.py +++ b/cosmos_predict2/_src/predict2/interactive/configs/data.py @@ -30,6 +30,12 @@ cr1_embeddings_path="cr1_empty_string_text_embeddings.pt", ) +# Long-horizon Phase 0 cache used by the reference tabletop warmup and SF stages. +dataset_jhu_dvrk_mono_warmup_h73 = L(ActionDatasetSFWarmup)( + data_path="datasets/jhu_dvrk_mono_warmup_4step_h73_tabletop", + cr1_embeddings_path="cr1_empty_string_text_embeddings.pt", +) + # ----------- Dataloaders ----------- @@ -69,6 +75,12 @@ def register_interactive_data(): name="gr00t_g1_warmup", node=make_dataloader(dataset_gr00t_g1_warmup), ) + cs.store( + group="data_train", + package="dataloader_train", + name="jhu_dvrk_mono_warmup_h73", + node=make_dataloader(dataset_jhu_dvrk_mono_warmup_h73), + ) cs.store( group="data_val", package="dataloader_val", @@ -81,3 +93,9 @@ def register_interactive_data(): name="gr00t_g1_warmup", node=make_dataloader(dataset_gr00t_g1_warmup), ) + cs.store( + group="data_val", + package="dataloader_val", + name="jhu_dvrk_mono_warmup_h73", + node=make_dataloader(dataset_jhu_dvrk_mono_warmup_h73), + ) diff --git a/cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_self_forcing.py b/cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_self_forcing.py index aaece14..4f2668b 100644 --- a/cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_self_forcing.py +++ b/cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_self_forcing.py @@ -14,14 +14,30 @@ # limitations under the License. import math +import os from hydra.core.config_store import ConfigStore from cosmos_predict2._src.imaginaire.lazy_config import LazyDict -from cosmos_predict2._src.predict2.distill.utils.config_helper import build_no_s3_run, deep_update_config_dict +from cosmos_predict2._src.predict2.distill.utils.config_helper import ( + build_no_s3_run, + build_no_s3_run_v2, + deep_update_config_dict, +) from cosmos_predict2._src.predict2.models.video2world_model import HighSigmaStrategy from cosmos_predict2._src.predict2.text_encoders.text_encoder import EmbeddingConcatStrategy +_TABLETOP_OUTPUT_ROOT = os.environ.get("IMAGINAIRE_OUTPUT_ROOT", "imaginaire/output") +_JHU_H73_WARMUP_CHECKPOINT = ( + f"{_TABLETOP_OUTPUT_ROOT}/cosmos_predict2_action_conditioned/interactive_warmup/" + "jhu_dvrk_mono_i4_lr3e-5_h73_tabletop_no_s3_resumable/checkpoints/iter_000018000" +) +_JHU_H73_TEACHER_MODEL = ( + f"{_TABLETOP_OUTPUT_ROOT}/cosmos_predict2_action_conditioned/official_runs_vid2vid/" + "cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_" + "8nodes_release_oss_h73_tabletop/checkpoints/iter_000005000/model" +) + def make_experiment( name: str, @@ -277,6 +293,35 @@ def make_experiment( ), ) +ACTION_JHU_DVRK_MONO_TABLETOP_H73_SELF_FORCING = make_experiment( + name="jhu_dvrk_mono_i4-sf_h73_tabletop", + data="jhu_dvrk_mono_warmup_h73", + overrides=dict( + job=dict( + project="cosmos_predict2_action_conditioned", + group="interactive_self_forcing", + ), + checkpoint=dict(load_path=_JHU_H73_WARMUP_CHECKPOINT), + trainer=dict(max_iter=3000), + optimizer=dict(lr=5e-8), + model=dict( + config=dict( + state_t=1 + 72 // 4, + net=dict(action_dim=44), + net_fake_score=dict(action_dim=44), + net_teacher=dict(action_dim=44), + optimizer_discriminator_config=dict(lr=5e-6), + optimizer_fake_score_config=dict(lr=5e-6), + resolution="288", + teacher_load_from=dict( + load_path=_JHU_H73_TEACHER_MODEL, + credentials="", + ), + ), + ), + ), +) + cs = ConfigStore.instance() cs.store( @@ -297,3 +342,15 @@ def make_experiment( name="cosmos_predict2p5_2B_action_gr00t_gr1_self_forcing_no_s3", node=build_no_s3_run(ACTION_GR00T_GR1_SELF_FORCING), ) +cs.store( + group="experiment", + package="_global_", + name="cosmos_predict2p5_2B_action_jhu_dvrk_mono_tabletop_h73_self_forcing_no_s3_resumable", + node=build_no_s3_run_v2( + ACTION_JHU_DVRK_MONO_TABLETOP_H73_SELF_FORCING, + local_path=True, + resumable=True, + load_training_state=False, + wandb_mode="offline", + ), +) diff --git a/cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_warmup.py b/cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_warmup.py index 6500508..b085ca2 100644 --- a/cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_warmup.py +++ b/cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_warmup.py @@ -13,10 +13,23 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os + from hydra.core.config_store import ConfigStore from cosmos_predict2._src.imaginaire.lazy_config import LazyDict -from cosmos_predict2._src.predict2.distill.utils.config_helper import build_no_s3_run, deep_update_config_dict +from cosmos_predict2._src.predict2.distill.utils.config_helper import ( + build_no_s3_run, + build_no_s3_run_v2, + deep_update_config_dict, +) + +_TABLETOP_OUTPUT_ROOT = os.environ.get("IMAGINAIRE_OUTPUT_ROOT", "imaginaire/output") +_JHU_H13_FINE_ANNEAL_CHECKPOINT = ( + f"{_TABLETOP_OUTPUT_ROOT}/cosmos_predict2_action_conditioned/official_runs_vid2vid/" + "cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_" + "8nodes_release_oss_fine_anneal_4k/checkpoints/iter_000004000" +) def make_experiment( @@ -146,6 +159,24 @@ def make_experiment( ), ) +ACTION_JHU_DVRK_MONO_TABLETOP_H73_WARMUP = make_experiment( + name="jhu_dvrk_mono_i4_lr3e-5_h73_tabletop", + data="jhu_dvrk_mono_warmup_h73", + overrides=dict( + checkpoint=dict(load_path=_JHU_H13_FINE_ANNEAL_CHECKPOINT), + model=dict( + config=dict( + state_t=1 + 72 // 4, + net=dict(action_dim=44), + resolution="288", + ), + ), + optimizer=dict(lr=3e-5), + dataloader_train=dict(batch_size=2), + trainer=dict(max_iter=20000), + ), +) + """ torchrun --nproc_per_node=1 --master_port=12341 -m scripts.train --config=cosmos_predict2/_src/predict2/interactive/configs/config_warmup.py -- experiment=cosmos_predict2p5_2B_action_gr00t_gr1_warmup """ @@ -170,3 +201,15 @@ def make_experiment( name="cosmos_predict2p5_2B_action_gr00t_gr1_warmup_no_s3", node=build_no_s3_run(ACTION_GR00T_WARMUP_GR1), ) +cs.store( + group="experiment", + package="_global_", + name="cosmos_predict2p5_2B_action_jhu_dvrk_mono_tabletop_h73_warmup_no_s3_resumable", + node=build_no_s3_run_v2( + ACTION_JHU_DVRK_MONO_TABLETOP_H73_WARMUP, + local_path=True, + resumable=True, + load_training_state=False, + wandb_mode="offline", + ), +) diff --git a/docs/tutorial_teacher_training_and_self_forcing.md b/docs/tutorial_teacher_training_and_self_forcing.md index c9abc01..7010692 100644 --- a/docs/tutorial_teacher_training_and_self_forcing.md +++ b/docs/tutorial_teacher_training_and_self_forcing.md @@ -1,20 +1,1150 @@ # From a Cosmos-H-Surgical-Simulator Teacher to a Real-Time Causal Student -This tutorial describes how to adapt the bidirectional -Cosmos-H-Surgical-Simulator (C-H-S-S) model to a new action-conditioned video -dataset and then distill it into a causal, streaming student with Self Forcing. -The final student can be deployed with Cosmos-H-Dreams for low-latency, -closed-loop generation. - -The recipe is based on the JHU dVRK tabletop experiment used to train a -73-frame causal student: - -1. Fine-tune a short-horizon, bidirectional teacher from C-H-S-S. -2. Fine-tune a long-horizon, bidirectional teacher from the short-horizon - checkpoint. -3. Generate a Phase 0 cache of teacher denoising trajectories. +This recipe describes how to adapt the bidirectional Cosmos-H-Surgical-Simulator (C-H-S-S) model to a new paired kinematics-video dataset and then distill it into a causal, streaming student with [Self Forcing](https://arxiv.org/abs/2506.08009). The resulting student can be deployed with [Cosmos-H-Dreams](https://github.com/isaac-for-healthcare/Cosmos-H-Dreams) for low-latency, real-time closed-loop generation. + +The reference run used the SutureBot tabletop dataset, which is part of Open-H-Embodiment, to train a 73-frame causal student that ran in real time with the Cosmos-H-Dreams library. The main steps to follow are: + +1. Fine-tune a short-horizon, bidirectional teacher with 13-frame samples from C-H-S-S. +2. Fine-tune a long-horizon, bidirectional teacher with 73-frame samples from the short-horizon checkpoint. +3. Generate a cache of teacher denoising trajectories (Phase 0). 4. Warm up a causal student against the cached trajectories. 5. Run Self Forcing distillation. 6. Convert and deploy the causal student with Cosmos-H-Dreams. +This tutorial first explains what you will train, then provides a reference recipe you can adapt to your dataset. + +## Table of contents + +- [1. What is being trained?](#1-what-is-being-trained) +- [2. Reference recipe](#2-reference-recipe) +- [3. Prerequisites](#3-prerequisites) +- [4. Prepare a custom action-conditioned dataset](#4-prepare-a-custom-action-conditioned-dataset) +- [5. Stage 1: short-horizon bidirectional teacher](#5-stage-1-short-horizon-bidirectional-teacher) +- [6. Stage 2: long-horizon bidirectional teacher](#6-stage-2-long-horizon-bidirectional-teacher) +- [7. Stage 3: Phase 0 teacher-trajectory cache](#7-stage-3-phase-0-teacher-trajectory-cache) +- [8. Stage 4: causal-student warmup](#8-stage-4-causal-student-warmup) +- [9. Stage 5: Self Forcing distillation](#9-stage-5-self-forcing-distillation) +- [10. Convert and validate the distilled student](#10-convert-and-validate-the-distilled-student) +- [11. Deploy with Cosmos-H-Dreams](#11-deploy-with-cosmos-h-dreams) +- [12. Common failure modes](#12-common-failure-modes) +- [13. Further reading and resources](#13-further-reading-and-resources) + +## 1. What is being trained? + +Each stage serves a different purpose: + + +| Stage | Network | Attention | Training signal | Output | Initialization / checkpoint source | +| ------------- | ---------------------------------------------------- | ---------------------------- | ---------------------------------------------- | ---------------------------------- | ---------------------------------------------------- | +| Short teacher | `cosmos_v1_2B_action_chunk_conditioned` | Bidirectional | Rectified-flow teacher objective on real clips | Domain-adapted 13-frame teacher | C-H-S-S | +| Long teacher | Same teacher network | Bidirectional | Same objective on longer clips | 73-frame teacher | Short teacher | +| Phase 0 | Long teacher, frozen | Bidirectional | Inference only | Cached teacher latent trajectories | Long teacher | +| Warmup | `action_causal_cosmos_v1_2B` | Causal | Regression to cached teacher trajectories | Initialized causal student | Short-teacher; Phase 0 supplies long teacher targets | +| Self Forcing | Causal student + frozen teacher + fake-score network | Causal student rollout | DMD/fake-score/adversarial objectives | Streaming causal student | Warmup student plus frozen long teacher | +| Deployment | Causal student, frozen | Causal with rolling KV cache | Inference only | Real-time generated video | Self Forcing student | + + +The key design decision is to learn the target domain first with a bidirectional model. Causal distillation is then performed only after the teacher produces useful long-horizon rollouts. + +## 2. Reference recipe + +The reference tabletop run used the following values. Treat them as a validated starting point, not universal constants. + + +| Quantity | Short teacher | Long teacher | Phase 0 | Warmup | Self Forcing | +| --------------------------------- | -------------------------------------- | ----------------------------- | ------------------------ | ----------------------------- | ------------------------------------- | +| Pixel frames | 13 | 73 | 73 | 73 | 73 | +| Predicted/action frames | 12 | 72 | 72 | 72 | 72 | +| Latent temporal length, `state_t` | 4 | 19 | 19 | 19 | 19 | +| Per-GPU batch, 8 nodes × 8 GPUs | 16 | 4 | N/A | 2 | 1 | +| Global batch | 1024 | 256 | N/A | 128 | 64 | +| Main learning rate | `1.6e-4` | `4e-5` | N/A | `3e-5` | `5e-8` | +| Iteration budget | selected at 16k, then 4k cosine annealing | 5k | 10k cached samples | 20k ceiling; 18k selected | 3k | +| LR schedule | base training, then cosine annealing | 1k warmup + cosine | N/A | Constant | Constant | +| Resolution | 288 × 512 | 288 × 512 | 288 × 512 | 288 × 512 | 288 × 512 | +| Model-facing action width | 44 | 44 | 44 | 44 | 44 | +| Stage warm-start/input | C-H-S-S DCP | annealed h13 teacher, iter 4k | h73 teacher EMA, iter 5k | annealed h13 teacher, iter 4k | warmup iter 18k + h73 teacher iter 5k | + + + + +### 2.1 Reference implementation + +This repository contains the complete reference tabletop implementation used to obtain the released [checkpoint on Hugging Face](https://huggingface.co/nvidia/Cosmos-H-Dreams). Use these files as the concrete reference while adapting the `my_robot_*` examples below: + +- dataset mixture: [`JHU_DVRK_MONO_FINETUNE_*_DATASET_SPECS`](../cosmos_predict2/_src/predict2/action/datasets/gr00t_dreams/groot_configs.py); +- 13- and 73-frame loaders: [`action_conditioned/data.py`](../cosmos_predict2/_src/predict2/action/configs/action_conditioned/data.py); +- short teacher, cosine-annealing stage, and long teacher: [`exp_2B_action_conditioned_rectify_flow_gr00t.py`](../cosmos_predict2/_src/predict2/action/configs/action_conditioned/experiment/exp_2B_action_conditioned_rectify_flow_gr00t.py); +- Phase 0 cache generator: [`inference_jhu_dvrk_warmup.py`](../cosmos_predict2/_src/predict2/action/inference/inference_jhu_dvrk_warmup.py); +- cache loader: [`interactive/configs/data.py`](../cosmos_predict2/_src/predict2/interactive/configs/data.py); +- causal warmup: [`exp_action_warmup.py`](../cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_warmup.py); +- Self Forcing: [`exp_action_self_forcing.py`](../cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_self_forcing.py); +- launch templates and environment contract: [`train_scripts/tabletop/`](../train_scripts/tabletop/README.md). + +The configs read the environment variables `JHU_TABLETOP_DATA_ROOT`, `CHSS_CHECKPOINT_DIR`, and `IMAGINAIRE_OUTPUT_ROOT`. + +Cosmos-Predict2.5's video tokenizer has a temporal compression ratio of four. The pixel-frame and latent-frame relationship used throughout this recipe is: + +```text +num_frames = 1 + num_actions +state_t = 1 + num_actions // 4 +``` + +Therefore: + +```text +13 frames = 1 context + 12 predicted frames -> state_t = 4 +73 frames = 1 context + 72 predicted frames -> state_t = 19 +``` + +Valid horizons satisfy: + +```python +assert (num_frames - 1) % 4 == 0 +``` + +Examples are 13, 25, 49, and 73 frames. + +## 3. Prerequisites + +Follow [setup.md](setup.md), authenticate with Hugging Face, and download the C-H-S-S checkpoint. + +```bash +git clone https://github.com/nvidia-cosmos/Cosmos-H-Surgical-Simulator.git +cd Cosmos-H-Surgical-Simulator +git lfs pull + +curl -LsSf https://astral.sh/uv/install.sh | sh +source "$HOME/.local/bin/env" +uv sync --extra=cu128 +source .venv/bin/activate + +uv tool install -U "huggingface_hub[cli]" +hf auth login + +export IMAGINAIRE_OUTPUT_ROOT=/persistent/path/imaginaire/output +export IMAGINAIRE_CACHE_DIR=/persistent/path/imaginaire/cache +mkdir -p "$IMAGINAIRE_OUTPUT_ROOT" "$IMAGINAIRE_CACHE_DIR" +``` + +The complete workflow is GPU-heavy. The reference run used 8 nodes, each with eight 80 GB GPUs. CPU-only validation is still useful for syntax, config registration, dataset metadata, and cache-integrity checks, but it does not validate model memory or throughput. + +### 3.1 Always use persistent output storage + +Set `IMAGINAIRE_OUTPUT_ROOT` *inside the training container or process*, not just in the login shell. Otherwise, the experiment directory may be created in the container's overlay storage. A checkpoint save may appear to succeed but disappear when the container exits. + +For an Enroot/Pyxis-style launch: + +```bash +OUTPUT_DIR=/persistent/path/imaginaire/output +mkdir -p "$OUTPUT_DIR" +export OUTPUT_DIR + +srun \ + --container-mounts="${OUTPUT_DIR}:${OUTPUT_DIR},${PWD}:/workspace" \ + --container-workdir=/workspace \ + bash -c ' + export IMAGINAIRE_OUTPUT_ROOT="$OUTPUT_DIR" + python -m scripts.train ... + ' +``` + +Verify persistence after the first save: + +```bash +test -f \ + "$IMAGINAIRE_OUTPUT_ROOT////checkpoints/latest_checkpoint.txt" +``` + + + +## 4. Prepare a custom action-conditioned dataset + +The existing pipeline expects LeRobot-format datasets. Before training, each dataset needs: + +- video referenced by `meta/modality.json`; +- robot state at the context/reference timestep; +- per-timestep robot actions; +- `meta/info.json`, episode metadata, and parquet data in the normal LeRobot layout; +- post-transform normalization statistics; +- an embodiment entry describing action keys, state keys, transforms, frame stride, and target video resolution. + +Read [README_ACTION_SPACE.md](../scripts/README_ACTION_SPACE.md) before defining a new action space. + +### 4.1 Define the model-facing action representation + +C-H-S-S uses a unified 44D action input. For non-CMR embodiments, actions are transformed into their native model-facing representation and then zero-padded to 44D. For a dual-arm dVRK: + +```text +PSM1 relative xyz + rot6d 9D +PSM1 absolute gripper 1D +PSM2 relative xyz + rot6d 9D +PSM2 absolute gripper 1D + = 20D native +20D native + 24 trailing zeros = 44D model input +``` + +Do not pad raw parquet columns by hand if you use `MixedLeRobotDataset`; it pads the transformed action tensor. Make sure deployment uses the same transform, normalization, concatenation order, and padding convention as training. + +For an existing embodiment, reuse its registry entry. For a new dual-arm embodiment, add a registry entry in `cosmos_predict2/_src/predict2/action/datasets/gr00t_dreams/groot_configs.py`: + +```python +EMBODIMENT_REGISTRY["my_dual_arm_robot"] = { + "timestep_interval": 3, # 30 Hz raw -> 10 Hz model rate + "video_keys": ["video.endoscope_left"], + "state_keys": [ + "state.psm1_pose", + "state.psm1_gripper", + "state.psm2_pose", + "state.psm2_gripper", + ], + "action_keys": [ + "action.psm1_pose", + "action.psm1_gripper", + "action.psm2_pose", + "action.psm2_gripper", + ], + "action_key_configs": _dual_arm_eef_configs( + "action.psm1_pose", + "action.psm1_gripper", + "action.psm2_pose", + "action.psm2_gripper", + "state.psm1_pose", + "state.psm2_pose", + input_rot="quat", + ref_rot="quat", + input_quat="xyzw", + ref_quat="xyzw", + ), + "video_width": 512, + "video_height": 288, + "modality_filename": "meta/modality.json", + "normalization_mode": "mean_std", +} +``` + +Also add an `EmbodimentTag` value in +`cosmos_predict2/_src/predict2/action/datasets/gr00t_dreams/data/embodiment_tags.py`. + +### 4.2 Create the dataset mixture + +Define one spec per dataset. Setting `mix_ratio` in proportion to frame count makes sampling roughly frame-proportional: + +The reference nine-subset configuration is defined in +[`groot_configs.py`](../cosmos_predict2/_src/predict2/action/datasets/gr00t_dreams/groot_configs.py), +with all paths resolved relative to `JHU_TABLETOP_DATA_ROOT`. + +```python +MY_TRAIN_SPECS = [ + { + "path": "/datasets/my_robot/success", + "embodiment": EmbodimentTag.MY_DUAL_ARM_ROBOT, + "mix_ratio": 500_000.0, + "test_split_ratio_override": 0.01, + }, + { + "path": "/datasets/my_robot/failures", + "embodiment": EmbodimentTag.MY_DUAL_ARM_ROBOT, + "mix_ratio": 100_000.0, + "test_split_ratio_override": 0.02, + }, + { + "path": "/datasets/my_robot/ood", + "embodiment": EmbodimentTag.MY_DUAL_ARM_ROBOT, + "mix_ratio": 50_000.0, + "data_split_override": "full", + }, +] + +MY_VAL_SPECS = MY_TRAIN_SPECS[:2] +``` + +Equal `mix_ratio=1.0` values weight subsets equally, not individual samples. This can oversample a small failure set hundreds of times, so use equal ratios only when that is your intent. + +As noted in [Cosmos-Surg-dVRK](https://arxiv.org/abs/2510.16240), failure episodes are essential when training a simulator. A simulator trained only on successful episodes will be biased toward successful outcomes and may not handle failures well. The reference tabletop run included both failure episodes and out-of-distribution (OOD) data consisting of random trajectories. + +### 4.3 Generate post-transform statistics + +Raw `meta/stats.json` is insufficient when transforms change dimensionality, for example, when converting a 7D quaternion representation to a 9D xyz+rot6d representation. Generate statistics for the exact post-transform representation using [`scripts/compute_openh_action_stats.py`](../scripts/compute_openh_action_stats.py): + +```bash +python scripts/compute_openh_action_stats.py \ + --dataset-path /datasets/my_robot/success \ + --embodiment my_dual_arm_robot + +python scripts/compute_openh_action_stats.py \ + --dataset-path /datasets/my_robot/failures \ + --embodiment my_dual_arm_robot +``` + +Each dataset should then contain: + +```text +meta/stats_cosmos.json +``` + +CMR Versius uses its specialized `meta/stats_cosmos-44D.json` path instead. If you change `timestep_interval`, rotation convention, key order, or action representation, regenerate statistics. + +### 4.4 Register 13-frame and 73-frame dataloaders + +In `cosmos_predict2/_src/predict2/action/configs/action_conditioned/data.py`: + +The reference `jhu_dvrk_mono_finetune_{train,val}` and +`jhu_dvrk_mono_finetune_h73_{train,val}` loaders are already registered in +[`action_conditioned/data.py`](../cosmos_predict2/_src/predict2/action/configs/action_conditioned/data.py). +The following generic form shows what to change for another robot: + +```python +my_robot_h13_train_dataset = L(MixedLeRobotDataset)( + dataset_specs=MY_TRAIN_SPECS, + num_frames=13, + data_split="train", + max_action_dim=MAX_ACTION_DIM, + downscaled_res=False, +) +my_robot_h13_val_dataset = L(MixedLeRobotDataset)( + dataset_specs=MY_VAL_SPECS, + num_frames=13, + data_split="test", + max_action_dim=MAX_ACTION_DIM, + downscaled_res=False, +) + +my_robot_h73_train_dataset = L(MixedLeRobotDataset)( + dataset_specs=MY_TRAIN_SPECS, + num_frames=73, + data_split="train", + max_action_dim=MAX_ACTION_DIM, + downscaled_res=False, +) +my_robot_h73_val_dataset = L(MixedLeRobotDataset)( + dataset_specs=MY_VAL_SPECS, + num_frames=73, + data_split="test", + max_action_dim=MAX_ACTION_DIM, + downscaled_res=False, +) +``` + +Register each with a `DataLoader` and Hydra `ConfigStore`, following the existing `suturebot_train` and `suturebot_val` registrations: + +```python +cs.store( + group="data_train", + package="dataloader_train", + name="my_robot_h13_train", + node=L(DataLoader)( + dataset=my_robot_h13_train_dataset, + sampler=L(get_sampler)(dataset=my_robot_h13_train_dataset), + batch_size=1, + drop_last=True, + ), +) +``` + +Repeat for `my_robot_h13_val`, `my_robot_h73_train`, and `my_robot_h73_val`. + +### 4.5 Dataset preflight + +Before allocating many GPUs, instantiate the 13-frame and 73-frame datasets in a single process and check: + +```python +sample = dataset[0] +assert sample["video"].shape[1] == NUM_FRAMES +assert sample["action"].shape == (NUM_FRAMES - 1, 44) +assert sample["video"].dtype == torch.uint8 +assert torch.isfinite(sample["action"]).all() +``` + +Confirm that: + +- the context and actions come from the same episode; +- the 73-frame window does not cross episode boundaries; +- video and action timestamps use the same stride; +- transformed actions have plausible means and scales; +- train and validation episode splits do not overlap. + +Unlike the CMR procedure-filtering path, a standard `MixedLeRobotDataset` does not require a separate filter manifest for each horizon. It does require a `stats_cosmos.json` file for every dataset. + +## 5. Stage 1: short-horizon bidirectional teacher + +Warm-start the fine-tuning from the [C-H-S-S checkpoint](https://huggingface.co/nvidia/Cosmos-H-Surgical-Simulator) rather than from the generic Cosmos-Predict2.5 checkpoint to benefit from the surgical visual prior and the trained 44D action embedder. + +Add an experiment to +`cosmos_predict2/_src/predict2/action/configs/action_conditioned/experiment/exp_2B_action_conditioned_rectify_flow_gr00t.py`: + +The committed reference symbol is +`AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_FINETUNE_13FRAME_8NODES_OSS` in +[`exp_2B_action_conditioned_rectify_flow_gr00t.py`](../cosmos_predict2/_src/predict2/action/configs/action_conditioned/experiment/exp_2B_action_conditioned_rectify_flow_gr00t.py). + +```python +MY_ROBOT_H13_TEACHER = LazyDict( + dict( + defaults=[ + "/experiment/2b_bridge_action_conditioned_oss", + {"override /net": "cosmos_v1_2B_action_chunk_conditioned"}, + {"override /data_train": "my_robot_h13_train"}, + {"override /data_val": "my_robot_h13_val"}, + "_self_", + ], + job=dict( + project="cosmos_predict2_action_conditioned", + group="official_runs_vid2vid", + name="my_robot_h13_teacher", + ), + checkpoint=dict( + # DCP directory, not model_ema_bf16.pt. + load_path="/checkpoints/cosmos_h_surgical_simulator/iter_NNNNN", + load_training_state=False, + strict_resume=False, + ), + model=dict( + config=dict( + state_t=1 + 12 // 4, + net=dict(action_dim=44), + ), + ), + dataloader_train=dict(batch_size=16), + optimizer=dict(lr=1.6e-4, weight_decay=0.1), + ), + flags={"allow_objects": True}, +) +``` + +Register it: + +```python +cs.store( + group="experiment", + package="_global_", + name=MY_ROBOT_H13_TEACHER["job"]["name"], + node=MY_ROBOT_H13_TEACHER, +) +``` + +Important checkpoint rules: + +- Training initialization expects the DCP iteration directory. +- `load_training_state=False` resets the source optimizer, scheduler, and iteration count. +- `strict_resume=False` allows nonessential source-run state to differ. +- The model-facing action width must remain 44 if the C-H-S-S action embedder is to load without shape changes. +- A consolidated `model_ema_bf16.pt` is primarily for inference, not this DCP warm start. + +Launch the committed reference config directly, or use the sanitized +[`01_train_short_teacher_h13.sh`](../train_scripts/tabletop/01_train_short_teacher_h13.sh) +template: + +```bash +torchrun --nproc_per_node=8 --master_port=12341 -m scripts.train \ + --config=cosmos_predict2/_src/predict2/action/configs/action_conditioned/config.py \ + -- \ + experiment=cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss \ + checkpoint.save_iter=200 \ + ~dataloader_train.dataloaders +``` + +For multiple nodes, add the usual `torchrun --nnodes`, `--node_rank`, and `--master_addr` arguments under your scheduler. + +### 5.1 Fine-tune the short teacher with cosine annealing + +In the reference tabletop run, iteration 16,000 was selected as the base checkpoint and refined in a separate 4,000-iteration cosine-annealing run. This made the phase boundary explicit and avoided carrying over the optimizer and scheduler state: + +The exact reference config is +`AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_FINETUNE_13FRAME_8NODES_OSS_FINE_ANNEAL_4K`; +the portable launcher is +[`02_fine_anneal_short_teacher_h13.sh`](../train_scripts/tabletop/02_fine_anneal_short_teacher_h13.sh). + +```python +MY_ROBOT_H13_FINE_ANNEAL = LazyDict( + dict( + defaults=["/experiment/my_robot_h13_teacher", "_self_"], + job=dict( + project="cosmos_predict2_action_conditioned", + group="official_runs_vid2vid", + name="my_robot_h13_teacher_fine_anneal_4k", + ), + checkpoint=dict( + load_path="/output/.../my_robot_h13_teacher/checkpoints/iter_000016000", + load_training_state=False, + strict_resume=False, + ), + scheduler=L(LambdaWarmUpCosineScheduler)( + warm_up_steps=[100], + f_start=[0.10], + f_max=[1.00], + f_min=[0.05], + cycle_lengths=[4000], + ), + trainer=dict(max_iter=4000), + ), + flags={"allow_objects": True}, +) +``` + +Do not automatically select iteration 16,000 for a different dataset. Use validation rollouts and smoothed loss to select the short-teacher checkpoint, then run the cosine-annealing stage from that checkpoint. + +## 6. Stage 2: long-horizon bidirectional teacher + +This recipe applies the progressive temporal post-training approach from [OmniDreams](https://arxiv.org/abs/2606.03159): first learn the target domain over a manageable short horizon, then extend the same teacher to a longer horizon. The network weights are compatible across horizons; the longer run changes the training clip length and latent temporal length. + +For the 73-frame reference tabletop run, the committed symbol is +`AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_TABLETOP_H73_8NODES_OSS`; see +[`exp_2B_action_conditioned_rectify_flow_gr00t.py`](../cosmos_predict2/_src/predict2/action/configs/action_conditioned/experiment/exp_2B_action_conditioned_rectify_flow_gr00t.py) +and [`03_train_long_teacher_h73.sh`](../train_scripts/tabletop/03_train_long_teacher_h73.sh). + +```python +MY_ROBOT_H73_TEACHER = LazyDict( + dict( + defaults=[ + "/experiment/my_robot_h13_teacher", + {"override /data_train": "my_robot_h73_train"}, + {"override /data_val": "my_robot_h73_val"}, + "_self_", + ], + job=dict( + project="cosmos_predict2_action_conditioned", + group="official_runs_vid2vid", + name="my_robot_h73_teacher", + ), + checkpoint=dict( + load_path=( + "/output/.../my_robot_h13_teacher_fine_anneal_4k/" + "checkpoints/iter_000004000" + ), + load_training_state=False, + strict_resume=False, + ), + model=dict( + config=dict( + state_t=1 + 72 // 4, # 19 + net=dict(action_dim=44), + ), + ), + dataloader_train=dict(batch_size=4), + optimizer=dict(lr=4e-5, weight_decay=0.1), + scheduler=L(LambdaWarmUpCosineScheduler)( + warm_up_steps=[1000], + f_start=[0.10], + f_max=[1.00], + f_min=[0.05], + cycle_lengths=[5000], + ), + trainer=dict(max_iter=5000), + ), + flags={"allow_objects": True}, +) +``` + +The reference schedule was: + +```text +iterations 0-1000: linear warmup, 0.1x -> 1.0x base LR +iterations 1000-5000: cosine decay, 1.0x -> 0.05x base LR +``` + +With a base LR of `4e-5`, the peak is `4e-5` and the final LR is `2e-6`. +Keep `cycle_lengths[0] == trainer.max_iter` so the cosine reaches its endpoint at the checkpoint consumed by Phase 0. + +Launch it with the same training entry point: + +```bash +torchrun --nproc_per_node=8 --master_port=12341 -m scripts.train \ + --config=cosmos_predict2/_src/predict2/action/configs/action_conditioned/config.py \ + -- \ + experiment=cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss_h73_tabletop \ + checkpoint.save_iter=200 \ + ~dataloader_train.dataloaders +``` + + + +### 6.1 Batch and LR scaling + +The reference 8-node run used: + +```text +8 nodes x 8 GPUs x batch 4 = global batch 256, LR 4e-5 +``` + +If 73 frames at batch 4 OOMs on 80-GB GPUs, use batch 2 and, if the node count is unchanged, reduce the LR to `2e-5`. For other cluster sizes, preserve the global batch and learning rate unless you intentionally want a different optimization run. + +### 6.2 Decide whether the long teacher is ready + +Raw rectified-flow loss depends strongly on sampled diffusion time and can look like a sawtooth. Prefer: + +- a moving mean over at least 100 iterations; +- like-for-like loss buckets at the same logging offset; +- held-out 73-frame rollouts; +- FDS (frame decay score) or task-specific metrics; +- visual action responsiveness and temporal consistency. + +In the reference tabletop run, the 100-step moving mean was flat for roughly the final 1,100 iterations, and the cosine had reached its floor at iteration 5,000. That was stronger evidence of convergence than the jagged raw loss. + +Convert the selected teacher checkpoint for cache-generation inference: + +```bash +TEACHER_DCP=/output/.../my_robot_h73_teacher/checkpoints/iter_000005000 +python scripts/convert_distcp_to_pt.py \ + "$TEACHER_DCP/model" \ + "$TEACHER_DCP" + +test -f "$TEACHER_DCP/model_ema_bf16.pt" +``` + + + +## 7. Stage 3: Phase 0 teacher-trajectory cache + +Phase 0 runs the frozen bidirectional teacher and stores its intermediate denoising latents at selected query steps. The causal warmup does not train directly on source videos; it trains on these teacher targets. + +The reference cache used: + +```text +cache size: 10,000 examples +num_frames: 73 +actions/example: 72 +query steps: 0, 9, 18, 27, 34 +sampling: random over the full training mixture +seed: 0 +guidance: 0 +``` + +Expected layout: + +```text +datasets/my_robot_warmup_4step_h73/ +├── actions/ +│ └── .json +├── images/ +│ └── .png +├── latents/ +│ └── .pt +├── videos/ +│ └── .mp4 +└── indices.json +``` + +Each latent file is a dictionary keyed by the queried denoising-step indices. Each action JSON file must contain an array with shape `(72, 44)`. + +### 7.1 Use the mixed-dataset cache generator + +The script `cosmos_predict2/_src/predict2/action/inference/inference_gr00t_warmup.py` shows the core call: + +```python +_, _, latents_to_save = video2world_cli.step_inference_with_latents( + img_array=first_frame, + action=action, + guidance=0, + seed=0, + num_latent_conditional_frames=1, + query_steps=[0, 9, 18, 27, 34], +) + +latents_to_save = { + step: latent.squeeze(0).cpu() + for step, latent in latents_to_save.items() +} +``` + +The script [`inference_jhu_dvrk_warmup.py`](../cosmos_predict2/_src/predict2/action/inference/inference_jhu_dvrk_warmup.py) implements the mixed-dataset cache generation used by the tabletop recipe. It does the following: + +1. constructs the same `MixedLeRobotDataset` and dataset specs used by the teacher; +2. accepts `--num_frames` rather than hardcoding 13; +3. asserts `(num_frames - 1) % 4 == 0`; +4. samples across the full virtual mixture rather than taking the first 10,000 sequential indices; +5. writes globally unique filenames based on the selected dataset index; +6. skips only examples for which all required artifacts already exist. + +Its seeded `build_global_index_list()` logic is equivalent to: + +```python +def build_global_index_list(dataset_len: int, total: int, seed: int) -> np.ndarray: + n = min(dataset_len, total) + return np.random.default_rng(seed).permutation(dataset_len)[:n] +``` + +Every rank must build the same seeded list, then consume a disjoint slice. + +### 7.2 Single-GPU command + +For the reference tabletop mixture, run the committed cache generator directly: + +```bash +CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. python \ + cosmos_predict2/_src/predict2/action/inference/inference_jhu_dvrk_warmup.py \ + --experiment cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss_h73_tabletop \ + --ckpt_path /output/.../iter_000005000/model_ema_bf16.pt \ + --save_root datasets/jhu_dvrk_mono_warmup_4step_h73_tabletop \ + --resolution 288,512 \ + --guidance 0 \ + --num_frames 73 \ + --chunk_size 72 \ + --sample_strategy random \ + --total_samples 10000 \ + --indices_seed 0 \ + --start 0 \ + --end 10000 \ + --query_steps 0,9,18,27,34 +``` + +The Phase 0 `.pt` input is the consolidated long-horizon teacher EMA checkpoint. This is different from the DCP directory used to initialize training. + +For SLURM, use the sanitized +[`04_phase0_teacher_cache_h73.sh`](../train_scripts/tabletop/04_phase0_teacher_cache_h73.sh) +template. + +### 7.3 Multi-GPU sharding + +Phase 0 performs inference only, and ranks can process their shards independently. With `N` total ranks: + +```bash +SAMPLES_PER_RANK=$(( (TOTAL_SAMPLES + N_RANKS - 1) / N_RANKS )) +START=$(( GLOBAL_RANK * SAMPLES_PER_RANK )) +END=$(( (GLOBAL_RANK + 1) * SAMPLES_PER_RANK )) +(( END > TOTAL_SAMPLES )) && END=$TOTAL_SAMPLES +``` + +On SLURM with eight tasks per node: + +```bash +GLOBAL_RANK=$(( SLURM_NODEID * 8 + SLURM_LOCALID )) +N_RANKS=$(( SLURM_NNODES * 8 )) +``` + +Do not shard only by `SLURM_LOCALID`; it repeats ranks 0-7 on every node and causes duplicate cache generation. + +### 7.4 Cache integrity gate + +Do not start warmup merely because the cache job exited. Verify complete artifact quartets: + +```bash +CACHE=datasets/jhu_dvrk_mono_warmup_4step_h73_tabletop + +python - "$CACHE" <<'PY' +import json +import sys +from pathlib import Path + +root = Path(sys.argv[1]) +latent_ids = {p.stem for p in (root / "latents").glob("*.pt")} +action_ids = {p.stem for p in (root / "actions").glob("*.json")} +image_ids = {p.stem for p in (root / "images").glob("*.png")} +video_ids = {p.stem for p in (root / "videos").glob("*.mp4")} +complete = latent_ids & action_ids & image_ids & video_ids + +print({ + "latents": len(latent_ids), + "actions": len(action_ids), + "images": len(image_ids), + "videos": len(video_ids), + "complete": len(complete), +}) +assert len(complete) == 10_000 + +example = next(iter(complete)) +actions = json.loads((root / "actions" / f"{example}.json").read_text()) +assert len(actions) == 72 +assert all(len(row) == 44 for row in actions) +PY +``` + +Use the same teacher checkpoint for Phase 0 and Self Forcing. + +## 8. Stage 4: causal-student warmup + + + +### 8.1 Register the cache + +In `cosmos_predict2/_src/predict2/interactive/configs/data.py`: + +The reference cache is already registered as `jhu_dvrk_mono_warmup_h73` in +[`interactive/configs/data.py`](../cosmos_predict2/_src/predict2/interactive/configs/data.py). +Use the following pattern for a differently named cache: + +```python +dataset_my_robot_warmup_h73 = L(ActionDatasetSFWarmup)( + data_path="datasets/my_robot_warmup_4step_h73", + cr1_embeddings_path="cr1_empty_string_text_embeddings.pt", +) + +cs.store( + group="data_train", + package="dataloader_train", + name="my_robot_warmup_h73", + node=make_dataloader(dataset_my_robot_warmup_h73), +) +cs.store( + group="data_val", + package="dataloader_val", + name="my_robot_warmup_h73", + node=make_dataloader(dataset_my_robot_warmup_h73), +) +``` + +`ActionDatasetSFWarmup` reads the five latent targets `[0, 9, 18, 27, 34]`, the context image, and the action sequence. + +### 8.2 Register the warmup experiment + +In `cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_warmup.py`: + +The committed reference symbol is `ACTION_JHU_DVRK_MONO_TABLETOP_H73_WARMUP` in +[`exp_action_warmup.py`](../cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_warmup.py). + +```python +MY_ROBOT_H73_WARMUP = make_experiment( + name="my_robot_h73_warmup", + data="my_robot_warmup_h73", + overrides=dict( + checkpoint=dict( + # DCP directory. This initializes the causal network weights. + load_path=( + "/output/.../my_robot_h13_teacher_fine_anneal_4k/" + "checkpoints/iter_000004000" + ), + ), + model=dict( + config=dict( + state_t=1 + 72 // 4, + net=dict(action_dim=44), + resolution=288, + ), + ), + dataloader_train=dict(batch_size=2), + optimizer=dict(lr=3e-5), + trainer=dict(max_iter=20000), + ), +) +``` + +The validated 73-frame experiment deliberately kept the causal network's inherited `num_action_per_chunk=12`. This field controls the local causal generation chunk, not the full 72-action training horizon. The full horizon is defined by `state_t=19` and the cache tensor shape. Do not change the local chunk to 72 unless the network and runtime are explicitly designed and validated for that attention and chunking behavior. + +Register a local/no-object-store variant with the override-preserving, resumable helper: + +```python +cs.store( + group="experiment", + package="_global_", + name="my_robot_h73_warmup_no_s3_resumable", + node=build_no_s3_run_v2( + MY_ROBOT_H73_WARMUP, + local_path=True, + resumable=True, + load_training_state=False, + wandb_mode="offline", + ), +) +``` + +The tabletop experiment used a no-S3 resumable helper that preserved the custom overrides and loaded the EMA teacher weights into the fresh student's regular network. + +### 8.3 Launch warmup + +The portable reference launcher is +[`05_warmup_student_h73.sh`](../train_scripts/tabletop/05_warmup_student_h73.sh). + +```bash +torchrun --nproc_per_node=8 --master_port=12342 -m scripts.train \ + --config=cosmos_predict2/_src/predict2/interactive/configs/config_warmup.py \ + -- \ + experiment=cosmos_predict2p5_2B_action_jhu_dvrk_mono_tabletop_h73_warmup_no_s3_resumable \ + checkpoint.save_iter=200 +``` + +The reference scaling was: + +```text +8 nodes x 8 GPUs x batch 2 = global batch 128, LR 3e-5 +``` + + + +### 8.4 Select the warmup checkpoint + +The reference warmup used 20,000 iterations as a ceiling, not a requirement. In the tabletop run: + +- loss fell steeply through the first several thousand iterations; +- continued improving slowly after 10,000; +- became effectively flat around 16,000-19,000; +- iteration 18,000 was selected because it was a complete checkpoint inside the plateau. + +Use a moving mean rather than a single batch loss. Once the moving mean has no trend for several thousand iterations, a checkpoint in that plateau is a reasonable Self Forcing initialization. + +The causal warmup and Self Forcing networks require the NATTEN multidimensional attention backend used by `action_causal_cosmos_v1_2B`. Verify `import natten` and the configured attention backend before allocating a multi-node job. + +## 9. Stage 5: Self Forcing distillation + +Self Forcing closes the train-test gap by rolling out the causal student during training, conditioning future predictions on its own generated history, and matching the resulting distribution to the teacher's distribution. The training model holds: + +- `net`: causal student; +- `net_teacher`: frozen bidirectional teacher; +- `net_fake_score`: auxiliary score/critic network; +- optional discriminator components used by the configured DMD/GAN losses. + +The loss is adversarial and is not expected to decrease monotonically. Judge health by losses that remain finite and bounded, successful checkpointing, and rollout quality. + +### 9.1 Register the Self Forcing experiment + +In `cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_self_forcing.py`: + +The committed reference symbol is +`ACTION_JHU_DVRK_MONO_TABLETOP_H73_SELF_FORCING` in +[`exp_action_self_forcing.py`](../cosmos_predict2/_src/predict2/interactive/configs/experiment/exp_action_self_forcing.py). + +```python +MY_ROBOT_H73_SELF_FORCING = make_experiment( + name="my_robot_h73_self_forcing", + data="my_robot_warmup_h73", + overrides=dict( + job=dict( + project="cosmos_predict2_action_conditioned", + group="interactive_self_forcing", + ), + checkpoint=dict( + # Selected causal warmup DCP. + load_path=( + "/output/.../interactive_warmup/my_robot_h73_warmup/" + "checkpoints/iter_000018000" + ), + ), + trainer=dict(max_iter=3000), + optimizer=dict(lr=5e-8), + model=dict( + config=dict( + state_t=1 + 72 // 4, + net=dict(action_dim=44), + net_fake_score=dict(action_dim=44), + net_teacher=dict(action_dim=44), + optimizer_discriminator_config=dict(lr=5e-6), + optimizer_fake_score_config=dict(lr=5e-6), + resolution="288", + teacher_load_from=dict( + # Must match the teacher used to create Phase 0. + load_path=( + "/output/.../my_robot_h73_teacher/" + "checkpoints/iter_000005000/model" + ), + credentials="", + ), + ), + ), + ), +) +``` + +The three path types are deliberately different: + +```text +warmup student initialization: .../iter_000018000 (DCP iteration dir) +frozen teacher initialization: .../iter_000005000/model (DCP model shard dir) +eventual deployment checkpoint: .../model_ema_bf16.pt (consolidated file) +``` + +Use `config_distill.py`, which selects the distillation-aware checkpointer: + +```python +cs.store( + group="experiment", + package="_global_", + name="my_robot_h73_self_forcing_no_s3_resumable", + node=build_no_s3_run_v2( + MY_ROBOT_H73_SELF_FORCING, + local_path=True, + resumable=True, + load_training_state=False, + wandb_mode="offline", + ), +) +``` + + + +### 9.2 Launch Self Forcing + +Use [`06_self_forcing_h73.sh`](../train_scripts/tabletop/06_self_forcing_h73.sh) +for the validated 8-node reference. + +```bash +torchrun --nproc_per_node=8 --master_port=12343 -m scripts.train \ + --config=cosmos_predict2/_src/predict2/interactive/configs/config_distill.py \ + -- \ + experiment=cosmos_predict2p5_2B_action_jhu_dvrk_mono_tabletop_h73_self_forcing_no_s3_resumable \ + checkpoint.save_iter=200 +``` + +The reference 8-node run used: + +```text +global batch: 8 x 8 x 1 = 64 +student/DMD LR: 5e-8 +discriminator LR: 5e-6 +fake-score LR: 5e-6 +iterations: 3000 +``` + + + +### 9.3 Self Forcing health checks + +Stop and investigate if you see: + +- NaN or infinite losses; +- rapidly exploding loss magnitudes; +- repeated CUDA OOM before the first optimizer step; +- missing student, fake-score, optimizer, or trainer shards; +- a cache loaded with substantially fewer examples than expected; +- a teacher path different from the Phase 0 teacher; +- repeated restarts from iteration zero. + +Do not stop merely because the composite training loss rises. Evaluate checkpoints with open-loop action sequences and, ideally, closed-loop policy rollouts. + +## 10. Convert and validate the distilled student + +The Self Forcing checkpointer stores multiple networks and training state in DCP format. Convert the selected iteration to `.pt` format: + +```bash +SF_DCP=/output/.../interactive_self_forcing/my_robot_h73_self_forcing/checkpoints/iter_000003000 + +python scripts/convert_distcp_to_pt.py \ + "$SF_DCP/model" \ + "$SF_DCP" + +ls -lh \ + "$SF_DCP/model.pt" \ + "$SF_DCP/model_ema_fp32.pt" \ + "$SF_DCP/model_ema_bf16.pt" +``` + +Use `model_ema_bf16.pt` for deployment unless the runtime's model card says otherwise. + +Before moving to an optimized runtime, smoke-test the DCP checkpoint as follows: + +1. Create the streaming manifest and its corresponding ground-truth MP4 files and normalized action files using the committed [`extract_jhu_inference_manifest.py`](../scripts/extract_jhu_inference_manifest.py): + +```bash +python scripts/extract_jhu_inference_manifest.py \ + --subset hf_suturebot \ + --episode-ids 1440,1441,1442 \ + --num-frames 73 \ + --output-dir sf_inference_data/jhu_tabletop_test_h73 +``` + +Choose test episode IDs that exist in your dataset. The extractor uses the same JHU action transform, 44D padding, and `stats_cosmos.json` normalization as training. + +2. Generate rollouts for inspection using the generated streaming manifest: + +```bash +python -m cosmos_predict2._src.predict2.interactive.inference.action_video2world_streaming \ + --config cosmos_predict2/_src/predict2/interactive/configs/config_distill.py \ + --experiment cosmos_predict2p5_2B_action_jhu_dvrk_mono_tabletop_h73_self_forcing_no_s3_resumable \ + --ckpt_path "$SF_DCP" \ + --input_json sf_inference_data/jhu_tabletop_test_h73/jhu_tabletop_inference_manifest.json \ + --resolution 288,512 \ + --num_steps 4 \ + --max_frames 73 +``` + +Note: Parts of the public streaming script's rollout loop assume 12-action chunks. Audit and parameterize the affected fields before using it as the definitive 73-frame evaluation path. + +## 11. Deploy with Cosmos-H-Dreams + +[Cosmos-H-Dreams](https://github.com/isaac-for-healthcare/Cosmos-H-Dreams) runs the distilled surgical causal student in real time. Use its C-H-S-S/Cosmos-H model integration and add a checkpoint preset for the converted student. See the [documentation](https://github.com/isaac-for-healthcare/Cosmos-H-Dreams/blob/main/README.md) for instructions on running the student with your chosen controller. + +Keep the following training and distillation requirements in mind during deployment. + +### 11.1 Action preprocessing at runtime + +The model does not accept an arbitrary vector of 44 raw robot values. Deployment must reproduce the training pipeline: + +1. sample actions at the training rate; +2. use the same context/reference robot state; +3. convert translations and rotations to the configured relative representation; +4. normalize with the same `stats_cosmos.json`; +5. concatenate keys in the same order; +6. pad the transformed native action from 20D to 44D; +7. provide one 44D vector per generated pixel-frame transition. + +For a dVRK-like 20D transformed vector: + +```python +def pad_action_to_44(action_20d: np.ndarray) -> np.ndarray: + assert action_20d.shape[-1] == 20 + out = np.zeros((*action_20d.shape[:-1], 44), dtype=np.float32) + out[..., :20] = action_20d.astype(np.float32) + return out +``` + +This helper is only the final padding operation. It does not perform pose conversion or normalization. + +### 11.2 Real-time validation sequence + +Validate in this order: + +1. Load the checkpoint and verify that there are no missing or unexpected model keys. +2. Generate a rollout from a recorded first frame and normalized action trace. +3. Compare the first short rollout against the training-repository streaming implementation. +4. Verify action direction with simple isolated motions. +5. Measure first-chunk latency separately from steady-state latency. +6. Discard compilation/autotuning warmup chunks before reporting throughput. +7. Run a long open-loop trace and inspect rolling-cache degradation. +8. Connect a live controller or policy only after confirming parity with recorded traces. + +Measure both generation throughput and end-to-end control-loop latency. A model can generate faster than real time while still having an unacceptable first-chunk or action-ingestion delay. + +## 12. Common failure modes + +The following failure modes can occur during distillation. + +### Checkpoints “saved” but are absent + +`IMAGINAIRE_OUTPUT_ROOT` points to container-local storage. Bind persistent storage and export the variable inside the container. + +### The long teacher OOMs + +Reduce the per-GPU batch size. If the number of GPUs is unchanged, scale the LR with the global batch. Context parallelism or activation checkpointing are the next options when a batch size of 1 still fails. + +### A long-horizon run silently trains on 13 frames + +Check all three values together: + +```text +dataset num_frames = 73 +model state_t = 19 +action sequence = 72 +``` + +Also inspect cache-generation and streaming scripts for hardcoded 12-action loops. + +Run a separate horizon-aware evaluation job for 49- or 73-frame teachers. + +### Warmup starts on a partial cache + +Verify that every sample has all required artifacts; do not rely only on job completion or the number of latent files. + +### Multi-node Phase 0 creates duplicates + +Use global rank, not local rank, for shard boundaries. + +### The student checkpoint fails to load in the runtime + +Check `action_dim`, `num_action_per_latent_frame`, and `hidden_dim_in_action_embedder` first. In the reference tabletop run, `action_dim` was 44 and `num_action_per_latent_frame` was 4; derive the hidden width from the converted checkpoint. The current default is 8192. + +### Self Forcing loss rises + +That alone is not a divergence signal. Check that losses remain finite and bounded, then evaluate rollouts. The objective alternates student and critic/fake-score updates. + +### Hugging Face returns HTTP 429 at distributed startup + +Pre-download shared text embeddings and checkpoint assets to a mounted cache. Avoid having every rank independently fetch the same files. + +## 13. Further reading and resources + +- Cosmos-H-Surgical-Simulator: [Hugging Face model](https://huggingface.co/nvidia/Cosmos-H-Surgical-Simulator) +- Open-H-Embodiment: [Hugging Face dataset](https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-Open-H-Embodiment) +- Cosmos-H-Dreams code and examples: [GitHub repository](https://github.com/isaac-for-healthcare/Cosmos-H-Dreams) +- Cosmos-H-Dreams model: [Hugging Face checkpoint](https://huggingface.co/nvidia/Cosmos-H-Dreams) +- NVIDIA Cosmos-Predict2.5: [GitHub repository](https://github.com/nvidia-cosmos/cosmos-predict2.5) +- Cosmos-Surg-dVRK: *World Foundation Model-based Automated Online Evaluation of Surgical Robot Policy Learning*: [https://arxiv.org/abs/2510.16240](https://arxiv.org/abs/2510.16240) +- Self Forcing: *Bridging the Train-Test Gap in Autoregressive Video Diffusion*: [https://arxiv.org/abs/2506.08009](https://arxiv.org/abs/2506.08009) +- NVIDIA OmniDreams: *Real-Time Generative World Model for Closed-Loop Autonomous Vehicle Simulation*: [https://arxiv.org/abs/2606.03159](https://arxiv.org/abs/2606.03159) diff --git a/scripts/compute_openh_action_stats.py b/scripts/compute_openh_action_stats.py new file mode 100644 index 0000000..df8460f --- /dev/null +++ b/scripts/compute_openh_action_stats.py @@ -0,0 +1,881 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Compute per-key normalization statistics for any Open-H embodiment. + +Unlike compute_cmr_action_stats.py (which is CMR-specific with hardcoded raw +indices, clutch filtering, and motion scaling), this script is GENERIC: it +instantiates the real transform pipeline (GenericRelativeActionTransform) for +any embodiment registered in EMBODIMENT_REGISTRY and collects statistics on +the TRANSFORMED action/state output. + +This guarantees that the statistics exactly match what the training pipeline +produces, regardless of the embodiment's delta conversion (rel_xyz_rot6d, +relative, delta, or absolute). + +Output: + meta/stats_cosmos.json — per-key statistics in the same format as + stats_cosmos-44D.json (action.psm1_pose → {mean, std, min, max, q01, q99}) + +For CMR Versius, continue using compute_cmr_action_stats.py (it handles +the additional clutch filtering and motion scaling that are CMR-specific). + +Usage: + # All Open-H datasets at once (reads OPEN_H_DATASET_SPECS, includes exclude_splits): + python compute_openh_action_stats.py --all + + # Quick test (10 samples per dataset): + python compute_openh_action_stats.py --all --max-samples 10 + + # Single dataset: + python compute_openh_action_stats.py \\ + --dataset-path /path/to/lerobot/dataset \\ + --embodiment dvrk + + # Single dataset with episode filtering: + python compute_openh_action_stats.py \\ + --dataset-path /path/to/stanford/Needle_Transfer \\ + --embodiment dvrk_stanford_real \\ + --exclude-splits fail bad_frames + + # All dVRK datasets under a root: + python compute_openh_action_stats.py \\ + --dataset-path-root /path/to/jhu \\ + --embodiment jhu_dvrk_mono +""" + +import argparse +import json +import os +import time +import warnings +from concurrent.futures import ProcessPoolExecutor, as_completed +from functools import partial +from pathlib import Path + +# Suppress noisy torchvision video deprecation warnings +warnings.filterwarnings("ignore", message=".*video decoding and encoding capabilities of torchvision.*") + +import numpy as np +import pandas as pd +from tqdm import tqdm + +from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.embodiment_tags import EmbodimentTag +from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.dataset import ( + LE_ROBOT_INFO_FILENAME, + resolve_excluded_episode_indices, +) +from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.schema import ( + LeRobotModalityMetadata, + LeRobotStateActionMetadata, +) +from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.groot_configs import ( + EMBODIMENT_REGISTRY, + OPEN_H_DATASET_SPECS, +) +from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.transform.state_action import ( + convert_to_hybrid_relative, +) + +# Maximum parallel workers +MAX_WORKERS = 64 + + +# ============================================================================ +# Statistics helpers (streaming, memory-efficient) +# ============================================================================ + +class StreamingStats: + """Memory-efficient streaming statistics using Welford's algorithm + reservoir sampling.""" + + def __init__(self, num_dims: int, reservoir_size: int = 2_000_000): + self.num_dims = num_dims + self.reservoir_size = reservoir_size + self.count = 0 + self.mean = np.zeros(num_dims, dtype=np.float64) + self.M2 = np.zeros(num_dims, dtype=np.float64) + self.min_vals = np.full(num_dims, np.inf, dtype=np.float64) + self.max_vals = np.full(num_dims, -np.inf, dtype=np.float64) + self.reservoir = None + self.reservoir_count = 0 + + def update(self, batch: np.ndarray): + if batch.shape[0] == 0: + return + n = batch.shape[0] + batch_mean = np.mean(batch, axis=0) + batch_var = np.var(batch, axis=0, ddof=0) + batch_M2 = batch_var * n + self.min_vals = np.minimum(self.min_vals, np.min(batch, axis=0)) + self.max_vals = np.maximum(self.max_vals, np.max(batch, axis=0)) + if self.count == 0: + self.mean = batch_mean + self.M2 = batch_M2 + self.count = n + else: + n_total = self.count + n + delta = batch_mean - self.mean + self.mean = (self.count * self.mean + n * batch_mean) / n_total + self.M2 = self.M2 + batch_M2 + delta ** 2 * self.count * n / n_total + self.count = n_total + # Reservoir sampling + if self.reservoir is None: + self.reservoir = batch[:self.reservoir_size].copy() + self.reservoir_count = n + elif len(self.reservoir) < self.reservoir_size: + space = self.reservoir_size - len(self.reservoir) + self.reservoir = np.vstack([self.reservoir, batch[:space]]) + self.reservoir_count += n + else: + # Algorithm R replacement + for row in batch: + self.reservoir_count += 1 + j = np.random.randint(0, self.reservoir_count) + if j < self.reservoir_size: + self.reservoir[j] = row + + def get_stats(self) -> dict: + if self.count == 0: + z = [0.0] * self.num_dims + return {"mean": z, "std": z, "min": z, "max": z, "q01": z, "q99": z} + std = np.sqrt(self.M2 / self.count) + if self.reservoir is not None and len(self.reservoir) > 0: + q01 = np.quantile(self.reservoir, 0.01, axis=0) + q99 = np.quantile(self.reservoir, 0.99, axis=0) + else: + q01 = self.min_vals + q99 = self.max_vals + return { + "mean": self.mean.tolist(), + "std": std.tolist(), + "min": self.min_vals.tolist(), + "max": self.max_vals.tolist(), + "q01": q01.tolist(), + "q99": q99.tolist(), + } + + +# ============================================================================ +# Main processing +# ============================================================================ + +def _is_lerobot_dataset(path: Path) -> bool: + return (path / "data").is_dir() and (path / "meta").is_dir() and any((path / "data").rglob("*.parquet")) + + +def _discover_datasets(root: Path) -> list[Path]: + datasets = [] + for child in sorted(root.iterdir()): + if child.is_dir() and _is_lerobot_dataset(child): + datasets.append(child) + return datasets + + +def _process_episode_parquet( + parquet_path: Path, + modality_meta: dict, + action_key_configs: dict, + action_delta_indices: list[int], + state_delta_indices: list[int], + action_keys: list[str], + state_keys: list[str], +) -> tuple[np.ndarray, np.ndarray, str | None]: + """Worker function: process one parquet file (episode) without video decoding. + + Reads state/action arrays from parquet, applies delta conversion, concatenates + per-key arrays, and returns the result for streaming stats collection. + + This runs in a subprocess via ProcessPoolExecutor for parallelism. + + Returns: + (action_array, state_array, warning_or_None) + action_array shape: (N_samples * T_action, action_dim) + state_array shape: (N_samples, state_dim) + """ + try: + df = pd.read_parquet(parquet_path) + T = len(df) + + # Maximum delta for action horizon + max_action_delta = max(action_delta_indices) if action_delta_indices else 0 + effective_length = max(0, T - max_action_delta) + if effective_length == 0: + return np.empty((0, 0)), np.empty((0, 0)), None + + # Extract per-key arrays from the flat parquet columns using modality metadata + def extract_key_data(key: str, df: pd.DataFrame) -> np.ndarray: + """Extract a named key's data from parquet using modality.json metadata.""" + modality, subkey = key.split(".", 1) + meta = modality_meta.get(modality, {}).get(subkey) + if meta is None: + raise KeyError(f"{key} config not found") + original_col = meta.get("original_key") + if original_col is None: + # Default: observation.state for state, action for action + original_col = "observation.state" if modality == "state" else "action" + start, end = meta["start"], meta["end"] + col_data = np.stack(df[original_col].values) # (T, D_flat) + return col_data[:, start:end].astype(np.float32) # (T, key_dim) + + # Process all valid starting indices + all_action_rows = [] + all_state_rows = [] + n_skipped = 0 + + for base_idx in range(effective_length): + # --- State at t=0 --- + state_parts = [] + for key in state_keys: + s_idx = base_idx + state_delta_indices[0] # delta_indices=[0] + s_idx = max(0, min(s_idx, T - 1)) + key_data = extract_key_data(key, df) + state_parts.append(key_data[s_idx]) + if state_parts: + all_state_rows.append(np.concatenate(state_parts)) + + # --- Action over horizon --- + action_timestep_rows = [] + for a_delta in action_delta_indices: + a_idx = base_idx + a_delta + a_idx = max(0, min(a_idx, T - 1)) + action_parts = [] + for key in action_keys: + key_data = extract_key_data(key, df) + action_parts.append(key_data[a_idx]) + action_timestep_rows.append(np.concatenate(action_parts)) + action_horizon = np.stack(action_timestep_rows) # (T_action, action_dim_raw) + + # --- Apply delta conversion per key --- + # We need to apply the same GenericRelativeActionTransform logic + # but directly on numpy arrays (no torch, no dataset wrapper). + offset = 0 + converted_parts = [] + state_row = all_state_rows[-1] if all_state_rows else None + skip_sample = False + + for key in action_keys: + cfg = action_key_configs.get(key) + key_data_raw = extract_key_data(key, df) + raw_dim = key_data_raw.shape[1] + + # Extract this key's action horizon data + key_action = action_horizon[:, offset:offset + raw_dim] + + if cfg is not None and cfg.rep == "rel_xyz_rot6d": + # Get reference state pose + ref_key = cfg.state_key + if ref_key and state_row is not None: + # Find the state key's slice in the concatenated state + s_off = 0 + ref_pose = None + for sk in state_keys: + sk_data = extract_key_data(sk, df) + sk_dim = sk_data.shape[1] + if sk == ref_key: + ref_pose = state_row[s_off:s_off + sk_dim] + break + s_off += sk_dim + + if ref_pose is not None: + # Guard: check for zero-norm quaternions (invalid data / + # padding at episode boundaries). Scipy's Rotation.from_quat() + # crashes on [0,0,0,0]. + if cfg.input_rotation_format == "quat": + quat_slice = key_action[:, 3:7] + norms = np.linalg.norm(quat_slice, axis=-1) + if np.any(norms < 1e-8): + skip_sample = True + break + if cfg.reference_rotation_format == "quat": + ref_quat = ref_pose[3:7] if len(ref_pose) >= 7 else ref_pose[3:] + if np.linalg.norm(ref_quat) < 1e-8: + skip_sample = True + break + + key_action = convert_to_hybrid_relative( + action_data=key_action, + eef_pose=ref_pose, + input_rotation_format=cfg.input_rotation_format, + reference_rotation_format=cfg.reference_rotation_format, + input_quat_order=cfg.input_quat_order, + reference_quat_order=cfg.reference_quat_order, + ) # (T_action, 9) + + elif cfg is not None and cfg.rep == "relative": + # Joint-space subtraction + ref_key = cfg.state_key + if ref_key and state_row is not None: + s_off = 0 + ref_val = None + for sk in state_keys: + sk_data = extract_key_data(sk, df) + sk_dim = sk_data.shape[1] + if sk == ref_key: + ref_val = state_row[s_off:s_off + sk_dim] + break + s_off += sk_dim + if ref_val is not None: + key_action = key_action - ref_val + + # delta / absolute: pass through unchanged + converted_parts.append(key_action) + offset += raw_dim + + if skip_sample: + # Remove the state row we just added (it corresponds to this skipped sample) + if all_state_rows: + all_state_rows.pop() + n_skipped += 1 + continue + + converted_action = np.concatenate(converted_parts, axis=-1) # (T_action, action_dim) + all_action_rows.append(converted_action) + + if not all_action_rows: + warn = None + if n_skipped > 0: + warn = f"{parquet_path.name}: all {effective_length} samples skipped ({n_skipped} had zero-norm quaternions)" + return np.empty((0, 0)), np.empty((0, 0)), warn + + actions = np.concatenate(all_action_rows, axis=0) # (N*T_action, action_dim) + states = np.stack(all_state_rows) if all_state_rows else np.empty((0, 0)) + + warn = None + if n_skipped > 0: + warn = (f"{parquet_path.name}: {n_skipped}/{effective_length} samples skipped " + f"(zero-norm quaternions), {len(all_action_rows)} valid") + return actions, states, warn + + except Exception as e: + import traceback + return np.empty((0, 0)), np.empty((0, 0)), f"Error processing {parquet_path.name}: {type(e).__name__}: {e}" + + +def _load_modality_meta(dataset_path: Path, modality_filename: str) -> dict: + """Load modality.json and return a simplified {modality: {subkey: {start, end, original_key}}} dict.""" + modality_path = dataset_path / modality_filename + if not modality_path.exists(): + raise FileNotFoundError(f"Modality file not found: {modality_path}") + + with open(modality_path, "r") as f: + raw = json.load(f) + + result: dict[str, dict] = {} + for modality in ["state", "action"]: + result[modality] = {} + if modality not in raw: + continue + for subkey, meta in raw[modality].items(): + result[modality][subkey] = { + "start": meta.get("start", 0), + "end": meta.get("end", 1), + "original_key": meta.get("original_key"), + } + return result + + +def process_single_dataset( + dataset_path: Path, + embodiment: str, + num_frames: int, + max_samples: int | None, + output_filename: str, + exclude_splits: list[str] | None = None, + num_workers: int | None = None, + timestep_interval_override: int | None = None, +): + """Process one dataset using parallel episode processing (no video decoding). + + Each parquet file (episode) is processed by a worker subprocess that: + 1. Reads state/action arrays from parquet (fast, no video) + 2. Applies the delta conversion (rel_xyz_rot6d, relative, etc.) + 3. Returns the transformed arrays for streaming stats collection + + This is ~100x faster than the sequential video-decoding approach. + + Args: + dataset_path: Path to the LeRobot dataset. + embodiment: Embodiment tag string. + num_frames: Number of video frames (e.g. 13). + max_samples: Max episodes to process (for testing). None = all. + output_filename: Output filename in meta/ dir. + exclude_splits: Split names from info.json to exclude. + num_workers: Parallel workers (default: min(cpu_count, MAX_WORKERS)). + timestep_interval_override: If given (int > 0), override the + ``timestep_interval`` value read from EMBODIMENT_REGISTRY for + this run. Useful to experiment with a different effective training + rate *without* editing ``groot_configs.py``. IMPORTANT: if you use + this, training must run with the same stride, otherwise + ``stats_cosmos.json`` will not match the distribution the model + sees. + """ + if num_workers is None: + num_workers = min(os.cpu_count() or 8, MAX_WORKERS) + + reg = EMBODIMENT_REGISTRY.get(embodiment) + if reg is None: + raise ValueError(f"Unknown embodiment '{embodiment}'. Available: {list(EMBODIMENT_REGISTRY.keys())}") + + registry_timestep_interval = reg["timestep_interval"] + if timestep_interval_override is not None: + if timestep_interval_override < 1: + raise ValueError( + f"--timestep-interval must be >= 1, got {timestep_interval_override}" + ) + timestep_interval = int(timestep_interval_override) + else: + timestep_interval = registry_timestep_interval + + print("=" * 80) + print(f"COMPUTING OPEN-H ACTION STATS — {dataset_path.name}") + print(f" embodiment: {embodiment}") + print(f" num_frames: {num_frames}") + if timestep_interval_override is not None and timestep_interval_override != registry_timestep_interval: + print( + f" timestep_interval: {timestep_interval} " + f"(OVERRIDE; registry default for '{embodiment}' is {registry_timestep_interval})" + ) + else: + print(f" timestep_interval: {timestep_interval} (from EMBODIMENT_REGISTRY)") + print(f" workers: {num_workers}") + if exclude_splits: + print(f" exclude_splits: {exclude_splits}") + print("=" * 80) + + num_action_frames = num_frames - 1 + action_delta_indices = list(range(0, num_action_frames * timestep_interval, timestep_interval)) + state_delta_indices = [0] + + action_keys = reg["action_keys"] + state_keys = reg["state_keys"] + action_key_configs = reg.get("action_key_configs", {}) + modality_filename = reg.get("modality_filename", "meta/modality.json") + + # Load modality metadata (maps key names to parquet column indices) + modality_meta = _load_modality_meta(dataset_path, modality_filename) + + # Discover parquet files (one per episode) + parquet_files = sorted(dataset_path.glob("data/*/*.parquet")) + if not parquet_files: + print(f" ERROR: No parquet files found in {dataset_path / 'data'}") + return False + + # Apply exclude_splits filtering at episode level + if exclude_splits: + excluded_ids = resolve_excluded_episode_indices(dataset_path, exclude_splits) + # Parse episode index from filename (e.g., episode_000042.parquet → 42) + filtered = [] + for pf in parquet_files: + try: + ep_idx = int(pf.stem.split("_")[-1]) + except ValueError: + filtered.append(pf) # Can't parse → keep + continue + if ep_idx not in excluded_ids: + filtered.append(pf) + n_excluded = len(parquet_files) - len(filtered) + print(f" exclude_splits: removed {n_excluded} episodes, {len(filtered)} remaining") + parquet_files = filtered + + if max_samples is not None: + parquet_files = parquet_files[:max_samples] + + print(f" Processing {len(parquet_files)} episodes with {num_workers} workers...") + + # Create partial function with fixed arguments for the worker + worker_fn = partial( + _process_episode_parquet, + modality_meta=modality_meta, + action_key_configs=action_key_configs, + action_delta_indices=action_delta_indices, + state_delta_indices=state_delta_indices, + action_keys=action_keys, + state_keys=state_keys, + ) + + # Process episodes in parallel + action_tracker = None + state_tracker = None + total_action_samples = 0 + total_state_samples = 0 + episodes_with_warnings = 0 + episodes_empty = 0 + + with ProcessPoolExecutor(max_workers=num_workers) as executor: + futures = {executor.submit(worker_fn, pf): pf for pf in parquet_files} + + for future in tqdm(as_completed(futures), total=len(futures), desc=f"Stats for {dataset_path.name}"): + actions, states, warning = future.result() + + if warning: + episodes_with_warnings += 1 + if episodes_with_warnings <= 10: + print(f" [WARN] {warning}") + + if actions.size == 0: + episodes_empty += 1 + continue + + # Lazy-init trackers based on first result's dimensions + if action_tracker is None: + action_dim = actions.shape[1] + action_tracker = StreamingStats(action_dim) + print(f" Action dim (post-transform): {action_dim}") + if state_tracker is None and states.size > 0: + state_dim = states.shape[1] + state_tracker = StreamingStats(state_dim) + print(f" State dim (post-transform): {state_dim}") + + action_tracker.update(actions.astype(np.float64)) + total_action_samples += len(actions) + + if state_tracker is not None and states.size > 0: + state_tracker.update(states.astype(np.float64)) + total_state_samples += len(states) + + if episodes_with_warnings > 10: + print(f" [WARN] {episodes_with_warnings} episodes had warnings (showing first 10)") + if episodes_empty > 0: + print(f" [INFO] {episodes_empty}/{len(parquet_files)} episodes produced no valid samples " + f"(zero-norm quaternions or empty episodes)") + + if action_tracker is None or action_tracker.count == 0: + print(f" ERROR: No valid samples found! ({episodes_empty} empty, {episodes_with_warnings} warnings)") + return False + + valid_episodes = len(parquet_files) - episodes_empty + print(f"\nValid episodes: {valid_episodes}/{len(parquet_files)}") + print(f"Total action timesteps: {total_action_samples:,}") + print(f"Total state samples: {total_state_samples:,}") + + # ---------------------------------------------------------------- + # Build per-key stats by slicing the concatenated action vector + # ---------------------------------------------------------------- + action_global = action_tracker.get_stats() + stats: dict = {} + + # Determine per-key dimensions from modality metadata + action configs + # For rel_xyz_rot6d keys, output is 9D (not raw 7D) + action_key_dims: dict[str, tuple[int, int]] = {} + offset = 0 + for key in action_keys: + modality, subkey = key.split(".", 1) + meta = modality_meta.get(modality, {}).get(subkey) + if meta is None: + continue + raw_dim = meta["end"] - meta["start"] + + cfg = action_key_configs.get(key) + if cfg is not None and cfg.rep == "rel_xyz_rot6d": + out_dim = 9 # xyz(3) + rot6d(6) + else: + out_dim = raw_dim + + action_key_dims[key] = (offset, offset + out_dim) + offset += out_dim + + state_key_dims: dict[str, tuple[int, int]] = {} + offset = 0 + for key in state_keys: + modality, subkey = key.split(".", 1) + meta = modality_meta.get(modality, {}).get(subkey) + if meta is None: + continue + dim = meta["end"] - meta["start"] + state_key_dims[key] = (offset, offset + dim) + offset += dim + + # Extract per-key stats + for key, (s, e) in action_key_dims.items(): + key_stats = { + "mean": action_global["mean"][s:e], + "std": action_global["std"][s:e], + "min": action_global["min"][s:e], + "max": action_global["max"][s:e], + } + if action_tracker.reservoir is not None: + res = action_tracker.reservoir[:, s:e] + key_stats["q01"] = np.quantile(res, 0.01, axis=0).tolist() + key_stats["q99"] = np.quantile(res, 0.99, axis=0).tolist() + else: + key_stats["q01"] = key_stats["min"] + key_stats["q99"] = key_stats["max"] + stats[key] = key_stats + + if state_tracker is not None: + state_global = state_tracker.get_stats() + for key, (s, e) in state_key_dims.items(): + key_stats = { + "mean": state_global["mean"][s:e], + "std": state_global["std"][s:e], + "min": state_global["min"][s:e], + "max": state_global["max"][s:e], + } + if state_tracker.reservoir is not None: + res = state_tracker.reservoir[:, s:e] + key_stats["q01"] = np.quantile(res, 0.01, axis=0).tolist() + key_stats["q99"] = np.quantile(res, 0.99, axis=0).tolist() + else: + key_stats["q01"] = key_stats["min"] + key_stats["q99"] = key_stats["max"] + stats[key] = key_stats + + # Global concatenated stats for convenience + stats["action"] = action_global + if state_tracker is not None: + stats["state"] = state_tracker.get_stats() + + # ---------------------------------------------------------------- + # Stamp provenance metadata for downstream guards. + # ---------------------------------------------------------------- + # We record the exact ``timestep_interval`` used to compute these + # statistics as a top-level integer so ``LeRobotSingleDataset`` can + # assert it still matches EMBODIMENT_REGISTRY at training time. + # Since all real stat entries are per-key dicts (``"action.psm1_pose"`` + # etc.), an int at the top level never collides with a stat, and the + # existing validation loop in ``dataset.py`` (``if isinstance(stat, int): + # continue``) already skips it. + stats["timestep_interval"] = int(timestep_interval) + + # ---------------------------------------------------------------- + # Write to disk + # ---------------------------------------------------------------- + out_path = dataset_path / "meta" / output_filename + out_path.parent.mkdir(parents=True, exist_ok=True) + with open(out_path, "w") as f: + json.dump(stats, f, indent=2) + + print(f"\nSaved stats to {out_path}") + print(f"Per-key action stats: {list(action_key_dims.keys())}") + if state_key_dims: + print(f"Per-key state stats: {list(state_key_dims.keys())}") + + for key, (s, e) in action_key_dims.items(): + dim = e - s + print(f" {key} ({dim}D): mean_abs={np.mean(np.abs(action_global['mean'][s:e])):.6f}, " + f"std_mean={np.mean(action_global['std'][s:e]):.6f}") + + return True + + +def _resolve_embodiment_string(spec_embodiment) -> str: + """Normalise the 'embodiment' field from OPEN_H_DATASET_SPECS to a plain string.""" + if isinstance(spec_embodiment, EmbodimentTag): + return spec_embodiment.value + return str(spec_embodiment) + + +def run_all(args): + """Process every dataset listed in OPEN_H_DATASET_SPECS (--all mode). + + Skips CMR Versius entries (they use compute_cmr_action_stats.py and + stats_cosmos-44D.json instead). + + Each spec's ``exclude_splits`` is forwarded so that the same episodes + excluded during training are also excluded from statistics computation. + """ + # Deduplicate: multiple specs may share the same (path, embodiment). + # Keep the first occurrence's exclude_splits. + seen: set[tuple[str, str]] = set() + jobs: list[tuple[Path, str, list[str] | None]] = [] + + for spec in OPEN_H_DATASET_SPECS: + emb = _resolve_embodiment_string(spec["embodiment"]) + dp = Path(spec["path"]) + + # CMR has its own dedicated script → skip + if emb == EmbodimentTag.CMR_VERSIUS.value: + continue + + key = (str(dp), emb) + if key in seen: + continue + seen.add(key) + jobs.append((dp, emb, spec.get("exclude_splits", None))) + + if not jobs: + print("No non-CMR datasets found in OPEN_H_DATASET_SPECS.") + return + + print("#" * 80) + print(f"OPEN-H BATCH MODE: {len(jobs)} dataset(s) to process") + print("#" * 80) + for i, (dp, emb, excl) in enumerate(jobs, 1): + excl_str = f" exclude={excl}" if excl else "" + print(f" [{i:2d}] [{emb:<22s}] {dp.name}{excl_str}") + print("#" * 80) + + total_start = time.time() + results: dict[str, str] = {} + + for i, (dp, emb, excl) in enumerate(jobs, 1): + print(f"\n{'#' * 80}") + print(f"# [{i}/{len(jobs)}] embodiment={emb} path={dp.name}") + if excl: + print(f"# exclude_splits={excl}") + print(f"{'#' * 80}") + + if not dp.exists(): + print(f" SKIPPED — path does not exist: {dp}") + results[f"{emb}/{dp.name}"] = "SKIPPED (path missing)" + continue + + try: + ok = process_single_dataset( + dataset_path=dp, + embodiment=emb, + num_frames=args.num_frames, + max_samples=args.max_samples, + output_filename=args.output_filename, + exclude_splits=excl, + num_workers=args.num_workers, + timestep_interval_override=args.timestep_interval, + ) + results[f"{emb}/{dp.name}"] = "OK" if ok else "FAILED" + except Exception as e: + print(f" ERROR: {e}") + results[f"{emb}/{dp.name}"] = f"ERROR ({e})" + + elapsed = time.time() - total_start + print(f"\n{'#' * 80}") + print(f"ALL DONE — {elapsed:.1f}s total, {len(results)} dataset(s)") + print(f"{'#' * 80}") + for name, status in results.items(): + print(f" {name}: {status}") + print(f"{'#' * 80}") + + +def run_single(args): + """Process a single dataset or auto-discovered datasets (original mode).""" + if args.dataset_path: + dataset_paths = [Path(args.dataset_path)] + else: + root = Path(args.dataset_path_root) + if _is_lerobot_dataset(root): + dataset_paths = [root] + else: + dataset_paths = _discover_datasets(root) + + if not dataset_paths: + print("ERROR: No datasets found!") + return + + exclude_splits = args.exclude_splits if args.exclude_splits else None + print(f"Found {len(dataset_paths)} dataset(s) for embodiment '{args.embodiment}'") + if exclude_splits: + print(f" exclude_splits: {exclude_splits}") + + total_start = time.time() + results = {} + for i, dp in enumerate(dataset_paths, 1): + if len(dataset_paths) > 1: + print(f"\n{'#' * 80}") + print(f"# DATASET {i}/{len(dataset_paths)}: {dp.name}") + print(f"{'#' * 80}") + + ok = process_single_dataset( + dataset_path=dp, + embodiment=args.embodiment, + num_frames=args.num_frames, + max_samples=args.max_samples, + output_filename=args.output_filename, + exclude_splits=exclude_splits, + num_workers=args.num_workers, + timestep_interval_override=args.timestep_interval, + ) + results[dp.name] = "OK" if ok else "FAILED" + + elapsed = time.time() - total_start + if len(dataset_paths) > 1: + print(f"\n{'#' * 80}") + print(f"ALL DONE — {elapsed:.1f}s") + for name, status in results.items(): + print(f" {name}: {status}") + print(f"{'#' * 80}") + + +def main(): + parser = argparse.ArgumentParser( + description="Compute normalization stats for Open-H embodiments (post-transform)", + epilog=( + "Modes:\n" + " --all Process ALL datasets in OPEN_H_DATASET_SPECS\n" + " (skips CMR Versius — use compute_cmr_action_stats.py).\n" + " No --dataset-path or --embodiment needed.\n\n" + " --dataset-path + --embodiment\n" + " Process a single dataset with a specific embodiment.\n\n" + " --dataset-path-root + --embodiment\n" + " Auto-discover datasets under a root directory.\n\n" + "Examples:\n" + " # All Open-H datasets at once:\n" + " python compute_openh_action_stats.py --all\n\n" + " # Quick test (10 samples per dataset):\n" + " python compute_openh_action_stats.py --all --max-samples 10\n\n" + " # Single dataset:\n" + " python compute_openh_action_stats.py \\\n" + " --dataset-path /path/to/suturebot_2 --embodiment dvrk\n" + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + # --- mutually exclusive: --all vs --dataset-path / --dataset-path-root --- + path_group = parser.add_mutually_exclusive_group(required=True) + path_group.add_argument( + "--all", action="store_true", + help="Process every dataset in OPEN_H_DATASET_SPECS (skips CMR Versius)", + ) + path_group.add_argument("--dataset-path", type=str, + help="Path to a single LeRobot dataset") + path_group.add_argument("--dataset-path-root", type=str, + help="Root directory containing multiple LeRobot datasets") + + parser.add_argument("--embodiment", type=str, default=None, + choices=list(EMBODIMENT_REGISTRY.keys()), + help="Embodiment tag (required for --dataset-path / --dataset-path-root; " + "ignored for --all)") + parser.add_argument("--exclude-splits", type=str, nargs="+", default=None, + help="Split names from info.json to exclude (e.g., --exclude-splits fail bad_frames). " + "For --all mode, exclude_splits from OPEN_H_DATASET_SPECS are used automatically.") + parser.add_argument("--num-frames", type=int, default=13, + help="Number of video frames (default: 13 = 1 context + 12 prediction)") + parser.add_argument("--max-samples", type=int, default=None, + help="Max episodes per dataset (for quick testing)") + parser.add_argument("--num-workers", type=int, default=None, + help=f"Number of parallel workers (default: min(cpu_count, {MAX_WORKERS}))") + parser.add_argument("--output-filename", type=str, default="stats_cosmos.json", + help="Output filename in meta/ dir (default: stats_cosmos.json)") + parser.add_argument("--timestep-interval", type=int, default=None, + help="Override the 'timestep_interval' (action stride) that is otherwise " + "read from EMBODIMENT_REGISTRY in groot_configs.py. Useful when you " + "want to compute stats for a different effective training rate than " + "the registry default (e.g. use 3 instead of 5 on a 30Hz dVRK dataset " + "to target 10 Hz effective). IMPORTANT: if you override this here, " + "the training pipeline MUST use the same value or the resulting " + "stats_cosmos.json will not match the distribution the model sees.") + args = parser.parse_args() + + # Validate: --dataset-path / --dataset-path-root require --embodiment + if not args.all and args.embodiment is None: + parser.error("--embodiment is required when using --dataset-path or --dataset-path-root") + + if args.all: + run_all(args) + else: + run_single(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/extract_jhu_inference_manifest.py b/scripts/extract_jhu_inference_manifest.py new file mode 100644 index 0000000..2fbe5e2 --- /dev/null +++ b/scripts/extract_jhu_inference_manifest.py @@ -0,0 +1,315 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Create streaming-inference manifests from JHU dVRK tabletop episodes. + +The script writes one ground-truth MP4 and one normalized, 44D-padded action +array per selected episode, plus the JSON manifest consumed by +``action_video2world_streaming``. It uses the same ``MixedLeRobotDataset`` +transforms and ``meta/stats_cosmos.json`` files as the reference training +recipe, so actions must not be normalized again downstream. + +Example: + +.. code-block:: bash + + python scripts/extract_jhu_inference_manifest.py \ + --subset hf_suturebot \ + --episode-ids 1440,1441,1442 \ + --num-frames 73 \ + --output-dir sf_inference_data/jhu_tabletop_test_h73 + +For multiple subsets, pass ``--episodes-json`` with a list of objects: + +.. code-block:: json + + [ + { + "subset": "hf_suturebot", + "path": "/datasets/jhu_tabletop/hf_suturebot", + "ids": [1440, 1441, 1442] + } + ] + +The ``path`` field is optional when the subset basename is present in +``JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS``. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np + +DEFAULT_TIMESTEP_INTERVAL = 3 + + +def parse_arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Extract JHU dVRK episode MP4/action pairs and an input manifest " + "for reference-tabletop streaming inference." + ), + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--subset", help="Dataset basename, for example hf_suturebot.") + parser.add_argument("--episode-ids", help="Comma-separated LeRobot episode IDs.") + parser.add_argument("--dataset-path", help="Override the path resolved for --subset.") + parser.add_argument( + "--episodes-json", + help="JSON list of {subset, path?, ids} objects; overrides single-subset arguments.", + ) + parser.add_argument( + "--num-frames", + type=int, + default=73, + help="Pixel-frame horizon per window; (num_frames-1) must be divisible by four.", + ) + parser.add_argument( + "--num-chunks", + type=int, + default=1, + help="Maximum consecutive non-overlapping windows per episode.", + ) + parser.add_argument( + "--start-margin", + type=int, + default=0, + help="Raw source-frame offset for the first window.", + ) + parser.add_argument( + "--split", + choices=("train", "test", "full"), + default="test", + help="Select the augmented or deterministic modality transform.", + ) + parser.add_argument( + "--timestep-interval", + type=int, + default=DEFAULT_TIMESTEP_INTERVAL, + help="Raw-frame stride used by the JHU embodiment.", + ) + parser.add_argument( + "--test-split-ratio", + type=float, + default=0.02, + help="Held-out ratio used while constructing the dataset.", + ) + parser.add_argument("--output-dir", required=True, help="Manifest and episode output directory.") + parser.add_argument( + "--tag", + default="jhu_tabletop", + help="Filename, manifest, and prediction-directory prefix.", + ) + parser.add_argument( + "--predicted-output-root", + default="interactive-output", + help="Manifest output_video directory prefix.", + ) + parser.add_argument("--fps", type=float, default=10.0, help="Output MP4 frame rate.") + return parser.parse_args() + + +def _load_episode_specs(args: argparse.Namespace) -> list[dict]: + if args.episodes_json: + with open(args.episodes_json) as file: + raw = json.load(file) + if not isinstance(raw, list): + raise ValueError("--episodes-json must contain a JSON list") + specs = [] + for index, entry in enumerate(raw): + if "subset" not in entry or "ids" not in entry: + raise ValueError(f"--episodes-json entry {index} requires 'subset' and 'ids'") + specs.append( + { + "subset": entry["subset"], + "path": entry.get("path"), + "ids": [int(value) for value in entry["ids"]], + } + ) + return specs + + if not args.subset or not args.episode_ids: + raise ValueError("provide --episodes-json or both --subset and --episode-ids") + ids = [int(value) for value in args.episode_ids.split(",") if value.strip()] + return [{"subset": args.subset, "path": args.dataset_path, "ids": ids}] + + +def _video_to_numpy(sample_video) -> np.ndarray: + """Convert a ``(C,T,H,W)`` sample to ``(T,H,W,C)`` uint8.""" + array = sample_video.permute(1, 2, 3, 0).cpu().numpy() + if array.dtype == np.uint8: + return array + return (np.clip(array, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8) + + +def main() -> int: + args = parse_arguments() + if args.num_frames < 2 or (args.num_frames - 1) % 4 != 0: + print( + f"ERROR: --num-frames must be >= 2 and satisfy " + f"(num_frames-1) % 4 == 0; got {args.num_frames}", + file=sys.stderr, + ) + return 2 + if args.num_chunks < 1: + print("ERROR: --num-chunks must be >= 1", file=sys.stderr) + return 2 + + try: + episode_specs = _load_episode_specs(args) + except ValueError as error: + print(f"ERROR: {error}", file=sys.stderr) + return 2 + + try: + import mediapy + + from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.dataset import ( + MixedLeRobotDataset, + ) + from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.embodiment_tags import ( + EmbodimentTag, + ) + from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.groot_configs import ( + JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS, + MAX_ACTION_DIM, + ) + except ImportError as error: + print( + f"ERROR: failed to import runtime dependencies: {error}\n" + "Run inside the Cosmos-Predict2.5 environment.", + file=sys.stderr, + ) + return 2 + + subset_to_path = { + Path(spec["path"]).name: spec["path"] + for spec in JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS + } + num_actions = args.num_frames - 1 + raw_chunk_stride = num_actions * args.timestep_interval + output_dir = Path(args.output_dir).expanduser().resolve() + episodes_dir = output_dir / "episodes" + episodes_dir.mkdir(parents=True, exist_ok=True) + predicted_dir = Path(args.predicted_output_root) / args.tag + manifest_entries: list[dict] = [] + + for spec in episode_specs: + subset = spec["subset"] + dataset_path = spec["path"] or subset_to_path.get(subset) + if dataset_path is None: + print(f"ERROR: no path found for subset '{subset}'", file=sys.stderr) + return 2 + dataset_path = os.path.abspath(os.path.expanduser(dataset_path)) + if not os.path.isdir(dataset_path): + print(f"ERROR: dataset path is missing: {dataset_path}", file=sys.stderr) + return 2 + + dataset = MixedLeRobotDataset( + dataset_specs=[ + { + "path": dataset_path, + "embodiment": EmbodimentTag.JHU_DVRK_MONO, + "mix_ratio": 1.0, + } + ], + num_frames=args.num_frames, + data_split=args.split, + max_action_dim=MAX_ACTION_DIM, + downscaled_res=False, + test_split_ratio=args.test_split_ratio, + ) + sub_dataset = dataset.sub_datasets[0] + trajectory_ids = np.asarray(sub_dataset.trajectory_ids).astype(int) + trajectory_lengths = np.asarray(sub_dataset.trajectory_lengths).astype(int) + length_by_id = dict(zip(trajectory_ids.tolist(), trajectory_lengths.tolist())) + + def fetch_window(episode_id: int, base_index: int): + saved_steps = sub_dataset._all_steps + try: + sub_dataset._all_steps = [(episode_id, base_index)] + return dataset[0] + finally: + sub_dataset._all_steps = saved_steps + + for episode_id in spec["ids"]: + if episode_id not in length_by_id: + print(f"WARNING: episode {episode_id} is absent from '{subset}'; skipping") + continue + usable_frames = length_by_id[episode_id] - 1 - args.start_margin + chunks_that_fit = max(0, usable_frames // raw_chunk_stride) + chunks_to_use = min(args.num_chunks, chunks_that_fit) + if chunks_to_use < 1: + print(f"WARNING: episode {episode_id} is shorter than one full window; skipping") + continue + + action_chunks = [] + video_chunks = [] + for chunk_index in range(chunks_to_use): + base_index = args.start_margin + chunk_index * raw_chunk_stride + sample = fetch_window(episode_id, base_index) + actions = sample["action"].cpu().numpy().astype(np.float32) + if actions.shape != (num_actions, MAX_ACTION_DIM): + raise ValueError( + f"episode {episode_id} produced action shape {actions.shape}; " + f"expected {(num_actions, MAX_ACTION_DIM)}" + ) + action_chunks.append(actions) + video_chunks.append(_video_to_numpy(sample["video"])) + + actions = np.concatenate(action_chunks, axis=0) + video = np.concatenate( + [video_chunks[0], *(chunk[1:] for chunk in video_chunks[1:])], + axis=0, + ) + stem = f"{args.tag}_{subset}_ep{episode_id:06d}" + video_path = episodes_dir / f"{stem}.mp4" + actions_path = episodes_dir / f"{stem}_actions.npy" + mediapy.write_video(str(video_path), video, fps=args.fps) + np.save(actions_path, actions) + + manifest_entries.append( + { + "input_video": os.path.relpath(video_path, start=os.getcwd()), + "input_action": os.path.relpath(actions_path, start=os.getcwd()), + "output_video": str(predicted_dir / f"{stem}.mp4"), + "_subset": subset, + "_episode_id": episode_id, + "_num_chunks": chunks_to_use, + "_action_shape": list(actions.shape), + "_video_shape_thwc": list(video.shape), + } + ) + + if not manifest_entries: + print("ERROR: no episodes were extracted", file=sys.stderr) + return 1 + + manifest_path = output_dir / f"{args.tag}_inference_manifest.json" + manifest_path.write_text(json.dumps(manifest_entries, indent=2)) + provenance = { + "script": os.path.relpath(__file__, start=os.getcwd()), + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "embodiment": "jhu_dvrk_mono", + "num_frames": args.num_frames, + "state_t": 1 + num_actions // 4, + "num_action_per_chunk": num_actions, + "timestep_interval": args.timestep_interval, + "episode_specs": episode_specs, + "manifest_file": manifest_path.name, + } + (output_dir / f"{args.tag}_inference_manifest_provenance.json").write_text( + json.dumps(provenance, indent=2) + ) + print(f"Wrote {len(manifest_entries)} entries to {manifest_path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/train_scripts/tabletop/01_train_short_teacher_h13.sh b/train_scripts/tabletop/01_train_short_teacher_h13.sh new file mode 100755 index 0000000..8ded7ff --- /dev/null +++ b/train_scripts/tabletop/01_train_short_teacher_h13.sh @@ -0,0 +1,55 @@ +#!/bin/bash +# Reference stage 1: 13-frame bidirectional tabletop teacher. +# Add site-specific --account/--partition directives if required. +#SBATCH --job-name=tabletop-teacher-h13 +#SBATCH --nodes=8 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --time=4:00:00 +#SBATCH --output=tabletop-teacher-h13_%A_%a.out +#SBATCH --error=tabletop-teacher-h13_%A_%a.out +#SBATCH --array=0-9%1 +#SBATCH --dependency=singleton +#SBATCH --requeue + +set -euo pipefail + +: "${OUTPUT_ROOT:?Set OUTPUT_ROOT to persistent storage}" +: "${TABLETOP_DATA_ROOT:?Set TABLETOP_DATA_ROOT to the nine-dataset LeRobot root}" +: "${CHSS_CHECKPOINT_DIR:?Set CHSS_CHECKPOINT_DIR to the C-H-S-S DCP directory}" +: "${CONTAINER_COSMOS25:?Set CONTAINER_COSMOS25 to the teacher container image}" + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +mkdir -p "$OUTPUT_ROOT" +export MASTER_ADDR +MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | sed -n '1p')" + +MOUNTS="${REPO_ROOT}:/workspace,${OUTPUT_ROOT}:/imaginaire_output" +MOUNTS="${MOUNTS},${TABLETOP_DATA_ROOT}:/datasets/jhu_tabletop" +MOUNTS="${MOUNTS},${CHSS_CHECKPOINT_DIR}:/checkpoints/chss" + +srun --export=ALL \ + --container-image="$CONTAINER_COSMOS25" \ + --container-mounts="$MOUNTS" \ + --container-workdir=/workspace \ + bash -c ' + set -euo pipefail + source .venv/bin/activate + export IMAGINAIRE_OUTPUT_ROOT=/imaginaire_output + export JHU_TABLETOP_DATA_ROOT=/datasets/jhu_tabletop + export CHSS_CHECKPOINT_DIR=/checkpoints/chss + NODE_RANK=${SLURM_NODEID:-0} + NNODES=${SLURM_JOB_NUM_NODES:-1} + torchrun \ + --nnodes="$NNODES" \ + --nproc_per_node=8 \ + --master_port=25001 \ + --master_addr="$MASTER_ADDR" \ + --node_rank="$NODE_RANK" \ + -m scripts.train \ + --config=cosmos_predict2/_src/predict2/action/configs/action_conditioned/config.py \ + -- \ + experiment=cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss \ + checkpoint.save_iter=200 \ + ~dataloader_train.dataloaders + ' diff --git a/train_scripts/tabletop/02_fine_anneal_short_teacher_h13.sh b/train_scripts/tabletop/02_fine_anneal_short_teacher_h13.sh new file mode 100755 index 0000000..433f6f0 --- /dev/null +++ b/train_scripts/tabletop/02_fine_anneal_short_teacher_h13.sh @@ -0,0 +1,52 @@ +#!/bin/bash +# Reference stage 1b: 4k cosine fine anneal from short-teacher iter 16,000. +#SBATCH --job-name=tabletop-teacher-h13-anneal +#SBATCH --nodes=8 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --time=4:00:00 +#SBATCH --output=tabletop-teacher-h13-anneal_%A_%a.out +#SBATCH --error=tabletop-teacher-h13-anneal_%A_%a.out +#SBATCH --array=0-3%1 +#SBATCH --dependency=singleton +#SBATCH --requeue + +set -euo pipefail + +: "${OUTPUT_ROOT:?Set OUTPUT_ROOT to persistent storage}" +: "${TABLETOP_DATA_ROOT:?Set TABLETOP_DATA_ROOT to the LeRobot root}" +: "${CONTAINER_COSMOS25:?Set CONTAINER_COSMOS25 to the teacher container image}" + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +SHORT_DCP="${OUTPUT_ROOT}/cosmos_predict2_action_conditioned/official_runs_vid2vid/cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss/checkpoints/iter_000016000" +test -d "$SHORT_DCP" || { echo "Missing short-teacher DCP: $SHORT_DCP" >&2; exit 1; } + +export MASTER_ADDR +MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | sed -n '1p')" +MOUNTS="${REPO_ROOT}:/workspace,${OUTPUT_ROOT}:/imaginaire_output" +MOUNTS="${MOUNTS},${TABLETOP_DATA_ROOT}:/datasets/jhu_tabletop" + +srun --export=ALL \ + --container-image="$CONTAINER_COSMOS25" \ + --container-mounts="$MOUNTS" \ + --container-workdir=/workspace \ + bash -c ' + set -euo pipefail + source .venv/bin/activate + export IMAGINAIRE_OUTPUT_ROOT=/imaginaire_output + export JHU_TABLETOP_DATA_ROOT=/datasets/jhu_tabletop + NODE_RANK=${SLURM_NODEID:-0} + NNODES=${SLURM_JOB_NUM_NODES:-1} + torchrun \ + --nnodes="$NNODES" \ + --nproc_per_node=8 \ + --master_port=25001 \ + --master_addr="$MASTER_ADDR" \ + --node_rank="$NODE_RANK" \ + -m scripts.train \ + --config=cosmos_predict2/_src/predict2/action/configs/action_conditioned/config.py \ + -- \ + experiment=cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss_fine_anneal_4k \ + checkpoint.save_iter=200 \ + ~dataloader_train.dataloaders + ' diff --git a/train_scripts/tabletop/03_train_long_teacher_h73.sh b/train_scripts/tabletop/03_train_long_teacher_h73.sh new file mode 100755 index 0000000..fe5f75a --- /dev/null +++ b/train_scripts/tabletop/03_train_long_teacher_h73.sh @@ -0,0 +1,53 @@ +#!/bin/bash +# Reference stage 2: 73-frame teacher warm-started from the annealed h13 teacher. +#SBATCH --job-name=tabletop-teacher-h73 +#SBATCH --nodes=8 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --time=4:00:00 +#SBATCH --output=tabletop-teacher-h73_%A_%a.out +#SBATCH --error=tabletop-teacher-h73_%A_%a.out +#SBATCH --array=0-9%1 +#SBATCH --dependency=singleton +#SBATCH --requeue + +set -euo pipefail + +: "${OUTPUT_ROOT:?Set OUTPUT_ROOT to persistent storage}" +: "${TABLETOP_DATA_ROOT:?Set TABLETOP_DATA_ROOT to the LeRobot root}" +: "${CONTAINER_COSMOS25:?Set CONTAINER_COSMOS25 to the teacher container image}" + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +ANNEALED_DCP="${OUTPUT_ROOT}/cosmos_predict2_action_conditioned/official_runs_vid2vid/cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss_fine_anneal_4k/checkpoints/iter_000004000" +test -d "$ANNEALED_DCP" || { echo "Missing annealed teacher DCP: $ANNEALED_DCP" >&2; exit 1; } +mkdir -p "$OUTPUT_ROOT" + +export MASTER_ADDR +MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | sed -n '1p')" +MOUNTS="${REPO_ROOT}:/workspace,${OUTPUT_ROOT}:/imaginaire_output" +MOUNTS="${MOUNTS},${TABLETOP_DATA_ROOT}:/datasets/jhu_tabletop" + +srun --export=ALL \ + --container-image="$CONTAINER_COSMOS25" \ + --container-mounts="$MOUNTS" \ + --container-workdir=/workspace \ + bash -c ' + set -euo pipefail + source .venv/bin/activate + export IMAGINAIRE_OUTPUT_ROOT=/imaginaire_output + export JHU_TABLETOP_DATA_ROOT=/datasets/jhu_tabletop + NODE_RANK=${SLURM_NODEID:-0} + NNODES=${SLURM_JOB_NUM_NODES:-1} + torchrun \ + --nnodes="$NNODES" \ + --nproc_per_node=8 \ + --master_port=25001 \ + --master_addr="$MASTER_ADDR" \ + --node_rank="$NODE_RANK" \ + -m scripts.train \ + --config=cosmos_predict2/_src/predict2/action/configs/action_conditioned/config.py \ + -- \ + experiment=cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss_h73_tabletop \ + checkpoint.save_iter=200 \ + ~dataloader_train.dataloaders + ' diff --git a/train_scripts/tabletop/04_phase0_teacher_cache_h73.sh b/train_scripts/tabletop/04_phase0_teacher_cache_h73.sh new file mode 100755 index 0000000..b42260b --- /dev/null +++ b/train_scripts/tabletop/04_phase0_teacher_cache_h73.sh @@ -0,0 +1,71 @@ +#!/bin/bash +# Reference stage 3: one-node, eight-rank Phase 0 teacher cache generation. +#SBATCH --job-name=tabletop-phase0-h73 +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=8 +#SBATCH --gres=gpu:8 +#SBATCH --time=4:00:00 +#SBATCH --output=tabletop-phase0-h73_%A_%a.out +#SBATCH --error=tabletop-phase0-h73_%A_%a.out +#SBATCH --array=0-3%1 +#SBATCH --dependency=singleton +#SBATCH --requeue + +set -euo pipefail + +: "${OUTPUT_ROOT:?Set OUTPUT_ROOT to persistent storage}" +: "${CACHE_ROOT:?Set CACHE_ROOT to persistent cache storage}" +: "${TABLETOP_DATA_ROOT:?Set TABLETOP_DATA_ROOT to the LeRobot root}" +: "${CONTAINER_COSMOS25:?Set CONTAINER_COSMOS25 to the teacher container image}" + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +TEACHER_RELATIVE="cosmos_predict2_action_conditioned/official_runs_vid2vid/cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss_h73_tabletop/checkpoints/iter_000005000/model_ema_bf16.pt" +test -f "${OUTPUT_ROOT}/${TEACHER_RELATIVE}" || { + echo "Missing consolidated teacher checkpoint: ${OUTPUT_ROOT}/${TEACHER_RELATIVE}" >&2 + exit 1 +} + +export TOTAL_SAMPLES="${TOTAL_SAMPLES:-10000}" +export SAMPLE_STRATEGY="${SAMPLE_STRATEGY:-random}" +export INDICES_SEED="${INDICES_SEED:-0}" +export TEACHER_CKPT="/imaginaire_output/${TEACHER_RELATIVE}" +export EXPERIMENT="cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss_h73_tabletop" +export SAVE_ROOT="datasets/jhu_dvrk_mono_warmup_4step_h73_tabletop" + +MOUNTS="${REPO_ROOT}:/workspace,${OUTPUT_ROOT}:/imaginaire_output" +MOUNTS="${MOUNTS},${CACHE_ROOT}:/imaginaire_cache" +MOUNTS="${MOUNTS},${TABLETOP_DATA_ROOT}:/datasets/jhu_tabletop" + +srun --export=ALL \ + --container-image="$CONTAINER_COSMOS25" \ + --container-mounts="$MOUNTS" \ + --container-workdir=/workspace \ + bash -c ' + set -euo pipefail + source .venv/bin/activate + export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID + export JHU_TABLETOP_DATA_ROOT=/datasets/jhu_tabletop + export IMAGINAIRE_CACHE_DIR=/imaginaire_cache + export HF_HOME=/imaginaire_cache/huggingface + + N_RANKS=8 + SAMPLES_PER_RANK=$(( (TOTAL_SAMPLES + N_RANKS - 1) / N_RANKS )) + START=$(( SLURM_LOCALID * SAMPLES_PER_RANK )) + END=$(( (SLURM_LOCALID + 1) * SAMPLES_PER_RANK )) + (( END > TOTAL_SAMPLES )) && END=$TOTAL_SAMPLES + + python cosmos_predict2/_src/predict2/action/inference/inference_jhu_dvrk_warmup.py \ + --experiment "$EXPERIMENT" \ + --ckpt_path "$TEACHER_CKPT" \ + --save_root "$SAVE_ROOT" \ + --resolution 288,512 \ + --guidance 0 \ + --num_frames 73 \ + --chunk_size 72 \ + --sample_strategy "$SAMPLE_STRATEGY" \ + --total_samples "$TOTAL_SAMPLES" \ + --indices_seed "$INDICES_SEED" \ + --start "$START" \ + --end "$END" \ + --query_steps 0,9,18,27,34 + ' diff --git a/train_scripts/tabletop/05_warmup_student_h73.sh b/train_scripts/tabletop/05_warmup_student_h73.sh new file mode 100755 index 0000000..4d93e7c --- /dev/null +++ b/train_scripts/tabletop/05_warmup_student_h73.sh @@ -0,0 +1,58 @@ +#!/bin/bash +# Reference stage 4: 8-node causal-student warmup, state_t=19. +#SBATCH --job-name=tabletop-warmup-h73 +#SBATCH --nodes=8 +#SBATCH --ntasks-per-node=8 +#SBATCH --gres=gpu:8 +#SBATCH --time=4:00:00 +#SBATCH --output=tabletop-warmup-h73_%A_%a.out +#SBATCH --error=tabletop-warmup-h73_%A_%a.out +#SBATCH --array=0-9%1 +#SBATCH --dependency=singleton +#SBATCH --requeue + +set -euo pipefail + +: "${OUTPUT_ROOT:?Set OUTPUT_ROOT to persistent storage}" +: "${CACHE_ROOT:?Set CACHE_ROOT to persistent cache storage}" +: "${CONTAINER_CAUSAL:?Set CONTAINER_CAUSAL to an image containing NATTEN}" + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +PHASE0_ROOT="${REPO_ROOT}/datasets/jhu_dvrk_mono_warmup_4step_h73_tabletop" +for artifact in latents images actions videos; do + test -d "${PHASE0_ROOT}/${artifact}" || { + echo "Missing Phase 0 artifact directory: ${PHASE0_ROOT}/${artifact}" >&2 + exit 1 + } +done +mkdir -p "$OUTPUT_ROOT" "$CACHE_ROOT" + +export MASTER_ADDR +MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | sed -n '1p')" +export MASTER_PORT=25002 +export WORLD_SIZE=$SLURM_NTASKS + +MOUNTS="${REPO_ROOT}:/workspace,${OUTPUT_ROOT}:/imaginaire_output" +MOUNTS="${MOUNTS},${CACHE_ROOT}:/imaginaire_cache" + +srun --export=ALL \ + --container-image="$CONTAINER_CAUSAL" \ + --container-mounts="$MOUNTS" \ + --container-workdir=/workspace \ + bash -c ' + set -euo pipefail + source .venv/bin/activate + export RANK=$SLURM_PROCID + export LOCAL_RANK=$SLURM_LOCALID + export IMAGINAIRE_OUTPUT_ROOT=/imaginaire_output + export IMAGINAIRE_CACHE_DIR=/imaginaire_cache + export HF_HOME=/imaginaire_cache/huggingface + export TORCH_HOME=/imaginaire_cache/torch + export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + python -c "import natten" + python -m scripts.train \ + --config=cosmos_predict2/_src/predict2/interactive/configs/config_warmup.py \ + -- \ + experiment=cosmos_predict2p5_2B_action_jhu_dvrk_mono_tabletop_h73_warmup_no_s3_resumable \ + checkpoint.save_iter=200 + ' diff --git a/train_scripts/tabletop/06_self_forcing_h73.sh b/train_scripts/tabletop/06_self_forcing_h73.sh new file mode 100755 index 0000000..94d821d --- /dev/null +++ b/train_scripts/tabletop/06_self_forcing_h73.sh @@ -0,0 +1,54 @@ +#!/bin/bash +# Reference stage 5: 8-node Self Forcing distillation, 3,000 iterations. +#SBATCH --job-name=tabletop-self-forcing-h73 +#SBATCH --nodes=8 +#SBATCH --ntasks-per-node=8 +#SBATCH --gres=gpu:8 +#SBATCH --time=4:00:00 +#SBATCH --output=tabletop-self-forcing-h73_%A_%a.out +#SBATCH --error=tabletop-self-forcing-h73_%A_%a.out +#SBATCH --array=0-9%1 +#SBATCH --dependency=singleton +#SBATCH --requeue + +set -euo pipefail + +: "${OUTPUT_ROOT:?Set OUTPUT_ROOT to persistent storage}" +: "${CACHE_ROOT:?Set CACHE_ROOT to persistent cache storage}" +: "${CONTAINER_CAUSAL:?Set CONTAINER_CAUSAL to an image containing NATTEN}" + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +WARMUP_DCP="${OUTPUT_ROOT}/cosmos_predict2_action_conditioned/interactive_warmup/jhu_dvrk_mono_i4_lr3e-5_h73_tabletop_no_s3_resumable/checkpoints/iter_000018000" +TEACHER_MODEL="${OUTPUT_ROOT}/cosmos_predict2_action_conditioned/official_runs_vid2vid/cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss_h73_tabletop/checkpoints/iter_000005000/model" +test -d "$WARMUP_DCP" || { echo "Missing warmup DCP: $WARMUP_DCP" >&2; exit 1; } +test -d "$TEACHER_MODEL" || { echo "Missing teacher model DCP: $TEACHER_MODEL" >&2; exit 1; } +mkdir -p "$OUTPUT_ROOT" "$CACHE_ROOT" + +export MASTER_ADDR +MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | sed -n '1p')" +export MASTER_PORT=25003 +export WORLD_SIZE=$SLURM_NTASKS +MOUNTS="${REPO_ROOT}:/workspace,${OUTPUT_ROOT}:/imaginaire_output" +MOUNTS="${MOUNTS},${CACHE_ROOT}:/imaginaire_cache" + +srun --export=ALL \ + --container-image="$CONTAINER_CAUSAL" \ + --container-mounts="$MOUNTS" \ + --container-workdir=/workspace \ + bash -c ' + set -euo pipefail + source .venv/bin/activate + export RANK=$SLURM_PROCID + export LOCAL_RANK=$SLURM_LOCALID + export IMAGINAIRE_OUTPUT_ROOT=/imaginaire_output + export IMAGINAIRE_CACHE_DIR=/imaginaire_cache + export HF_HOME=/imaginaire_cache/huggingface + export TORCH_HOME=/imaginaire_cache/torch + export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + python -c "import natten" + python -m scripts.train \ + --config=cosmos_predict2/_src/predict2/interactive/configs/config_distill.py \ + -- \ + experiment=cosmos_predict2p5_2B_action_jhu_dvrk_mono_tabletop_h73_self_forcing_no_s3_resumable \ + checkpoint.save_iter=200 + ' diff --git a/train_scripts/tabletop/README.md b/train_scripts/tabletop/README.md new file mode 100644 index 0000000..956b627 --- /dev/null +++ b/train_scripts/tabletop/README.md @@ -0,0 +1,43 @@ +# Tabletop teacher-to-student SLURM templates + +These templates reproduce the tabletop reference pipeline described in [`docs/tutorial_teacher_training_and_self_forcing.md`](../../docs/tutorial_teacher_training_and_self_forcing.md). +They use Pyxis/Enroot-style `srun --container-*` options; adapt those options if your cluster uses another container runtime. + +Before submitting, create the output directories and export host paths: + +```bash +export OUTPUT_ROOT=/persistent/path/imaginaire/output +export CACHE_ROOT=/persistent/path/imaginaire/cache +export TABLETOP_DATA_ROOT=/persistent/path/jhu_tabletop_lerobot +export CHSS_CHECKPOINT_DIR=/persistent/path/cosmos_h_surgical_simulator_dcp +export CONTAINER_COSMOS25=/path/to/cosmos-predict-2.5.sqsh +export CONTAINER_CAUSAL=/path/to/image-with-natten.sqsh +mkdir -p "$OUTPUT_ROOT" "$CACHE_ROOT" +``` + +`TABLETOP_DATA_ROOT` must contain the nine directories listed by `JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS`. Add your cluster's `#SBATCH --account` and `#SBATCH --partition` directives locally if required. + +Run order: + +```bash +sbatch train_scripts/tabletop/01_train_short_teacher_h13.sh +sbatch train_scripts/tabletop/02_fine_anneal_short_teacher_h13.sh +sbatch train_scripts/tabletop/03_train_long_teacher_h73.sh + +# Convert long-teacher iter_000005000/model to model_ema_bf16.pt first. +sbatch train_scripts/tabletop/04_phase0_teacher_cache_h73.sh + +sbatch train_scripts/tabletop/05_warmup_student_h73.sh +sbatch train_scripts/tabletop/06_self_forcing_h73.sh +``` + +Keep these invariants synchronized: + +- short teacher: 13 frames, `state_t=4`, selected iteration 16,000; +- fine anneal: 4,000 iterations; +- long teacher: 73 frames, `state_t=19`, selected iteration 5,000; +- Phase 0: 10,000 randomly selected samples, 72 actions, 73 frames; +- warmup: 20,000-iteration ceiling; reference SF initialization at 18,000; +- Self Forcing: 3,000 iterations. + +The Phase 0 script is resumable because existing complete artifact quartets are skipped. Before warmup, verify that `latents/`, `images/`, `actions/`, and `videos/` contain the same index set; do not rely only on SLURM completion. \ No newline at end of file From 484abf1ce139c82fba6a86a328b4235febfa2214 Mon Sep 17 00:00:00 2001 From: Lukas Zbinden Date: Sun, 26 Jul 2026 21:51:02 +0200 Subject: [PATCH 2/2] Fix pre-commit formatting and license checks --- ...B_action_conditioned_rectify_flow_gr00t.py | 6 +- ...orial_teacher_training_and_self_forcing.md | 1 - scripts/compute_openh_action_stats.py | 1781 +++++++++-------- scripts/extract_jhu_inference_manifest.py | 15 +- .../tabletop/01_train_short_teacher_h13.sh | 15 + .../02_fine_anneal_short_teacher_h13.sh | 15 + .../tabletop/03_train_long_teacher_h73.sh | 15 + .../tabletop/04_phase0_teacher_cache_h73.sh | 15 + .../tabletop/05_warmup_student_h73.sh | 15 + train_scripts/tabletop/06_self_forcing_h73.sh | 15 + train_scripts/tabletop/README.md | 2 +- 11 files changed, 997 insertions(+), 898 deletions(-) diff --git a/cosmos_predict2/_src/predict2/action/configs/action_conditioned/experiment/exp_2B_action_conditioned_rectify_flow_gr00t.py b/cosmos_predict2/_src/predict2/action/configs/action_conditioned/experiment/exp_2B_action_conditioned_rectify_flow_gr00t.py index 8cfb0ba..ad799a3 100644 --- a/cosmos_predict2/_src/predict2/action/configs/action_conditioned/experiment/exp_2B_action_conditioned_rectify_flow_gr00t.py +++ b/cosmos_predict2/_src/predict2/action/configs/action_conditioned/experiment/exp_2B_action_conditioned_rectify_flow_gr00t.py @@ -52,6 +52,7 @@ def _tabletop_teacher_checkpoint(run_name: str, iteration: int) -> str: f"official_runs_vid2vid/{run_name}/checkpoints/iter_{iteration:09d}" ) + _TRAINER_DEBUG_CONFIG = dict( max_iter=1000, logging_iter=50, @@ -1001,10 +1002,7 @@ def build_debug_runs(job): flags={"allow_objects": True}, ) -_JHU_H13_TEACHER_RUN = ( - "cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_" - "finetune_13frame_8nodes_release_oss" -) +_JHU_H13_TEACHER_RUN = "cosmos_predict2p5_2B_action_conditioned_jhu_dvrk_mono_finetune_13frame_8nodes_release_oss" AC_CHUNK_SINGLE_VIEW_2B_JHU_DVRK_MONO_FINETUNE_13FRAME_8NODES_OSS_FINE_ANNEAL_4K = LazyDict( dict( defaults=[ diff --git a/docs/tutorial_teacher_training_and_self_forcing.md b/docs/tutorial_teacher_training_and_self_forcing.md index 7010692..f5b4c46 100644 --- a/docs/tutorial_teacher_training_and_self_forcing.md +++ b/docs/tutorial_teacher_training_and_self_forcing.md @@ -1147,4 +1147,3 @@ Pre-download shared text embeddings and checkpoint assets to a mounted cache. Av - Cosmos-Surg-dVRK: *World Foundation Model-based Automated Online Evaluation of Surgical Robot Policy Learning*: [https://arxiv.org/abs/2510.16240](https://arxiv.org/abs/2510.16240) - Self Forcing: *Bridging the Train-Test Gap in Autoregressive Video Diffusion*: [https://arxiv.org/abs/2506.08009](https://arxiv.org/abs/2506.08009) - NVIDIA OmniDreams: *Real-Time Generative World Model for Closed-Loop Autonomous Vehicle Simulation*: [https://arxiv.org/abs/2606.03159](https://arxiv.org/abs/2606.03159) - diff --git a/scripts/compute_openh_action_stats.py b/scripts/compute_openh_action_stats.py index df8460f..4f3a5b7 100644 --- a/scripts/compute_openh_action_stats.py +++ b/scripts/compute_openh_action_stats.py @@ -1,881 +1,900 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Compute per-key normalization statistics for any Open-H embodiment. - -Unlike compute_cmr_action_stats.py (which is CMR-specific with hardcoded raw -indices, clutch filtering, and motion scaling), this script is GENERIC: it -instantiates the real transform pipeline (GenericRelativeActionTransform) for -any embodiment registered in EMBODIMENT_REGISTRY and collects statistics on -the TRANSFORMED action/state output. - -This guarantees that the statistics exactly match what the training pipeline -produces, regardless of the embodiment's delta conversion (rel_xyz_rot6d, -relative, delta, or absolute). - -Output: - meta/stats_cosmos.json — per-key statistics in the same format as - stats_cosmos-44D.json (action.psm1_pose → {mean, std, min, max, q01, q99}) - -For CMR Versius, continue using compute_cmr_action_stats.py (it handles -the additional clutch filtering and motion scaling that are CMR-specific). - -Usage: - # All Open-H datasets at once (reads OPEN_H_DATASET_SPECS, includes exclude_splits): - python compute_openh_action_stats.py --all - - # Quick test (10 samples per dataset): - python compute_openh_action_stats.py --all --max-samples 10 - - # Single dataset: - python compute_openh_action_stats.py \\ - --dataset-path /path/to/lerobot/dataset \\ - --embodiment dvrk - - # Single dataset with episode filtering: - python compute_openh_action_stats.py \\ - --dataset-path /path/to/stanford/Needle_Transfer \\ - --embodiment dvrk_stanford_real \\ - --exclude-splits fail bad_frames - - # All dVRK datasets under a root: - python compute_openh_action_stats.py \\ - --dataset-path-root /path/to/jhu \\ - --embodiment jhu_dvrk_mono -""" - -import argparse -import json -import os -import time -import warnings -from concurrent.futures import ProcessPoolExecutor, as_completed -from functools import partial -from pathlib import Path - -# Suppress noisy torchvision video deprecation warnings -warnings.filterwarnings("ignore", message=".*video decoding and encoding capabilities of torchvision.*") - -import numpy as np -import pandas as pd -from tqdm import tqdm - -from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.embodiment_tags import EmbodimentTag -from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.dataset import ( - LE_ROBOT_INFO_FILENAME, - resolve_excluded_episode_indices, -) -from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.schema import ( - LeRobotModalityMetadata, - LeRobotStateActionMetadata, -) -from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.groot_configs import ( - EMBODIMENT_REGISTRY, - OPEN_H_DATASET_SPECS, -) -from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.transform.state_action import ( - convert_to_hybrid_relative, -) - -# Maximum parallel workers -MAX_WORKERS = 64 - - -# ============================================================================ -# Statistics helpers (streaming, memory-efficient) -# ============================================================================ - -class StreamingStats: - """Memory-efficient streaming statistics using Welford's algorithm + reservoir sampling.""" - - def __init__(self, num_dims: int, reservoir_size: int = 2_000_000): - self.num_dims = num_dims - self.reservoir_size = reservoir_size - self.count = 0 - self.mean = np.zeros(num_dims, dtype=np.float64) - self.M2 = np.zeros(num_dims, dtype=np.float64) - self.min_vals = np.full(num_dims, np.inf, dtype=np.float64) - self.max_vals = np.full(num_dims, -np.inf, dtype=np.float64) - self.reservoir = None - self.reservoir_count = 0 - - def update(self, batch: np.ndarray): - if batch.shape[0] == 0: - return - n = batch.shape[0] - batch_mean = np.mean(batch, axis=0) - batch_var = np.var(batch, axis=0, ddof=0) - batch_M2 = batch_var * n - self.min_vals = np.minimum(self.min_vals, np.min(batch, axis=0)) - self.max_vals = np.maximum(self.max_vals, np.max(batch, axis=0)) - if self.count == 0: - self.mean = batch_mean - self.M2 = batch_M2 - self.count = n - else: - n_total = self.count + n - delta = batch_mean - self.mean - self.mean = (self.count * self.mean + n * batch_mean) / n_total - self.M2 = self.M2 + batch_M2 + delta ** 2 * self.count * n / n_total - self.count = n_total - # Reservoir sampling - if self.reservoir is None: - self.reservoir = batch[:self.reservoir_size].copy() - self.reservoir_count = n - elif len(self.reservoir) < self.reservoir_size: - space = self.reservoir_size - len(self.reservoir) - self.reservoir = np.vstack([self.reservoir, batch[:space]]) - self.reservoir_count += n - else: - # Algorithm R replacement - for row in batch: - self.reservoir_count += 1 - j = np.random.randint(0, self.reservoir_count) - if j < self.reservoir_size: - self.reservoir[j] = row - - def get_stats(self) -> dict: - if self.count == 0: - z = [0.0] * self.num_dims - return {"mean": z, "std": z, "min": z, "max": z, "q01": z, "q99": z} - std = np.sqrt(self.M2 / self.count) - if self.reservoir is not None and len(self.reservoir) > 0: - q01 = np.quantile(self.reservoir, 0.01, axis=0) - q99 = np.quantile(self.reservoir, 0.99, axis=0) - else: - q01 = self.min_vals - q99 = self.max_vals - return { - "mean": self.mean.tolist(), - "std": std.tolist(), - "min": self.min_vals.tolist(), - "max": self.max_vals.tolist(), - "q01": q01.tolist(), - "q99": q99.tolist(), - } - - -# ============================================================================ -# Main processing -# ============================================================================ - -def _is_lerobot_dataset(path: Path) -> bool: - return (path / "data").is_dir() and (path / "meta").is_dir() and any((path / "data").rglob("*.parquet")) - - -def _discover_datasets(root: Path) -> list[Path]: - datasets = [] - for child in sorted(root.iterdir()): - if child.is_dir() and _is_lerobot_dataset(child): - datasets.append(child) - return datasets - - -def _process_episode_parquet( - parquet_path: Path, - modality_meta: dict, - action_key_configs: dict, - action_delta_indices: list[int], - state_delta_indices: list[int], - action_keys: list[str], - state_keys: list[str], -) -> tuple[np.ndarray, np.ndarray, str | None]: - """Worker function: process one parquet file (episode) without video decoding. - - Reads state/action arrays from parquet, applies delta conversion, concatenates - per-key arrays, and returns the result for streaming stats collection. - - This runs in a subprocess via ProcessPoolExecutor for parallelism. - - Returns: - (action_array, state_array, warning_or_None) - action_array shape: (N_samples * T_action, action_dim) - state_array shape: (N_samples, state_dim) - """ - try: - df = pd.read_parquet(parquet_path) - T = len(df) - - # Maximum delta for action horizon - max_action_delta = max(action_delta_indices) if action_delta_indices else 0 - effective_length = max(0, T - max_action_delta) - if effective_length == 0: - return np.empty((0, 0)), np.empty((0, 0)), None - - # Extract per-key arrays from the flat parquet columns using modality metadata - def extract_key_data(key: str, df: pd.DataFrame) -> np.ndarray: - """Extract a named key's data from parquet using modality.json metadata.""" - modality, subkey = key.split(".", 1) - meta = modality_meta.get(modality, {}).get(subkey) - if meta is None: - raise KeyError(f"{key} config not found") - original_col = meta.get("original_key") - if original_col is None: - # Default: observation.state for state, action for action - original_col = "observation.state" if modality == "state" else "action" - start, end = meta["start"], meta["end"] - col_data = np.stack(df[original_col].values) # (T, D_flat) - return col_data[:, start:end].astype(np.float32) # (T, key_dim) - - # Process all valid starting indices - all_action_rows = [] - all_state_rows = [] - n_skipped = 0 - - for base_idx in range(effective_length): - # --- State at t=0 --- - state_parts = [] - for key in state_keys: - s_idx = base_idx + state_delta_indices[0] # delta_indices=[0] - s_idx = max(0, min(s_idx, T - 1)) - key_data = extract_key_data(key, df) - state_parts.append(key_data[s_idx]) - if state_parts: - all_state_rows.append(np.concatenate(state_parts)) - - # --- Action over horizon --- - action_timestep_rows = [] - for a_delta in action_delta_indices: - a_idx = base_idx + a_delta - a_idx = max(0, min(a_idx, T - 1)) - action_parts = [] - for key in action_keys: - key_data = extract_key_data(key, df) - action_parts.append(key_data[a_idx]) - action_timestep_rows.append(np.concatenate(action_parts)) - action_horizon = np.stack(action_timestep_rows) # (T_action, action_dim_raw) - - # --- Apply delta conversion per key --- - # We need to apply the same GenericRelativeActionTransform logic - # but directly on numpy arrays (no torch, no dataset wrapper). - offset = 0 - converted_parts = [] - state_row = all_state_rows[-1] if all_state_rows else None - skip_sample = False - - for key in action_keys: - cfg = action_key_configs.get(key) - key_data_raw = extract_key_data(key, df) - raw_dim = key_data_raw.shape[1] - - # Extract this key's action horizon data - key_action = action_horizon[:, offset:offset + raw_dim] - - if cfg is not None and cfg.rep == "rel_xyz_rot6d": - # Get reference state pose - ref_key = cfg.state_key - if ref_key and state_row is not None: - # Find the state key's slice in the concatenated state - s_off = 0 - ref_pose = None - for sk in state_keys: - sk_data = extract_key_data(sk, df) - sk_dim = sk_data.shape[1] - if sk == ref_key: - ref_pose = state_row[s_off:s_off + sk_dim] - break - s_off += sk_dim - - if ref_pose is not None: - # Guard: check for zero-norm quaternions (invalid data / - # padding at episode boundaries). Scipy's Rotation.from_quat() - # crashes on [0,0,0,0]. - if cfg.input_rotation_format == "quat": - quat_slice = key_action[:, 3:7] - norms = np.linalg.norm(quat_slice, axis=-1) - if np.any(norms < 1e-8): - skip_sample = True - break - if cfg.reference_rotation_format == "quat": - ref_quat = ref_pose[3:7] if len(ref_pose) >= 7 else ref_pose[3:] - if np.linalg.norm(ref_quat) < 1e-8: - skip_sample = True - break - - key_action = convert_to_hybrid_relative( - action_data=key_action, - eef_pose=ref_pose, - input_rotation_format=cfg.input_rotation_format, - reference_rotation_format=cfg.reference_rotation_format, - input_quat_order=cfg.input_quat_order, - reference_quat_order=cfg.reference_quat_order, - ) # (T_action, 9) - - elif cfg is not None and cfg.rep == "relative": - # Joint-space subtraction - ref_key = cfg.state_key - if ref_key and state_row is not None: - s_off = 0 - ref_val = None - for sk in state_keys: - sk_data = extract_key_data(sk, df) - sk_dim = sk_data.shape[1] - if sk == ref_key: - ref_val = state_row[s_off:s_off + sk_dim] - break - s_off += sk_dim - if ref_val is not None: - key_action = key_action - ref_val - - # delta / absolute: pass through unchanged - converted_parts.append(key_action) - offset += raw_dim - - if skip_sample: - # Remove the state row we just added (it corresponds to this skipped sample) - if all_state_rows: - all_state_rows.pop() - n_skipped += 1 - continue - - converted_action = np.concatenate(converted_parts, axis=-1) # (T_action, action_dim) - all_action_rows.append(converted_action) - - if not all_action_rows: - warn = None - if n_skipped > 0: - warn = f"{parquet_path.name}: all {effective_length} samples skipped ({n_skipped} had zero-norm quaternions)" - return np.empty((0, 0)), np.empty((0, 0)), warn - - actions = np.concatenate(all_action_rows, axis=0) # (N*T_action, action_dim) - states = np.stack(all_state_rows) if all_state_rows else np.empty((0, 0)) - - warn = None - if n_skipped > 0: - warn = (f"{parquet_path.name}: {n_skipped}/{effective_length} samples skipped " - f"(zero-norm quaternions), {len(all_action_rows)} valid") - return actions, states, warn - - except Exception as e: - import traceback - return np.empty((0, 0)), np.empty((0, 0)), f"Error processing {parquet_path.name}: {type(e).__name__}: {e}" - - -def _load_modality_meta(dataset_path: Path, modality_filename: str) -> dict: - """Load modality.json and return a simplified {modality: {subkey: {start, end, original_key}}} dict.""" - modality_path = dataset_path / modality_filename - if not modality_path.exists(): - raise FileNotFoundError(f"Modality file not found: {modality_path}") - - with open(modality_path, "r") as f: - raw = json.load(f) - - result: dict[str, dict] = {} - for modality in ["state", "action"]: - result[modality] = {} - if modality not in raw: - continue - for subkey, meta in raw[modality].items(): - result[modality][subkey] = { - "start": meta.get("start", 0), - "end": meta.get("end", 1), - "original_key": meta.get("original_key"), - } - return result - - -def process_single_dataset( - dataset_path: Path, - embodiment: str, - num_frames: int, - max_samples: int | None, - output_filename: str, - exclude_splits: list[str] | None = None, - num_workers: int | None = None, - timestep_interval_override: int | None = None, -): - """Process one dataset using parallel episode processing (no video decoding). - - Each parquet file (episode) is processed by a worker subprocess that: - 1. Reads state/action arrays from parquet (fast, no video) - 2. Applies the delta conversion (rel_xyz_rot6d, relative, etc.) - 3. Returns the transformed arrays for streaming stats collection - - This is ~100x faster than the sequential video-decoding approach. - - Args: - dataset_path: Path to the LeRobot dataset. - embodiment: Embodiment tag string. - num_frames: Number of video frames (e.g. 13). - max_samples: Max episodes to process (for testing). None = all. - output_filename: Output filename in meta/ dir. - exclude_splits: Split names from info.json to exclude. - num_workers: Parallel workers (default: min(cpu_count, MAX_WORKERS)). - timestep_interval_override: If given (int > 0), override the - ``timestep_interval`` value read from EMBODIMENT_REGISTRY for - this run. Useful to experiment with a different effective training - rate *without* editing ``groot_configs.py``. IMPORTANT: if you use - this, training must run with the same stride, otherwise - ``stats_cosmos.json`` will not match the distribution the model - sees. - """ - if num_workers is None: - num_workers = min(os.cpu_count() or 8, MAX_WORKERS) - - reg = EMBODIMENT_REGISTRY.get(embodiment) - if reg is None: - raise ValueError(f"Unknown embodiment '{embodiment}'. Available: {list(EMBODIMENT_REGISTRY.keys())}") - - registry_timestep_interval = reg["timestep_interval"] - if timestep_interval_override is not None: - if timestep_interval_override < 1: - raise ValueError( - f"--timestep-interval must be >= 1, got {timestep_interval_override}" - ) - timestep_interval = int(timestep_interval_override) - else: - timestep_interval = registry_timestep_interval - - print("=" * 80) - print(f"COMPUTING OPEN-H ACTION STATS — {dataset_path.name}") - print(f" embodiment: {embodiment}") - print(f" num_frames: {num_frames}") - if timestep_interval_override is not None and timestep_interval_override != registry_timestep_interval: - print( - f" timestep_interval: {timestep_interval} " - f"(OVERRIDE; registry default for '{embodiment}' is {registry_timestep_interval})" - ) - else: - print(f" timestep_interval: {timestep_interval} (from EMBODIMENT_REGISTRY)") - print(f" workers: {num_workers}") - if exclude_splits: - print(f" exclude_splits: {exclude_splits}") - print("=" * 80) - - num_action_frames = num_frames - 1 - action_delta_indices = list(range(0, num_action_frames * timestep_interval, timestep_interval)) - state_delta_indices = [0] - - action_keys = reg["action_keys"] - state_keys = reg["state_keys"] - action_key_configs = reg.get("action_key_configs", {}) - modality_filename = reg.get("modality_filename", "meta/modality.json") - - # Load modality metadata (maps key names to parquet column indices) - modality_meta = _load_modality_meta(dataset_path, modality_filename) - - # Discover parquet files (one per episode) - parquet_files = sorted(dataset_path.glob("data/*/*.parquet")) - if not parquet_files: - print(f" ERROR: No parquet files found in {dataset_path / 'data'}") - return False - - # Apply exclude_splits filtering at episode level - if exclude_splits: - excluded_ids = resolve_excluded_episode_indices(dataset_path, exclude_splits) - # Parse episode index from filename (e.g., episode_000042.parquet → 42) - filtered = [] - for pf in parquet_files: - try: - ep_idx = int(pf.stem.split("_")[-1]) - except ValueError: - filtered.append(pf) # Can't parse → keep - continue - if ep_idx not in excluded_ids: - filtered.append(pf) - n_excluded = len(parquet_files) - len(filtered) - print(f" exclude_splits: removed {n_excluded} episodes, {len(filtered)} remaining") - parquet_files = filtered - - if max_samples is not None: - parquet_files = parquet_files[:max_samples] - - print(f" Processing {len(parquet_files)} episodes with {num_workers} workers...") - - # Create partial function with fixed arguments for the worker - worker_fn = partial( - _process_episode_parquet, - modality_meta=modality_meta, - action_key_configs=action_key_configs, - action_delta_indices=action_delta_indices, - state_delta_indices=state_delta_indices, - action_keys=action_keys, - state_keys=state_keys, - ) - - # Process episodes in parallel - action_tracker = None - state_tracker = None - total_action_samples = 0 - total_state_samples = 0 - episodes_with_warnings = 0 - episodes_empty = 0 - - with ProcessPoolExecutor(max_workers=num_workers) as executor: - futures = {executor.submit(worker_fn, pf): pf for pf in parquet_files} - - for future in tqdm(as_completed(futures), total=len(futures), desc=f"Stats for {dataset_path.name}"): - actions, states, warning = future.result() - - if warning: - episodes_with_warnings += 1 - if episodes_with_warnings <= 10: - print(f" [WARN] {warning}") - - if actions.size == 0: - episodes_empty += 1 - continue - - # Lazy-init trackers based on first result's dimensions - if action_tracker is None: - action_dim = actions.shape[1] - action_tracker = StreamingStats(action_dim) - print(f" Action dim (post-transform): {action_dim}") - if state_tracker is None and states.size > 0: - state_dim = states.shape[1] - state_tracker = StreamingStats(state_dim) - print(f" State dim (post-transform): {state_dim}") - - action_tracker.update(actions.astype(np.float64)) - total_action_samples += len(actions) - - if state_tracker is not None and states.size > 0: - state_tracker.update(states.astype(np.float64)) - total_state_samples += len(states) - - if episodes_with_warnings > 10: - print(f" [WARN] {episodes_with_warnings} episodes had warnings (showing first 10)") - if episodes_empty > 0: - print(f" [INFO] {episodes_empty}/{len(parquet_files)} episodes produced no valid samples " - f"(zero-norm quaternions or empty episodes)") - - if action_tracker is None or action_tracker.count == 0: - print(f" ERROR: No valid samples found! ({episodes_empty} empty, {episodes_with_warnings} warnings)") - return False - - valid_episodes = len(parquet_files) - episodes_empty - print(f"\nValid episodes: {valid_episodes}/{len(parquet_files)}") - print(f"Total action timesteps: {total_action_samples:,}") - print(f"Total state samples: {total_state_samples:,}") - - # ---------------------------------------------------------------- - # Build per-key stats by slicing the concatenated action vector - # ---------------------------------------------------------------- - action_global = action_tracker.get_stats() - stats: dict = {} - - # Determine per-key dimensions from modality metadata + action configs - # For rel_xyz_rot6d keys, output is 9D (not raw 7D) - action_key_dims: dict[str, tuple[int, int]] = {} - offset = 0 - for key in action_keys: - modality, subkey = key.split(".", 1) - meta = modality_meta.get(modality, {}).get(subkey) - if meta is None: - continue - raw_dim = meta["end"] - meta["start"] - - cfg = action_key_configs.get(key) - if cfg is not None and cfg.rep == "rel_xyz_rot6d": - out_dim = 9 # xyz(3) + rot6d(6) - else: - out_dim = raw_dim - - action_key_dims[key] = (offset, offset + out_dim) - offset += out_dim - - state_key_dims: dict[str, tuple[int, int]] = {} - offset = 0 - for key in state_keys: - modality, subkey = key.split(".", 1) - meta = modality_meta.get(modality, {}).get(subkey) - if meta is None: - continue - dim = meta["end"] - meta["start"] - state_key_dims[key] = (offset, offset + dim) - offset += dim - - # Extract per-key stats - for key, (s, e) in action_key_dims.items(): - key_stats = { - "mean": action_global["mean"][s:e], - "std": action_global["std"][s:e], - "min": action_global["min"][s:e], - "max": action_global["max"][s:e], - } - if action_tracker.reservoir is not None: - res = action_tracker.reservoir[:, s:e] - key_stats["q01"] = np.quantile(res, 0.01, axis=0).tolist() - key_stats["q99"] = np.quantile(res, 0.99, axis=0).tolist() - else: - key_stats["q01"] = key_stats["min"] - key_stats["q99"] = key_stats["max"] - stats[key] = key_stats - - if state_tracker is not None: - state_global = state_tracker.get_stats() - for key, (s, e) in state_key_dims.items(): - key_stats = { - "mean": state_global["mean"][s:e], - "std": state_global["std"][s:e], - "min": state_global["min"][s:e], - "max": state_global["max"][s:e], - } - if state_tracker.reservoir is not None: - res = state_tracker.reservoir[:, s:e] - key_stats["q01"] = np.quantile(res, 0.01, axis=0).tolist() - key_stats["q99"] = np.quantile(res, 0.99, axis=0).tolist() - else: - key_stats["q01"] = key_stats["min"] - key_stats["q99"] = key_stats["max"] - stats[key] = key_stats - - # Global concatenated stats for convenience - stats["action"] = action_global - if state_tracker is not None: - stats["state"] = state_tracker.get_stats() - - # ---------------------------------------------------------------- - # Stamp provenance metadata for downstream guards. - # ---------------------------------------------------------------- - # We record the exact ``timestep_interval`` used to compute these - # statistics as a top-level integer so ``LeRobotSingleDataset`` can - # assert it still matches EMBODIMENT_REGISTRY at training time. - # Since all real stat entries are per-key dicts (``"action.psm1_pose"`` - # etc.), an int at the top level never collides with a stat, and the - # existing validation loop in ``dataset.py`` (``if isinstance(stat, int): - # continue``) already skips it. - stats["timestep_interval"] = int(timestep_interval) - - # ---------------------------------------------------------------- - # Write to disk - # ---------------------------------------------------------------- - out_path = dataset_path / "meta" / output_filename - out_path.parent.mkdir(parents=True, exist_ok=True) - with open(out_path, "w") as f: - json.dump(stats, f, indent=2) - - print(f"\nSaved stats to {out_path}") - print(f"Per-key action stats: {list(action_key_dims.keys())}") - if state_key_dims: - print(f"Per-key state stats: {list(state_key_dims.keys())}") - - for key, (s, e) in action_key_dims.items(): - dim = e - s - print(f" {key} ({dim}D): mean_abs={np.mean(np.abs(action_global['mean'][s:e])):.6f}, " - f"std_mean={np.mean(action_global['std'][s:e]):.6f}") - - return True - - -def _resolve_embodiment_string(spec_embodiment) -> str: - """Normalise the 'embodiment' field from OPEN_H_DATASET_SPECS to a plain string.""" - if isinstance(spec_embodiment, EmbodimentTag): - return spec_embodiment.value - return str(spec_embodiment) - - -def run_all(args): - """Process every dataset listed in OPEN_H_DATASET_SPECS (--all mode). - - Skips CMR Versius entries (they use compute_cmr_action_stats.py and - stats_cosmos-44D.json instead). - - Each spec's ``exclude_splits`` is forwarded so that the same episodes - excluded during training are also excluded from statistics computation. - """ - # Deduplicate: multiple specs may share the same (path, embodiment). - # Keep the first occurrence's exclude_splits. - seen: set[tuple[str, str]] = set() - jobs: list[tuple[Path, str, list[str] | None]] = [] - - for spec in OPEN_H_DATASET_SPECS: - emb = _resolve_embodiment_string(spec["embodiment"]) - dp = Path(spec["path"]) - - # CMR has its own dedicated script → skip - if emb == EmbodimentTag.CMR_VERSIUS.value: - continue - - key = (str(dp), emb) - if key in seen: - continue - seen.add(key) - jobs.append((dp, emb, spec.get("exclude_splits", None))) - - if not jobs: - print("No non-CMR datasets found in OPEN_H_DATASET_SPECS.") - return - - print("#" * 80) - print(f"OPEN-H BATCH MODE: {len(jobs)} dataset(s) to process") - print("#" * 80) - for i, (dp, emb, excl) in enumerate(jobs, 1): - excl_str = f" exclude={excl}" if excl else "" - print(f" [{i:2d}] [{emb:<22s}] {dp.name}{excl_str}") - print("#" * 80) - - total_start = time.time() - results: dict[str, str] = {} - - for i, (dp, emb, excl) in enumerate(jobs, 1): - print(f"\n{'#' * 80}") - print(f"# [{i}/{len(jobs)}] embodiment={emb} path={dp.name}") - if excl: - print(f"# exclude_splits={excl}") - print(f"{'#' * 80}") - - if not dp.exists(): - print(f" SKIPPED — path does not exist: {dp}") - results[f"{emb}/{dp.name}"] = "SKIPPED (path missing)" - continue - - try: - ok = process_single_dataset( - dataset_path=dp, - embodiment=emb, - num_frames=args.num_frames, - max_samples=args.max_samples, - output_filename=args.output_filename, - exclude_splits=excl, - num_workers=args.num_workers, - timestep_interval_override=args.timestep_interval, - ) - results[f"{emb}/{dp.name}"] = "OK" if ok else "FAILED" - except Exception as e: - print(f" ERROR: {e}") - results[f"{emb}/{dp.name}"] = f"ERROR ({e})" - - elapsed = time.time() - total_start - print(f"\n{'#' * 80}") - print(f"ALL DONE — {elapsed:.1f}s total, {len(results)} dataset(s)") - print(f"{'#' * 80}") - for name, status in results.items(): - print(f" {name}: {status}") - print(f"{'#' * 80}") - - -def run_single(args): - """Process a single dataset or auto-discovered datasets (original mode).""" - if args.dataset_path: - dataset_paths = [Path(args.dataset_path)] - else: - root = Path(args.dataset_path_root) - if _is_lerobot_dataset(root): - dataset_paths = [root] - else: - dataset_paths = _discover_datasets(root) - - if not dataset_paths: - print("ERROR: No datasets found!") - return - - exclude_splits = args.exclude_splits if args.exclude_splits else None - print(f"Found {len(dataset_paths)} dataset(s) for embodiment '{args.embodiment}'") - if exclude_splits: - print(f" exclude_splits: {exclude_splits}") - - total_start = time.time() - results = {} - for i, dp in enumerate(dataset_paths, 1): - if len(dataset_paths) > 1: - print(f"\n{'#' * 80}") - print(f"# DATASET {i}/{len(dataset_paths)}: {dp.name}") - print(f"{'#' * 80}") - - ok = process_single_dataset( - dataset_path=dp, - embodiment=args.embodiment, - num_frames=args.num_frames, - max_samples=args.max_samples, - output_filename=args.output_filename, - exclude_splits=exclude_splits, - num_workers=args.num_workers, - timestep_interval_override=args.timestep_interval, - ) - results[dp.name] = "OK" if ok else "FAILED" - - elapsed = time.time() - total_start - if len(dataset_paths) > 1: - print(f"\n{'#' * 80}") - print(f"ALL DONE — {elapsed:.1f}s") - for name, status in results.items(): - print(f" {name}: {status}") - print(f"{'#' * 80}") - - -def main(): - parser = argparse.ArgumentParser( - description="Compute normalization stats for Open-H embodiments (post-transform)", - epilog=( - "Modes:\n" - " --all Process ALL datasets in OPEN_H_DATASET_SPECS\n" - " (skips CMR Versius — use compute_cmr_action_stats.py).\n" - " No --dataset-path or --embodiment needed.\n\n" - " --dataset-path + --embodiment\n" - " Process a single dataset with a specific embodiment.\n\n" - " --dataset-path-root + --embodiment\n" - " Auto-discover datasets under a root directory.\n\n" - "Examples:\n" - " # All Open-H datasets at once:\n" - " python compute_openh_action_stats.py --all\n\n" - " # Quick test (10 samples per dataset):\n" - " python compute_openh_action_stats.py --all --max-samples 10\n\n" - " # Single dataset:\n" - " python compute_openh_action_stats.py \\\n" - " --dataset-path /path/to/suturebot_2 --embodiment dvrk\n" - ), - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - - # --- mutually exclusive: --all vs --dataset-path / --dataset-path-root --- - path_group = parser.add_mutually_exclusive_group(required=True) - path_group.add_argument( - "--all", action="store_true", - help="Process every dataset in OPEN_H_DATASET_SPECS (skips CMR Versius)", - ) - path_group.add_argument("--dataset-path", type=str, - help="Path to a single LeRobot dataset") - path_group.add_argument("--dataset-path-root", type=str, - help="Root directory containing multiple LeRobot datasets") - - parser.add_argument("--embodiment", type=str, default=None, - choices=list(EMBODIMENT_REGISTRY.keys()), - help="Embodiment tag (required for --dataset-path / --dataset-path-root; " - "ignored for --all)") - parser.add_argument("--exclude-splits", type=str, nargs="+", default=None, - help="Split names from info.json to exclude (e.g., --exclude-splits fail bad_frames). " - "For --all mode, exclude_splits from OPEN_H_DATASET_SPECS are used automatically.") - parser.add_argument("--num-frames", type=int, default=13, - help="Number of video frames (default: 13 = 1 context + 12 prediction)") - parser.add_argument("--max-samples", type=int, default=None, - help="Max episodes per dataset (for quick testing)") - parser.add_argument("--num-workers", type=int, default=None, - help=f"Number of parallel workers (default: min(cpu_count, {MAX_WORKERS}))") - parser.add_argument("--output-filename", type=str, default="stats_cosmos.json", - help="Output filename in meta/ dir (default: stats_cosmos.json)") - parser.add_argument("--timestep-interval", type=int, default=None, - help="Override the 'timestep_interval' (action stride) that is otherwise " - "read from EMBODIMENT_REGISTRY in groot_configs.py. Useful when you " - "want to compute stats for a different effective training rate than " - "the registry default (e.g. use 3 instead of 5 on a 30Hz dVRK dataset " - "to target 10 Hz effective). IMPORTANT: if you override this here, " - "the training pipeline MUST use the same value or the resulting " - "stats_cosmos.json will not match the distribution the model sees.") - args = parser.parse_args() - - # Validate: --dataset-path / --dataset-path-root require --embodiment - if not args.all and args.embodiment is None: - parser.error("--embodiment is required when using --dataset-path or --dataset-path-root") - - if args.all: - run_all(args) - else: - run_single(args) - - -if __name__ == "__main__": - main() +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Compute per-key normalization statistics for any Open-H embodiment. + +Unlike compute_cmr_action_stats.py (which is CMR-specific with hardcoded raw +indices, clutch filtering, and motion scaling), this script is GENERIC: it +instantiates the real transform pipeline (GenericRelativeActionTransform) for +any embodiment registered in EMBODIMENT_REGISTRY and collects statistics on +the TRANSFORMED action/state output. + +This guarantees that the statistics exactly match what the training pipeline +produces, regardless of the embodiment's delta conversion (rel_xyz_rot6d, +relative, delta, or absolute). + +Output: + meta/stats_cosmos.json — per-key statistics in the same format as + stats_cosmos-44D.json (action.psm1_pose → {mean, std, min, max, q01, q99}) + +For CMR Versius, continue using compute_cmr_action_stats.py (it handles +the additional clutch filtering and motion scaling that are CMR-specific). + +Usage: + # All Open-H datasets at once (reads OPEN_H_DATASET_SPECS, includes exclude_splits): + python compute_openh_action_stats.py --all + + # Quick test (10 samples per dataset): + python compute_openh_action_stats.py --all --max-samples 10 + + # Single dataset: + python compute_openh_action_stats.py \\ + --dataset-path /path/to/lerobot/dataset \\ + --embodiment dvrk + + # Single dataset with episode filtering: + python compute_openh_action_stats.py \\ + --dataset-path /path/to/stanford/Needle_Transfer \\ + --embodiment dvrk_stanford_real \\ + --exclude-splits fail bad_frames + + # All dVRK datasets under a root: + python compute_openh_action_stats.py \\ + --dataset-path-root /path/to/jhu \\ + --embodiment jhu_dvrk_mono +""" + +import argparse +import json +import os +import time +import warnings +from concurrent.futures import ProcessPoolExecutor, as_completed +from functools import partial +from pathlib import Path + +# Suppress noisy torchvision video deprecation warnings +warnings.filterwarnings("ignore", message=".*video decoding and encoding capabilities of torchvision.*") + +import numpy as np +import pandas as pd +from tqdm import tqdm + +from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.dataset import ( + resolve_excluded_episode_indices, +) +from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.embodiment_tags import EmbodimentTag +from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.data.transform.state_action import ( + convert_to_hybrid_relative, +) +from cosmos_predict2._src.predict2.action.datasets.gr00t_dreams.groot_configs import ( + EMBODIMENT_REGISTRY, + OPEN_H_DATASET_SPECS, +) + +# Maximum parallel workers +MAX_WORKERS = 64 + + +# ============================================================================ +# Statistics helpers (streaming, memory-efficient) +# ============================================================================ + + +class StreamingStats: + """Memory-efficient streaming statistics using Welford's algorithm + reservoir sampling.""" + + def __init__(self, num_dims: int, reservoir_size: int = 2_000_000): + self.num_dims = num_dims + self.reservoir_size = reservoir_size + self.count = 0 + self.mean = np.zeros(num_dims, dtype=np.float64) + self.M2 = np.zeros(num_dims, dtype=np.float64) + self.min_vals = np.full(num_dims, np.inf, dtype=np.float64) + self.max_vals = np.full(num_dims, -np.inf, dtype=np.float64) + self.reservoir = None + self.reservoir_count = 0 + + def update(self, batch: np.ndarray): + if batch.shape[0] == 0: + return + n = batch.shape[0] + batch_mean = np.mean(batch, axis=0) + batch_var = np.var(batch, axis=0, ddof=0) + batch_M2 = batch_var * n + self.min_vals = np.minimum(self.min_vals, np.min(batch, axis=0)) + self.max_vals = np.maximum(self.max_vals, np.max(batch, axis=0)) + if self.count == 0: + self.mean = batch_mean + self.M2 = batch_M2 + self.count = n + else: + n_total = self.count + n + delta = batch_mean - self.mean + self.mean = (self.count * self.mean + n * batch_mean) / n_total + self.M2 = self.M2 + batch_M2 + delta**2 * self.count * n / n_total + self.count = n_total + # Reservoir sampling + if self.reservoir is None: + self.reservoir = batch[: self.reservoir_size].copy() + self.reservoir_count = n + elif len(self.reservoir) < self.reservoir_size: + space = self.reservoir_size - len(self.reservoir) + self.reservoir = np.vstack([self.reservoir, batch[:space]]) + self.reservoir_count += n + else: + # Algorithm R replacement + for row in batch: + self.reservoir_count += 1 + j = np.random.randint(0, self.reservoir_count) + if j < self.reservoir_size: + self.reservoir[j] = row + + def get_stats(self) -> dict: + if self.count == 0: + z = [0.0] * self.num_dims + return {"mean": z, "std": z, "min": z, "max": z, "q01": z, "q99": z} + std = np.sqrt(self.M2 / self.count) + if self.reservoir is not None and len(self.reservoir) > 0: + q01 = np.quantile(self.reservoir, 0.01, axis=0) + q99 = np.quantile(self.reservoir, 0.99, axis=0) + else: + q01 = self.min_vals + q99 = self.max_vals + return { + "mean": self.mean.tolist(), + "std": std.tolist(), + "min": self.min_vals.tolist(), + "max": self.max_vals.tolist(), + "q01": q01.tolist(), + "q99": q99.tolist(), + } + + +# ============================================================================ +# Main processing +# ============================================================================ + + +def _is_lerobot_dataset(path: Path) -> bool: + return (path / "data").is_dir() and (path / "meta").is_dir() and any((path / "data").rglob("*.parquet")) + + +def _discover_datasets(root: Path) -> list[Path]: + datasets = [] + for child in sorted(root.iterdir()): + if child.is_dir() and _is_lerobot_dataset(child): + datasets.append(child) + return datasets + + +def _process_episode_parquet( + parquet_path: Path, + modality_meta: dict, + action_key_configs: dict, + action_delta_indices: list[int], + state_delta_indices: list[int], + action_keys: list[str], + state_keys: list[str], +) -> tuple[np.ndarray, np.ndarray, str | None]: + """Worker function: process one parquet file (episode) without video decoding. + + Reads state/action arrays from parquet, applies delta conversion, concatenates + per-key arrays, and returns the result for streaming stats collection. + + This runs in a subprocess via ProcessPoolExecutor for parallelism. + + Returns: + (action_array, state_array, warning_or_None) + action_array shape: (N_samples * T_action, action_dim) + state_array shape: (N_samples, state_dim) + """ + try: + df = pd.read_parquet(parquet_path) + T = len(df) + + # Maximum delta for action horizon + max_action_delta = max(action_delta_indices) if action_delta_indices else 0 + effective_length = max(0, T - max_action_delta) + if effective_length == 0: + return np.empty((0, 0)), np.empty((0, 0)), None + + # Extract per-key arrays from the flat parquet columns using modality metadata + def extract_key_data(key: str, df: pd.DataFrame) -> np.ndarray: + """Extract a named key's data from parquet using modality.json metadata.""" + modality, subkey = key.split(".", 1) + meta = modality_meta.get(modality, {}).get(subkey) + if meta is None: + raise KeyError(f"{key} config not found") + original_col = meta.get("original_key") + if original_col is None: + # Default: observation.state for state, action for action + original_col = "observation.state" if modality == "state" else "action" + start, end = meta["start"], meta["end"] + col_data = np.stack(df[original_col].values) # (T, D_flat) + return col_data[:, start:end].astype(np.float32) # (T, key_dim) + + # Process all valid starting indices + all_action_rows = [] + all_state_rows = [] + n_skipped = 0 + + for base_idx in range(effective_length): + # --- State at t=0 --- + state_parts = [] + for key in state_keys: + s_idx = base_idx + state_delta_indices[0] # delta_indices=[0] + s_idx = max(0, min(s_idx, T - 1)) + key_data = extract_key_data(key, df) + state_parts.append(key_data[s_idx]) + if state_parts: + all_state_rows.append(np.concatenate(state_parts)) + + # --- Action over horizon --- + action_timestep_rows = [] + for a_delta in action_delta_indices: + a_idx = base_idx + a_delta + a_idx = max(0, min(a_idx, T - 1)) + action_parts = [] + for key in action_keys: + key_data = extract_key_data(key, df) + action_parts.append(key_data[a_idx]) + action_timestep_rows.append(np.concatenate(action_parts)) + action_horizon = np.stack(action_timestep_rows) # (T_action, action_dim_raw) + + # --- Apply delta conversion per key --- + # We need to apply the same GenericRelativeActionTransform logic + # but directly on numpy arrays (no torch, no dataset wrapper). + offset = 0 + converted_parts = [] + state_row = all_state_rows[-1] if all_state_rows else None + skip_sample = False + + for key in action_keys: + cfg = action_key_configs.get(key) + key_data_raw = extract_key_data(key, df) + raw_dim = key_data_raw.shape[1] + + # Extract this key's action horizon data + key_action = action_horizon[:, offset : offset + raw_dim] + + if cfg is not None and cfg.rep == "rel_xyz_rot6d": + # Get reference state pose + ref_key = cfg.state_key + if ref_key and state_row is not None: + # Find the state key's slice in the concatenated state + s_off = 0 + ref_pose = None + for sk in state_keys: + sk_data = extract_key_data(sk, df) + sk_dim = sk_data.shape[1] + if sk == ref_key: + ref_pose = state_row[s_off : s_off + sk_dim] + break + s_off += sk_dim + + if ref_pose is not None: + # Guard: check for zero-norm quaternions (invalid data / + # padding at episode boundaries). Scipy's Rotation.from_quat() + # crashes on [0,0,0,0]. + if cfg.input_rotation_format == "quat": + quat_slice = key_action[:, 3:7] + norms = np.linalg.norm(quat_slice, axis=-1) + if np.any(norms < 1e-8): + skip_sample = True + break + if cfg.reference_rotation_format == "quat": + ref_quat = ref_pose[3:7] if len(ref_pose) >= 7 else ref_pose[3:] + if np.linalg.norm(ref_quat) < 1e-8: + skip_sample = True + break + + key_action = convert_to_hybrid_relative( + action_data=key_action, + eef_pose=ref_pose, + input_rotation_format=cfg.input_rotation_format, + reference_rotation_format=cfg.reference_rotation_format, + input_quat_order=cfg.input_quat_order, + reference_quat_order=cfg.reference_quat_order, + ) # (T_action, 9) + + elif cfg is not None and cfg.rep == "relative": + # Joint-space subtraction + ref_key = cfg.state_key + if ref_key and state_row is not None: + s_off = 0 + ref_val = None + for sk in state_keys: + sk_data = extract_key_data(sk, df) + sk_dim = sk_data.shape[1] + if sk == ref_key: + ref_val = state_row[s_off : s_off + sk_dim] + break + s_off += sk_dim + if ref_val is not None: + key_action = key_action - ref_val + + # delta / absolute: pass through unchanged + converted_parts.append(key_action) + offset += raw_dim + + if skip_sample: + # Remove the state row we just added (it corresponds to this skipped sample) + if all_state_rows: + all_state_rows.pop() + n_skipped += 1 + continue + + converted_action = np.concatenate(converted_parts, axis=-1) # (T_action, action_dim) + all_action_rows.append(converted_action) + + if not all_action_rows: + warn = None + if n_skipped > 0: + warn = f"{parquet_path.name}: all {effective_length} samples skipped ({n_skipped} had zero-norm quaternions)" + return np.empty((0, 0)), np.empty((0, 0)), warn + + actions = np.concatenate(all_action_rows, axis=0) # (N*T_action, action_dim) + states = np.stack(all_state_rows) if all_state_rows else np.empty((0, 0)) + + warn = None + if n_skipped > 0: + warn = ( + f"{parquet_path.name}: {n_skipped}/{effective_length} samples skipped " + f"(zero-norm quaternions), {len(all_action_rows)} valid" + ) + return actions, states, warn + + except Exception as e: + return np.empty((0, 0)), np.empty((0, 0)), f"Error processing {parquet_path.name}: {type(e).__name__}: {e}" + + +def _load_modality_meta(dataset_path: Path, modality_filename: str) -> dict: + """Load modality.json and return a simplified {modality: {subkey: {start, end, original_key}}} dict.""" + modality_path = dataset_path / modality_filename + if not modality_path.exists(): + raise FileNotFoundError(f"Modality file not found: {modality_path}") + + with open(modality_path, "r") as f: + raw = json.load(f) + + result: dict[str, dict] = {} + for modality in ["state", "action"]: + result[modality] = {} + if modality not in raw: + continue + for subkey, meta in raw[modality].items(): + result[modality][subkey] = { + "start": meta.get("start", 0), + "end": meta.get("end", 1), + "original_key": meta.get("original_key"), + } + return result + + +def process_single_dataset( + dataset_path: Path, + embodiment: str, + num_frames: int, + max_samples: int | None, + output_filename: str, + exclude_splits: list[str] | None = None, + num_workers: int | None = None, + timestep_interval_override: int | None = None, +): + """Process one dataset using parallel episode processing (no video decoding). + + Each parquet file (episode) is processed by a worker subprocess that: + 1. Reads state/action arrays from parquet (fast, no video) + 2. Applies the delta conversion (rel_xyz_rot6d, relative, etc.) + 3. Returns the transformed arrays for streaming stats collection + + This is ~100x faster than the sequential video-decoding approach. + + Args: + dataset_path: Path to the LeRobot dataset. + embodiment: Embodiment tag string. + num_frames: Number of video frames (e.g. 13). + max_samples: Max episodes to process (for testing). None = all. + output_filename: Output filename in meta/ dir. + exclude_splits: Split names from info.json to exclude. + num_workers: Parallel workers (default: min(cpu_count, MAX_WORKERS)). + timestep_interval_override: If given (int > 0), override the + ``timestep_interval`` value read from EMBODIMENT_REGISTRY for + this run. Useful to experiment with a different effective training + rate *without* editing ``groot_configs.py``. IMPORTANT: if you use + this, training must run with the same stride, otherwise + ``stats_cosmos.json`` will not match the distribution the model + sees. + """ + if num_workers is None: + num_workers = min(os.cpu_count() or 8, MAX_WORKERS) + + reg = EMBODIMENT_REGISTRY.get(embodiment) + if reg is None: + raise ValueError(f"Unknown embodiment '{embodiment}'. Available: {list(EMBODIMENT_REGISTRY.keys())}") + + registry_timestep_interval = reg["timestep_interval"] + if timestep_interval_override is not None: + if timestep_interval_override < 1: + raise ValueError(f"--timestep-interval must be >= 1, got {timestep_interval_override}") + timestep_interval = int(timestep_interval_override) + else: + timestep_interval = registry_timestep_interval + + print("=" * 80) + print(f"COMPUTING OPEN-H ACTION STATS — {dataset_path.name}") + print(f" embodiment: {embodiment}") + print(f" num_frames: {num_frames}") + if timestep_interval_override is not None and timestep_interval_override != registry_timestep_interval: + print( + f" timestep_interval: {timestep_interval} " + f"(OVERRIDE; registry default for '{embodiment}' is {registry_timestep_interval})" + ) + else: + print(f" timestep_interval: {timestep_interval} (from EMBODIMENT_REGISTRY)") + print(f" workers: {num_workers}") + if exclude_splits: + print(f" exclude_splits: {exclude_splits}") + print("=" * 80) + + num_action_frames = num_frames - 1 + action_delta_indices = list(range(0, num_action_frames * timestep_interval, timestep_interval)) + state_delta_indices = [0] + + action_keys = reg["action_keys"] + state_keys = reg["state_keys"] + action_key_configs = reg.get("action_key_configs", {}) + modality_filename = reg.get("modality_filename", "meta/modality.json") + + # Load modality metadata (maps key names to parquet column indices) + modality_meta = _load_modality_meta(dataset_path, modality_filename) + + # Discover parquet files (one per episode) + parquet_files = sorted(dataset_path.glob("data/*/*.parquet")) + if not parquet_files: + print(f" ERROR: No parquet files found in {dataset_path / 'data'}") + return False + + # Apply exclude_splits filtering at episode level + if exclude_splits: + excluded_ids = resolve_excluded_episode_indices(dataset_path, exclude_splits) + # Parse episode index from filename (e.g., episode_000042.parquet → 42) + filtered = [] + for pf in parquet_files: + try: + ep_idx = int(pf.stem.split("_")[-1]) + except ValueError: + filtered.append(pf) # Can't parse → keep + continue + if ep_idx not in excluded_ids: + filtered.append(pf) + n_excluded = len(parquet_files) - len(filtered) + print(f" exclude_splits: removed {n_excluded} episodes, {len(filtered)} remaining") + parquet_files = filtered + + if max_samples is not None: + parquet_files = parquet_files[:max_samples] + + print(f" Processing {len(parquet_files)} episodes with {num_workers} workers...") + + # Create partial function with fixed arguments for the worker + worker_fn = partial( + _process_episode_parquet, + modality_meta=modality_meta, + action_key_configs=action_key_configs, + action_delta_indices=action_delta_indices, + state_delta_indices=state_delta_indices, + action_keys=action_keys, + state_keys=state_keys, + ) + + # Process episodes in parallel + action_tracker = None + state_tracker = None + total_action_samples = 0 + total_state_samples = 0 + episodes_with_warnings = 0 + episodes_empty = 0 + + with ProcessPoolExecutor(max_workers=num_workers) as executor: + futures = {executor.submit(worker_fn, pf): pf for pf in parquet_files} + + for future in tqdm(as_completed(futures), total=len(futures), desc=f"Stats for {dataset_path.name}"): + actions, states, warning = future.result() + + if warning: + episodes_with_warnings += 1 + if episodes_with_warnings <= 10: + print(f" [WARN] {warning}") + + if actions.size == 0: + episodes_empty += 1 + continue + + # Lazy-init trackers based on first result's dimensions + if action_tracker is None: + action_dim = actions.shape[1] + action_tracker = StreamingStats(action_dim) + print(f" Action dim (post-transform): {action_dim}") + if state_tracker is None and states.size > 0: + state_dim = states.shape[1] + state_tracker = StreamingStats(state_dim) + print(f" State dim (post-transform): {state_dim}") + + action_tracker.update(actions.astype(np.float64)) + total_action_samples += len(actions) + + if state_tracker is not None and states.size > 0: + state_tracker.update(states.astype(np.float64)) + total_state_samples += len(states) + + if episodes_with_warnings > 10: + print(f" [WARN] {episodes_with_warnings} episodes had warnings (showing first 10)") + if episodes_empty > 0: + print( + f" [INFO] {episodes_empty}/{len(parquet_files)} episodes produced no valid samples " + f"(zero-norm quaternions or empty episodes)" + ) + + if action_tracker is None or action_tracker.count == 0: + print(f" ERROR: No valid samples found! ({episodes_empty} empty, {episodes_with_warnings} warnings)") + return False + + valid_episodes = len(parquet_files) - episodes_empty + print(f"\nValid episodes: {valid_episodes}/{len(parquet_files)}") + print(f"Total action timesteps: {total_action_samples:,}") + print(f"Total state samples: {total_state_samples:,}") + + # ---------------------------------------------------------------- + # Build per-key stats by slicing the concatenated action vector + # ---------------------------------------------------------------- + action_global = action_tracker.get_stats() + stats: dict = {} + + # Determine per-key dimensions from modality metadata + action configs + # For rel_xyz_rot6d keys, output is 9D (not raw 7D) + action_key_dims: dict[str, tuple[int, int]] = {} + offset = 0 + for key in action_keys: + modality, subkey = key.split(".", 1) + meta = modality_meta.get(modality, {}).get(subkey) + if meta is None: + continue + raw_dim = meta["end"] - meta["start"] + + cfg = action_key_configs.get(key) + if cfg is not None and cfg.rep == "rel_xyz_rot6d": + out_dim = 9 # xyz(3) + rot6d(6) + else: + out_dim = raw_dim + + action_key_dims[key] = (offset, offset + out_dim) + offset += out_dim + + state_key_dims: dict[str, tuple[int, int]] = {} + offset = 0 + for key in state_keys: + modality, subkey = key.split(".", 1) + meta = modality_meta.get(modality, {}).get(subkey) + if meta is None: + continue + dim = meta["end"] - meta["start"] + state_key_dims[key] = (offset, offset + dim) + offset += dim + + # Extract per-key stats + for key, (s, e) in action_key_dims.items(): + key_stats = { + "mean": action_global["mean"][s:e], + "std": action_global["std"][s:e], + "min": action_global["min"][s:e], + "max": action_global["max"][s:e], + } + if action_tracker.reservoir is not None: + res = action_tracker.reservoir[:, s:e] + key_stats["q01"] = np.quantile(res, 0.01, axis=0).tolist() + key_stats["q99"] = np.quantile(res, 0.99, axis=0).tolist() + else: + key_stats["q01"] = key_stats["min"] + key_stats["q99"] = key_stats["max"] + stats[key] = key_stats + + if state_tracker is not None: + state_global = state_tracker.get_stats() + for key, (s, e) in state_key_dims.items(): + key_stats = { + "mean": state_global["mean"][s:e], + "std": state_global["std"][s:e], + "min": state_global["min"][s:e], + "max": state_global["max"][s:e], + } + if state_tracker.reservoir is not None: + res = state_tracker.reservoir[:, s:e] + key_stats["q01"] = np.quantile(res, 0.01, axis=0).tolist() + key_stats["q99"] = np.quantile(res, 0.99, axis=0).tolist() + else: + key_stats["q01"] = key_stats["min"] + key_stats["q99"] = key_stats["max"] + stats[key] = key_stats + + # Global concatenated stats for convenience + stats["action"] = action_global + if state_tracker is not None: + stats["state"] = state_tracker.get_stats() + + # ---------------------------------------------------------------- + # Stamp provenance metadata for downstream guards. + # ---------------------------------------------------------------- + # We record the exact ``timestep_interval`` used to compute these + # statistics as a top-level integer so ``LeRobotSingleDataset`` can + # assert it still matches EMBODIMENT_REGISTRY at training time. + # Since all real stat entries are per-key dicts (``"action.psm1_pose"`` + # etc.), an int at the top level never collides with a stat, and the + # existing validation loop in ``dataset.py`` (``if isinstance(stat, int): + # continue``) already skips it. + stats["timestep_interval"] = int(timestep_interval) + + # ---------------------------------------------------------------- + # Write to disk + # ---------------------------------------------------------------- + out_path = dataset_path / "meta" / output_filename + out_path.parent.mkdir(parents=True, exist_ok=True) + with open(out_path, "w") as f: + json.dump(stats, f, indent=2) + + print(f"\nSaved stats to {out_path}") + print(f"Per-key action stats: {list(action_key_dims.keys())}") + if state_key_dims: + print(f"Per-key state stats: {list(state_key_dims.keys())}") + + for key, (s, e) in action_key_dims.items(): + dim = e - s + print( + f" {key} ({dim}D): mean_abs={np.mean(np.abs(action_global['mean'][s:e])):.6f}, " + f"std_mean={np.mean(action_global['std'][s:e]):.6f}" + ) + + return True + + +def _resolve_embodiment_string(spec_embodiment) -> str: + """Normalise the 'embodiment' field from OPEN_H_DATASET_SPECS to a plain string.""" + if isinstance(spec_embodiment, EmbodimentTag): + return spec_embodiment.value + return str(spec_embodiment) + + +def run_all(args): + """Process every dataset listed in OPEN_H_DATASET_SPECS (--all mode). + + Skips CMR Versius entries (they use compute_cmr_action_stats.py and + stats_cosmos-44D.json instead). + + Each spec's ``exclude_splits`` is forwarded so that the same episodes + excluded during training are also excluded from statistics computation. + """ + # Deduplicate: multiple specs may share the same (path, embodiment). + # Keep the first occurrence's exclude_splits. + seen: set[tuple[str, str]] = set() + jobs: list[tuple[Path, str, list[str] | None]] = [] + + for spec in OPEN_H_DATASET_SPECS: + emb = _resolve_embodiment_string(spec["embodiment"]) + dp = Path(spec["path"]) + + # CMR has its own dedicated script → skip + if emb == EmbodimentTag.CMR_VERSIUS.value: + continue + + key = (str(dp), emb) + if key in seen: + continue + seen.add(key) + jobs.append((dp, emb, spec.get("exclude_splits", None))) + + if not jobs: + print("No non-CMR datasets found in OPEN_H_DATASET_SPECS.") + return + + print("#" * 80) + print(f"OPEN-H BATCH MODE: {len(jobs)} dataset(s) to process") + print("#" * 80) + for i, (dp, emb, excl) in enumerate(jobs, 1): + excl_str = f" exclude={excl}" if excl else "" + print(f" [{i:2d}] [{emb:<22s}] {dp.name}{excl_str}") + print("#" * 80) + + total_start = time.time() + results: dict[str, str] = {} + + for i, (dp, emb, excl) in enumerate(jobs, 1): + print(f"\n{'#' * 80}") + print(f"# [{i}/{len(jobs)}] embodiment={emb} path={dp.name}") + if excl: + print(f"# exclude_splits={excl}") + print(f"{'#' * 80}") + + if not dp.exists(): + print(f" SKIPPED — path does not exist: {dp}") + results[f"{emb}/{dp.name}"] = "SKIPPED (path missing)" + continue + + try: + ok = process_single_dataset( + dataset_path=dp, + embodiment=emb, + num_frames=args.num_frames, + max_samples=args.max_samples, + output_filename=args.output_filename, + exclude_splits=excl, + num_workers=args.num_workers, + timestep_interval_override=args.timestep_interval, + ) + results[f"{emb}/{dp.name}"] = "OK" if ok else "FAILED" + except Exception as e: + print(f" ERROR: {e}") + results[f"{emb}/{dp.name}"] = f"ERROR ({e})" + + elapsed = time.time() - total_start + print(f"\n{'#' * 80}") + print(f"ALL DONE — {elapsed:.1f}s total, {len(results)} dataset(s)") + print(f"{'#' * 80}") + for name, status in results.items(): + print(f" {name}: {status}") + print(f"{'#' * 80}") + + +def run_single(args): + """Process a single dataset or auto-discovered datasets (original mode).""" + if args.dataset_path: + dataset_paths = [Path(args.dataset_path)] + else: + root = Path(args.dataset_path_root) + if _is_lerobot_dataset(root): + dataset_paths = [root] + else: + dataset_paths = _discover_datasets(root) + + if not dataset_paths: + print("ERROR: No datasets found!") + return + + exclude_splits = args.exclude_splits if args.exclude_splits else None + print(f"Found {len(dataset_paths)} dataset(s) for embodiment '{args.embodiment}'") + if exclude_splits: + print(f" exclude_splits: {exclude_splits}") + + total_start = time.time() + results = {} + for i, dp in enumerate(dataset_paths, 1): + if len(dataset_paths) > 1: + print(f"\n{'#' * 80}") + print(f"# DATASET {i}/{len(dataset_paths)}: {dp.name}") + print(f"{'#' * 80}") + + ok = process_single_dataset( + dataset_path=dp, + embodiment=args.embodiment, + num_frames=args.num_frames, + max_samples=args.max_samples, + output_filename=args.output_filename, + exclude_splits=exclude_splits, + num_workers=args.num_workers, + timestep_interval_override=args.timestep_interval, + ) + results[dp.name] = "OK" if ok else "FAILED" + + elapsed = time.time() - total_start + if len(dataset_paths) > 1: + print(f"\n{'#' * 80}") + print(f"ALL DONE — {elapsed:.1f}s") + for name, status in results.items(): + print(f" {name}: {status}") + print(f"{'#' * 80}") + + +def main(): + parser = argparse.ArgumentParser( + description="Compute normalization stats for Open-H embodiments (post-transform)", + epilog=( + "Modes:\n" + " --all Process ALL datasets in OPEN_H_DATASET_SPECS\n" + " (skips CMR Versius — use compute_cmr_action_stats.py).\n" + " No --dataset-path or --embodiment needed.\n\n" + " --dataset-path + --embodiment\n" + " Process a single dataset with a specific embodiment.\n\n" + " --dataset-path-root + --embodiment\n" + " Auto-discover datasets under a root directory.\n\n" + "Examples:\n" + " # All Open-H datasets at once:\n" + " python compute_openh_action_stats.py --all\n\n" + " # Quick test (10 samples per dataset):\n" + " python compute_openh_action_stats.py --all --max-samples 10\n\n" + " # Single dataset:\n" + " python compute_openh_action_stats.py \\\n" + " --dataset-path /path/to/suturebot_2 --embodiment dvrk\n" + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + # --- mutually exclusive: --all vs --dataset-path / --dataset-path-root --- + path_group = parser.add_mutually_exclusive_group(required=True) + path_group.add_argument( + "--all", + action="store_true", + help="Process every dataset in OPEN_H_DATASET_SPECS (skips CMR Versius)", + ) + path_group.add_argument("--dataset-path", type=str, help="Path to a single LeRobot dataset") + path_group.add_argument("--dataset-path-root", type=str, help="Root directory containing multiple LeRobot datasets") + + parser.add_argument( + "--embodiment", + type=str, + default=None, + choices=list(EMBODIMENT_REGISTRY.keys()), + help="Embodiment tag (required for --dataset-path / --dataset-path-root; ignored for --all)", + ) + parser.add_argument( + "--exclude-splits", + type=str, + nargs="+", + default=None, + help="Split names from info.json to exclude (e.g., --exclude-splits fail bad_frames). " + "For --all mode, exclude_splits from OPEN_H_DATASET_SPECS are used automatically.", + ) + parser.add_argument( + "--num-frames", type=int, default=13, help="Number of video frames (default: 13 = 1 context + 12 prediction)" + ) + parser.add_argument("--max-samples", type=int, default=None, help="Max episodes per dataset (for quick testing)") + parser.add_argument( + "--num-workers", + type=int, + default=None, + help=f"Number of parallel workers (default: min(cpu_count, {MAX_WORKERS}))", + ) + parser.add_argument( + "--output-filename", + type=str, + default="stats_cosmos.json", + help="Output filename in meta/ dir (default: stats_cosmos.json)", + ) + parser.add_argument( + "--timestep-interval", + type=int, + default=None, + help="Override the 'timestep_interval' (action stride) that is otherwise " + "read from EMBODIMENT_REGISTRY in groot_configs.py. Useful when you " + "want to compute stats for a different effective training rate than " + "the registry default (e.g. use 3 instead of 5 on a 30Hz dVRK dataset " + "to target 10 Hz effective). IMPORTANT: if you override this here, " + "the training pipeline MUST use the same value or the resulting " + "stats_cosmos.json will not match the distribution the model sees.", + ) + args = parser.parse_args() + + # Validate: --dataset-path / --dataset-path-root require --embodiment + if not args.all and args.embodiment is None: + parser.error("--embodiment is required when using --dataset-path or --dataset-path-root") + + if args.all: + run_all(args) + else: + run_single(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/extract_jhu_inference_manifest.py b/scripts/extract_jhu_inference_manifest.py index 2fbe5e2..fdd3b70 100644 --- a/scripts/extract_jhu_inference_manifest.py +++ b/scripts/extract_jhu_inference_manifest.py @@ -152,8 +152,7 @@ def main() -> int: args = parse_arguments() if args.num_frames < 2 or (args.num_frames - 1) % 4 != 0: print( - f"ERROR: --num-frames must be >= 2 and satisfy " - f"(num_frames-1) % 4 == 0; got {args.num_frames}", + f"ERROR: --num-frames must be >= 2 and satisfy (num_frames-1) % 4 == 0; got {args.num_frames}", file=sys.stderr, ) return 2 @@ -182,16 +181,12 @@ def main() -> int: ) except ImportError as error: print( - f"ERROR: failed to import runtime dependencies: {error}\n" - "Run inside the Cosmos-Predict2.5 environment.", + f"ERROR: failed to import runtime dependencies: {error}\nRun inside the Cosmos-Predict2.5 environment.", file=sys.stderr, ) return 2 - subset_to_path = { - Path(spec["path"]).name: spec["path"] - for spec in JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS - } + subset_to_path = {Path(spec["path"]).name: spec["path"] for spec in JHU_DVRK_MONO_FINETUNE_TRAIN_DATASET_SPECS} num_actions = args.num_frames - 1 raw_chunk_stride = num_actions * args.timestep_interval output_dir = Path(args.output_dir).expanduser().resolve() @@ -304,9 +299,7 @@ def fetch_window(episode_id: int, base_index: int): "episode_specs": episode_specs, "manifest_file": manifest_path.name, } - (output_dir / f"{args.tag}_inference_manifest_provenance.json").write_text( - json.dumps(provenance, indent=2) - ) + (output_dir / f"{args.tag}_inference_manifest_provenance.json").write_text(json.dumps(provenance, indent=2)) print(f"Wrote {len(manifest_entries)} entries to {manifest_path}") return 0 diff --git a/train_scripts/tabletop/01_train_short_teacher_h13.sh b/train_scripts/tabletop/01_train_short_teacher_h13.sh index 8ded7ff..117cb1a 100755 --- a/train_scripts/tabletop/01_train_short_teacher_h13.sh +++ b/train_scripts/tabletop/01_train_short_teacher_h13.sh @@ -1,4 +1,19 @@ #!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Reference stage 1: 13-frame bidirectional tabletop teacher. # Add site-specific --account/--partition directives if required. #SBATCH --job-name=tabletop-teacher-h13 diff --git a/train_scripts/tabletop/02_fine_anneal_short_teacher_h13.sh b/train_scripts/tabletop/02_fine_anneal_short_teacher_h13.sh index 433f6f0..5262ffa 100755 --- a/train_scripts/tabletop/02_fine_anneal_short_teacher_h13.sh +++ b/train_scripts/tabletop/02_fine_anneal_short_teacher_h13.sh @@ -1,4 +1,19 @@ #!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Reference stage 1b: 4k cosine fine anneal from short-teacher iter 16,000. #SBATCH --job-name=tabletop-teacher-h13-anneal #SBATCH --nodes=8 diff --git a/train_scripts/tabletop/03_train_long_teacher_h73.sh b/train_scripts/tabletop/03_train_long_teacher_h73.sh index fe5f75a..7d8b1d2 100755 --- a/train_scripts/tabletop/03_train_long_teacher_h73.sh +++ b/train_scripts/tabletop/03_train_long_teacher_h73.sh @@ -1,4 +1,19 @@ #!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Reference stage 2: 73-frame teacher warm-started from the annealed h13 teacher. #SBATCH --job-name=tabletop-teacher-h73 #SBATCH --nodes=8 diff --git a/train_scripts/tabletop/04_phase0_teacher_cache_h73.sh b/train_scripts/tabletop/04_phase0_teacher_cache_h73.sh index b42260b..a941ccc 100755 --- a/train_scripts/tabletop/04_phase0_teacher_cache_h73.sh +++ b/train_scripts/tabletop/04_phase0_teacher_cache_h73.sh @@ -1,4 +1,19 @@ #!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Reference stage 3: one-node, eight-rank Phase 0 teacher cache generation. #SBATCH --job-name=tabletop-phase0-h73 #SBATCH --nodes=1 diff --git a/train_scripts/tabletop/05_warmup_student_h73.sh b/train_scripts/tabletop/05_warmup_student_h73.sh index 4d93e7c..d6f2948 100755 --- a/train_scripts/tabletop/05_warmup_student_h73.sh +++ b/train_scripts/tabletop/05_warmup_student_h73.sh @@ -1,4 +1,19 @@ #!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Reference stage 4: 8-node causal-student warmup, state_t=19. #SBATCH --job-name=tabletop-warmup-h73 #SBATCH --nodes=8 diff --git a/train_scripts/tabletop/06_self_forcing_h73.sh b/train_scripts/tabletop/06_self_forcing_h73.sh index 94d821d..19ec5b4 100755 --- a/train_scripts/tabletop/06_self_forcing_h73.sh +++ b/train_scripts/tabletop/06_self_forcing_h73.sh @@ -1,4 +1,19 @@ #!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Reference stage 5: 8-node Self Forcing distillation, 3,000 iterations. #SBATCH --job-name=tabletop-self-forcing-h73 #SBATCH --nodes=8 diff --git a/train_scripts/tabletop/README.md b/train_scripts/tabletop/README.md index 956b627..4798db0 100644 --- a/train_scripts/tabletop/README.md +++ b/train_scripts/tabletop/README.md @@ -40,4 +40,4 @@ Keep these invariants synchronized: - warmup: 20,000-iteration ceiling; reference SF initialization at 18,000; - Self Forcing: 3,000 iterations. -The Phase 0 script is resumable because existing complete artifact quartets are skipped. Before warmup, verify that `latents/`, `images/`, `actions/`, and `videos/` contain the same index set; do not rely only on SLURM completion. \ No newline at end of file +The Phase 0 script is resumable because existing complete artifact quartets are skipped. Before warmup, verify that `latents/`, `images/`, `actions/`, and `videos/` contain the same index set; do not rely only on SLURM completion.