periphanes/cosmos3-gr1-difforce-statefix

Cosmos3-Nano GR-1 — Diffusion-Forcing + StateFix checkpoints

PyTorch Distributed Checkpoint (DCP) checkpoints from joint WM + Action SFT of NVIDIA Cosmos3-Nano (Omni-MoT World Foundation Model) on GR-1, using a causal diffusion-forcing schedule with a fixed proprio-state prefix ("statefix").

Each variant is provided at iteration 20,000 (the latest checkpoint):

Subfolder	Horizon	chunk_length	latent_t	Training
h65/	65-frame	64	17	8×B200, FSDP, 45,056-token packing
h129/	129-frame	128	33	8×B200, FSDP, 45,056-token packing
h65_joint_1000/	65-frame	64	17	4×B200, FSDP, 45,056-token packing, joint diffusion-forcing (WM + action jointly supervised)
h65_arroll_joint_1000/	65-frame	64	17	as h65_joint_1000 + AR rollout loss (autoregressive rollout on own predictions: horizon 4, loss decay 0.8, weight 1.0, 2,000-iter warmup)

Layout

Each subfolder mirrors the framework's checkpoint layout (full DCP — model + optimizer + scheduler + trainer state), so it can be resumed or evaluated directly:

<variant>/
  config.yaml, config.pkl, job_env.yaml, launch_info.yaml
  checkpoints/
    latest_checkpoint.txt
    iter_000020000/
      model/      # FSDP-sharded DCP (.distcp shards + .metadata)
      optim/      # optimizer state (enables training resume)
      scheduler/
      trainer/

Note: these are sharded DCP checkpoints (.distcp + .metadata), not consolidated safetensors. Load them with torch.distributed.checkpoint via the Cosmos3 framework, or consolidate to HF format with cosmos_framework.scripts.export_model.