When does pixel reconstruction help a JEPA world model plan?

Code

TL;DR: JEPA world models can lose track of a moving object when the background barely changes. In our experiments, adding a strongly weighted reconstruction loss fixed this, raising planning success on frozen backgrounds from 11% to 71%, while costing about 4 percentage points on flickering ones.

While LLMs have proven their worth in the digital realm, there exists a considerable gap in how they adapt to perform real tasks in the physical world. VLA (Vision Language Action) models - a derivative of the VLM (Vision Language Model) - built upon the LLM architecture to meet the physical world, are currently the front runner for most robotics startups (e.g., Figure AI, Physical Intelligence). While VLAs are making progress, some argue they have their limitations. Namely, they map what they see to an action and are reactive in nature. They do not predict what their actions cause and primarily rely on large amounts of robot demonstration data to “learn”. A compelling alternative, one rising in popularity, is world models. Unlike VLAs, they predict what happens after an action, so a robot can imagine outcomes and plan its actions. General Intuition, AMI Labs, and Google DeepMind are a few companies focused on this paradigm.

Such models have a wide range of applications, ranging from robotics and interactive games to self-driving and biology. Today, “world models” are used as a blanket term across the industry. For the ones we care about, they can be loosely grouped into two camps: generative and latent world models. While both learn to predict what happens next after an action, they differ in what they predict. Generative models predict the next entire observation, like an image (pixels), while latent models predict a compressed representation of it.

A prominent example of the latter is JEPA (Joint Embedding Predictive Architecture), proposed by Yann LeCun (2022). The underlying premise is that the model predicts in representation space instead of reconstructing the entire image, and can learn to ignore details that are random or irrelevant to the task, such as sensor noise. For instance, a world model trying to predict where a car will be next does not need to exactly predict how the leaves on the trees behind it will sway. In contrast, a generative model has to render them anyway, leading to more effort on irrelevant details. The bet is that much of the future is unpredictable, and a model focusing on predicting only what’s necessary in a compact representation learns more useful features while being cheaper to plan with.

However, this freedom to ignore details poses a risk of ignoring the wrong thing. Sobal et al. (2022) showed that JEPA models tend to fixate on aspects of a scene that change slowly, like the background, a tendency known as the slow-feature bias. This happens because the model is rewarded for predicting its next representation accurately, and anything in the scene that’s fixed is the easiest to predict.

Generative models are much less prone to this problem, since they must redraw the entire scene, including the moving object, to predict the next frame. This raises a natural question: can we get the best of both? Specifically, does adding pixel reconstruction to a JEPA model stop it from losing track of the object, and when does that actually help it plan? To test this, we built a small environment of a single dot moving over a noisy background, varying how quickly the background changes, from frozen to flickering. We found that on stable backgrounds, a sufficiently weighted reconstruction loss improved planning from 11–19% to 71–79% (Fig. 1), with a slight reduction on flickering backgrounds. We repeat each main comparison with three training seeds and evaluate on unseen data.

One thing to acknowledge before proceeding is that reconstruction partly goes against JEPA’s core idea of ignoring irrelevant detail, a trade-off we discuss later on. Our goal is simply to see how close we can get to the best of both.

How a JEPA model works

A JEPA model consists of an encoder and a predictor. At the simplest level, the encoder converts an image, or state, $o_t$ into an embedding $z_t$ in latent space. The predictor takes that embedding and an action $a_t$ to predict the embedding of the next state, $\hat{z}_{t+1}$. During training, the model learns to make this prediction match the encoder’s embedding of the actual next state, $z_{t+1}$ (Fig. 2). The model can be prone to cheating, commonly known as representation collapse, where the encoder outputs the same embedding for every image, making the prediction trivially perfect. LeWorldModel (LeWM; Maes et al., 2026), a model based on JEPA, prevents this issue with SIGReg (Sketched Isotropic Gaussian Regularization; Balestriero & LeCun, 2025), which pushes the embeddings of different images to spread out evenly in every direction, so they cannot all collapse to a single point. However, SIGReg does not address the slow-feature bias described earlier, where the model focuses on the wrong features.

Where existing fixes fall short

This slow-feature bias has received growing attention. A recent study by Gulati & Nemenman (2026) showed that the issue worsens gradually as the background changes more slowly. However, they only studied contrastive JEPA models. In such approaches, the models avoid collapse by comparing clips against each other, eventually learning to discern different clips. LeWM, on the other hand, uses SIGReg instead to avoid collapse. Recent JEPA work has been moving towards LeWM-style approaches, partly because contrastive learning requires a large number of negative samples, or mismatched clips to compare against.

LeWM is not immune to this either. MotionJEPA (Karmann et al., 2026) showed that static background distractors can severely degrade LeWM’s planning. They went on to propose a training objective that encourages the model to retain what changes between frames by predicting an embedding of the difference between consecutive frames, where anything static cancels out. However, their planning experiments kept these distractors fixed within each trajectory, without varying how quickly the background changes. In reality, backgrounds are rarely perfectly still and a slowly changing background may be just as bad as a static one. It is also unclear whether MotionJEPA’s fix holds up in that case, since a drifting background no longer fully cancels out between frames.

