---
license: mit
---
# commaai/worldmodel-4B
4B-parameter driving world model trained using [vit-ae-2x-f8c32](https://huggingface.co/commaai/vit-ae-2x-f8c32). Predicts paired-camera latents and a driving plan from future anchors, frame history, and ego motion. The video dense [DiT](https://arxiv.org/abs/2212.09748) backbone uses [parameter efficient adaptive normalization](https://arxiv.org/abs/2603.06741) (RMSNorm) and an [inverse dynamics](https://arxiv.org/abs/2206.11795) plan head.
Learning to Drive from a World Model ([paper](https://arxiv.org/abs/2504.19077), [blog](https://blog.comma.ai/mlsim/)) | [Inference and RL examples](https://github.com/YassineYousfi/openpilot.distill/tree/master/rl)
The checkpoint uses FP8 attention weights, NVFP4 MLP weights, and FP8 KV caches.
```text
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.
```text
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).
```text
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