jere-mybao/cosmos3-nano-droid-forward-dynamics

Cosmos3-Nano DROID Forward Dynamics

An action-conditioned forward dynamics (video world) model: given the first frame of a Franka DROID scene and a chunk of 16 end-effector actions, it predicts the remaining 16 frames — the video those actions produce. The first frame and all 16 action steps are clean conditioning; only the future frames are denoised ("forward_dynamics": Predict video given first frame and all actions). It is a post-train of nvidia/Cosmos3-Nano on nvidia/Cosmos3-DROID.

• Base model: nvidia/Cosmos3-Nano (15.75 B parameters, bf16, 1165 tensors)
• This model: 15.17 B parameters, bf16, 809 tensors
• Training: 20,000 iterations, data_parallel_shard_degree=8, bf16, single node
• Weights published: the EMA weights (net_ema), which is what the framework loads by default at inference (use_ema_weights=True)

On naming: Cosmos3-Nano is described by NVIDIA as an "8B" model, which counts one tower. The network is a mixture-of-transformers with an understanding (und) tower and a generation (gen) tower, so the full parameter count is ~15.2 B. Both numbers describe the same model.

The 356-tensor difference from the base is not missing weights. This recipe runs with vlm_config.include_visual = False and no audio branch, so the base model's Qwen vision tower (language_model.visual.*, 351 tensors) and its audio bridge (llm2sound, sound2llm, sound_modality_embed) are simply not part of this model — the checkpoint contains zero tensors absent from the base. Image conditioning reaches the model through the Wan VAE and the vae2llm projection, not through a visual encoder.

What was actually trained

Measured by diffing this checkpoint against the base weights tensor by tensor:

Component	Relative change vs base	Status
Generation tower (*_moe_gen, 397 tensors)	~1e-2 to 5e-2	trained
Understanding tower + shared (412 tensors)	~1e-6 (bf16 rounding level)	effectively unchanged

The optimizer state corroborates this from the other direction: it covers 6.98 B parameters (generation tower 6.946 B + adapters 0.036 B), not the full 15.17 B. So this post-train adapted the generation tower and the action/video adapters while leaving the language understanding tower at its pretrained values. The published weights contain both towers, so the model is complete and standalone.

Read the Trained window, Action format, and Training-data caveat sections before you evaluate this model. Each one is a way to get plausible-looking output that is quietly wrong.

---

Trained window

The model was trained on a 17-frame window at 15 fps — about 1.07 seconds.

This is read off the training configuration, not assumed:

