DROID action-policy post-training: lazy LeRobot dataset + Cosmos3-Nano recipe + numerical regression test by fwd4 · Pull Request #86 · NVIDIA/cosmos-framework

fwd4 · 2026-07-06T05:47:09Z

Summary

Adds DROID action-policy post-training on Cosmos3-Nano and a numerical regression test for the action-policy launches — the reference reproduction recipe for the DROID policy result.

What's included

Lazy DROID LeRobot dataset — cosmos_framework/data/generator/action/datasets/cosmos3_action_lerobot.py (streaming BaseActionLeRobotDataset) + a rewritten droid_lerobot_dataset.py on top of it, plus droid_lerobot_dataset_config.py. Keys the versioned merged root; use_success_only resolves the success/ split; eager _register_sources().
DROID Nano recipe — configs/base/experiment/action/posttrain_config/action_policy_droid_nano.py + examples/toml/sft_config/action_policy_droid_repro.toml: res480, joint_pos 8D + use_state, JSON action prompt (format_prompt_as_json=True), CPU-side color jitter. The TOML pins the GB200 reference shape — HSDP 32×8 (256 ranks), global batch 8192, lr 2e-4, 10000 iters — and trains the generation + action heads from the public Cosmos3-Nano base.
net_ema warm-start — checkpoint/dcp.py: seed net_ema from net when net_ema is skipped on load (fresh action heads).
Action-policy numerical regression test — tests/action_policy_regression_test.py: the action-policy analogue of tests/launch_regression_test.py. Deterministic 10-iter re-run of the LIBERO + DROID launches (single-node, --deterministic, seed 42), asserting per-arch rank-0 loss goldens with a tolerance. LIBERO golden captured on the H200 CI arch; the DROID spec skips unless its (large, out-of-CI) dataset is supplied via DROID_ROOT.

Reproduction

The recipe reproduces the DROID action-policy result from the public Cosmos3-Nano base. The exact training code was validated by a from-scratch 64-GPU (GB200) run to 10k iterations, and the recipe TOML now encodes that run's configuration directly.

Companion

Cookbook PR: NVIDIA/cosmos#261 — the runnable DROID/LIBERO finetune cookbook that drives this recipe.

🤖 Generated with Claude Code

