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

Lucas Maes; Quentin Le Lidec; Damien Scieur; Yann LeCun; Randall Balestriero

cs.LGMar 13, 2026

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

Authors: Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun, Randall Balestriero

Organizations: Mila & Université de Montréal · New York University · Samsung SAIL · Brown University

Abstract

Joint Embedding Predictive Architectures (JEPAs) offer a compelling framework for learning world models in compact latent spaces, yet existing methods remain fragile, relying on complex multi-term losses, exponential moving averages, pre-trained encoders, or auxiliary supervision to avoid representation collapse. In this work, we introduce LeWorldModel (LeWM), the first JEPA that trains stably end-to-end from raw pixels using only two loss terms: a next-embedding prediction loss and a regularizer enforcing Gaussian-distributed latent embeddings. This reduces tunable loss hyperparameters from six to one compared to the only existing end-to-end alternative. With ~15M parameters trainable on a single GPU in a few hours, LeWM plans up to 48x faster than foundation-model-based world models while remaining competitive across diverse 2D and 3D control tasks. Beyond control, we show that LeWM's latent space encodes meaningful physical structure through probing of physical quantities. Surprise evaluation confirms that the model reliably detects physically implausible events.

Figures & tables

Figure 1: LeWorldModel Training Pipeline. Given frame observations ${\bm{o}}_{1:T}$ and actions ${\bm{a}}_{1:T}$ , the encoder maps frames into low-dimensional latent representations ${\bm{z}}_{1:T}$ . The predictor models the environment dynamics by autoregressively predicting the next latent state ${\bm{z}}_{t+1}$ from the current latent state ${\bm{z}}_{t}$ and action ${\bm{a}}_{t}$ . The encoder and predictor are jointly optimized using a mean-squared error (MSE) prediction loss. LeWM does not rely on any training heuristics, such as stop-gradient, exponential moving averages, or pre-trained representations. To prevent trivial collapse, the SIGReg regularization term enforces Gaussian-distributed latent embeddings, promoting feature diversity. For tractability, latent embeddings are projected onto multiple random directions, and a normality test is applied to each one-dimensional projection. Aggregating these statistics encourages the full embedding distribution to match an isotropic Gaussian.

Figure 2: Characteristics of latent world model approaches. Methods are grouped by training paradigm. End-to-end methods (PLDM) learn the encoder and predictor jointly from pixels without pre-trained representations or heuristics like stop-gradient or EMAs, but require many hyperparameters and lack collapse guarantees. Foundation-based methods (DINO-WM) avoid collapse by freezing a pre-trained vision encoder, forgoing end-to-end learning. Task-specific methods (Dreamer, TD-MPC) require reward signals or privileged state access. LeWM combines the strengths of all three: end-to-end, task-agnostic, pixel-based, reconstruction- and reward-free, with a single hyperparameter with provable anti-collapse guarantees.

Figure 3: Planning time and performance under fixed compute. Left: Planning time comparison averaged over 50 runs. Encoding observations with

\sim 200\times

fewer tokens than DINO-WM allows LeWM to achieve planning speeds comparable to PLDM while being up to

\sim 50\times

faster than DINO-WM. Center–Right: Planning performance under the same computational budget (fixed FLOPs). LeWM significantly outperforms DINO-WM on Push-T (center) and OGBench-Cube (right). See App. D for planning setup details.

Figure 4: Latent Planning with LeWorldModel . Given an initial observation ${\bm{o}}_{1}$ and a goal ${\bm{o}}_{g}$ , the world model learned in Fig. 2 performs planning in the LeWM latent space. The initial state embedding ${\bm{z}}_{1}$ and the goal embedding ${\bm{z}}_{g}$ are obtained from the encoder. The predictor then rolls out future latent states up to a horizon $H$ . A latent cost between the final predicted state and the goal embedding guides a solver to optimize the action sequence. This prediction–optimization loop is repeated until convergence to a good plan candidate.

Figure 5: Environments used for evaluation. Left: Push-T, a 2D manipulation task where the agent must push a block toward a target configuration, commonly used as a robotics benchmark. Center (1): OGBench-Cube, a visually richer 3D manipulation environment where a robotic arm interacts with a cube to reach a target position. Center (2): Two-Room, a simple 2D navigation environment where an agent moves between rooms to reach target positions. Right: Reacher, a task where a 2-joint arm needs to reach a target configuration in a 2D plane. All environments have a continuous action space. More details on environment and datasets are available in appendix E .

