Cosmos3 Edge Generator SO101 Forward Dynamics — 15k
This repository contains the iteration-15,000 EMA checkpoint exported from
Distributed Checkpoint (DCP) to sharded Hugging Face safetensors. It is a
Cosmos3 4B Edge Generator forward-dynamics model conditioned on one observation
frame, a 16-step action chunk, and the task-language caption, predicting a
17-frame video at 20 Hz (0.8 seconds) and resolution class 480.
The model was post-trained jointly on single-arm and bimanual ArmNetBench SO101:
| Domain | Domain ID | Raw action width |
|---|---|---|
armnet_so101_single | 23 | 6 |
armnet_so101_bimanual | 23 | 12 |
The default policy manifest uses armnet_so101_single; both domains share the
same learned domain ID. Actions are conditioned inputs in forward-dynamics mode,
so action-generation loss is zero by design.
Evaluation
Single-arm native validation split: 75 episodes and 32,264 retained samples.
| Metric | Result |
|---|---|
| Vision rectified-flow loss (all 32,264 samples) | 0.1080426808 |
| Total weighted rectified-flow loss | 1.0804268077 |
| Action rectified-flow loss | 0.0 |
| Mean decoded-video PSNR (diverse 30) | 18.2413917612 dB |
| Pooled-pixel PSNR (diverse 30) | 17.8811561069 dB |
Bimanual diverse-30 mean PSNR is 18.2390195243 dB; pooled-pixel PSNR is 17.8093258752 dB.
Inference and evaluation code
Use NVIDIA's Cosmos repository and the forward-dynamics Action notebook:
The native export can be loaded by Cosmos3OmniModel.from_pretrained_dcp in the
matching Cosmos3 checkout. The checkpoint includes its pinned Edge Reasoner
processor and vision assets under edge_reasoner/; no raw DCP optimizer,
scheduler, or trainer state is included.
The SO101 dataset adapter/domain registration used for this model is from the
training checkout and is not yet on the public Cosmos main branch. For direct
SO101 runs, register domain ID 23 with raw action width 6 (single-arm) or 12
(bimanual), matching checkpoint.json.
Provenance
Exported from the iteration-15,000 training checkpoint using EMA weights.
No training or checkpoint modification was performed during evaluation or export. The source DCP was loaded with EMA-to-regular weight mapping under inference-only workflows.