PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2603.19312v3cs.LGcs.AI

TL;DR

Existing JEPA world models face representation collapse and often depend on complex stabilization strategies. LeWM trains a stable JEPA end-to-end from raw pixels with a two-term objective, achieving competitive control and faster planning while encoding physical structure in its latent space.

  • Problem

    Existing JEPA methods are prone to representation collapse and typically require heuristic regularization, multi-objective losses, external information, or pretrained encoders.

  • Method

    LeWM jointly trains an encoder and predictor from raw pixels with prediction and Gaussian-distribution regularization terms, using offline trajectories and latent-space planning.

  • Results

    LeWM achieves competitive control across diverse 2D and 3D tasks, with a compact 15M-parameter model and planning up to 48× faster than foundation-model-based world models.

  • Takeaways & Limitations

    Encouraging an isotropic Gaussian latent distribution provides a scalable and principled alternative for preventing collapse while supporting efficient latent world modeling.

  • Takeaways & Limitations

    Planning remains restricted to short horizons, and performance depends on offline datasets with sufficient coverage; low data diversity can weaken SIGReg.

Abstract

from arXiv · show

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.

1 Introduction

LeWM addresses JEPA collapse and training complexity with a stable end-to-end approach from raw pixels. It achieves competitive control and faster planning while probing physical structure in its latent space.

  • World Models learn action consequences from sensory inputs, enabling planning and improvement in imagination space.
  • JEPAs encode observations into compact latent spaces and predict future representations, but existing methods are prone to collapse and often require heuristics or extra supervision.
  • LeWM introduces a stable end-to-end JEPA trained from raw pixels with a simple two-term objective on a single GPU.
  • 15M parameters yield competitive control across diverse 2D and 3D tasks, including competitiveness with foundation-model world models at substantially lower cost.
  • 48× faster planning is achieved while preserving competitive performance, and latent-space evaluations probe physical quantities and unphysical trajectories.

2 Related Work

Related work spans generative pixel-space world models, compact-latent JEPA methods, and latent planning approaches. These lines differ in representation learning, stabilization, and whether policies or planners operate in imagination.

  • Generative world models simulate environment dynamics by producing future observations conditioned on past states and actions.
  • JEPA methods predict dynamic evolution in compact latent spaces, with representation-learning variants commonly using EMA and stop-gradient stabilization.
  • Action-conditioned JEPA world models may use pretrained encoders, which avoid collapse but limit representation expressivity.
  • Latent-dynamics approaches can train policies through imagined rollouts, after which the policy executes directly without requiring the world model at test time.
  • LeWM instead performs test-time planning directly in latent space using model predictive control.

3 Method: LeWorldModel

LeWM learns a latent world model end-to-end from offline pixel trajectories by jointly encoding observations and predicting future embeddings conditioned on actions. Its two-term objective combines next-embedding prediction with SIGReg, which promotes isotropic Gaussian latent embeddings to prevent collapse.

  • Training setup: The model is trained on unannotated offline trajectories of raw pixel observations and actions without rewards or task specifications.Behavior trajectories may be pseudo-expert or exploratory if they sufficiently cover environment dynamics.
  • Model architecture: LeWM encodes each frame into a compact latent representation and autoregressively predicts the next embedding from latent history and actions.The encoder and predictor are jointly optimized, with causal masking preventing access to future embeddings.
  • Training objective: LeWM minimizes two terms: a teacher-forced next-embedding prediction loss and a regularization loss.The prediction loss encourages representations that are predictable for the latent dynamics predictor.
  • Training objective: SIGReg prevents trivial collapse by encouraging latent embeddings to match an isotropic Gaussian distribution.It estimates high-dimensional normality through random one-dimensional projections and the Epps–Pulley test statistic.
  • Hyperparameter selection: Only the regularization weight λ is effectively tuned because the number of SIGReg projections has negligible downstream impact.The default settings are M = 1024 projections and λ = 0.1; bisection enables logarithmic-time search.
  • Latent planning: At inference, CEM optimizes action sequences against terminal latent goal matching, while MPC executes only an initial action segment before replanning.The fixed world model rolls out predicted latent states over horizon H, trading longer lookahead against computational cost and accumulated model bias.

