anshks/cosmos3-nano-bridge-fd-wm

Cosmos3-Nano Bridge Forward-Dynamics World Model

Action-conditioned forward-dynamics (FD) world model for the BridgeData / WidowX setup, post-trained from Cosmos3-Nano. Given a single conditioning frame plus a chunk of robot actions, it predicts the resulting multi-view video — used as the rollout environment in a WorldGym-style policy ranking over our bridge WM.

• Base model: Cosmos3-Nano (Qwen3-VL-8B MoT backbone + Wan2.2 VAE tokenizer)
• Mode: forward_dynamics — condition on frame 0 + all actions → predict video
• Action space: bridge 7D [Δxyz, Δrpy, gripper(abs)]
• Views: multi-view concat (3-camera WidowX)
• Chunk length: 16 · operating point: resolution 480, shift 5
• Precision: bf16

Checkpoints

Each subfolder is a self-contained HF safetensors export (7 shards + config.json + checkpoint.json).

subfolder	iter	held-out motion-PSNR	PSNR	SSIM	note
best-iter14000/	14,000	20.99 dB	25.43	0.862	best by held-out motion-PSNR
latest-iter30400/	30,400	20.36 dB	24.93	0.857	latest; metric plateaued

Held-out eval = seeded 1% split (never trained on), n=6 episodes, motion-PSNR masked to moving regions. The single-chunk metric plateaus / slightly regresses past ~14k iters; best is the recommended default, latest is provided for completeness.

Training

• g512: 16× H200 (2 nodes), global batch 512, bridge_multimodal (WidowX, multi-view).
• Post-trained from the base Cosmos3-Nano DCP checkpoint via the cosmos-framework action SFT recipe.
• Note: the bridge multi-node dataloader occasionally stalled on NFS during training (recovered by resubmit).

Usage

These are cosmos-framework model exports (not transformers-native). Load with the cosmos-framework inference stack (the exported config.json / checkpoint.json describe the OmniMoTModel). Point the bridge multiview FD rollout tooling at a subfolder as the checkpoint path.

Caveats

Research checkpoints. Bridge episodes are short (~≤98 frames); FD conditions on a single history frame so long autoregressive rollouts drift. Best for short-horizon dreaming and policy-ranking use.