commaai/worldmodel-4B

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