🤖 machine learning

XP-JEPA: Cross-Predictive Physics Grounding for Forecastable Latent Dynamics

The paper introduces XP-JEPA, a cross-predictive learning framework that grounds visual latent dynamics in privileged physical trajectories during training to significantly improve forecastability and control performance, while discarding the physical branch for visual-only deployment.

Original authors: Kehan Wen, Ziming Li, Siyuan Luo, Fan Shi

Published 2026-08-26

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Kehan Wen, Ziming Li, Siyuan Luo, Fan Shi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Robots that can plan ahead often rely on a mental shortcut: they build an internal map of the world, a kind of simulation where they can test out different moves before actually making them. To do this, the robot takes a picture of its surroundings, turns that image into a simplified set of numbers, and then asks a computer program to guess what those numbers will look like a moment later if the robot pushes, lifts, or turns something. If the guess matches the goal, the robot executes the move. This process works well when the robot's internal map accurately reflects how the physical world changes. However, a common problem arises when the robot learns to predict the future based only on pictures. It can become too clever for its own good, finding patterns in the visual data that are easy to predict but have nothing to do with real physics. For instance, the robot might learn that a certain shade of gray usually follows a certain shade of blue, even if the actual object hasn't moved at all. This creates a fragile understanding where the robot's predictions drift away from reality, causing it to fail when it tries to interact with the world.

A team of researchers at the National University of Singapore has developed a new method to fix this flaw, ensuring that the robot's internal predictions stay grounded in the laws of physics. They call their approach XP-JEPA. Instead of letting the robot learn solely from watching video, they gave it a second stream of information during training: a direct readout of the physical state of the objects, such as their exact position, size, and orientation in space. This physical data is like a hidden layer of truth that the robot cannot see with its camera but can access while it is learning. The researchers designed a system where the robot's visual brain and its physical brain work together. Both sides receive the same history of actions and try to predict the future. The visual side predicts what the next picture will look like, while the physical side predicts the next set of coordinates. Crucially, the system forces these two predictions to agree with each other. If the visual side predicts that a block will move up, but the physical side predicts it will stay still, the system learns that the visual prediction is wrong. This cross-checking process teaches the visual system to pay attention to the real physical changes that matter, rather than just guessing based on superficial visual patterns.

The researchers tested this method on a wide variety of robotic tasks, including pushing blocks, stacking objects, and tossing items into containers. They compared their new system against standard models that learn only from pictures. The results were striking. When the researchers let the new system run a simulation of a task, the error in its prediction grew very slowly over time. In contrast, the standard model's prediction drifted significantly, becoming unreliable after just a few steps. Specifically, the new method reduced the prediction error from a value of 0.361 down to 0.104. This improvement in accuracy translated directly into better performance. When the robots were asked to actually perform the tasks, the new system succeeded 78.2 percent of the time, a substantial jump from the 53.6 percent success rate of the standard visual-only model. The improvement held true across six different types of interactions, showing that the method works for a broad range of physical challenges.

One of the most important findings of the study is what happens when the robot stops using the physical data. The physical information is only available while the robot is in the training phase; it is not something a real-world robot can measure during a mission. The researchers found that once the robot finished learning with the help of the physical data, they could throw away the physical sensors entirely. The robot then operated using only its camera, just like the standard models. Despite losing access to the physical data, the robot retained the improved ability to predict the future. This suggests that the robot had successfully internalized the rules of physics into its visual understanding. It learned to see the world in a way that naturally respects how objects move and interact, rather than just memorizing visual tricks.

The study also explored whether simply teaching the robot to recognize where objects are located would be enough to solve the problem. They tried a different approach where the robot was forced to guess the position of an object directly from a single snapshot of the image. While this method made the robot very good at identifying where an object was at a specific moment, it did not help the robot predict how that object would move in the future. The robot could still guess the position accurately, but its predictions about the future path of the object were just as unreliable as the standard model. This distinction is vital: knowing where something is right now is not the same as understanding how it will change. The new method succeeded because it focused on the relationship between actions and future states, using the physical data to guide the learning of that relationship, rather than just teaching the robot to label the current scene.

