Add DROID action-policy SFT recipe (Cosmos3-Nano, joint_pos) - #24
Merged
Conversation
fwd4
force-pushed
the
action-policy-droid-sft
branch
from
June 8, 2026 09:33
e20fbb6 to
b8a1cad
Compare
Adds a DROID action-policy post-training recipe for nvidia/Cosmos3-Nano: - DROID LeRobot dataset: compact columnar load + episode-aware windowing, joint_pos (8D) actions + use_state proprioception. - ActionTransformPipeline wrapper (get_action_droid_sft_dataset). - Registered action_policy_droid_nano experiment + res480 repro TOML (lr 2e-4, lambdalinear, grad_clip 1.0, count-based batch, fresh action heads). - EMA warm-start in checkpoint/dcp.py (net_ema := net on warm-start init). - Docs + example launcher. Validated end-to-end on H200 (1 node/8 GPU and 2 nodes/16 ranks, HSDP). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Hao Liang <haolia@nvidia.com>
fwd4
force-pushed
the
action-policy-droid-sft
branch
from
June 8, 2026 10:10
b8a1cad to
5a9d716
Compare
lfengad
previously approved these changes
Jun 8, 2026
Per-sample random crop+rescale (5% spatial jitter) + color jitter, applied to all camera views with shared params (temporally and cross-view consistent) before the concat-view assembly. Off by default; enabled in the action_policy_droid_nano experiment to match the reference recipe. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The warm-start branch that copies net.* -> net_ema.* (so EMA starts from the loaded weights instead of random construction-time values) was gated on net_ema.* being present in the model state dict, which is always true — so it fired on every warm start and overwrote a genuinely loaded, trained net_ema (e.g. when initializing from a training checkpoint that carries its own EMA). Gate the copy on net_ema.* being explicitly listed in keys_to_skip_loading instead: only reset net_ema = net when net_ema was actually skipped on load (the HF->DCP init case). When net_ema.* is loaded, leave it intact. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
mli0603
reviewed
Jun 9, 2026
| fps=fps, | ||
| chunk_length=chunk_length, | ||
| viewpoint=viewpoint, | ||
| action_space=action_space, |
Collaborator
There was a problem hiding this comment.
I see that the released Droid dataset doesn't accept action_space as an arg. Have you tried running this script with cosmos-framework?
Collaborator
Author
There was a problem hiding this comment.
yes, it's runnable, dataset is updated, check droid_lerobot_dataset.py
Co-authored-by: Yu-Wei Chao <82182961+ychao-nvidia@users.noreply.github.com>
Co-authored-by: Yu-Wei Chao <82182961+ychao-nvidia@users.noreply.github.com>
Co-authored-by: Yu-Wei Chao <82182961+ychao-nvidia@users.noreply.github.com>
Co-authored-by: Yu-Wei Chao <82182961+ychao-nvidia@users.noreply.github.com>
Signed-off-by: Hao Liang <haolia@nvidia.com>
…on large video batches) Signed-off-by: Hao Liang <haolia@nvidia.com>
Ports the internal use_filter_dict from the DROID policy recipe: restrict training windows to curated keep-ranges (drops idle/non-task frames), matching the reference distribution (kept 74.42% = 12.54M of 16.85M windows). Off by default; enabled via filter_dict_path (an internal data artifact, not shipped). Additive branch — the unfiltered path is unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Hao Liang <haolia@nvidia.com>
…ternal The OSS recipe used torch AdamW (bf16 params) + eps 1e-6, which under-steps the small 5x-lr action heads (sub-ULP updates swamped in bf16) and leaves the action loss on a noisy high plateau. Switch to FusedAdam with fp32 master_weights + eps 1e-8 to match the internal droid_lerobot_8b_policy optimizer. An exact-match forward+optimizer test (identical batch/seed, 1-GPU eager) confirmed the forward is byte-identical across repos and the convergence gap was solely the optimizer. Signed-off-by: Hao Liang <haolia@nvidia.com>
ychao-nvidia
approved these changes
Jun 11, 2026
lfengad
approved these changes
Jun 11, 2026
rahul-steiger-nv
pushed a commit
to rahul-steiger-nv/cosmos-framework
that referenced
this pull request
Jun 15, 2026
## 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>
mli0603
pushed a commit
that referenced
this pull request
Jun 22, 2026
…ility) (#37) ## Problem The DROID action SFT dataloader trained with an **unstable, slow-settling grad-norm** (and a noisy action-loss plateau) vs the internal reference. Root cause: the DROID action dataset is **map-style** and — unlike the iterable vision `SFTDataset`, which self-shuffles — does **not** shuffle, and `RankPartitionedDataLoader` wraps it in a `DataLoader` with no `shuffle`, i.e. a `SequentialSampler`. Every rank then iterates the *same* consecutive, overlapping windows, so the all-reduced global batch is effectively **~1 episode** → high gradient variance. (Forward + gradients were verified numerically equivalent to the internal model on identical input, so this was a data-path issue, not the model/loss/optimizer.) ## Fix `ActionIterableShuffleDataset` (`iterable_shuffle=True`): an `IterableDataset` view of the map-style dataset that streams **rank × worker-sharded, episode-order-shuffled, sequential-within-episode** — decorrelated batches with sequential reads (preserves I/O locality + copy-on-write; a plain `shuffle=True`/`RandomSampler` instead does random-access I/O → ~11 min/iter and OOM from broken COW). Mirrors the internal iterable dataset's per-worker episode assignment. - Adds `DROIDLeRobotDataset.get_shuffle_blocks()` (per-episode/segment flat-index blocks the iterable streams). - No `DataLoader`/sampler change needed — `IterableDataset` is handled natively (`sampler=None`). ## Validation (8192 global batch) | iter | this fix | internal ref | no-shuffle | |------|---------|------|------| | 100 | grad-norm 2.9 | 4.7 | 21 | | 450 | grad-norm 1.7 | 1.9 | — | Per-component **action loss converges to ~0.0055** (matches internal ~0.005; the no-shuffle run plateaued noisily at 0.03–0.07). Builds on #24 (recipe + FusedAdam optimizer). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: Hao Liang <haolia@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a DROID action-policy SFT recipe for
nvidia/Cosmos3-Nano, mirroring the internaldroid_lerobot_8bpolicy 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), plusjoint_pos(8D: 7 joints + gripper) anduse_statesupport.data/vfm/action/datasets/action_sft_dataset.py(new) —get_action_droid_sft_dataset(...)wrapping the dataset throughActionTransformPipeline.configs/.../action/posttrain_config/action_policy_droid_nano.py(new) — registeredaction_policy_droid_nanoexperiment (Cosmos3-Nano / 8B MoT): optimizer trains gen+action heads (5× LR on action heads),LambdaLinearschedule, count-based batch, res480,encode_exact_durations=[33](chunk 32 → 33 frames).checkpoint/dcp.py— EMA warm-start: whenkeys_to_skip_loadingexcludesnet_ema., initializenet_ema = netfrom 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:
max_samples_per_batch=32(64 OOMs at 139 GiB; internal used 128 on GB200).shard 8 × replicate 2,TRAIN_EXIT=0.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
max_samples_per_batch,max_sequence_length=None) lives in the experiment Python — TOML cannot expressnull, and the loader only overrides keys present in the TOML.nvidia/Cosmos3-Nano→ DCP and pass viaBASE_CHECKPOINT_PATH; action heads init fresh (skipped on load).