File size: 3,237 Bytes
71d3f6f 2e1041d 71d3f6f 2e1041d 71d3f6f 2e1041d 91bcebe 168e1d7 289abf8 d08c251 58c3bff db10190 4831182 71d3f6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | ---
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.
|