Add DROID action-policy SFT recipe (Cosmos3-Nano, joint_pos) (NVIDIA#24)

## Summary

Adds a **DROID action-policy SFT recipe** for `nvidia/Cosmos3-Nano`,
mirroring the internal `droid_lerobot_8b` policy run, so users can
post-train the action-generation + action heads on DROID (LeRobot v3.0)
data.

## What's included

- **`data/vfm/action/datasets/droid_lerobot_dataset.py`** — DROID
LeRobot dataset: compact columnar load + episode-aware windowing
(replaces an eager full-table materialization), plus `joint_pos` (8D: 7
joints + gripper) and `use_state` support.
- **`data/vfm/action/datasets/action_sft_dataset.py`** (new) —
`get_action_droid_sft_dataset(...)` wrapping the dataset through
`ActionTransformPipeline`.
- **`configs/.../action/posttrain_config/action_policy_droid_nano.py`**
(new) — registered `action_policy_droid_nano` experiment (Cosmos3-Nano /
8B MoT): optimizer trains gen+action heads (5× LR on action heads),
`LambdaLinear` schedule, count-based batch, res480,
`encode_exact_durations=[33]` (chunk 32 → 33 frames).
- **`checkpoint/dcp.py`** — EMA warm-start: when `keys_to_skip_loading`
excludes `net_ema.`, initialize `net_ema = net` from the base weights so
EMA starts from the init rather than zeros.
- **`examples/toml/sft_config/action_policy_droid_{nano,repro}.toml`** —
1-GPU smoke + scaled (res480) configs.
- **`examples/launch_sft_action_policy_droid.sh`** +
**`docs/action_policy_droid_posttraining.md`** — runnable launcher and
walkthrough.

## Validation

End-to-end on H200:
- **1 node / 8×H200** — dry-run + training at res480,
`max_samples_per_batch=32` (64 OOMs at 139 GiB; internal used 128 on
GB200).
- **2 nodes / 16 ranks** — HSDP `shard 8 × replicate 2`, `TRAIN_EXIT=0`.
- Recipe faithful to internal `droid_lerobot_8b`: lr 1e-4 / betas / wd,
5× action-head LR, `LambdaLinear`, shift `{256:3,480:5,720:10}`,
`concat_view`, `chunk_length=32`.

## Notes

- Count-based batch (`max_samples_per_batch`,
`max_sequence_length=None`) lives in the experiment Python — TOML cannot
express `null`, and the loader only overrides keys present in the TOML.
- Base checkpoint: convert `nvidia/Cosmos3-Nano` → DCP and pass via
`BASE_CHECKPOINT_PATH`; action heads init fresh (skipped on load).

---------

Signed-off-by: Hao Liang <haolia@nvidia.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: lfengad <liangf@nvidia.com>
Co-authored-by: Yu-Wei Chao <82182961+ychao-nvidia@users.noreply.github.com>

Author: 4 people

cosmos_framework/checkpoint/dcp.py

@@ -539,6 +539,20 @@ def load(
                                     "Ensure the model has net_ema submodule."
                                 )
                                 _state_dict[sd_key] = _state_dict[key_ema]
+                    elif warm_start and any(str(s).startswith("net_ema") for s in self.keys_to_skip_loading):
+                        # Only when net_ema.* is explicitly skipped on load (e.g. an HF->DCP
+                        # init from convert_model_to_dcp that has only net.*): the skipped
+                        # net_ema.* keep build_net() construction values (random init when
+                        # vlm_config.pretrained_weights.enabled=False), which would seed EMA
+                        # from random weights -> copy net.* -> net_ema.* so EMA starts from the
+                        # freshly-loaded init. When net_ema.* IS loaded (e.g. a training DCP
+                        # that carries a trained EMA), do NOT clobber it.
+                        log.info("Warm start: net_ema. skipped on load -> resetting net_ema = net.")
+                        for sd_key in list(_state_dict.keys()):
+                            if sd_key.startswith("net."):
+                                key_ema = "net_ema." + sd_key.removeprefix("net.")
+                                if key_ema in _state_dict:
+                                    _state_dict[key_ema] = _state_dict[sd_key]
                     results = _model_wrapper.load_state_dict(_state_dict)
                     if results is not None:
                         if len(results.missing_keys) > 0:

cosmos_framework/configs/base/config.py

@@ -98,4 +98,5 @@ def make_config() -> Config:
     # Register shipped experiments explicitly.
     import cosmos_framework.configs.base.experiment.sft.vision_sft_nano  # noqa: F401
     import cosmos_framework.configs.base.experiment.sft.vision_sft_super  # noqa: F401
+    import cosmos_framework.configs.base.experiment.action.posttrain_config.action_policy_droid_nano  # noqa: F401
     return c

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

@@ -0,0 +1,229 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""``action_policy_droid_nano`` — Cosmos3-Nano DROID action policy SFT recipe.
+
+Mirrors the vision SFT stack (PackingDataLoader + RankPartitionedDataLoader),
+but feeds the DROID action dataset (``joint_pos`` 8D + ``use_state``, raw/
+un-normalized — same as the internal ``droid_lerobot_8b_policy`` run) through
+``ActionTransformPipeline``, and trains the generation + action heads from the
+public ``nvidia/Cosmos3-Nano`` base.
+
+Usage (1 node, 8 GPU)::
+
+    DROID_ROOT=/path/to/droid_lerobot_640x360/success \\
+    BASE_CHECKPOINT_PATH=<Cosmos3-Nano DCP dir> \\
+    WAN_VAE_PATH=<Wan2.2_VAE.pth> \\
+    torchrun --nproc_per_node=8 -m cosmos_framework.scripts.train \\
+        --sft-toml examples/toml/sft_config/action_policy_droid_repro.toml
+"""
+
+import copy
+
+from hydra.core.config_store import ConfigStore
+
+from cosmos_framework.utils.lazy_config import LazyCall as L
+from cosmos_framework.utils.lazy_config import LazyDict
+
+from cosmos_framework.configs.base.experiment.sft.models.nano_model_config import NANO_MODEL_CONFIG
+from cosmos_framework.data.vfm.joint_dataloader import (
+    PackingDataLoader,
+    RankPartitionedDataLoader,
+)
+from cosmos_framework.data.vfm.action.datasets.action_sft_dataset import get_action_droid_sft_dataset
+
+cs = ConfigStore.instance()
+
+
+action_policy_droid_nano = LazyDict(
+    dict(
+        defaults=[
+            {"override /model": "mot_fsdp"},
+            {"override /data_train": None},
+            {"override /data_val": None},
+            # Match internal droid_lerobot_8b_policy: apex FusedAdam with fp32
+            # master_weights + eps 1e-8. adamw + fused + eps 1e-6 (bf16, no fp32
+            # master) under-steps the small 5x-lr action heads and leaves the action
+            # loss on a noisy high plateau; an exact-match forward/optimizer test
+            # confirmed the convergence gap was the optimizer, not the model.
+            {"override /optimizer": "fusedadamw"},
+            {"override /scheduler": "lambdalinear"},  # matches internal droid_lerobot_8b (was lambdacosine)
+            {"override /checkpoint": "s3"},
+            {
+                "override /callbacks": [
+                    "basic",
+                    "optimization",
+                    "job_monitor",
+                ]
+            },
+            {"override /ema": "power"},
+            {"override /tokenizer": "wan2pt2_tokenizer"},
+            {"override /sound_tokenizer": None},
+            {"override /cluster": None},
+            {"override /vlm_config": None},
+            {"override /ckpt_type": "dcp"},
+            "_self_",
+        ],
+        job=dict(
+            project="cosmos3",
+            group="action_sft",
+            name="action_policy_droid_nano",
+            wandb_mode="disabled",
+        ),
+        model=dict(
+            config=copy.deepcopy(NANO_MODEL_CONFIG),  # action_gen=True, max_action_dim=64
+        ),
+        optimizer=dict(
+            betas=[0.9, 0.99],
+            eps=1.0e-08,
+            fused=True,  # popped by build_optimizer for FusedAdam (fused by construction)
+            # Generation + action heads (mirrors internal droid_lerobot_8b_policy).
+            keys_to_select=[
+                "moe_gen",
+                "time_embedder",
+                "vae2llm",
+                "llm2vae",
+                "action2llm",
+                "llm2action",
+                "action_modality_embed",
+            ],
+            lr=2.0e-04,  # matches internal droid_lerobot_8b_policy submit (--lr 2e-4)
+            lr_multipliers={
+                "action2llm": 5.0,
+                "llm2action": 5.0,
+                "action_modality_embed": 5.0,
+            },
+            optimizer_type="FusedAdam",
+            weight_decay=0.05,
+        ),
+        scheduler=dict(
+            lr_scheduler_type="LambdaLinear",  # matches internal droid_lerobot_8b (was LambdaCosine)
+            cycle_lengths=[100],  # smoke: 100 iters (real run sets via TOML)
+            f_max=[0.4],
+            f_min=[0.0],
+            f_start=[0.0],
+            verbosity_interval=0,
+            warm_up_steps=[0],
+        ),
+        trainer=dict(
+            distributed_parallelism="fsdp",
+            grad_accum_iter=1,
+            logging_iter=1,
+            max_iter=100,  # smoke
+            max_val_iter=None,
+            run_validation=False,
+            run_validation_on_start=False,
+            save_zero_checkpoint=False,
+            seed=42,
+            timeout_period=999999999,
+            validation_iter=100,
+            compile_config=dict(recompile_limit=8, use_duck_shape=False),
+            cudnn=dict(benchmark=True, deterministic=False),
+            ddp=dict(broadcast_buffers=True, find_unused_parameters=False, static_graph=True),
+            grad_scaler_args=dict(enabled=False),
+            callbacks=dict(
+                dataloader_speed=dict(every_n=100, save_s3=False, step_size=1),
+                device_monitor=dict(
+                    every_n=200, log_memory_detail=True, save_s3=False, step_size=1, upload_every_n_mul=5
+                ),
+                grad_clip=dict(clip_norm=1.0, force_finite=True),  # matches internal make_8b
+                heart_beat=dict(every_n=200, save_s3=False, step_size=1, update_interval_in_minute=20),
+                iter_speed=dict(every_n=1, hit_thres=50, save_s3=False, save_s3_every_log_n=500),
+                low_precision=dict(update_iter=1),
+                manual_gc=dict(every_n=5, gc_level=1, warm_up=1),
+                param_count=dict(save_s3=False),
+                skip_nan_step=dict(max_consecutive_nan=100),
+                training_stats=dict(log_freq=100),
+            ),
+        ),
+        checkpoint=dict(
+            broadcast_via_filesystem=False,
+            dcp_async_mode_enabled=False,
+            enable_gcs_patch_in_boto3=True,
+            keys_not_to_resume=[],
+            # Skip net_ema. (→ EMA warm-start copies net→net_ema, see dcp.py) AND the
+            # action heads, so they init fresh from the base — matches internal
+            # make_8b _DEFAULT_KEYS_TO_SKIP (Cosmos3-Nano's action heads are not
+            # DROID-policy-trained).
+            keys_to_skip_loading=[
+                "net_ema.",
+                "action2llm",
+                "llm2action",
+                "action_modality_embed",
+                "action_pos_embed",
+            ],
+            load_ema_to_reg=False,
+            load_path="???",  # Cosmos3-Nano DCP dir; supply via TOML/env
+            load_training_state=False,
+            only_load_scheduler_state=False,
+            save_iter=100,
+            strict_resume=False,  # base init: tolerate key set differences
+            verbose=True,
+            hf_export=dict(
+                enabled=False,
+                export_every_n=1,
+                hf_repo_id=None,
+                upload_to_object_store=dict(bucket="", credentials="", enabled=False),
+            ),
+            jit=dict(device="cuda", dtype="bfloat16", enabled=False, input_shape=None, strict=True),
+            load_from_object_store=dict(bucket="", credentials="", enabled=False),
+            save_to_object_store=dict(bucket="", credentials="", enabled=False),
+        ),
+        dataloader_train=L(PackingDataLoader)(
+            audio_sample_rate=48000,
+            dataset_name="action_droid",
+            max_samples_per_batch=128,  # count-based batch (matches internal res480 8B)
+            max_sequence_length=None,  # None disables token packing (TOML can't express null)
+            patch_spatial=2,
+            sound_latent_fps=0,
+            tokenizer_spatial_compression_factor=16,
+            tokenizer_temporal_compression_factor=4,
+            dataloader=L(RankPartitionedDataLoader)(
+                batch_size=1,
+                in_order=False,
+                num_workers=4,
+                persistent_workers=True,
+                pin_memory=True,
+                prefetch_factor=4,
+                sampler=None,
+                datasets=dict(
+                    droid=dict(
+                        ratio=1,
+                        dataset=L(get_action_droid_sft_dataset)(
+                            root="${oc.env:DROID_ROOT}",
+                            fps=15.0,
+                            chunk_length=32,
+                            action_space="joint_pos",
+                            use_state=True,
+                            use_image_augmentation=True,  # SR boost (random crop+rescale + color jitter)
+                            # Keep-ranges window filter (drops idle/non-task frames). Off by default;
+                            # the launcher sets use_filter_dict=True + filter_dict_path for internal parity.
+                            use_filter_dict=False,
+                            filter_dict_path=None,
+                            action_normalization=None,
+                            viewpoint="concat_view",  # wrist 480p (top) + L/R shoulder 320x180 (bottom)
+                            resolution="480",  # 640x360 data @ 480p (matches internal res480 run)
+                            max_action_dim="${model.config.max_action_dim}",
+                            cfg_dropout_rate=0.1,
+                            tokenizer_config="${model.config.vlm_config.tokenizer}",
+                        ),
+                    ),
+                ),
+            ),
+        ),
+        dataloader_val=None,
+        upload_reproducible_setup=False,
+    ),
+    flags={"allow_objects": True},
+)
+
+
+# chunk_length=32 → 33 observation frames; pin the VAE encode duration to match
+# (internal used [17] for chunk_length=16). Set post-construction so it lands on
+# the deep-copied NANO_MODEL_CONFIG.tokenizer.
+action_policy_droid_nano["model"]["config"]["tokenizer"]["encode_exact_durations"] = [33]
+
+
+for _item in [action_policy_droid_nano]:
+    _name = [k for k, v in globals().items() if v is _item][0]
+    cs.store(group="experiment", package="_global_", name=_name, node=_item)

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

@@ -0,0 +1,86 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Map-style action SFT dataset: ``DROIDLeRobotDataset`` → ``ActionTransformPipeline``.
+
+The base ``DROIDLeRobotDataset.__getitem__`` returns the raw sample
+(``video``/``action``/``ai_caption``/``viewpoint``/``mode``/``domain_id``/
+``idle_frames``). The model expects each sample to be passed through
+``ActionTransformPipeline`` (spatial resize/pad, text tokenization, action
+padding to ``max_action_dim``, and ``sequence_plan`` construction). This thin
+wrapper composes the two so the experiment can hand a single map-style dataset
+to ``RankPartitionedDataLoader`` (mirroring how the vision recipe uses
+``get_sft_dataset``).
+"""
+from __future__ import annotations
+
+from typing import Any
+
+from torch.utils.data import Dataset
+
+from cosmos_framework.data.vfm.action.datasets.droid_lerobot_dataset import DROIDLeRobotDataset
+from cosmos_framework.data.vfm.action.transforms import ActionTransformPipeline
+
+
+class ActionSFTDataset(Dataset):
+    """Wraps a map-style action dataset and applies ``ActionTransformPipeline`` per sample."""
+
+    def __init__(self, dataset: Dataset, transform: ActionTransformPipeline, resolution: str | int | None):
+        super().__init__()
+        self._dataset = dataset
+        self._transform = transform
+        self._resolution = resolution
+
+    def __len__(self) -> int:
+        return len(self._dataset)
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        return self._transform(self._dataset[idx], self._resolution)
+
+
+def get_action_droid_sft_dataset(
+    *,
+    root: str,
+    fps: float = 15.0,
+    chunk_length: int = 32,
+    action_space: str = "joint_pos",
+    use_state: bool = True,
+    action_normalization: str | None = None,
+    viewpoint: str = "concat_view",
+    use_image_augmentation: bool = False,
+    use_filter_dict: bool = False,
+    filter_dict_path: str | None = None,
+    resolution: str | int = "256",
+    max_action_dim: int = 64,
+    tokenizer_config: dict | None = None,
+    cfg_dropout_rate: float = 0.1,
+    append_viewpoint_info: bool = True,
+    append_duration_fps_timestamps: bool = True,
+    append_resolution_info: bool = True,
+    append_idle_frames: bool = False,
+) -> ActionSFTDataset:
+    """Build the DROID action SFT dataset (joint_pos 8D by default), matching the
+    internal ``droid_lerobot_8b_policy`` data: ``action_space='joint_pos'`` +
+    ``use_state`` (8D, raw/un-normalized), concat_view, chunk_length 32."""
+    dataset = DROIDLeRobotDataset(
+        root=root,
+        fps=fps,
+        chunk_length=chunk_length,
+        viewpoint=viewpoint,
+        action_space=action_space,
+        use_state=use_state,
+        action_normalization=action_normalization,
+        use_image_augmentation=use_image_augmentation,
+        use_filter_dict=use_filter_dict,
+        filter_dict_path=filter_dict_path,
+    )
+    transform = ActionTransformPipeline(
+        tokenizer_config=tokenizer_config,
+        cfg_dropout_rate=cfg_dropout_rate,
+        max_action_dim=max_action_dim,
+        append_viewpoint_info=append_viewpoint_info,
+        append_duration_fps_timestamps=append_duration_fps_timestamps,
+        append_resolution_info=append_resolution_info,
+        append_idle_frames=append_idle_frames,
+    )
+    return ActionSFTDataset(dataset, transform, resolution)