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

1. Why I Care About World Models

When I talk about world models, I am not simply referring to another fashionable phrase in artificial intelligence. I mean something more fundamental: a learned representation of how the world changes. In my formulation, a world model takes the current state or observation of a system, combines it with a candidate action, and predicts what will happen next. Symbolically, I can think of this as taking a situation sss, applying an action aaa, and predicting a future situation s′s’s′.

This capability matters because intelligence is not just reactive pattern matching. An intelligent agent should be able to ask: “What will happen if I do this?” Before acting, it should be able to simulate, anticipate, and evaluate possible futures. That is why world models are powerful. They allow an agent to imagine candidate outcomes, perform model-based control, and estimate surprise or uncertainty when reality diverges from expectation.

In practical terms, this means that a robot, a game-playing agent, or a physical AI system can use an internal predictive model to reason before it acts. Instead of merely mapping observations directly to actions, the agent can internally roll out possible futures and choose actions with better foresight. This is the conceptual foundation of the work I am presenting.

2. World Models Are Not a New Idea

Although world models are currently attracting enormous attention, the core idea is old. I emphasize this because scientific progress is often packaged as novelty when it is actually a refinement of older principles. Richard Sutton’s Dyna architecture, proposed in 1990, already described something strikingly similar to what we now call a world model: a black-box action model that takes a situation and an action as input and predicts the next situation.

The historical continuity is important. Dyna integrated learning, planning, and reacting. It allowed an agent to use real experience, learn a predictive model from that experience, and then use imagined experience for planning. That is exactly the intellectual lineage behind modern world-model research.

What has changed today is not the conceptual ambition but the computational substrate. Instead of small tabular models or hand-designed representations, we now train neural systems on high-dimensional sensory observations such as images, LiDAR, proprioception, and video. The question is no longer whether world models are useful in principle. The question is whether we can train them stably, efficiently, and scalably from raw perceptual data.

3. From State Prediction to Observation Prediction

In idealized reinforcement learning, we often talk about states. But in real systems, the agent rarely has access to the true underlying state of the world. A robot does not directly perceive “the state”; it receives observations through sensors. These observations may include camera images, joint angles, velocities, tactile readings, or depth scans.

That is why I shift the notation from sss to ooo. A world model in realistic embodied AI should predict what an agent will observe after executing an action. Given a current observation ooo and an action aaa, it predicts a future observation o′o’o′. This is a much harder problem than simple state transition modeling because observations can be extremely high-dimensional and partially observable.

For example, if a quadrotor sees a room through a forward-facing camera and executes a yaw command, the model should predict how the camera view will change. It should understand, implicitly, that rotating left changes the visual field, that nearby objects move across the frame, and that the geometry of the environment remains coherent. This is not merely image generation. It is action-conditioned physical prediction.

4. Model-Free and Model-Based Policies

The central strategic question is whether intelligent agents should be model-free or model-based. In a model-free policy, a neural network maps observations directly to actions. The internal mechanism may be sophisticated, but there is no explicit step in which the agent predicts future outcomes. The policy simply learns: “Given this observation, take this action.”

In a model-based policy, by contrast, the agent explicitly uses an internal world model. It asks what will happen under possible actions, evaluates those imagined futures, and then selects an action. This makes the policy more modular and potentially more interpretable. The action is not just emitted by an opaque network; it is chosen after internal forecasting.

However, model-based control also has costs. The agent needs a mechanism for proposing candidate actions, rolling them out, and evaluating them. This can be computationally expensive. Yet the benefit is substantial: if the world model is accurate, the agent can reason counterfactually. It can consider actions it has not yet taken and estimate their consequences before committing to them.

5. General Agents Need Predictive Models

I view world models as more than engineering tools; I see them as a prerequisite for generalization. A system that can only react to familiar observations may perform well inside its training distribution, but it becomes brittle when the world changes. A system with a predictive model can do something more abstract: it can infer how actions transform the environment.

The claim that general agents need world models is not merely rhetorical. If an agent must solve multi-step, goal-directed tasks, it needs some way to evaluate future consequences. It must relate present observations, possible actions, intermediate states, and goals. Without a predictive structure, it is forced to rely on memorized action patterns or implicit correlations.

