Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI 👨🏻‍💻🤖

LeWorldModel: Learning Action-Conditioned Dynamics from Pixels

3 minute read

Published:

TL;DR: LeWorldModel learns visual representations and action-conditioned latent dynamics together. A next-embedding prediction loss is paired with SIGReg to discourage collapse. A planner can then search for action sequences whose predicted final representation approaches an image goal.
First page of LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels
Paper preview: LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels.
Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun, and Randall Balestriero
Read the paper (PDF)

A world model must connect actions to consequences

Intuition first. A robot sees a block and considers pushing it left or right. A useful model must predict different consequences for those actions, while retaining the information needed to judge progress towards the goal.

Maes et al. introduce LeWorldModel, or LeWM, which learns an encoder and predictor jointly from pixel observations and actions. It uses prediction plus Gaussian-distribution regularisation rather than a pretrained visual encoder, a pixel-reconstruction objective, or a moving-average target encoder. The method connects the representation geometry of LeJEPA to a sequential control problem.

Two loss terms, jointly differentiated

Let \(z_t=f_\theta(o_t)\) encode observation \(o_t\). Suppressing the predictor’s history window for readability, the objective is

\[ \begin{aligned} \hat z_{t+1}&=g_\phi(z_t,a_t),\\ \mathcal L_{\mathrm{pred}}&=\|\hat z_{t+1}-z_{t+1}\|_2^2,\\ \mathcal L&=\mathcal L_{\mathrm{pred}}+ \lambda\operatorname{SIGReg}(Z). \end{aligned} \]

\(Z\) collects training embeddings. SIGReg compares random projections with a Gaussian target. Gradients flow through both predicted and target representations; the complete predictor uses a causal history and action conditioning. Training formulation.

Prediction alone permits an uninformative constant representation. Regularisation constrains its distribution, while the action-conditioned prediction task supplies the relationship between successive observations.

Learn representations that predict action consequencesThe prediction loss shapes dynamics and representations jointly. SIGReg separately constrains the embedding distribution. Encode observationsPredict with actionsCompare next features
The prediction loss shapes dynamics and representations jointly. SIGReg separately constrains the embedding distribution.

Worked example: choosing an action sequence in latent space

Consider a deliberately simple scalar latent model \(\hat z_{t+1}=\hat z_t+a_t\). Start at \(z_0=0.4\) and set the goal to \(z_g=1.0\). Over three steps, candidate A applies \(0.2\) each time; candidate B applies \(0.1\) each time.

\[ \begin{aligned} \hat z_3(A)&=0.4+3(0.2)=1.0, & J(A)&=0,\\ \hat z_3(B)&=0.4+3(0.1)=0.7, & J(B)&=0.09, \end{aligned} \]

where \(J=(\hat z_3-z_g)^2\). The model prefers A. This arithmetic illustrates planning, not LeWM’s learned architecture or a physical unit conversion.

If the real latent transition adds an unmodelled 0.02 per step, A finishes at 1.06 instead. Its goal error becomes \(0.06^2=0.0036\). Even a small systematic one-step error can change a planned outcome over a longer horizon.

Teacher forcing and rollout are different conditions

During one-step training, the predictor receives representations of observed history. In a multi-step imagined rollout, its predictions become part of the context for later predictions. Errors can therefore move it away from the inputs used during training.

Evaluate rollouts at the horizons used by the controller, and test whether candidate trajectories are ranked correctly. Low average one-step error does not establish either property. The authors’ project page describes image-goal planning with the cross-entropy method.

Representation quality and task success need separate checks

A latent feature can encode position while ignoring whether an object is securely grasped. If that distinction matters for the task, a low goal distance can be misleading. Probe relevant physical quantities and evaluate full executions rather than treating a smooth latent trajectory as a success certificate.

The paper’s control and physical-prediction experiments support results within their protocols. Model size and planning speed comparisons also depend on the chosen baseline, hardware, and horizon.

Key Insight: Joint representation and dynamics learning makes predictability part of the encoder's objective. Planning still requires accurate action consequences and a goal representation that preserves task-relevant distinctions.

Read learned MPC for the control loop, and compare V-JEPA 2 for a pipeline using large-scale visual pretraining before action-conditioned post-training.

References

  1. Maes, L., et al. LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels. 2026.
  2. Authors’ project page and implementation.

You May Also Enjoy

The Standard Library Modules Worth Knowing Before You pip install

8 minute read

Published:

Python ships with a counter, a queue, a memoiser, a CLI parser, a test framework’s smaller cousin, and a JSON-safe date library already installed, most of what people reach for a package to solve is one import away.

Files and Context Managers: Why with Is Not Optional

9 minute read

Published:

A file object that falls out of scope without with does get closed eventually, but ‘eventually’ means whenever the garbage collector gets around to it, which is not a promise any program handling more than a handful of files can live with.

Errors and Exceptions: EAFP, the Hierarchy, and Reading a Traceback

9 minute read

Published:

Exceptions in Python are not exceptional. They are a normal control-flow mechanism that the language leans on so heavily that the idiomatic style is to try the operation and handle the failure, rather than check first, which is faster, shorter, and free of race conditions.