| --- |
| license: other |
| license_name: openmdw1.1-license |
| license_link: https://openmdw.ai/license/1-1/ |
| library_name: cosmos |
| base_model: nvidia/Cosmos3-Nano-Policy-DROID |
| tags: |
| - robotics |
| - droid |
| - policy |
| - world-model |
| - cosmos |
| --- |
| |
| # Cosmos-droid |
|
|
| Intermediate checkpoints from a full fine-tune of |
| [nvidia/Cosmos3-Nano-Policy-DROID](https://huggingface.co/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: |
|
|
| ```shell |
| 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](https://openmdw.ai/license/1-1/) license. See the base model card for the |
| bias, explainability, privacy and safety notes, which apply to this model as well. |
|
|
