| --- |
| license: other |
| license_name: nvidia-open-model-license |
| base_model: nvidia/Cosmos3-Edge-Policy-DROID |
| tags: |
| - lerobot |
| - cosmos3 |
| - robotics |
| - so101 |
| --- |
| |
| # Cosmos3-Edge-Policy-SO101-init |
|
|
| `nvidia/Cosmos3-Edge-Policy-DROID` (4B Mixture-of-Transformers world action model, |
| DROID 8-D joint policy) converted to the **LeRobot `cosmos3` policy format** and |
| retargeted to the **SO-101** embodiment (5 joints + gripper, absolute joint position). |
| This is the **un-finetuned initialization checkpoint**: weights are bit-identical to |
| the NVIDIA release; only the config/normalization contract changed. No weight surgery |
| is needed because the Cosmos3 action head is a `DomainAwareLinear` over a zero-padded |
| 64-dim action vector with a per-sample raw-dim mask. |
|
|
| Converted with `scripts/convert_cosmos3_edge_policy_to_lerobot.py` from |
| [nota-github/xpu-lerobot#6](https://github.com/nota-github/xpu-lerobot/pull/6). |
|
|
| ## SO-101 retargeting (vs the DROID contract) |
|
|
| | field | DROID | this repo | |
| |---|---|---| |
| | `raw_action_dim` | 8 (7 joints + gripper) | **6 (5 joints + gripper)** | |
| | cameras | wrist + left/right over-shoulder | **wrist + top** (3rd view zero-padded) | |
| | normalization | — | **QUANTILES (q01/q99 → [-1,1])**, stats injected from the training dataset at train time | |
| | `conditioning_fps` | 15 | **30** | |
| | `invert_gripper` | true | **false** | |
| | `chunk_size` / `n_action_steps` | 32 / 32 | 32 / 32 | |
| | `domain_id` | 8 (`droid_lerobot`) | 8 (reused, adapted by fine-tuning) | |
|
|
| ## Usage |
|
|
| Requires the `cosmos3` policy from the PR above and `diffusers` git main (>= 0.40.0.dev0 |
| schema — 0.39 silently drops `use_und_k_norm_for_gen`/`qk_norm_for_text`/`hidden_act` |
| and builds a different architecture), `huggingface_hub>=1.23`, `safetensors>=0.8`. |
|
|
| ```bash |
| lerobot-train \ |
| --policy.path=geonmin-kim/Cosmos3-Edge-Policy-SO101-init \ |
| --dataset.repo_id=<your SO-101 LeRobot dataset> ... |
| ``` |
|
|
| Note: `config.json`'s `text_processor_name_or_path` points at the checkpoint |
| directory; when loading from the Hub on a fresh machine, download the full snapshot |
| (`snapshot_download`) and pass the local path, or override |
| `--policy.text_processor_name_or_path=<local snapshot dir>` (the Qwen tokenizer files |
| `tokenizer.json` / `tokenizer_config.json` / `special_tokens_map.json` / |
| `chat_template.jinja` ship in this repo). |
|
|
| ## Fine-tuned descendants |
|
|
| - Flow matching (4-step): `geonmin-kim/Cosmos3-Edge-SO101-Morig-D6231-step{3000..30000}` — |
| use `guidance_scale=1.0` at inference (no prompt dropout during fine-tuning, CFG hurts). |
| - Drift 1-step (`num_inference_steps=1`): `geonmin-kim/Cosmos3-Edge-SO101-Drift-D6231-step{5000..100000}` — |
| recommended for deployment (open-loop MAE 5.45 vs 26.96, latency 0.25s vs 0.62s per chunk on B200). |
|
|
| Training data of the descendants: `geonmin-kim/SO101-large-small-cube-11dataset-ratio6231-0729`. |
|
|
