加载中…
加载中…
4B-parameter driving world model trained using vit-ae-2x-f8c32. Predicts paired-camera latents and a driving plan from future anchors, frame history, and ego motion. The video dense DiT backbone uses parameter efficient adaptive normalization (RMSNorm) and an inverse dynamics plan head.
Learning to Drive from a World Model (paper, blog) | Inference and RL examples
The checkpoint uses FP8 attention weights, NVFP4 MLP weights, and FP8 KV caches.
model.fp8_nvfp4.torchpackage (ZIP; abridged)
`-- archive/
|-- torchtitan/experiments/worldmodel/
| |-- model.py # architecture
| |-- model_for_inference.py # KV cache and sampling
| `-- schedulers.py # rectified flow
|-- model/model.pkl # serialized model
|-- assets/state_dict.pt # quantized state dict
`-- meta/meta.pkl # I/O shapes, dtypes, export metadata
Inference inputs: latents (B, 15, 32, 16, 32), relative translation/Euler angles, pose mask, and frame indices.
f f f f f h h h h h h h h h x (h9)
<------ prefill: t=0 -------> decode: n steps, t=1 -> 0
f0-f4: five future anchors; h0-h8: nine history frames; x (h9): target initialized with Gaussian noise. Each frame contains 128 spatial tokens (2x2 latent patches).
Block-causal mask in packed order (rows: queries; columns: keys)
q\k f0 f1 f2 f3 f4 h0 h1 h2 h3 h4 h5 h6 h7 h8 x (h9)
f0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0
f1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0
f2 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0
f3 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0
f4 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0
h0 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0
h1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0
h2 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0
h3 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0
h4 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0
h5 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0
h6 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0
h7 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0
h8 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0
x (h9) 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
1: dense attention within a 128x128 token block; 0: masked. Diagonal blocks allow all tokens within a frame to attend to each other.
Generation is autoregressive. For each prediction, prefill caches 14 clean frames at t=0; decode reuses their K/Vs for n rectified-flow Euler steps. The generated latent is decoded by the autoencoder and appended to the rolling history for the next prediction.
The model was not post-trained using forcing techniques, training used sufficient augmentation to make it robust to corrupted history and autoregressive drift.
Examples
Euler sampling, cfg=2.0, num_denoising_steps=15