| --- |
| 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> |
|
|
