hdf-world-model-ogbench-cube-single

Heterogeneous Diffusion Forcing World Model — OGBench cube-single (1 cube)

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-single (1 cube) — 28 state dims, absolute raw states and absolute actions, checkpoint-owned min-max normalizers.
• Data: cube-single-play-v0 (official OGBench play, 1,000 trajectories x 1001 steps); 901,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_Learning commit cc88389e9c06b2a8f3dca103718cce248210d468 (algorithms/world_model/heterogeneous_diffusion_forcing.py).

Validation (H=100 rollout, mean reference-normalized MSE)

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 4).
• 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.