A simpler alternative is to borrow from generative models. As discussed earlier, generative models are much less prone to the slow-feature bias because they must redraw the entire frame, moving object included, so adding a reconstruction loss to a JEPA model could have a similar effect. This idea is not new to latent world models. Dreamer (Hafner et al., 2020) learns its latent space partly by redrawing images, and LRC-JEPA (Huang et al., 2026) explores a more elaborate version, routing slowly changing information into a separate branch used only for reconstruction. However, reconstruction comes with its own trade-off: in reinforcement learning, reconstruction-based world models are known to struggle when backgrounds contain complex moving distractors (Deng et al., 2022), since they waste effort redrawing them. The LeWM authors also found that adding a reconstruction loss did not improve planning on clean scenes. That being said, we do not know if a simple reconstruction loss helps LeWM when a distracting background is present, and if it does, how much weight it needs. So we look at two things: how LeWM behaves as the background changes more slowly, and whether reconstruction or MotionJEPA can fix it. For this, we need an environment where we have full control over the background.

Experimental Setup

We use the moving-dot environment from Sobal et al. (2022). The environment consists of a 28×28 image with a moving dot where an action moves it a small amount. The model is then trained on 20,000 clips of the dot randomly moving on top of a noisy background. The noise is relatively faint with a per-pixel standard deviation  (σ) of 0.03 in relation to the dot’s brightness.

The speed with which the background changes is controlled with the ρ parameter. We test five speeds: frozen (ρ = 1), slow drift (ρ = 0.99), moderate drift (ρ = 0.9), fast drift (ρ = 0.5), and flicker (ρ = 0), where the noise is new every frame (Fig. 3). Each clip has its own random background. As a reference, we also train and test plain LeWM with no background noise—just the dot on a blank image. The chart labels this “No background.”

During testing, we give the model a goal image from later in the same clip, and use a planner (CEM) to move the dot towards it. If it ends up within 1.5 pixels of the goal within 20 moves, we count it as a success. For reference, random actions succeed 16.4% of the time, while a controller given the true dot position succeeds on all 500 goals.

We then compare four models: plain LeWM, a pixel model, LeWM + reconstruction, and MotionJEPA, all using the same encoder and predictor, and the same training settings. Each model is trained three times with different random seeds and tested on the same set of 500 unseen goals. Unless stated otherwise, the percentages below are averages across those three runs, rounded to the nearest whole percent; ranges span background conditions.

Results

Where plain LeWM fails

On a clean background with no noise, plain LeWM plans well, reaching the goal 85% of the time. However, once a background is present, its performance relies heavily on how quickly the background changes. For instance, with frozen backgrounds, slow drift, or moderate drift (ρ = 0.9), planning drops to 11–19%: below random actions on frozen backgrounds and roughly at random with slow or moderate drift. It only retains performance when the background changes quickly, reaching 83% with fast drift (ρ = 0.5) and 86% with flicker (Fig. 4).

To investigate this, we checked if the dot’s position could still be read from the model’s embeddings. On frozen backgrounds, the position readout was poor, with a mean error of 8.82 pixels. Despite this, the model could identify the source clip 100% of the time in a check that fixed the dot positions across clips, leaving only the background to distinguish them. In other words, the model had learned to recognize backgrounds instead of tracking the dot.

This follows from how the model is trained. Since the background stays the same throughout a clip, so does the embedding, making the next one trivial to predict. At the same time, each clip has a different background, so the embeddings still differ across clips, which is enough to satisfy SIGReg. As a result, the model satisfies its training goals without needing to track the dot, essentially using the background as a name tag for each clip. With a flickering background, this shortcut no longer works, since the random noise is hard to predict frame to frame.

The pixel model, on the other hand, holds up much better on stable backgrounds, reaching 72% on frozen. However, it falls behind plain LeWM on flicker (77% vs. 86%), since it wastes effort trying to predict the noise. In other words, each approach fails where the other succeeds, leading back to our question of whether a JEPA model can get the best of both.

Fixing it with reconstruction

The pixel model largely avoids this shortcut by redrawing the entire frame. So, we add a similar pressure to LeWM. Alongside its usual prediction task, the model also redraws the current image from its embedding. We use the β parameter to control the extent to which this redraw influences training (Fig. 5).

With small weights, it essentially had no impact. In the single-seed weight sweep, β = 0.1, 1, and 30 reached 12.8%, 9.0%, and 13.0% planning success on frozen backgrounds, respectively, compared with 9.8% for plain LeWM using that seed. One likely reason is the size of the dot in relation to its background, only taking up a few of the 784 pixels. So a redrawn image can leave the dot out entirely and still have a low average pixel error. The gradient diagnostic supports this: at β = 1, the reconstruction gradient norm on the encoder was just 0.39% of the JEPA gradient norm at initialization and 0.08% at the end of training. Of the weights tested, only β = 300 produced a large improvement, reaching 69.4% for that seed and 71% averaged across all three seeds (Figs. 6 and 7).