Panel of Figure 6 — Figure 6: Planning performance across environments. Results are shown for Two-Room (left), Reacher (center 1), PushT (center-2) and OGBench-Cube (right). LeWM consistently outperforms PLDM and DINO-WM on Push-T and Reacher. On OGBench-Cube, DINO-WM slightly outperforms LeWM, possibly due to the higher visual complexity and the 3D nature of the environment, which makes encoder training more challenging. In the simpler Two-Room environment, PLDM and DINO-WM outperform LeWM, which may be explained by the SIGReg regularization encouraging a Gaussian distribution in a high-dimensional latent space, while the intrinsic dimensionality of the environment is much lower.

		Linear		MLP
Property	Model	MSE $\downarrow$	r $\uparrow$	MSE $\downarrow$	r $\uparrow$
Agent Location	DINO-WM	$1.888\pm 0.500$	$\mathbf{0.977}$	$\mathbf{0.003\pm 0.022}$	$\mathbf{0.999}$
	PLDM	$0.090\pm 0.311$	$0.955$	$0.014\pm 0.119$	$0.993$
	LeWM	$\mathbf{0.052\pm 0.149}$	$0.974$	$0.004\pm 0.056$	$0.998$
Block Location	DINO-WM	$\mathbf{0.006\pm 0.007}$	$\mathbf{0.997}$	$0.002\pm 0.003$	$\mathbf{0.999}$
	PLDM	$0.122\pm 0.341$	$0.938$	$0.011\pm 0.066$	$0.994$

Table 1: Physical latent probing results on Push-T. LeWM consistently outperforms PLDM while remaining competitive with DINO-WM. The strong probing performance of DINO-WM on certain properties may stem from its foundation-model pretraining: the DINOv2 encoder is trained on two orders of magnitude more data (

\sim

124M images) spanning a far more diverse distribution, which likely allows it to capture some physical properties in its embeddings by default.

Figure 7: Predictor rollout on OGBench-Cube . We visualize decoded latent plans from LeWM given context and action sequences. Each rollout encodes three image observations as context, then autoregressively generates future latents conditioned on the actions in an open-loop fashion. Latents are decoded via a decoder trained a posteriori . The imagined rollout suggests that latent representations capture the global scene structure while finer details like end-effector angle are not fully preserved. Additional rollouts for Push-T and OGBench-Cube are provided in Fig. 9 .

Figure 8: Violation-of-expectation evaluation across three environments. Each plot shows LeWM surprise along three trajectories: an unperturbed reference, a visually perturbed trajectory with abrupt object color change, and a physically perturbed trajectory where objects teleport to random positions. Teleportation violates physical continuity and produces a pronounced surprise spike, while the unperturbed trajectory stays at a low baseline. The increase is significant for teleportation across all environments (paired t-test,

p<0.01

) but weaker and non-significant for color changes, indicating greater sensitivity to physical than visual perturbations. Environments, left to right: TwoRoom , PushT , OGBench Cube .

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

Loss coefficient	Initial value
$\alpha$	18.0
$\beta$	12
$\gamma$	0.2
$\zeta$	0.7
$\nu$	0.0
$\mu$	0.0

Appendix

Table 2: Best coefficient found from grid search.

Panel of Figure 9 — Figure 9: Algorithm 9 . Pseudo-code for the training procedure of LeWorldModel. Pixel observations are encoded into latent embeddings, and a predictor estimates the dynamics by predicting the next-step embedding conditioned on actions. The model is optimized end-to-end using a next-embedding prediction loss together with a step-wise SIGReg regularization term to prevent representation collapse.

Figure 10: Decoder visualization during training. As training progresses, the latent representation increasingly captures the information required to reconstruct the visual scene, even though no reconstruction loss is used during training. Early in training, the decoded images correspond to slow features, a phenomenon previously reported [ 46 ] .

	Agent Position
	Linear		MLP
Model	MSE $\downarrow$	r $\uparrow$	MSE $\downarrow$	r $\uparrow$
DINO-WM	$0.488\pm 0.451$	$0.824$	$0.000\pm 0.000$	$0.999$
PLDM	$0.008\pm 0.041$	$0.996$	$0.000\pm 0.000$	$1.000$
LeWM	$0.008\pm 0.018$	$0.996$	$0.000\pm 0.000$	$1.000$

Appendix

Table 3: Physical Latent Probing results on TwoRoom. Although LeWM underperforms PLDM in downstream planning on this environment, it matches or outperforms PLDM across all probing metrics, and both methods substantially outperform DINO-WM on the linear probe. This suggests that the learned latent space captures the underlying physical state equally well and that the planning gap is not due to a less informative representation but rather to other factors such as the dynamics model or the planning procedure itself.

		Linear		MLP