Move the per-sample ColorJitter from the CPU dataloader workers (float32 rgb<->hsv on ~1 core/worker -- the dominant per-sample cost) to the GPU training step, and stop the base action dataset from eagerly materializing every frame row at construction. - model_config: add OmniMoTModelConfig.train_color_jitter (dict|None, default None) - omni_mot_model: build a torchvision ColorJitter from it; apply per-sample on the GPU float tensor in _normalize_video_databatch_inplace during training only (rearrange [B,C,T,H,W]->[(B T),C,H,W] so channels are at dim -3; one sampled param set per sample = temporally consistent, matching the dataset's behavior) - droid_lerobot_dataset / action_sft_dataset: add apply_color_jitter flag; when False the CPU augmentor keeps only the cheap spatial crop+resize - action_policy_droid_nano: enable the GPU path (dataset apply_color_jitter=False + model train_color_jitter={brightness:0.3,contrast:0.4,saturation:0.5,hue:0.08}) - base_dataset: build self._rows lazily (cached property); datasets with their own compact index (DROIDLeRobotDataset) skip the ~13min/tens-of-GB full-row build Semantically equivalent (color jitter post-compose with one shared param set == per-view jitter with shared params; spatial crop stays pre-compose). Gated by train_color_jitter default None -> other experiments unaffected. Measured (GB300): dataloader 2.06 -> 21.8 samples/s/rank; 1-node smoke step ~42s -> ~22s (was dataloader-bound); dataset build ~800s -> ~185s; loss trajectory unchanged (26->19->13). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

When warm-starting from an EMA-only export (e.g. cosmos3_nano_hf_dcp, which has only net.* = published EMA weights and no net_ema.*), keys_to_skip_loading includes "net_ema." so the EMA shadow is never loaded. Combined with init_moe being skipped whenever a checkpoint is present, net_ema's generation pathway is left at its random build-time init while net gets the loaded weights -> the EMA shadow is random at iter 0 (only the backbone is pretrained). After the model load, if a warm start skipped net_ema, copy net -> net_ema via the model's EMA worker so net_ema starts equal to net ("EMA warm-starts from net"). No effect on full-checkpoint resumes (net_ema not skipped) or non-EMA models. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…aset selection - action_sft_dataset: add format_prompt_as_json + sharded/flat (ShardedDROIDLeRobotDataset) selection + use_success_only, forwarded into ActionTransformPipeline. - recipe: CPU-side ColorJitter (train_color_jitter=None) to match i4's pipeline exactly; sharded=False (non-sharded merged) + format_prompt_as_json=True. - droid_lerobot_dataset: ShardedDROIDLeRobotDataset + LRU approximate-seek decode cache. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

… lazy dataset port (cosmos3_action_lerobot base + bespoke action_sft_dataset.py). This is the code that produced iter_10000 (40% SR json-prompt).

Deterministic 10-iter re-run of the action-policy SFT launches (action_policy_{libero,droid}_repro.toml on the Cosmos3-Nano base), asserting the rank-0 per-step loss series against per-arch goldens with a tolerance — the action-policy analogue of tests/launch_regression_test.py. Runs single-node shard=4/gpus(4) (like that file's nano spec) so the FSDP reduction stays intra-node and reproduces bit-exact. LIBERO data auto-downloads (libero_10 of nvidia/LIBERO_LeRobot_v3); DROID uses a local DROID_ROOT and skips if unset. Goldens captured on hardware via COSMOS_ACTION_REGRESSION_UPDATE_GOLDENS=1 (table empty until captured; test skips per arch/spec without an entry). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Captured on 4xGB200 single-node (shard=4, --deterministic, seed 42); the rank-0 10-iter loss series reproduced bit-exact across two independent capture jobs, so the default 1e-3 tolerance holds. DROID (gb200) determinism re-check + h100/H200 goldens land in follow-ups. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…s for eager recapture model.config.compile.enabled=false: torch.compile is a determinism hazard and its compiled Triton kernel fails on the DROID res480 sequence (CUDA driver error: invalid argument from static_triton_launcher). Eager fixes DROID and makes the loss bit-exact. The prior compiled gb200 LIBERO golden is superseded; goldens re-captured eager (gb200 + h100/H200) land next. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

model.config.compile.enabled=false left a hardcoded @torch.compile in the DROID forward whose Triton kernel fails on res480 (CUDA driver error: invalid argument). Set TORCHDYNAMO_DISABLE=1/TORCH_COMPILE_DISABLE=1 in the run env so every @torch.compile is a no-op -> true eager, which runs clean + bit-exact. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Both captured on 4xGB200 single-node (shard=4, --deterministic, seed 42, fully eager). LIBERO reproduced bit-exact across independent jobs; DROID matches its one compiled run bit-exact (loss is compile-invariant). DROID only ran cleanly once compiled because a hardcoded @torch.compile in the GradClip callback (_fused_nan_to_num) crashed on gb200 via torch's static Triton launcher (CUDA driver error: invalid argument) — eager (TORCHDYNAMO_DISABLE) avoids that launcher path and is numerically identical. h100/H200 golden pending (Lepton). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

… golden Compile stays ON (matches launch_regression_test.py's nano spec); loss is bit-exact under compile (grad-norm not asserted). Goldens are captured on the H200 CI arch — LIBERO is the primary numerical golden; DROID skips unless its (gb200-only) dataset is provided via DROID_ROOT. Drops the earlier gb200/eager goldens (the gb200 small-config DROID crash was a Blackwell static-Triton- launcher edge case, not a numeric issue). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Captured on H200 (Hopper) with torch.compile on, --deterministic, seed 42, single-node shard=4. DROID spec still skips without DROID_ROOT. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

fwd4 · 2026-07-06T06:00:52Z

Companion cookbook PR: NVIDIA/cosmos#261 — documents how to run this regression test to verify numerical reproducibility, with the reference H200 LIBERO-10 loss series.

The three ported DROID/action LeRobot dataset files carried imaginaire4's 'NVIDIA proprietary technology ... strictly confidential' banner, and memprofile.py had no header. Swap all four to the repo-standard SPDX OpenMDW-1.1 header so the addlicense pre-commit hook passes and no proprietary/confidential notice ships in the OSS repo. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Encode the actual Cosmos3-Nano-Policy-DROID reference shape in the run TOML (matching the LIBERO recipe's convention): HSDP 32x8 = 256 ranks (replicate=32/shard=8), lr=2e-4, max_samples_per_batch=32, grad_accum_iter=1 -> global batch 8192, 10000 iters. Previously these lived only in the launcher overrides. Document the single-node scale-down (replicate=1 + grad_accum=32). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Make the repro TOML self-contained instead of relying on launcher --opts: - model.rectified_flow_training_config.loss_scale = 10.0 - scheduler.cycle_lengths = [100000] (was [10000]; the reference run uses a long cosine cycle so lr decays slowly across the 10000-iter run) - trainer.callbacks.compile_tokenizer {enabled=true, warmup_resolutions=[480]} - dataloader_train.dataloader {num_workers=16, batch_size=16, prefetch_factor=2} Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

ychao-nvidia · 2026-07-06T14:58:32Z

LGTM!

@ychao-nvidia

Address @ychao-nvidia's review on PR NVIDIA#86 — drop unused/no-op knobs (behavior-preserving; ColorJitter still runs via use_image_augmentation): - model_config.train_color_jitter field (declared, set to None, never read) - action_policy_droid_nano: apply_color_jitter=True, sharded=False, and the train_color_jitter=None setter (all dead) + stale comments - get_action_droid_sft_dataset: remove apply_color_jitter, sharded (only ever raised NotImplementedError), and lerobot_roots params + the sharded str-parse Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

fwd4 · 2026-07-06T15:15:22Z

Thanks for the review! Addressed the dead-code comments in 6829e84:

train_color_jitter (model_config + the =None setter): removed — it was declared, set to None, and never read. ColorJitter still runs via use_image_augmentation=True in DROIDLeRobotDataset (the augmentor at droid_lerobot_dataset.py:310), so this is behavior-preserving.
apply_color_jitter: removed from both get_action_droid_sft_dataset and the experiment — it was received but never propagated (redundant with use_image_augmentation).
sharded / lerobot_roots: removed from get_action_droid_sft_dataset (+ the experiment's sharded=False) — sharded only ever raised NotImplementedError, and lerobot_roots was unused in that function (the real logic lives in droid_lerobot_dataset.py).

For base_dataset.py and memprofile.py — both are actually in use; replied inline with details and options.

Companion cookbook suggestion also applied on NVIDIA/cosmos#261.

Per review on PR NVIDIA#86: remove the separate memprofile.py stub and inline its five no-op profiling helpers (_memprofile_enabled/_deep_size/_fmt_mb/rss_tracker/ log_worker_memory_breakdown) at module scope in cosmos3_action_lerobot.py, its only consumer. Behavior-identical (profiling stays disabled); call sites unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

lfengad · 2026-07-06T15:41:33Z

+                    # Warm start that skipped net_ema (e.g. loading an EMA-only HF export
+                    # with no net_ema.* keys): the EMA shadow would otherwise keep its random
+                    # build-time generation pathway (init_moe is skipped when a checkpoint is
+                    # present). Seed net_ema from the freshly loaded net so the EMA starts equal


Would this need to be synced with internal?

For posttrain I think not, this is mainly for loading the HF converted dcp, which does not have ema weights. Internal ckpts will have this.

…duction (#261) ## What Aligns the DROID action-policy finetune cookbook (`cookbooks/cosmos3/generator/action/finetune/`) with the actual **Cosmos3-Nano-Policy-DROID GB200 reference reproduction run**. The DROID recipe files already exist on `main`; this PR pins the recipe's run-level scalars to the real reference shape instead of the single-node placeholder, mirroring how the LIBERO recipe already encodes its reference config. ## Changes - **`action_policy_droid_repro.toml`** — encode the GB200 reference config: - `data_parallel_replicate_degree=32` / `data_parallel_shard_degree=8` → **HSDP 32×8 = 256 ranks** (GB200: 64 nodes × 4) - `lr=2.0e-04`, `max_samples_per_batch=32`, `grad_accum_iter=1` → **global batch 8192** - `max_iter=10000`, `cycle_lengths=[10000]` - **README** — document the reference shape (HSDP 32×8, GBS 8192, lr 2e-4, 10k iters) and the **JSON action-prompt format** (`format_prompt_as_json=True`); note that eval must use the matching JSON prompt. Add the single-node scale-down override (`replicate=1` + `grad_accum_iter=32` keeps GBS 8192). - **`launch_sft_action_policy_droid.sh`** — fix the stale "8× H100" / "8-GPU FSDP by default" comments to reflect the 256-rank reference shape. The TOML is kept byte-identical to the cosmos-framework copy (`examples/toml/sft_config/action_policy_droid_repro.toml`). ## Notes - Companion to cosmos-framework PR [#86](NVIDIA/cosmos-framework#86), which carries the identical recipe TOML and the DROID/LIBERO stack. - The closed-loop RoboLab reference SR number will follow once the clean from-scratch reference run + eval land. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

fwd4 and others added 11 commits July 4, 2026 16:26

Capture the exact 40%-SR working-tree state: 7d9b443 + uncommitted i4…

b5827c0

… lazy dataset port (cosmos3_action_lerobot base + bespoke action_sft_dataset.py). This is the code that produced iter_10000 (40% SR json-prompt).

tests: add H200 LIBERO action-policy loss goldens

58f7462

Captured on H200 (Hopper) with torch.compile on, --deterministic, seed 42, single-node shard=4. DROID spec still skips without DROID_ROOT. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

fwd4 mentioned this pull request Jul 6, 2026

cookbook(action/droid): align DROID recipe with GB200 reference reproduction NVIDIA/cosmos#261

Merged

fwd4 and others added 3 commits July 6, 2026 15:37

recipe(action/droid): drop non-GB200 node-count mapping in TOML comment

19de528

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

fwd4 changed the title ~~[Draft] DROID action-policy post-training: lazy LeRobot dataset + Cosmos3-Nano recipe + numerical regression test~~ DROID action-policy post-training: lazy LeRobot dataset + Cosmos3-Nano recipe + numerical regression test Jul 6, 2026

fwd4 marked this pull request as ready for review July 6, 2026 09:32

fwd4 requested a review from ychao-nvidia July 6, 2026 13:04

ychao-nvidia reviewed Jul 6, 2026

View reviewed changes

ychao-nvidia approved these changes Jul 6, 2026

View reviewed changes

fwd4 requested review from lfengad, mli0603 and ychao-nvidia and removed request for lfengad July 6, 2026 15:33

fwd4 removed the request for review from ychao-nvidia July 6, 2026 15:39

lfengad reviewed Jul 6, 2026

View reviewed changes

Merge branch 'main' into haolia/droid-oss-40pct-full-rebased

a8537cf

lfengad approved these changes Jul 7, 2026

View reviewed changes

Merge branch 'main' into haolia/droid-oss-40pct-full-rebased

c0f45f7

fwd4 enabled auto-merge (squash) July 7, 2026 11:52

fwd4 merged commit 8ea4318 into NVIDIA:main Jul 7, 2026
8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

DROID action-policy post-training: lazy LeRobot dataset + Cosmos3-Nano recipe + numerical regression test#86

DROID action-policy post-training: lazy LeRobot dataset + Cosmos3-Nano recipe + numerical regression test#86
fwd4 merged 19 commits into
NVIDIA:mainfrom
fwd4:haolia/droid-oss-40pct-full-rebased

fwd4 commented Jul 6, 2026 •

edited

Loading

Uh oh!

fwd4 commented Jul 6, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ychao-nvidia commented Jul 6, 2026

Uh oh!

fwd4 commented Jul 6, 2026

Uh oh!

lfengad Jul 6, 2026

Uh oh!

fwd4 Jul 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

fwd4 commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What's included

Reproduction

Companion

Uh oh!

fwd4 commented Jul 6, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ychao-nvidia commented Jul 6, 2026

Uh oh!

fwd4 commented Jul 6, 2026

Uh oh!

lfengad Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

fwd4 Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

fwd4 commented Jul 6, 2026 •

edited

Loading