4 Latent Planning Performance

LeWM is evaluated for planning across diverse 2D and 3D navigation, motion-planning, and manipulation environments against JEPA and policy-based baselines. It performs strongly on challenging tasks and offers substantially faster planning, while performance depends on environmental complexity and latent-space regularization.

  • Evaluation setup: LeWM is compared with DINO-WM, PLDM, GCBC, GCIVL, and GCIQL, using fixed hyperparameters across environments.PLDM is the closest end-to-end pixel-based baseline but uses a seven-term VICReg-derived objective; DINO-WM uses a pretrained DINOv2 encoder.
  • Planning performance: 18% higher success rate on PushT is achieved by LeWM over PLDM, while LeWM remains competitive with DINO-WM.On PushT, pixels-only LeWM surpasses DINO-WM even when DINO-WM receives additional proprioceptive information.
  • Planning performance: LeWM outperforms PLDM and DINO-WM on Push-T and Reacher, but DINO-WM slightly leads on OGBench-Cube and both baselines outperform LeWM on Two-Room.The authors associate these differences with 3D visual complexity and the difficulty of matching a high-dimensional Gaussian prior in a low-complexity environment.
  • Planning efficiency: 48× faster planning is reported for LeWM, with full planning completing in under one second while preserving competitive performance across tasks.Planning time remains consistent across environments for a fixed planning setup.
  • Ablations and stability: SIGReg projection and integration-knot settings have little effect on performance, leaving λ as the only effective hyperparameter and enabling O(log n) bisection search.Performance saturates beyond a sufficient embedding dimension, and remains competitive when the ViT encoder is replaced with ResNet-18.
  • Ablations and stability: LeWM’s two-term objective converges smoothly and monotonically, unlike the noisy, non-monotonic behavior reported for PLDM’s seven-term objective.The prediction loss decreases steadily while SIGReg rapidly approaches a plateau.

5 Quantifying Physical Understanding in LeWM

LeWM’s latent space retains global scene and spatial structure while supporting recovery of physical quantities and detection of physically implausible events. Its temporal trajectories also become increasingly straight during training without explicit straightness regularization.

  • Latent-space structure: LeWM’s latent representations capture global scene structure, although finer details such as end-effector angle are not fully preserved in decoded rollouts.Rollouts use three context observations and autoregressively generate future latents conditioned on actions.
  • Physical probing: LeWM consistently outperforms PLDM on physical-quantity probing while remaining competitive with DINOv2-based representations.The probing evaluation trains linear and nonlinear probes on Push-T embeddings.
  • Latent-space structure: The latent space preserves sufficient information to reconstruct visual scenes from single 192-dimensional embeddings, despite reconstruction being absent from training.A decoder trained after the world-model training recovers images from latent embeddings.
  • Latent-space structure: t-SNE visualization suggests that LeWM preserves neighborhood relationships, spatial structure, and relative positions in the Push-T environment.
  • Temporal dynamics: LeWM’s latent trajectories become increasingly straight on Push-T during training and achieve higher temporal straightness than PLDM.This behavior emerges without explicit regularization for temporal straightness, whereas PLDM uses dedicated temporal smoothness regularization.
  • Violation-of-expectation evaluation: LeWM assigns higher surprise to physically perturbed trajectories than to unperturbed trajectories across TwoRoom, PushT, and OGBench Cube.Physical perturbations teleport objects, violating expected physical continuity; visual perturbations abruptly change object color.

6 Conclusion

