Cosmos-droid

Intermediate checkpoints from a full fine-tune of nvidia/Cosmos3-Nano-Policy-DROID (revision 805c0d6d46196ecdc789213a6de72303c383895b) on five DROID manipulation tasks. The model jointly generates video and actions (joint video/action "WAM" policy).

Work in progress: the run targets 5000 optimizer updates. These are intermediate checkpoints and have not been evaluated on a robot yet.

Checkpoints

Each folder is a self-contained checkpoint in the same layout as the base model (transformer/, vae/, vision_encoder/, text_tokenizer/, scheduler/, model_index.json, ...). They were produced with cosmos_framework.scripts.export_model using the EMA weights, followed by cosmos_framework.scripts.convert_model_to_diffusers.

Folder	Optimizer step	Samples seen	Training loss (median over the preceding 100 steps)
iter_000500/	500	256k	~1.19
iter_001000/	1000	512k	~1.04
iter_001500/	1500	768k	~0.97
iter_002000/	2000	1024k	~0.95
iter_002500/	2500	1280k	~0.87
iter_002629/	2629	1346k	~0.84
iter_002697/	2697	1380k	~0.89
iter_003000/	3000	1536k	~0.89
iter_003500/	3500	1792k	~0.81
iter_004000/	4000	2048k	~0.77
iter_004309/	4309	2206k	~0.76

Load one by downloading its folder and passing it as --checkpoint-path, for example:

hf download Cosmos-Data/Cosmos-droid --include "iter_001500/*" --local-dir Cosmos-droid
python -m cosmos_framework.scripts.inference --checkpoint-path Cosmos-droid/iter_001500 ...

Tasks and data

Five tabletop tasks: stack blue bowl on green bowl; serve bread on plate; put grapes into bowl; insert marker into mug; pick cherries and put on plate.

The training data was 96 successful, relabeled DROID-format episodes (19,627 training windows) at 640x360, 15 FPS. Each window has 33 frames and 32 action targets, plus three cameras: wrist on top, two exterior views below. Actions are 7 absolute joint commands plus gripper.

Training

• Full fine-tune (not LoRA), using the official generation/action optimizer selectors.
• Global batch 512 (8x H100 80GB, 16 samples/GPU x grad-accum 4), BF16, full activation checkpointing.
• FusedAdam, base LR 2e-4 with scheduler multiplier 0.4, action-head LR multiplier 5, weight decay 0.05, betas (0.9, 0.99), no warmup. EMA is enabled (power EMA, rate 0.1) and re-initialized from the base model's regular (non-EMA) weights.
• Recipe: examples/toml/sft_config/droid_five_task_8h100_gbs512.toml and docs/droid_five_task_policy.md in the Cosmos-Framework fork used for training.

License

Derived from nvidia/Cosmos3-Nano-Policy-DROID and released under the same OpenMDW-1.1 license. See the base model card for the bias, explainability, privacy and safety notes, which apply to this model as well.