File size: 3,934 Bytes
293401f 1db5557 1cc2ae9 1db5557 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | ---
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
<video src="https://huggingface.co/commaai/worldmodel-4B/resolve/main/assets/videos/0.mp4" width="320" style="max-width: 100%;" controls autoplay loop muted playsinline></video>
<video src="https://huggingface.co/commaai/worldmodel-4B/resolve/main/assets/videos/1.mp4" width="320" style="max-width: 100%;" controls autoplay loop muted playsinline></video>
<video src="https://huggingface.co/commaai/worldmodel-4B/resolve/main/assets/videos/2.mp4" width="320" style="max-width: 100%;" controls autoplay loop muted playsinline></video>
<video src="https://huggingface.co/commaai/worldmodel-4B/resolve/main/assets/videos/3.mp4" width="320" style="max-width: 100%;" controls autoplay loop muted playsinline></video>
|