This does not mean every useful agent must contain an explicit symbolic simulator. But it does suggest that some predictive latent structure is necessary. Whether that structure is explicit, modular, and accessible — or hidden inside the weights of a large model — is one of the most important open questions in AI.

6. A Toy Example: Push-T

To make the idea concrete, I use a simple Push-T environment. The task contains a small blue agent that must push a T-shaped object into a target slot. The observation includes an image of the scene and the position of the end effector. The action specifies where the end effector should move next.

This toy environment is useful because it distills the essence of world modeling. The agent must understand how its actions affect objects over time. It must learn that pushing at one point changes the pose of the T-shaped block, and that different actions produce different trajectories.

In my example, I train a compact neural architecture to predict the next state from the current observation and action. The predicted rollout is imperfect and somewhat unstable, but it captures the essential dynamics. Even in this small domain, we can see the core difficulty: the model must simultaneously learn a visual representation and the action-conditioned dynamics of that representation.

7. The Central Training Problem: Collapse

The hardest part of training world models is not merely making predictions. It is learning a useful representation while also learning dynamics. These two objectives are entangled. The model must compress high-dimensional sensory input into a latent space, but that latent space must preserve the information needed for future prediction and control.

This creates a dangerous optimization problem. A trivial solution may exist: the model can collapse all observations into the same representation. If every state looks identical in latent space, prediction becomes easy but useless. The model has minimized a superficial loss while destroying the information required for intelligent behavior.

Many world-model methods introduce tricks to avoid this collapse. Some use stop-gradient mechanisms, exponential moving averages, pretrained encoders, reconstruction losses, privileged state information, or carefully tuned regularizers. These methods may work empirically, but they often make the architecture more complicated and less elegant.

My motivation in LeWorldModel is to ask whether we can obtain stable end-to-end learning with a simpler and more principled regularization mechanism.

8. LeWorldModel: A Joint-Embedding Predictive Architecture

LeWorldModel is built around a joint-embedding predictive architecture. Instead of predicting future pixels directly, I predict future latent embeddings. The system encodes an observation oto_tot​ into a latent vector ztz_tzt​. Then, given an action atatat​, a predictor estimates the next latent vector z^t+1\hat{z}{t+1}z^t+1​. This prediction is compared against the encoded latent representation of the actual next observation zt+1z_{t+1}zt+1​.

This matters because pixel-space prediction is often wasteful. Pixels contain many details that are irrelevant for control, such as lighting variation, texture noise, and background appearance. A latent predictive model can focus on the structure that matters: object positions, agent configuration, contact dynamics, and task-relevant geometry.

The key contribution is the use of SIGReg regularization. SIGReg encourages the latent embeddings to remain Gaussian-distributed, isotropic, and feature-diverse. In simple terms, I want the latent space to stay healthy. It should not collapse into a point, stretch into a degenerate line, or encode everything redundantly. By regularizing the distribution of embeddings, I can stabilize end-to-end representation learning and dynamics learning.

9. Action-Conditioned Forecasting

The first capability I evaluate is open-loop action-conditioned forecasting. In open-loop prediction, the model receives an initial context and a sequence of actions, then predicts future latent states without being corrected by feedback at every step. This is a stringent test because errors can compound over time.

In Push-T and robotic manipulation settings, the model must predict how objects evolve as actions are executed. If the model merely memorizes static appearances, it will fail. It must encode causal consequences: pushing, moving, rotating, and contact-induced displacement.

The results show that LeWorldModel can generate coherent future predictions from pixels. The imagined trajectories are not perfect, but they retain the task-relevant structure. This is precisely what I want from a world model: not photorealistic fantasy, but useful counterfactual forecasting.

10. Model Predictive Control in Latent Space

The second capability is model predictive control. Given an initial observation and a goal observation, I encode both into latent space. Then I search for a sequence of actions that moves the predicted latent trajectory from the initial latent state toward the goal latent state.