Quantity	Value	Source
Action chunk length	16	chunk_length: 16 in the run's config.yaml
Frame rate	15 fps	fps: 15.0 (matches the dataset's meta/info.json)
Observation timestamps	17	observation_ts = [i*dt for i in range(0, chunk_length+1)]
Action timestamps	16	action_ts = [i*dt for i in range(0, chunk_length)]
Window duration	16/15 s ≈ 1.067 s	16 intervals at dt = 1/15 s
Video tokenizer exact duration	17 frames	tokenizer.encode_exact_durations: [17] in the exported config.json

The last row is an independent confirmation from a different part of the stack: the Wan VAE in this checkpoint is configured to encode exactly 17 frames, matching the 17 observation timestamps the dataloader produces. With a temporal compression factor of 4, those 17 frames become 5 latent frames.

So each training sample is 17 video frames and 16 actions: action t carries the scene from frame t to frame t+1.

Anyone rolling out beyond ~1.07 s of simulated time is extrapolating past the trained window. The model will still produce frames — autoregressive chaining of chunks is the obvious thing to do and it does not crash — but nothing in this checkpoint's training signal constrains that regime, and we have not characterised drift there. Treat longer horizons as an open question, not a supported mode.

Camera views

The model predicts three camera views jointly, as a single composite canvas — not three independent streams, and not one view at a time. The training recipe uses viewpoint: concat_view.

The canvas is assembled as:

+-----------------------------------------------+
|                                               |
|         wrist_image_left (full width)          |
|                                               |
+----------------------+------------------------+
| exterior_image_1_left| exterior_image_2_left   |   <- each downscaled to half H, half W
+----------------------+------------------------+

• top row: observation.image.wrist_image_left, full width
• bottom row: observation.image.exterior_image_1_left and observation.image.exterior_image_2_left, each bilinearly resized to (H/2, W/2) and concatenated horizontally

Source views are 360x640 each, giving a 540x640 composite, which the framework's resolution bucketing (tier "480", nearest aspect bucket 4,3) maps to a 544 x 736 (H x W) canvas. Because the three views share one canvas and one set of latents, they are denoised together and stay mutually consistent; you cannot request one view in isolation without changing the input geometry the model was trained on.

Action format

Get this wrong and the model will produce fluent, physically plausible video of the wrong motion. There is no runtime check that will catch it.

Dimensionality: 10 real channels, zero-padded to 64.

Index	Channels	Meaning
0-2	3	Position delta — end-effector translation, relative
3-8	6	Rotation delta, 6D rotation representation (first two columns of the rotation matrix), relative
9	1	Gripper
10-63	54	Zero padding to max_action_dim = 64

Ordering is [Pos(3), Rot6d(6), Gripper(1)], built by build_action_spec(Pos(), Rot("rot6d"), Gripper()). The action tensor handed to the model is [T=16, 64] per sample. The framework's own embodiment table agrees on the width: EMBODIMENT_TO_RAW_ACTION_DIM["droid_lerobot"] = 10.

You must also pass the right embodiment domain id

The action projection (action2llm) is a DomainAwareLinear: it holds a separate 64x4096 matrix per embodiment domain (num_embodiment_domains: 32) and selects one using a domain_id you supply alongside the actions.

For this model, domain_id = 8 (EMBODIMENT_TO_DOMAIN_ID["droid_lerobot"] = 8).

On the inference CLI you select this by name rather than by number — --domain-name droid_lerobot — which is resolved through get_domain_id().

Pass a different id and the actions are projected through a matrix this run never trained. The model will still generate video, and it will still look like a robot doing something. It will not be your actions.

Three more things that are easy to get wrong

1. Rotation is 6D, not Euler and not a quaternion. A 7-D [xyz, roll, pitch, yaw, gripper] vector — the most common DROID action convention — is not what this model expects. Channels 3-8 are the first two columns of the rotation matrix, flattened.
2. Poses are relative (deltas), not absolute. They come from pose_abs_to_rel(poses_abs, rotation_format="rot6d").
3. The gripper channel is inverted relative to the DROID field. This dataset version (droid_plus_lerobot_640x360_20260412) has IS_GRIPPER_ACTION_FLIPPED = True, so

   gripper_channel = 1.0 - action.gripper_position
   

   where action.gripper_position is the raw DROID/LeRobot feature. Feed the raw field unflipped and every grasp in your rollout is inverted.

Normalisation: none. The run sets action_normalization: null. Actions are consumed in raw metric units exactly as constructed above. There is no mean/std file, no min-max rescaling, and no quantile clipping to ship or to apply — if you normalise your actions before passing them in, you are off-distribution.

---

Action conditioning is verified

A forward-dynamics model can silently learn to ignore its action stream and predict plausible video unconditionally. That failure is real and documented for this base model and task (NVIDIA/cosmos-framework#175, where correct, shuffled, and zeroed actions all scored the same loss). It was tested here rather than assumed.

The probe holds the checkpoint fixed (optimizer.lr=0.0) and re-seeds the RNG identically before each arm, so all arms see the same noise and the same sigma draw, then scores the same batch three ways:

Arm	Manipulation	Result vs true (n = 40 paired samples)
true	batch unchanged	—
perm	action chunk permuted along time (identical marginals, wrong temporal correspondence)	degrades on 40 / 40, mean Δloss +0.289, Cohen's dz = 3.32, sign test p = 1.8e-12
zero	action chunk zeroed	degrades on 40 / 40, mean Δloss +0.344, dz = 3.16, p = 1.8e-12

perm is the load-bearing arm: a time-permuted chunk has the same marginal statistics as the real one, so beating it requires genuine temporal action-video correspondence, not just sensitivity to the action tensor's norm. zero is off-distribution and a model could separate it on norm alone.

Reproduce it with the framework's own probe:

cd $SWM_COSMOS
export DATASET_PATH=/path/to/Cosmos3-DROID
export BASE_CHECKPOINT_PATH=/path/to/Cosmos3-Nano
export WAN_VAE_PATH=/path/to/Wan2.2_VAE.pth
export IMAGINAIRE_OUTPUT_ROOT=/path/to/output_root

PYTHONPATH=. torchrun --nproc_per_node=4 --master_port=50113 \
  slurm_scripts/eval/ablate_actions.py \
  --sft-toml=examples/toml/sft_config/action_fd_droid_posttrain.toml \
  --out=ablation.jsonl \
  -- \
  job.name=action_fd_ablation \
  trainer.max_iter=40 \
  checkpoint.save_iter=100000 \
  checkpoint.load_path=/path/to/iter_000020000 \
  optimizer.lr=0.0 \
  model.config.compile.enabled=False \
  trainer.logging_iter=1

Then, over the 40 emitted rows:

import json, math
rows = [json.loads(l) for l in open("ablation.jsonl")]
d = [r["perm"] - r["true"] for r in rows]
n, pos = len(d), sum(x > 0 for x in d)
m = sum(d)/n
sd = math.sqrt(sum((x-m)**2 for x in d)/(n-1))
print(f"{pos}/{n} degraded, mean={m:.4f}, dz={m/sd:.3f}")

optimizer.lr=0.0 is what makes the comparison valid — the weights never move, so all 40 iterations score the same checkpoint.

---

Training-data caveat — read this before you report a number

97% of Cosmos3-DROID was in the training stream. The remaining 3% was held out by the data loader, and you can reconstruct exactly which episodes.

The dataset itself declares a single train split and the run defined no validation loader (dataloader_val: null), so the natural reading is that everything was trained on. That reading is wrong: the split happens a layer lower, inside the dataset class.

Subset	Episodes	Declared split	Trained on	Held out
success/	57,639	{"train": "0:57639"}	55,910	1,729
failure/	14,268	{"train": "0:14268"}	13,840	428
Total	71,907	single train, no val/test	69,750 (97.0%)	2,157 (3.0%)

DROIDMergedLeRobotDataset discovers success/ and failure/ as separate LeRobot roots and calls split_episode_ids(total_episodes, seed=42, val_ratio=0.03, split="train") once per root, dropping the first 3% of a seeded torch.randperm. The forward-dynamics experiment passes split="train" explicitly; split_seed and split_val_ratio are not exposed by the dataset factory, so the defaults 42 and 0.03 are what ran.

Reproduce the held-out ids without downloading anything:

import torch

def held_out(total_episodes, seed=42, val_ratio=0.03):
    g = torch.Generator().manual_seed(seed)
    ids = torch.randperm(total_episodes, generator=g).tolist()
    return ids[: int(round(total_episodes * val_ratio))]

held_out(57639)  # success -> 1729 ids, begins 48330, 22578, 51250, 19933, 16050
held_out(14268)  # failure ->  428 ids, begins  6150, 13604,  1028,  4732,  4582

What this means for your numbers. A loss, PSNR, SSIM, LPIPS, FVD or retrieval figure computed on an arbitrary Cosmos3-DROID episode is, with 97% probability, reconstruction on training data rather than generalization. Restrict to the 2,157 held-out ids above and it is a genuine held-out measurement — though still a weak one, because the split is by episode, not by scene or building. DROID episodes within a scene are highly correlated, so an episode-level split leaks scene identity. For a generalization estimate you can defend, hold out by scene or building yourself.

We are stating this because none of it is visible from the checkpoint or from the dataset's directory layout, and the dataset's own metadata actively suggests the opposite.

Other limitations

• Trained on a single robot embodiment: Franka Panda 7-DoF with a Robotiq 2F-85 gripper.
• The three-view composite geometry is part of the input contract; other camera counts or placements are out of distribution.
• Trained with use_state: false — the model does not consume proprioceptive state.
• Training used cfg_dropout_rate: 0.1, so classifier-free guidance is available.
• No systematic evaluation on held-out data exists for this checkpoint (see the caveat above). The action-conditioning probe is a diagnostic, not a benchmark score.

---

Usage

The published artifact is a consolidated Hugging Face safetensors directory, produced by the framework's own export_model converter. Inference runs from it directly:

python -m cosmos_framework.scripts.inference \
  --parallelism-preset=latency \
  --checkpoint-path /path/to/this/download \
  --action-path    your_actions.json \
  --domain-name    droid_lerobot \
  --action-chunk-size 16 \
  -i inputs/your_input.json \
  -o outputs/

No --config-file or --experiment is needed: the exported config.json carries them.

--action-path is required in forward_dynamics mode — it is the JSON holding the action chunk described above. --domain-name droid_lerobot selects the embodiment projection (id 8). Keep --action-chunk-size at 16 to stay inside the trained window.

The framework fetches the Wan 2.2 VAE from the Hub on first use, so the first run needs network access even though the weights here are self-contained.

How this artifact was produced, and how it was checked

Training wrote a PyTorch Distributed Checkpoint sharded 8 ways (iter_000020000/model/__0_0.distcp ... __7_0.distcp, 85 GiB) plus 53 GiB of optimizer state. The optimizer state is not published — it is useless for inference — and the shards were consolidated with the framework's own converter, not a hand-rolled merge:

python -m cosmos_framework.scripts.export_model \
  --checkpoint-path <run>/checkpoints/iter_000020000 \
  --config-file     <run>/config.yaml \
  --use-ema-weights --no-vit \
  -o <output>

export_model hands an unsharded target state dict to dcp.load, so PyTorch's own DCP resharding does the merge. It runs single-process on CPU. Framework commit 9cbd0841b50a1e667577292be1a4ad79cbc8e3d9 (recorded in export_manifest.json).

Equivalence check. Every one of the 809 published tensors was compared against the original shards, reading the reference values straight out of the .distcp files with PyTorch's low-level DCP key loader — deliberately not through the framework's load path, so the check validates the consolidation instead of re-running it.

matched bitwise : 809 / 809
mismatched      : 0
missing         : 0
export tensors with no checkpoint counterpart : 0
dtype: checkpoint float32 (net_ema) -> export bfloat16, for all 809

Equality is exact after the intended float32 -> bfloat16 cast (the EMA weights are kept in fp32 during training; the published model is bf16, matching the base model and the training precision). Because the published weights are bitwise identical to the source tensors, identical outputs on any batch follow from identical inputs and code — there is no numerical divergence to characterise.

Provenance and licence

• Base model: nvidia/Cosmos3-Nano (revision main), released by NVIDIA under the OpenMDW License Agreement v1.1.
• Video tokenizer: the Wan 2.2 VAE (Wan2.2_VAE.pth) from Wan-AI/Wan2.2-TI2V-5B, revision 921dbaf3f1674a56f47e83fb80a34bac8a8f203e. It is not bundled here; the framework fetches it. Spatial compression 16x, temporal compression 4x.
• Text tokenizer: Qwen/Qwen3-VL-8B-Instruct (tokenizer only; the vision tower is not used).
• Training data: nvidia/Cosmos3-DROID, NVIDIA's LeRobot v3.0 conversion of DROID, also released under OpenMDW-1.1. The underlying DROID dataset is released by its authors under CC-BY 4.0.
• These weights are a derivative of the above and are distributed under the same OpenMDW-1.1; the full text is in the LICENSE file in this repository, as that licence requires.

Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

OpenMDW-1.1 grants permission to deal in the Model Materials "without restriction", conditioned on retaining, in any distribution, a copy of the agreement and all notices of origin. Both are included here.

Citation

Please cite DROID:

@inproceedings{khazatsky2024droid,
  title     = {DROID: A Large-Scale In-The-Wild Robot Manipulation Dataset},
  author    = {Khazatsky, Alexander and Pertsch, Karl and Nair, Suraj and
               Balakrishna, Ashwin and Dasari, Sudeep and Karamcheti, Siddharth and
               others},
  booktitle = {Robotics: Science and Systems (RSS)},
  year      = {2024},
  url       = {https://arxiv.org/abs/2403.12945}
}

and NVIDIA Cosmos for the base model and the framework.