Property	Model	MSE $\downarrow$	r $\uparrow$	MSE $\downarrow$	r $\uparrow$
Joint Position	DINO-WM	$0.960\pm 1.150$	$0.808$	$0.200\pm 0.967$	$0.870$
	PLDM	$0.372\pm 1.172$	$0.695$	$0.340\pm 1.164$	$0.728$
	LeWM	$0.352\pm 1.173$	$0.706$	$0.330\pm 1.157$	$0.742$
Joint Velocity	DINO-WM	$0.792\pm 0.748$	$0.763$	$0.263\pm 0.683$	$0.852$
	PLDM	$1.016\pm 0.905$	$0.115$	$0.661\pm 0.830$	$0.536$

Appendix

Table 4: Physical latent probing results on OGBench-Cube. LeWM matches or outperforms PLDM on most properties and achieves the best results on positional quantities such as block position and end-effector position. DINO-WM retains a clear advantage on dynamic and rotational properties (joint velocity, end-effector yaw), likely because such quantities benefit from the richer visual priors learned during large-scale pretraining. All three methods struggle to recover block orientation (quaternion and yaw), suggesting that fine-grained rotational information remains difficult to encode in compact latent spaces regardless of the training strategy.

Panel of Figure 11 — Figure 11: Example of trajectories used for the Violation of Expectation experiments (Sec. 5.2 ). For each environment, the first row corresponds to the unperturbed trajectory, the second row corresponds to a trajectory where a visual perturbation occurs and the third row displays trajectories where the state of the system is randomly reset in the middle of the trajectory. The frame where the perturbation occurs is highlighted in red.

Figure 12: Violation-of-expectation evaluation with PLDM. From left to right: TwoRoom , PushT , and OGBench-Cube . Surprise is plotted over time for unperturbed, visually perturbed, and physically perturbed trajectories. In TwoRoom and PushT , the model assigns significantly higher surprise to both visual and physical perturbations. In OGBench-Cube , the increase in surprise is weaker and not consistently significant.

Figure 13: Visualization of the latent space obtained with LeWM for the PushT environment. On the left, the grid of states is obtained by moving the agent and the block in the x-y plane. On the right, the embeddings of these states are visualized using a t-SNE.

Figure 14: Violation-of-expectation evaluation with DINO-WM. From left to right: TwoRoom , PushT , and OGBench-Cube . Surprise is plotted over time for unperturbed, visually perturbed, and physically perturbed trajectories. While the model detects both perturbations in TwoRoom and PushT , surprise does not increase significantly for either perturbation in OGBench-Cube .

Figure 15: Ablation studies of key design choices in LeWM. Left: effect of the embedding dimension; performance improves with larger embeddings but quickly saturates beyond a certain threshold. Center: effect of the number of random projections used in SIGReg; performance remains stable, indicating that this parameter is not critical. Right: effect of the number of integration knots used to compute the SIGReg loss; results are similarly insensitive to this parameter.

Model	Push-T (SR $\uparrow$ )
DINO-WM	$92.0\pm 1.63$
PLDM	$78.0\pm 5.0$
LeWM (ours)	$\bm{96.0}\pm\bm{2.83}$

Appendix

Table 5: Training Variance. We report the mean success rate across three training seeds and the corresponding variance, evaluated over the same set of 50 trajectories on Push-T. The goal configuration is reachable within 25 steps, and we allow a planning budget of 50 steps. PLDM exhibits higher variance compared to DINO-WM and LeWM.

Figure 16: Effect of the SIGReg regularization weight

\lambda

on Push-T planning performance. Success rate remains above 80% across a wide range of values (

\lambda\in[0.01,0.2]

), peaking near

\lambda=0.09

. Performance degrades sharply only at

\lambda=0.5

, where the regularizer dominates the prediction loss and hinders dynamics modeling. Since

\lambda

is the only effective hyperparameter of LeWM, the SIGReg loss coefficient is easy to tune via a simple bisection search.

pred. size	Push-T (SR $\uparrow$ )
tiny	$80.67\pm 6.54$
small	$\bm{96.0}\pm\bm{2.83}$
base	$86.7\pm 3.06$

Appendix

Table 6: Effect of the predictor size on planning performance in the Push-T environment. We report the success rate (SR). The ViT-S predictor achieves the best performance.

	Push-T (SR $\uparrow$ )