This is an important shift. Instead of planning directly in pixel space, I plan in representation space. The world model rolls out future latent states under candidate actions, and an optimizer evaluates which action sequence best reduces the distance to the goal. In the implementation discussed here, methods such as the cross-entropy method can be used to iteratively refine candidate action sequences.

There is a caveat: not every real-world task conveniently provides a goal image. Many robotic tasks specify goals more abstractly, such as “place the object on the shelf” or “open the drawer.” Still, goal-conditioned latent planning is a valuable demonstration because it shows how a learned predictive model can be converted into a controller.

11. Empirical Control Performance

When I compare LeWorldModel with competing methods, the results are nuanced. LeWorldModel performs strongly on relatively low-complexity 2D tasks such as Push-T and Reacher. In more visually complex 3D environments, pretrained-encoder approaches such as DINO-based world models may outperform it because they benefit from large visual priors.

This distinction is important. LeWorldModel is not a claim that pretraining is useless. Rather, it demonstrates that stable end-to-end predictive learning from pixels can be achieved with a compact architecture and an elegant regularizer. The model is also computationally efficient. Because much of the work happens in latent space and the architecture remains small, it can be significantly faster than heavier alternatives.

The broader lesson is that world-model performance is not one-dimensional. We should evaluate accuracy, stability, compute cost, memory footprint, controllability, and robustness. A method that wins on raw performance but requires massive pretrained backbones may not be the right solution for every embodied system.

12. Surprise, Error, and Violation of Expectation

The third capability is the one I find especially compelling: surprise evaluation. A world model does not only predict the future; it can detect when the future violates its expectations.

In the experiment, I compare nominal trajectories with perturbed ones. For example, the object may suddenly change color, or it may be teleported to a new location. These interventions violate what the model expected to observe. If the world model is meaningful, its prediction error should increase sharply at the moment of perturbation.

This matters for real-world deployment. A robot should know when its internal model is failing. If a cup slips, a door jams, a human intervenes, or an object is unexpectedly displaced, the robot should not continue blindly. It should detect the anomaly and adapt. This makes world models relevant not only for planning but also for out-of-distribution detection, failure detection, and uncertainty estimation.

13. Quantifying Poor Predictions

The prediction-error plots show this phenomenon clearly. When a perturbation is applied, the model’s mean squared error spikes. This spike functions as a quantitative surprise signal. It tells the agent: “The world is no longer evolving according to my expectations.”

This is a major advantage over purely model-free policies. A model-free policy may still act confidently even when it has entered an unfamiliar regime. A world-model-based agent has a native mechanism for estimating mismatch between prediction and reality. That mismatch can be used to trigger replanning, request human intervention, slow down execution, or switch to a safer policy.

In this sense, prediction error is not merely a training metric. It becomes an operational signal. It allows the agent to monitor its own epistemic limitations. For embodied AI, that kind of self-monitoring is not optional; it is a prerequisite for reliable autonomy.

14. Broader Themes and Open Questions

The broader question is whether the future of intelligent agents will be model-free, model-based, or some hybrid of the two. I do not think the answer is obvious. Model-free systems are simple at inference time and can be extremely powerful when trained at scale. Model-based systems are more modular, more interpretable, and better suited to counterfactual reasoning, but they introduce additional computational and architectural complexity.

A second question concerns representation learning. Should an agent learn its representation and dynamics jointly, as LeWorldModel does, or should it use a frozen pretrained representation? Joint learning is elegant and task-adaptive, but it risks collapse. Pretraining provides strong visual priors, but it may encode features that are not optimized for control.

A third question is how to fight representational collapse elegantly. Many methods work, but not all are conceptually satisfying. My view is that SIGReg offers a clean contribution: it regularizes latent geometry directly, helping the model maintain a diverse and usable representation while still learning predictive dynamics end to end.

Ultimately, LeWorldModel is an argument for compact, stable, end-to-end world modeling from pixels. It does not solve every problem in model-based AI, but it clarifies an important path forward: learn useful latent dynamics, keep the representation healthy, plan in that latent space, and use prediction error as a signal for surprise. That combination — forecasting, control, and uncertainty estimation — is why world models remain one of the most consequential directions in embodied intelligence.