geonmin-kim/cosmos3-super-so101-fd-chunk32

Cosmos3-Super SO-101 Forward Dynamics — chunk_length=32 (iter 4500)

Action-conditioned world model for the SO-101 robot arm, post-trained from nvidia/Cosmos3-Super (64B). Feed one real camera frame and a recorded/candidate action sequence; the model generates the video of the robot executing it. Intended use is robot policy evaluation and failure analysis without running the physical robot.

This is the horizon-targeted chunk_length=32 variant: it resumed from the chunk_length=16 checkpoint (iter 4000) and continued training with 32-step action windows (33 frames at 30 fps, ~1.07 s per chunk).

Why this variant exists

The iter-4000 evaluation isolated the horizon failure mode: with real conditioning frames the model holds up across chunks (teacher-forced PSNR 31.1 → 34.4), but feeding it its own output collapses quality (31.1 → 18.0). Appearance is not the problem — trusting its own context is. Two interventions were planned; this checkpoint applies the first in isolation so its effect is attributable:

1. chunk_length 16 → 32 — training previously never showed the model a window longer than 17 frames, yet a rollout asks for 64. Roughly memory-neutral: the token-packing budget is fixed, so longer windows just mean fewer per pack (~1500 → ~2700 tokens each, ~16 → ~9 windows).
2. cond_noise_std > 0 (conditioning-frame noise) — deliberately still 0.0 here, to be added only after the chunk-length effect is measured.

Model details

Base model	nvidia/Cosmos3-Super (64B, MoT)
Resumed from	chunk16 variant @ iter 4000 (same run directory)
Adaptation	LoRA rank 16 / alpha 32 on q/k/v/o_proj_moe_gen + unfrozen action pathway (21.1M params)
Mode	forward_dynamics, action-conditioned video generation
Action space	10-D Cartesian EE (dx, dy, dz, 6-D rotation, gripper), quantile-normalized
Chunk length	32 action steps (33 frames @ 30 fps)
cond_noise_std	0.0 (intentionally — single-variable experiment)
Camera	single top view, ego_view, 480p
Training iteration	4500 (iters 4000–4500 at chunk_length=32)
Dataset	geonmin-kim/so101_merged_v2 — 1444 episodes / 486k frames / 33 tasks
Parallelism	FSDP 4-way shard, bf16 (fp32 master), 4× A100 80GB
Effective batch	grad_accum 4, max 24000 tokens after packing

Checkpoint format

PyTorch Distributed Checkpoint (DCP), saved from a 4-rank FSDP run:

model/      __0_0.distcp ... __3_0.distcp   (~120 GB total, full model + LoRA)
optim/      optimizer state (LoRA + action pathway only, ~250 MB)
scheduler/  LR scheduler state
trainer/    trainer bookkeeping (iteration counter, RNG)

This is not a safetensors/HF-format checkpoint. Load it through cosmos-framework (pinned to commit 5e67049) with the SO-101 overlay from nota-github/xpu-cosmos3-simulator (branch feat/so101-a100-port-and-action-pathway).

Usage

export COSMOS3_SIM_ROOT=/path/to/working/root

# 1. environment (see the simulator repo README for full setup)
git clone -b feat/so101-a100-port-and-action-pathway \
  https://github.com/nota-github/xpu-cosmos3-simulator.git
cd xpu-cosmos3-simulator
git clone https://github.com/NVIDIA/cosmos-framework.git "$COSMOS3_SIM_ROOT/packages/cosmos-framework"
git -C "$COSMOS3_SIM_ROOT/packages/cosmos-framework" checkout 5e67049
./overlay/apply_overlay.sh "$COSMOS3_SIM_ROOT/packages/cosmos-framework"
./scripts/setup_venv313.sh

# 2. this checkpoint
hf download geonmin-kim/cosmos3-super-so101-fd-chunk32 \
  --local-dir "$COSMOS3_SIM_ROOT/ckpt/chunk32_iter4500"

# 3. action normalization stats (REQUIRED — see warning below)
export SO101_ACTION_STATS="$COSMOS3_SIM_ROOT/ckpt/chunk32_iter4500/so101_stats_stride1_v2.json"

# 4. build conditioning inputs from a LeRobot episode, then roll out
source ./env.sh
python scripts/make_inputs.py --episodes 0 63 119
./scripts/run_rollout.sh 0,1 "$COSMOS3_SIM_ROOT/ckpt/chunk32_iter4500" \
    "$COSMOS3_SIM_ROOT/out/rollouts" \
    --input-dirs "$COSMOS3_SIM_ROOT"/out/inputs/ep*_droid_lerobot_s1_* \
    --modes autoregressive teacher_forced

# 5. score against the recorded episode
python scripts/evaluate.py --rollout-dir "$COSMOS3_SIM_ROOT/out/rollouts"
python scripts/make_comparison.py --rollout-dir "$COSMOS3_SIM_ROOT/out/rollouts"

To resume training from this checkpoint instead, point checkpoint.load_path at it and launch the super_horizon variant (train/run_train.sh super_horizon <gpus>), which sets SO101_TRAIN_ACTION_PATHWAY=1 and the 32-step TOML.

⚠️ Normalization stats are part of the model contract. This checkpoint was trained with so101_stats_stride1_v2.json (bundled in this repo). Using the pre-megamix v1 stats silently mis-scales actions by up to 19%. Do not mix.

Sampling configuration used at eval time: num_steps=30, guidance=1.0, shift=10.0, sigma_max=80.0, resolution 480, 16:9, fps 30.

Evaluation

Only 500 iterations of chunk-32 training separate this checkpoint from the chunk16 baseline; a dedicated eval sweep for it had not yet been run at export time. Baseline numbers (iter 4000, chunk_length=16) for reference:

Metric	Value
PSNR / SSIM (autoregressive)	18.8 dB / 0.844
Motion ratio (1.0 = matches reality)	0.87
Motion correlation	0.50
Axis separation (90° ≈ ceiling)	86° (base model: 32°)
Usable horizon	~0.5 s

Evaluate with scripts/evaluate.py / scripts/analyze_probes.py from the simulator repo and compare against the chunk16 card before drawing conclusions about the horizon intervention.

Training provenance

• Launched via train/run_train.sh super_horizon (TOML sets AC selective and compile.enabled=True directly; the TOML deliberately reuses [job].name = action_fd_so101_super_action so it resumes the same run directory and checkpoints)
• SO101_TRAIN_ACTION_PATHWAY=1 (required — without it the LoRA injector freezes the action pathway and training silently degenerates into first-frame video prediction)
• Base checkpoint: nvidia/Cosmos3-Super converted to DCP via cosmos_framework.scripts.convert_model_to_dcp

License

Base model and framework: OpenMDW-1.1 (NVIDIA Cosmos3). Fine-tuned weights released under the same terms. Training data: see geonmin-kim/so101_merged_v2 (Apache-2.0).