LeWM is a stable end-to-end JEPA for learning latent world models from pixels, using prediction and Gaussian-distribution regularization to avoid collapse. It performs competitively across control settings, but remains limited by short planning horizons and data-coverage requirements.

  • Conclusion: LeWM combines an image encoder and action-conditioned predictor to forecast future embeddings in a compact latent world model.
  • Conclusion: LeWM outperforms prior approaches in data efficiency, planning time, training time, and stability while remaining competitive in task performance across continuous-control environments.
  • Limitations: Planning remains restricted to short horizons, and the method relies on offline datasets with sufficient coverage.Low data diversity can weaken SIGReg in simple, low-dimensional environments.
  • Conclusion: LeWM’s training stability comes from explicitly encouraging latent embeddings toward an isotropic Gaussian distribution to prevent collapse.SIGReg uses random hyperspherical projections and univariate distribution matching; Cramér–Wold connects matching projections to the joint distribution.
  • SIGReg: SIGReg’s Gaussian matching is implemented through Cramér–Wold projections and Epps–Pulley statistics, with the integral approximated by numerical quadrature.The described implementation uses a quadrature scheme such as the trapezoid rule.

B Cross-Entropy Method

The appendix describes CEM-based action-sequence optimization and the training procedures and objectives used for LeWM and comparison policies. CEM iteratively refines a sampling distribution using low-cost elite plans, while the learning baselines use offline goal-conditioned objectives.

  • Cross-Entropy Method: CEM iteratively samples candidate action sequences, evaluates them with the world model, retains the lowest-cost elites, and updates the sampling distribution.The final action plan is obtained from the distribution mean at the last iteration.
  • Cross-Entropy Method: CEM has no guarantee of finding a global optimum in non-convex settings and becomes harder to apply as action-space dimensionality grows.
  • Cross-Entropy Method: In the reported experiments, CEM uses 300 sampled sequences per iteration, retains the top 30 elites, and runs up to 30 optimization steps.The experimental setup uses fewer iterations in some environments, as specified separately in the implementation details.
  • Compared objectives: PLDM uses a seven-term objective with six tunable loss weights, including prediction, variance, covariance, temporal, and inverse-dynamics terms.The paper notes that extensive environment- and dataset-specific tuning limits transferability.
  • Goal-conditioned baselines: GCIQL and GCIVL learn offline goal-conditioned value-based policies, while GCBC uses supervised imitation of expert actions from observation-goal pairs.Policy extraction for GCIQL and GCIVL uses advantage-weighted regression.
  • LeWM training: LeWM’s training procedure encodes pixels, predicts next embeddings from embeddings and actions, and minimizes prediction loss plus weighted SIGReg loss.

D Implementation details

LeWM uses a compact ViT-based encoder-predictor system and autoregressive latent rollouts for planning. The implementation combines frame skipping, short histories, and CEM action optimization.

  • Training setup: Frame-skip 5 groups consecutive actions into blocks, using 224 × 224 frames, batch size 128, and four-frame sub-trajectories.
  • Architecture: The encoder is a ViT-Tiny with patch size 14, while the predictor is a causally masked ViT-S with learned positional embeddings.
  • Architecture: Predictor history length is 3 for PushT and OGBench-Cube and 1 for TwoRoom, with autoregressive generation during planning.
  • Visualization: The visualization decoder reconstructs images from the 192-dimensional [CLS] embedding, but is used only for visualization.
  • Planning: CEM samples 300 action sequences, retains the top 30 trajectories, and plans over a five-step horizon corresponding to 25 environment timesteps.PushT uses up to 30 iterations, while other environments use up to 10.

E Environment & Dataset

LeWM is evaluated across diverse 2D and 3D control environments, using datasets of trajectories for navigation, manipulation, and reaching. The evaluation also examines predictive rollouts, physical quantities encoded in latent representations, and probing performance against PLDM and DINO-WM.

  • Environments: TwoRoom tests continuous 2D navigation through a doorway between two rooms.The dataset contains 10,000 episodes with an average trajectory length of 92 steps.
  • Environments: PushT tests continuous 2D manipulation in which an agent pushes a T-shaped block toward a target configuration.The dataset contains 20,000 expert episodes averaging 196 steps, while models are trained for 10 epochs.
  • Environments: OGBench-Cube tests continuous 3D robotic manipulation, requiring a robotic arm to place a cube at a target location.The dataset contains 10,000 episodes of 200 steps, and models are trained for 10 epochs.
  • Environments: Reacher tests two-joint robotic-arm control for reaching a target configuration in a 2D plane.Models are trained for 10 epochs on 10,000 episodes of 200 steps collected with a Soft Actor-Critic policy.
  • Predictive evaluation: Predictor rollouts on PushT closely track real agent and block motion, while OGBench-Cube predictions preserve the overall trajectory structure.Three context frames are encoded, future latent states are autoregressively predicted from actions, and predictions are decoded using a decoder not used during training.
  • Latent analysis: Physical probing evaluates whether latent embeddings encode agent, block, cube, and end-effector quantities using linear and nonlinear probes.Metrics include mean squared error and Pearson correlation; LeWM matches or outperforms PLDM on the reported TwoRoom and OGBench-Cube properties, with strengths varying by quantity.