The researchers also tested what would happen if the connection between the visual and physical data was broken. In one experiment, they paired the visual video of one task with the physical data of a completely different, unrelated task. In this scenario, the robot's predictions became chaotic, and its ability to control the robot collapsed to a success rate of just 16.6 percent. This result proves that the system relies on the correct pairing of visual and physical information. It is not enough to have access to physical data; the robot must learn the specific link between what it sees and how that scene physically evolves. The study confirms that the improvement comes from learning the correct dynamics of the world, not just from having more data.

In the end, the work demonstrates a path toward more reliable robotic planning. By using a temporary, privileged view of the physical world to train the robot's visual system, the researchers created a model that is far better at imagining the future. The robot learns to anticipate the consequences of its actions with greater precision, reducing the risk of failure when it tries to manipulate real objects. While the current experiments were conducted in a simulated environment, the approach offers a clear strategy for teaching machines to understand the physical world not just as a collection of images, but as a system of moving parts governed by consistent rules. The robot does not need to carry heavy physical sensors into the real world; it only needs to have learned the right way to see.

Technical Summary: XP-JEPA: Cross-Predictive Physics Grounding for Forecastable Latent Dynamics

1. Problem Statement

Latent world models enable planning by searching in representation space: an agent encodes the current scene, rolls forward candidate actions through learned dynamics, and selects actions whose predicted outcomes approach a goal. This approach relies on representations with reliable action-conditioned dynamics.

However, self-predictive models (where the encoder and predictor are jointly optimized) often suffer from a specific failure mode: the representation and predictor co-adapt to latent transitions that are internally predictable but only weakly constrained by the underlying physical evolution of the scene. As illustrated in the paper, two distinct physical states (e.g., a block held by a gripper vs. a block resting on a table along the same camera ray) may project to nearly identical images. A self-predictive objective does not require their latent transitions to separate, even if they evolve in opposite ways under the same action. This ambiguity leads to "rollout drift," where predicted futures diverge significantly from physical reality, undermining control performance.

While privileged physical states (e.g., object poses, velocities) provide an external reference to reduce this ambiguity, existing approaches often use them merely to supervise representation content (via distillation or state regression) or downstream behavior. The paper argues that making physical state decodable from a single latent snapshot does not necessarily constrain how those latents evolve over time under action. The core challenge is to use privileged trajectories to supervise the predictive transition itself, ensuring the learned dynamics are grounded in physical reality without requiring privileged inputs at test time.

2. Methodology: XP-JEPA

The authors introduce Cross-Predictive JEPA (XP-JEPA), a framework that treats visual observations and privileged physical states as two views of the same trajectory. Unlike methods that distill a fixed physical target into a visual representation, XP-JEPA couples the evolution of both modalities through a shared predictor.

Core Architecture and Objective

XP-JEPA builds upon the LeWM (Joint-Embedding Predictive Architecture) baseline, which jointly learns a visual encoder $f_\theta$ and an action-conditioned predictor $g_\psi$ . XP-JEPA augments this with a physical branch:

Dual Encoders: A visual encoder $f_\theta$ processes observations $o_t$ , and a physical encoder $h_\phi$ processes privileged states $s_t$ .
Shared Predictor: A single action-conditioned predictor $g_\psi$ advances histories from either modality ( $Z^o_t$ or $Z^s_t$ ) forward in time.
Cross-Predictive Loss: The model optimizes a loss function containing four terms:
- Within-modality terms: Visual self-prediction ( $g_\psi(Z^o_t, a_t) \to z^o_{t+1}$ ) and physical self-prediction ( $g_\psi(Z^s_t, a_t) \to z^s_{t+1}$ ).
- Cross-modal terms: A visual history predicts the future physical representation ( $g_\psi(Z^o_t, a_t) \to z^s_{t+1}$ ), and a physical history predicts the future visual representation ( $g_\psi(Z^s_t, a_t) \to z^o_{t+1}$ ).

The loss function is defined as:
$\mathcal{L}_{\text{XP-JEPA}} = \sum_{m,n \in \{o,s\}} \| g_\psi(Z^m_t, a_t) - z^n_{t+1} \|^2_2 + \beta [\Omega(z^o) + \Omega(z^s)]$
where $\Omega$ is an isotropy regularizer (SIGReg) to prevent representation collapse.

Unified Privileged-State Interface

To apply this across heterogeneous tasks (e.g., grasping, pushing, stacking), the authors design a unified rigid-body schema. Scenes are represented by a set of tokens:

