🤖 machine learning

Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency

This paper introduces Action-Conditioned Predictive Consistency (ACPC), a diagnostic framework grounded in bisimulation theory that quantifies the robustness of Joint-embedding predictive architecture (JEPA) world models by measuring the divergence between clean and perturbed state rollouts, thereby providing theoretical bounds on prediction errors and planner costs while demonstrating effectiveness across various visual control tasks.

Original authors: Guo An, Zijing Wu, Honghua Dong, Yuhao Yan, Zixuan Gui, Haochong Chen, Shanzhao Ruan, Xiang Wang, Yurong Ling, Qi Tian

Published 2026-08-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Guo An, Zijing Wu, Honghua Dong, Yuhao Yan, Zixuan Gui, Haochong Chen, Shanzhao Ruan, Xiang Wang, Yurong Ling, Qi Tian

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

Imagine you are teaching a robot to play a video game. You want the robot to learn the rules of the game—the physics, the goals, and how its actions change the world—so it can make smart decisions. But there's a catch: the robot sees the world through a camera, and cameras can be tricky. A little bit of static, a blur, or a sudden change in lighting shouldn't confuse the robot into thinking the game has changed completely. In the world of artificial intelligence, this is the challenge of building "world models." These are internal maps an AI builds to predict what will happen next if it takes a certain action.

Recently, a popular type of AI called a "Joint-embedding predictive architecture" (or JEPA for short) has become a favorite for building these maps. Instead of trying to redraw every single pixel of the game screen (which is like trying to memorize the exact color of every blade of grass), JEPAs learn a compact, abstract summary of the scene. Think of it like a robot learning that "the red ball is near the wall" rather than memorizing the exact shade of red. This is efficient, but it has a hidden flaw: sometimes, a tiny visual glitch (like a speck of dust on the lens) can trick the robot's summary into changing drastically, causing it to make terrible decisions. The big question researchers are asking is: How do we know if a robot's internal map is truly robust, or if it's just a house of cards waiting to collapse when the lighting changes?

This paper introduces a clever new way to test these robot brains, called Action-Conditioned Predictive Consistency (ACPC). Imagine you have a robot looking at a clean, perfect video of a game, and then you show it the exact same video but with a little bit of visual "noise" added, like static or blur. The researchers ask: If we tell the robot to perform the exact same sequence of moves in both scenarios, do its predictions of the future stay close together? If the robot is smart and robust, the two future paths it imagines should look very similar, even though the starting pictures were slightly different. If the robot is fragile, the two paths will drift far apart, leading to confusion.

The authors don't just look at this single comparison; they developed two specific "health checks" to summarize the robot's performance. The first is the Invariance Radius (IR). Think of this as measuring how much the robot's imagination wobbles when the input gets noisy. A low IR is good—it means the robot's predictions stay steady and close to each other, like a tightrope walker who doesn't sway much. The second check is the Separation Rate (SR). This is the safety net. We don't want the robot to be so steady that it ignores all differences. If the robot is broken, it might just predict the same boring future no matter what it sees (a "collapsed" representation). SR checks if the robot can still tell the difference between two genuinely different situations (like a ball being on the left vs. the right) even after the noise is added. A high SR means the robot keeps its eyes open and can still distinguish important details.

The researchers tested this idea on four different robotic control tasks, ranging from navigating a maze to pushing objects. They found that when a robot's model was trained to handle visual noise, its IR went down (it became steadier) and its SR went up (it stayed sharp). More importantly, they proved mathematically that if the robot's predictions stay close together (low ACPC), then the error in its future guesses and the cost of its planning decisions won't change wildly. In other words, if the robot's imagination doesn't drift apart when the camera gets dirty, it's much less likely to make a costly mistake.

They also showed that this diagnostic works across different types of robot brains, not just one specific design. When they tested it on a different architecture, the same pattern held: better performance under visual stress meant lower IR and higher SR. While they didn't claim to have solved the problem of robot vision forever, their results strongly suggest that checking how much a robot's predicted future drifts under the same actions is a powerful way to spot which models are truly ready for the messy, unpredictable real world. It's like checking if a ship's compass still points north even when the storm clouds roll in, rather than just hoping the ship doesn't sink.

Technical Summary: Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency

Problem Statement

Joint-embedding predictive architectures (JEPAs) learn world models by predicting target representations in a compact latent space rather than reconstructing raw pixels. While this design avoids modeling nuisance visual details, it offers no inherent guarantee that the learned representations are robust to visual perturbations. A critical gap exists in current diagnostics: encoder-level distances or single-transition metrics fail to capture how visual perturbations propagate through multi-step rollouts. Specifically, a perturbation might cause the predicted trajectory to diverge significantly from the clean trajectory, altering action-conditioned predictions and planner costs, even if the initial encoded representations appear similar. The core problem is determining how far apart the predicted trajectories of a clean history and its visually perturbed view become when rolled forward under the same action sequence.

Methodology

Action-Conditioned Predictive Consistency (ACPC)

The authors introduce Action-Conditioned Predictive Consistency (ACPC) as a diagnostic metric. ACPC measures the divergence between two predicted trajectories: one generated from a clean history ( $h$ ) and one from a visually perturbed view ( $\tilde{h}$ ), both rolled forward under an identical recorded action sequence ( $a$ ).