F.3 Violation-of-expectation

The violation-of-expectation evaluation tests whether LeWM assigns surprise to trajectories containing visual or physical perturbations. It compares unperturbed, visually perturbed, and physically perturbed trajectories across three environments.

  • Evaluation setup: Each environment contains an unperturbed reference, an abrupt visual color change, and a physical teleportation perturbation.The perturbation frame is highlighted in red in the trajectory examples.
  • TwoRoom: TwoRoom uses agent color changes and agent teleportation as its visual and physical violations.The agent follows an expert policy navigating toward a goal position.
  • PushT: PushT uses abrupt block-color changes and simultaneous agent-and-block teleportation as its perturbations.The agent is controlled by a random policy biased toward interacting with the block.
  • OGBench-Cube: OGBench-Cube uses cube-color changes and cube teleportation during expert pick-and-place trajectories.The expert policy picks up the cube and places it at a target position.

G Ablations.

The ablations examine training stability, representation size, SIGReg settings, predictor architecture, reconstruction, encoder choice, dropout, and planning solvers. Results indicate robust performance across several choices, with identifiable preferred settings and a single effective regularization hyperparameter.

  • Training stability: Training across three seeds produces consistently high success rates with low variance on PushT.Evaluation uses 50 trajectories, a reachable goal within 25 steps, and a planning budget of 50 steps.
  • Representation and SIGReg: Performance drops below an embedding dimension of about 184, then saturates as the dimension increases.Larger embeddings improve performance only up to a threshold.
  • Representation and SIGReg: Varying the number of SIGReg projections has little effect on downstream control, leaving λ as the only effective hyperparameter to optimize.The number of integration knots is similarly insensitive, while λ remains robust over a broad range.
  • Representation and SIGReg: λ values from 0.01 to 0.2 maintain success rates above 80%, whereas λ = 0.5 sharply degrades performance.The degradation occurs when the regularizer dominates the prediction loss; bisection search is suggested for tuning.
  • Architecture: The ViT-S predictor performs best, while ViT-T reduces performance and ViT-B provides no additional gains.The results identify ViT-S as the best trade-off between model capacity and optimization stability for this task.
  • Architecture: Adding a decoder and reconstruction loss slightly decreases downstream control performance compared with training without the decoder loss.The reported result suggests reconstruction is unnecessary for the planning objective.
  • Architecture: LeWM remains competitive with both ViT and ResNet-18 encoders, although ViT retains a modest advantage.This supports encoder-architecture flexibility in the reported PushT evaluation.
  • Regularization: A predictor dropout rate of 0.1 yields the highest success rate, while lower and higher rates perform worse.The ablation attributes the pattern to regularization and degradation from excessive dropout.

I Training Curves

Training curves compare LeWM with PLDM and examine temporal straightening on PushT. LeWM’s simpler objective makes its optimization dynamics easier to interpret, while straighter latent trajectories emerge during training without explicit temporal regularization.

  • Training curves: LeWM uses one regularization term alongside prediction loss, whereas PLDM contains multiple regularization terms.The simpler objective makes LeWM’s training dynamics easier to interpret and analyze.
  • Temporal latent straightening: LeWM’s temporal straightening emerges naturally during training without an explicit temporal regularization term.A higher S_straight indicates straighter latent trajectories because consecutive velocities are more nearly collinear.
  • Training curves: The training-curve comparison uses PushT curves for LeWM and PLDM.These curves provide the optimization-dynamics comparison referenced in the section.
Loading 2603.19312v3…