LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels
Authors: Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun, Randall Balestriero (2026)
arXiv: https://arxiv.org/abs/2603.19312
Abstract
The authors introduce LeWorldModel (LeWM), described as “the first JEPA that trains stably end-to-end from raw pixels using only two loss terms.” The method combines next-embedding prediction with a regularizer enforcing Gaussian-distributed embeddings (SIGReg — Sketched-Isotropic-Gaussian Regularizer). With 15M parameters trainable on a single GPU, LeWM achieves planning speeds 48x faster than foundation-model approaches while maintaining competitive control performance across 2D and 3D tasks.
Core Problem
Representation collapse in Joint Embedding Predictive Architectures (JEPAs). Existing approaches rely on complex multi-term losses, exponential moving averages, pre-trained encoders, or auxiliary supervision. LeWM reduces hyperparameters from six to one compared to existing end-to-end alternatives like PLDM.
Three prior categories it competes with:
- End-to-end methods (PLDM): lack formal collapse guarantees
- Foundation-based methods (DINO-WM): forgo true end-to-end learning
- Task-specific methods (Dreamer): require reward signals
Architecture
Encoder: Vision Transformer (~5M parameters) — maps observations to compact latent embeddings via [CLS] token + MLP projection with Batch Normalization
Predictor: Transformer (~10M parameters, 6 layers) — autoregressively predicts next-step embeddings conditioned on actions via Adaptive Layer Normalization
Training Objective (only two terms):
L_LeWM = L_pred + λ·SIGReg(Z)
- L_pred: MSE between predicted and actual next embeddings
- SIGReg: enforces Gaussian-distributed embeddings via one-dimensional projections using Epps-Pulley statistical test
No stop-gradient, no exponential moving averages, no pre-trained representations. Only λ as the effective hyperparameter.
Planning: Model Predictive Control using Cross-Entropy Method to optimize action sequences in latent space.
Results
- Push-T (2D manipulation): 18% higher success rate than PLDM
- OGBench-Cube (3D manipulation): competitive with DINO-WM
- Two-Room (navigation): PLDM and DINO-WM outperform LeWM in simpler settings
- Planning: under 1 second — up to 48x faster than DINO-WM (200x fewer latent tokens)
- Training curves: smooth and monotonic vs. PLDM’s noisy non-monotonic behavior
Physical understanding:
- Outperforms PLDM in probing physical quantities (block location, agent position, angles) from latent embeddings
- Post-hoc decoder recovers visual scenes from 192-dimensional embeddings (no reconstruction loss during training)
- Violation-of-expectation tests: assigns higher surprise to physically implausible events (teleportation) vs. visual perturbations
Limitations
- Short planning horizons
- Reliance on offline datasets with sufficient coverage
Future Work
- Hierarchical modeling
- Large-scale pretraining