armnet-so101-edge-generator-fd-15k

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:

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.

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:

https://github.com/NVIDIA/cosmos/blob/main/cookbooks/cosmos3/generator/action/run_fd_with_cosmos_framework.ipynb

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.