Results
Benchmarks and diagnostics
All results are from a single run: 20 epochs of encoder/predictor training on 4 chunks of Comma2k19 (789 segments, ~33 hours of driving), followed by 25 epochs of decoder training on the evaluation split.
Applying the LeWorldModel framework to Comma2k19 highway driving data. End-to-end JEPA training from raw pixels — no pretrained encoder, no reconstruction loss.
The encoder is a Vision Transformer (ViT, 384-dim hidden, 12 layers, patch size 14) that maps each frame to a 384-dimensional CLS-token embedding via a projection MLP. The predictor is a causal Transformer (8 layers, 16 heads) that autoregressively predicts the next embedding conditioned on vehicle actions embedded via Adaptive Layer Normalization.
The training objective is the sum of a next-embedding MSE prediction loss, the SIGReg regularisation term (with 2000-step linear warmup), and a cosine similarity smoothness term on temporally adjacent embeddings. No reconstruction decoder is used during training. A separate decoder is trained afterwards for visualisation only.
All results are from a single run: 20 epochs of encoder/predictor training on 4 chunks of Comma2k19 (789 segments, ~33 hours of driving), followed by 25 epochs of decoder training on the evaluation split.
The model follows the LeWM blueprint exactly, adapted for driving. Actions are (steering, speed, brake) multiplied across frameskip=5 timesteps, embedded via Conv1D + MLP, then injected into the predictor through Adaptive Layer Normalization at every Transformer layer.
# Encoder Encoder: ViT (hidden=384, layers=12, heads=6, patch=14) CLS token → ProjMLP(384→2048→384) → z_t ∈ R^384 # Predictor Predictor: CausalTransformer(depth=8, heads=16, mlp=2048, drop=0.1) Action conditioning via AdaLN at every layer ProjMLP(384→2048→384) output head # Actions Actions: (steer, speed, brake) × frameskip=5 → R^15 Conv1D + SiLU + Linear → R^384 # Loss L = MSE(z_hat_{t+1}, z_{t+1}) + λ · SIGReg(Z) # λ=0.025, warmup 2000 steps + 0.02 · ReprCosine(Z) # smoothness on nearby frames Total params: 57.7M
| Property | This work | PLDM | DINO-WM |
|---|---|---|---|
| End-to-end from pixels | Yes | Yes | No (frozen DINOv2) |
| Reconstruction-free | Yes | Yes | Yes |
| Reward-free | Yes | Yes | Yes |
| Anti-collapse guarantee | SIGReg (provable) | VICReg (heuristic) | Frozen encoder |
| Loss hyperparameters | 1 (lambda) | 6 | 0 |
| Real driving data | Comma2k19 | Sim only | Mixed |
Training was run on a single NVIDIA H100 80GB HBM3 GPU on Lightning.ai.
The VRAM utilisation shown in training logs (~1.7 GB) reflects allocated
memory as reported by torch.cuda.memory_allocated(). Total physical
VRAM consumption was ~84 GB due to PyTorch reserved buffers, the compiled model,
and the large batch of 128 sequences each of length 20 frames at 224×224 resolution.
Video decoding uses NVDEC (ffmpeg hevc_cuvid) for GPU-accelerated HEVC decode.
Frame normalisation uses ImageNet statistics. CAN bus signals (speed, steering)
and IMU data are loaded from comma2k19's processed_log directories
and interpolated onto the frame timeline.
# Hardware GPU : NVIDIA H100 80GB HBM3 Precision : bfloat16 + torch.compile (reduce-overhead) # Data Dataset : Comma2k19 chunks 1-4 (789 segments, ~33h driving) Resolution : 224×224 | frameskip=5 | seq_len=20 Train/eval : 711 / 78 segments | 78,665 training windows # Optimiser AdamW : lr=1e-4, weight_decay=1e-3, fused=True Sched : cosine decay with 2-epoch linear warmup Epochs: 20 (~19,660 steps) | Budget: 2.9h # Decoder (visualisation only) Arch : FC(384→1024→256×7×7) + 5× ConvTranspose2d Training : 25 epochs, MSE + 0.1 × edge loss, lr=3e-4
# Install dependencies pip install torch torchvision transformers einops matplotlib scipy scikit-learn av huggingface_hub # Train (Lightning.ai, 4 Comma2k19 chunks, 2.9h budget) python train.py --lightning --chunks 1 2 3 4 --budget-hours 2.9 # Train (custom data directory) python train.py --data-dir ./comma2k19 --chunks 1 2 --budget-hours 3.0 # Evaluate — produces all benchmark plots python tests.py --checkpoint ./checkpoints_v5/final.pt --lightning --output-dir ./results
The repository contains two files:
train.py — full model, data pipeline, and training loop;
and tests.py — the complete evaluation suite producing all benchmark plots above.
@article{maes2026leworldmodel,
title = {LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels},
author = {Maes, Lucas and Le Lidec, Quentin and Scieur, Damien and LeCun, Yann and Balestriero, Randall},
journal = {arXiv:2603.19312},
year = {2026}
}