Source Summary: LeWorldModel
Full title: LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels
Authors: Lucas Maes*, Quentin Le Lidec*, Damien Scieur, Yann LeCun, Randall Balestriero
Affiliations: Mila & Université de Montréal, NYU, Samsung SAIL, Brown University
arXiv: 2603.19312v2 — March 24, 2026
PDF: papers/maes-2026-leworldmodel-stable-end-to-end-joint-embedding-predictive.pdf
Contributions
- First stable end-to-end JEPA from raw pixels — trains without stop-gradient, EMA, or pretrained encoders.
- Two-term training objective — next-embedding MSE loss + SIGReg regularizer. Reduces tunable hyperparameters from 6 (PLDM) to 1 (λ).
- 48× faster planning than foundation-model-based world models (DINO-WM), completing full planning in under 1 second.
- Physical understanding — latent space encodes physical quantities (agent/block location, angles) and reliably detects physically implausible events via Violation-of-Expectation (VoE) framework.
- Single GPU trainable — 15M parameters, hours to train.
Methodology
Architecture
- Encoder: Vision Transformer (ViT-Tiny, ~5M params). Patch size 14, 12 layers, 3 attention heads, 192 hidden dim. [CLS] token → 1-layer MLP with BatchNorm projection.
- Predictor: Transformer, 6 layers, 16 attention heads, 10% dropout (~10M params). Action conditioning via Adaptive Layer Normalization (AdaLN) at each layer, initialized to zero for stable warm-up. Takes N history frames, predicts next frame autoregressively with causal masking.
Training Objective
L_LeWM = L_pred + λ · SIGReg(Z)
- L_pred (MSE):
||ẑ_{t+1} − z_{t+1}||²— next-embedding prediction (teacher-forcing) - SIGReg (Sketched-Isotropic-Gaussian Regularizer): projects embeddings onto M=1024 random unit-norm directions, applies Epps–Pulley normality test to each 1D projection, encourages isotropic Gaussian distribution. By Cramér–Wold theorem, matching all 1D marginals ≡ matching full joint distribution.
- λ = 0.1 (default). M has negligible effect → λ is the only hyperparameter. Binary search suffices: O(log n) vs. O(n⁶) for PLDM.
Latent Planning (Inference)
- Given initial obs o₁ and goal obs o_g, embed both via encoder.
- Predictor rolls out latent states to horizon H:
ẑ_{t+1} = pred(ẑ_t, a_t) - Minimize terminal cost:
C(ẑ_H) = ||ẑ_H − z_g||² - Solved with Cross-Entropy Method (CEM), iteratively.
- MPC strategy: execute first K actions, then replan.
Training Setup
- Offline, reward-free: unannotated trajectories of (o_{1:T}, a_{1:T}).
- Pseudo-expert or exploratory data acceptable.
Results
Planning Performance (4 environments)
| Environment | LeWM | PLDM | DINO-WM |
|---|---|---|---|
| Two-Room | 87% | 100% | 97% |
| Reacher | 86% | 78% | 79% |
| Push-T | 96% | 78% | 74% (92% w/ prop) |
| OGBench-Cube | 74% | 65% | 86% |
- Planning speed: LeWM 0.98s vs. DINO-WM 47s (48× faster).
- Under fixed FLOPs, LeWM significantly outperforms DINO-WM on Push-T (96% vs 13%) and OGBench-Cube (74% vs 48%).
- Underperforms on Two-Room — SIGReg struggles with very low intrinsic dimensionality environments.
Physical Understanding
- Linear probing on Push-T: LeWM consistently outperforms PLDM on agent location, block location, block angle.
- Latent straightening: trajectories become temporally straight as emergent phenomenon — higher than PLDM despite PLDM having explicit temporal smoothness regularization.
- VoE framework: significantly higher surprise for object teleportation (physical violation) than unperturbed; p < 0.01 paired t-test across all 3 environments.
Training Stability
- LeWM loss curves: smooth monotonic convergence (pred loss decreases steadily; SIGReg drops sharply early, then plateaus).
- PLDM: noisy, non-monotonic across several of its 7 loss components.
Limitations
- Short planning horizons — autoregressive rollouts accumulate errors; long-horizon planning not yet addressed. Hierarchical WM is suggested future direction.
- Offline data dependency — requires sufficient interaction coverage; SIGReg underperforms in very simple low-diversity environments (e.g., Two-Room).
- Action labels required — action-conditioned prediction needs action annotations. Inverse dynamics modeling is a proposed future direction.
- End-effector detail — fine-grained details (e.g., end-effector angle in OGBench-Cube) not fully captured in the 192-dim latent.
Key References
- LeCun 2022 — JEPA framework origin
- SIGReg: Balestriero & LeCun 2025 (LeJEPA, arXiv:2511.08544)
- PLDM: Sobal et al. 2022 — prior end-to-end JEPA baseline (7-term VICReg objective)
- DINO-WM: Zhou et al. 2025 — JEPA with frozen DINOv2 encoder
- VICReg: Bardes et al. 2022 — self-supervised regularization