Object Tokens: Contain canonical geometry (pooled Signed Distance Field), half-extents, pose (rotation and position).
Effector Token: End-effector rotation, position, and finger opening.
Table Token: A learned token completing the set.
Scenes are padded to a fixed number of object slots (e.g., six), with unused slots masked. This allows a single physical encoder to process diverse configurations without task-specific labels.

Training vs. Deployment

Training: Both encoders and the shared predictor are trained jointly. The physical branch acts as a second online predictive view, constraining the visual dynamics.
Deployment: The physical encoder $h_\phi$ and the privileged state $s_t$ are discarded. The deployed model consists only of the visual encoder and predictor. It receives the same visual observations as the baseline, with no privileged inputs during inference. Planning is performed by rolling out candidate action sequences in the learned latent dynamics and scoring them against an encoded goal.

3. Key Contributions

Cross-Predictive JEPA: A novel architecture that grounds latent transitions in privileged physical trajectories by advancing visual and physical histories through a shared predictor and matching each prediction to both modalities' futures. This couples pointwise correspondence with action-conditioned prediction.
Unified Privileged-State Interface: A rigid-body schema allowing a single physical encoder to operate across diverse object–interaction configurations without category labels or task-specific state fields.
Separation of Properties: The work explicitly distinguishes between decodability (recovering physical variables from a snapshot), forecastability (reliable action-conditioned prediction over time), and planning utility. It demonstrates that high decodability does not guarantee high forecastability or control success.

4. Experimental Results

The authors evaluate XP-JEPA on a multi-task suite (Meta-World based, 22 configurations) and four single-task environments (Push-T, OGBench-Block, Two-Room, Reacher).

Forecastability

Rollout Drift: XP-JEPA significantly reduces the relative rollout drift of a newly fitted predictor. On the multi-task suite, drift decreases from 0.361 (Visual baseline) to 0.104.
Generalization: This improvement is consistent across all single-task environments (e.g., Two-Room drift drops from 0.503 to 0.221).
Ablation: Removing cross-prediction (ALIGN-ONLY) reduces drift less effectively (0.158), while removing the shared predictor (CROSS-ONLY) maintains low drift (0.101), indicating that cross-modal prediction is the primary driver of forecastability.

Control Performance

Success Rates: On the multi-task suite, XP-JEPA increases mean control success from 53.6% to 78.2% (+24.6 percentage points).
Task Dependency: While forecastability improves universally, control gains vary by task. For example, XP-JEPA significantly outperforms the baseline in Two-Room and OGBench-Block but shows marginal gains in Push-T and Reacher.
Comparison to Regression: A baseline that regresses a comprehensive pose-based state target from the visual latent (REGRESS) achieves high position decodability ( $R^2 = 0.991$ ) but fails to improve forecastability (drift 0.373) or control (53.7% success). This confirms that snapshot content alone does not ensure reliable rollout prediction.

Ablation Insights

Correspondence vs. Prediction: Correct visual–physical correspondence (matching the right trajectory) is crucial for control. The SHUFFLE ablation (using correct architecture but mismatched trajectories) retains low drift (0.167) but collapses control to 16.6%, proving that predictable dynamics are insufficient if they do not preserve distinctions relevant to action selection.
Predictor Sharing: Sharing the predictor alone (SHARE-ONLY) yields little gain (56.7% success), suggesting that the coupling of modalities is more important than the shared weights.

5. Significance and Claims

The paper claims that cross-predictive physical grounding produces more forecastable latent dynamics for rollout-based control without requiring privileged inputs at test time.

Mechanism: The improvement arises not from simply encoding more physical information, but from constraining the predictive evolution of the representation using corresponding physical trajectories.
Distinction: The work clarifies that "decodability" (can we read the state?) and "forecastability" (can we predict the future?) are distinct properties. A model can be highly decodable but poorly forecastable, leading to poor control.
Limitations: The authors modestly note that results are limited to simulation, require paired privileged trajectories during training, and rely on a fixed rigid-body state schema. Extending this to real-world observations and broader physical interactions remains future work.

In summary, XP-JEPA demonstrates that by treating privileged state as a second predictive view and enforcing cross-modal consistency in the transition dynamics, one can learn visual world models that are robust to physical ambiguities, leading to significantly improved planning and control performance.