Heterogeneous Diffusion Forcing World Model — OGBench cube-triple (3 cubes)
State-based world model trained with the exhaustive-4 protocol: exactly four complete passes over every H=101 window (stride 1, episode-bounded) of the offline dataset, global batch 256 (4x H200 DDP, 64/rank).
- Architecture: semantic-slot factorized DiT (~9.6M params) — one token per state scalar, bidirectional within-frame attention + causal temporal attention, adaLN-Zero conditioning on the previous action and a per-frame diffusion noise level (Diffusion Forcing).
- Task: OGBench cube-triple (3 cubes) — 46 state dims, absolute raw states and absolute actions, checkpoint-owned min-max normalizers.
- Data: cube-triple-play-v0 (official OGBench play, 3,000 trajectories x 1001 steps); 2,703,000 training windows per epoch.
- Diffusion: 1000 train steps, squaredcos_cap_v2, sample-prediction, 20 DDIM inference steps, full-sequence sampling schedule.
- Source:
Robo_Continual_Learningcommitcc88389e9c06b2a8f3dca103718cce248210d468(algorithms/world_model/heterogeneous_diffusion_forcing.py).
Validation (H=100 rollout, mean reference-normalized MSE)
| epoch | step | val MSE |
|---|---|---|
| 1 | 10,559 | 0.01322 |
| 2 | 21,118 | 0.01001 |
| 3 | 31,677 | 0.00637 |
| 4 | 42,236 | 0.00714 |
Files
world_model_final.pt— EMA weights after exactly 4 exhaustive epochs (the protocol artifact).world_model_best.pt— EMA weights at the best validation epoch (epoch 3).metrics.jsonl— full training/validation log.source_commit.txt— exact training source revision.
Loading
import torch
payload = torch.load("world_model_final.pt", map_location="cpu", weights_only=False)
# payload["model"] holds the HeterogeneousDiffusionForcingModel state dict;
# payload["model_config"] / payload["schema"] rebuild the exact model.