LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels

TL;DR

LeWorldModel (LeWM) is the first JEPA-style latent world model that trains stably end-to-end from raw pixels using only two loss terms — a next-embedding prediction loss and a SIGReg regularizer that enforces isotropic-Gaussian latent embeddings — with no EMA, no stop-gradient, no pretrained encoders, and no proprioceptive auxiliaries. The recipe collapses from PLDM’s seven-term VICReg-style loss to two terms (one effective hyperparameter), trains 15M parameters on a single GPU, and plans up to 48× faster than DINOv2-backed world models while matching or beating them across 2D and 3D continuous-control tasks. A FAIR + Mila collaboration (Maes, Le Lidec, Scieur, LeCun, Balestriero).

Key claims

Two-loss objective trains stably end-to-end from raw pixels: a next-embedding prediction loss $\mathcal{L}_{\text{pred}}$ and a SIGReg regularizer that enforces isotropic Gaussian latent embeddings via Cramér-Wold projection to 1D Epps-Pulley tests — no EMA, no stop-gradient, no pretrained encoder, no auxiliary supervision [§3.1, Eqs. 1–3].
The full method has only one effective tunable hyperparameter (the regularization weight $\lambda$ ), versus PLDM’s seven-term VICReg-derived loss — enabling logarithmic-time bisection search rather than polynomial grid search [§3.1, §4.3].
LeWM (15M parameters, single GPU) outperforms PLDM on PushT by +18% success rate and remains competitive with DINO-WM even when DINO-WM is given proprioceptive inputs that LeWM does not see; on PushT, LeWM (pixels-only) surpasses DINO-WM (pixels + proprioception) [§4.2, Fig. 6].
Planning is 48× faster than DINOv2-backed world models with full planning completing in under one second, while preserving competitive control performance — narrowing the gap to real-time control [§4.2, Fig. 3].
Latent space encodes physical structure: linear and MLP probes recover agent location, block location, and block angle from LeWM embeddings at PLDM-level or better accuracy on Push-T, competitive with DINOv2 representations even though LeWM has 15M params vs DINOv2’s ~300M [§5.1, Table 1].
Violation-of-expectation surprise correctly detects physically implausible events (object teleport, color flip) across TwoRoom / PushT / OGBench Cube — LeWM assigns higher prediction-error surprise to perturbed frames [§5.2, Fig. 10].
Temporal latent-trajectory straightening emerges without an explicit regularizer for it — cosine similarity between consecutive latent velocities increases during training on PushT, and LeWM achieves higher straightness than PLDM despite PLDM having a dedicated temporal smoothness term [§5.1, Fig. 17, App. H].

Method

LeWM is a vanilla JEPA: a ViT-Tiny encoder (5M params) maps frame observation $o_t$ to latent $z_t = E_\theta(o_t)$ via the [CLS] token after a BatchNorm projector, and a transformer predictor (6 layers, 10M params) emits $\hat z_{t+1} = P_\psi(z_t, a_t)$ with actions injected via zero-initialized AdaLN. The predictor is autoregressive with temporal causal masking over a history of frame embeddings. The training loss is $\mathcal{L} = \mathcal{L}_{\text{pred}} + \lambda \mathcal{L}_{\text{SIGReg}}$ . The prediction loss is teacher-forced MSE between predicted and target embeddings of consecutive frames. SIGReg enforces an isotropic Gaussian shape on the batched embedding tensor: project onto random unit-norm directions, evaluate the 1D Epps-Pulley statistical test against $\mathcal{N}(0, 1)$ along each projection, and average. By Cramér-Wold, matching all 1D marginals matches the full joint distribution. The default is $K=128$ projections and $\lambda=1$ — both hyperparameters that have minimal downstream effect, leaving $\lambda$ as the single effective knob. At test time, planning is done via Cross-Entropy Method MPC over the action sequence with a terminal goal-matching loss $\|z_{t+H} - z_{\text{goal}}\|^2$ , re-planning every $h$ steps. The encoder is a from-scratch ViT, but the method is shown to work with ResNet-18 too (Tab. 8) — the recipe is encoder-agnostic.

Results

PushT planning success rate: LeWM beats PLDM by +18% and surpasses DINO-WM even when DINO-WM has access to proprioception [§4.2, Fig. 6]. Planning speed: 48× faster than DINOv2-backed world models, full planning < 1s [Fig. 3]. Two-Room (simplest environment, low intrinsic dimensionality) is the one setting where LeWM underperforms — the authors attribute it to SIGReg’s high-dimensional Gaussian target being mismatched against the very-low-dimensional underlying state, a real limitation [§4.2]. Probing on Push-T: LeWM linear/MLP probes recover Agent Location, Block Location, Block Angle at MSE / $r^2$ scores competitive with DINOv2 and consistently better than PLDM [Table 1]. Training stability: LeWM’s two loss curves are smooth and monotonic (prediction loss decreases steadily, SIGReg drops sharply then plateaus); PLDM’s seven-term objective is “noisy and non-monotonic across several of its loss components” [§4.3, Figs. 18–19]. Visualization: a decoder trained on the 192-dim latent reconstructs visual scenes faithfully despite reconstruction never being used during training, and t-SNE of latents on PushT preserves the environment’s spatial structure [§5.1, Figs. 8–9]. Author-named limitations: short planning horizon (compounding rollout error); reliance on offline action-labeled datasets; SIGReg degenerates in very-low-complexity environments; the chicken-and-egg of action-conditioned WFM training carries over.

Why it’s interesting

LeWM is the most direct successor on file to PLDM and slots into the wiki’s ongoing JEPA-vs-DINO-vs-pixel debate. Sitzmann’s The flavor of the bitter lesson for computer vision position essay argues video-generative pre-training should replace explicit 3D and camera-pose intermediates — LeWM is exactly the kind of “no explicit 3D, latent-predictive, pixel-conditioned” model Sitzmann predicts will win, but trained end-to-end at 15M params rather than scaled foundation-style. It complements V-JEPA 2.1: Unlocking Dense Features in Video Self-Supervised Learning (which extends V-JEPA 2’s loss to all tokens via a similar from-scratch end-to-end JEPA stance) and Back to the Features: DINO as a Foundation for Video World Models (DINO-world, which uses a frozen DINOv2 encoder and learns dynamics only — the position LeWM is arguing against). The single-GPU / 15M-param scale is the wiki’s first published-quality JEPA WFM that is reproducible on a research budget; useful to track as a baseline for the WFM-as-reward-model and WFM-as-planner threads currently dominated by V-JEPA 2 (Inference-time Physics Alignment of Video Generative Models with Latent World Models).