Mechanism: A frozen encoder maps $h$ and $\tilde{h}$ to latent states $z$ and $\tilde{z}$ . A frozen action-conditioned predictor then rolls these states forward for $H$ steps.
Metric: ACPC calculates the weighted Euclidean distance between the projected latent trajectories of the clean and perturbed rollouts.
Theoretical Bounds: The paper proves that ACPC bounds the perturbation-induced change in multi-step prediction error (Proposition 1). Furthermore, when evaluating candidate action sequences for planning, ACPC bounds the change in predicted costs (Proposition 2). These bounds hold for individual pairs without requiring distributional or smoothness assumptions.

Composite Diagnostics: IR and SR

To summarize ACPC across a model checkpoint, the authors define two complementary measures:

Invariance Radius (IR): This summarizes the sensitivity of paired rollouts to visual perturbations. It is computed as the quantile (e.g., 90th) of normalized ACPC values across many history anchors. Lower IR indicates that perturbed views remain close to their clean counterparts.
Separation Rate (SR): This checks whether the model preserves distinctions between different states. It measures the fraction of different-state pairs whose rollout distance exceeds the IR plus a margin. Higher SR indicates that distinct states remain distinguishable after rollout.

The authors emphasize that IR alone is insufficient because a collapsed representation (where all states map to the same point) would yield a low IR but fail to support control. Conversely, high SR alone does not guarantee robustness to perturbations. A robust checkpoint must exhibit both low IR and high SR.

Evaluation Protocol

The methodology is evaluated on the LeWM (Learning World Models) family and the PLDM architecture across four visual control tasks (TwoRoom, PushT, Reacher, OGBench-Cube).

Perturbations: Experiments utilize Gaussian noise, Gaussian blur, and image resizing.
Screening: A "checkpoint screen" is defined using thresholds for relative IR (compared to an unaugmented baseline) and SR. A checkpoint passes if it demonstrates reduced sensitivity (low relative IR) while maintaining state separation (high SR).

Key Contributions

Introduction of ACPC: The paper proposes ACPC, a pairwise diagnostic that compares predicted trajectories of clean and perturbed histories under shared actions. Unlike prior work focusing on encoder distances or single-step predictions, ACPC captures rollout-level effects.
Theoretical Bounds: The authors provide sample-wise proofs showing that ACPC bounds the change in multi-step prediction error and the change in predicted planning costs. These bounds explain why ACPC is informative about downstream planning failures.
Diagnostic Framework (IR and SR): The paper defines the Invariance Radius and Separation Rate to summarize ACPC across a checkpoint. This framework explicitly guards against representation collapse while measuring sensitivity to visual shifts.
Empirical Validation:
- Pair-Level: ACPC is shown to predict changes in prediction error and the "selection regret" (extra cost) incurred when a planner (CEM) selects a different plan due to perturbation.
- Checkpoint-Level: Across four tasks, checkpoints with lower IR and higher SR generally correspond to better planning success under visual noise.
- Generalization: The low-IR, high-SR pattern holds for a different architecture (PLDM) and extends to non-Gaussian perturbations (blur and resize).

Results

Prediction Error and Planning Cost: In cross-validated experiments, adding multi-step ACPC to regression models significantly reduced the error in predicting the drift in prediction error and the extra cost of CEM plan changes, outperforming baselines using only encoder shifts or one-step ACPC.
Recovery Identification: On LeWM, thresholds for IR and SR selected on a subset of tasks successfully identified checkpoints that recovered planning performance on held-out tasks under Gaussian noise.
Representation Collapse Detection: In a case study where representation collapse occurred (due to specific training weights), IR remained low (indicating low sensitivity), but SR dropped to near zero, correctly identifying the failure to maintain state distinctions.
Cross-Architecture and Cross-Perturbation: The diagnostic trends (lower IR, higher SR correlating with better performance) were observed on PLDM. Furthermore, the joint IR-SR score change ( $\Delta S$ ) correlated with planning success gains under blur and resize, suggesting the metric is not limited to Gaussian noise.

Significance and Claims

The paper claims that Action-Conditioned Predictive Consistency provides a necessary diagnostic for world models that goes beyond static representation analysis. By rolling paired observations forward under the same actions, ACPC reveals how visual perturbations propagate through the dynamics model, directly linking representation stability to planning performance.

The authors assert that:

Multi-step ACPC captures downstream effects that encoder-only and one-step comparisons miss.
The combination of IR and SR offers a practical "screen" to identify robust checkpoints without requiring retraining or access to ground-truth future states during the diagnostic phase (though observed futures are used for normalization in IR calculation).
The diagnostic is architecture-agnostic (valid on both LeWM and PLDM) and robust to different types of visual shifts.

The paper remains modest regarding its scope, noting that ACPC is a diagnostic tool and does not modify the planner itself. The theoretical bounds apply to specific pairs and fixed candidate pools, and the empirical screening results are specific to the evaluated visual shifts and state labels. The work does not claim to certify absolute robustness but rather to provide informative metrics that correlate strongly with planning success under perturbation.