LeWM w/o decoder loss	$\bm{96.0}\pm\bm{2.83}$
LeWM with decoder loss	$86.0\pm 7.54$

Appendix

Table 7: Effect of adding a reconstruction loss during training. We report the success rate (SR) on the Push-T planning task. The model trained without the decoder loss achieves higher performance.

	Push-T (SR $\uparrow$ )
LeWM ViT	$\bm{96.0}\pm\bm{2.83}$
LeWM ResNet-18	$94.0\pm 3.27$

Appendix

Table 8: Encoder Architecture Effect. We report the success rate (SR) on the Push-T planning task. LeWM achieves competitive performance across encoder architectures, with ViT holding a slight edge.

$p$	Push-T (SR $\uparrow$ )
$0.0$	$78\pm 6.54$
$0.1$	$\bm{96.0}\pm\bm{2.83}$
$0.2$	$85.33\pm 5.74$
$0.5$	$66.67\pm 4.11$

Appendix

Table 9: Effect of predictor dropout during training on Push-T planning performance. We report the success rate (SR). A small amount of dropout (

p=0.1

) yields the best results.

	LeWM	PLDM
CEM	$\bm{96.0}\pm\bm{2.83}$	$78.0\pm 5.0$
SGD ( $\beta_{1}$ )	$26\pm 4.32$	$4.67\pm 0.06$
RMSProp ( $\beta_{2}$ )	$67.33\pm 2.49$	$49.33\pm 8.26$
Adam ( $\beta_{1},\beta_{2}$ )	$84\pm 7.12$	$80\pm 3.27$

Appendix

Table 10: Planning Solver Performance. We report the success rate (SR) on the Push-T planning task.

Figure 17: Temporal Latent Straightening on Push-T. Mean cosine similarity between consecutive latent velocity vectors (Eq. 9 ) over training. Higher values indicate straighter latent trajectories. PLDM explicitly encourages temporal regularity through a dedicated temporal smoothness loss (

\mathcal{L}_{\text{time-sim}}

), yet LeWM achieves substantially straighter latent paths as a purely emergent phenomenon, without any temporal regularization term in its objective.

Figure 18: Push-T Training curves for LeWM.

Figure 19: Push-T Training curves for PLDM.

Explore similar work

Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models
May 10, 2026Kai Zhao, Dongliang Nie, Yuchen Lin +4Representation Learning Representation Collapse
1Shanghai University · 2The University of Manchester · 3Nankai University
UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling
Aug 7, 2026An Lanji, Dawei Liu, Jin Li +3World Model Learning Joint-Embedding Predictive Architecture
University of Electronic Science and Technology of China, Chengdu, China.
PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models
Aug 6, 2026Xi Zeng, Haojie Ren, Ziying Song +2Representation Learning World Model Learning
School of Mechanical and Aerospace Engineering, Nanyang Technological University, Singapore · The University of Sheﬃeld, Sheﬃeld, United Kingdom · School of Artiﬁcial Intelligence (School of Software), Yanshan University, Qinhuangdao, China

May 10, 2026J/K move · Enter open · S save

Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models

Kai Zhao, Dongliang Nie, Yuchen Lin, Zhehan Luo+3

1Shanghai University · 2The University of Manchester · 3Nankai University

Representation Learning Representation Collapse World Model Learning Joint-Embedding Predictive Architecture

Open paper AskPDF arXiv

Joint-Embedding Predictive Architectures (JEPAs) provide a simpleframework for learning world models by predicting future latent representations.However, JEPA training is subject to a bias-variance tradeoff.Without sufficient structural constraints, excessive representationalvariance causes the model to collapse to trivial solutions.The recent LeWorldModel (LeWM) shows that this issue can be alleviated bysimply constraining latent embeddings with an isotropic Gaussian prior.However, latent representations inherently lie on low-dimensional manifoldswithin a high-dimensional ambient space, and enforcing an isotropic Gaussianprior directly in this ambient space introduces an overly strong bias.In this work, we propose ame, which seeks a favorable operatingpoint on the bias-variance frontier by applying Gaussian constraints inmultiple random subspaces rather than in the originalembedding space.This design relaxes the global constraint while preserving itsanti-collapse effect, leading to a better balance between trainingstability and representation flexibility.Extensive experiments across fourcontinuous-control environments demonstrate that consistentlyoutperforms LeWM with very clear margins.Our method is simple yet effective, and serves as a strong baseline for future JEPA-based world model research.fdefinedeeemodeThe code is available at https://github.com/intcomp/Sub-JEPA.

Abstract

Explore similar work

Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models

UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models