Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency

Abstract: Joint-embedding predictive architectures (JEPAs) learn world models that predict in a compact latent space rather than in pixels, reducing the pressure to model nuisance appearance. Yet this provides no guarantee against visual perturbations: they can still alter the encoded representation and affect subsequent action-conditioned predictions. Bisimulation captures this requirement precisely: two observations should be treated as the same state only when their action-conditioned consequences agree. Guided by this criterion, we introduce Action-Conditioned Predictive Consistency (ACPC), a diagnostic that measures how far a clean history and a visually perturbed view of it diverge after being rolled forward under the same action sequence. We prove that this divergence bounds the perturbation-induced change in multi-step prediction error and planner cost. Building on pairwise ACPC, we define two complementary measures: the Invariance Radius (IR) summarizes clean-perturbed rollout spread, while the Separation Rate (SR) checks whether different states remain distinguishable after rollout. Experiments on four visual control tasks show that pairwise ACPC predicts perturbation-induced prediction and cost changes. On LeWM, the IR-SR screen transfers across tasks, and the joint diagnostic remains informative under blur and resize. PLDM exhibits similar diagnostic trends under a different architecture.

• The paper introduces Action-Conditioned Predictive Consistency (ACPC), which measures how visual perturbations diverge during identical multi-step latent rollouts and predicts error drift with 55.9% lower MAE than a base model.
• The paper combines the Invariance Radius (IR) with Separation Rate (SR) to distinguish robust representations from collapse, with SR exposing collapsed models that achieve low IR but sharply reduced task success.
• The paper shows that ACPC-based diagnostics transfer across tasks, visual stressors, and architectures, while providing conditional certificates for stable planner choices without directly improving planning or environment returns.

JEPA-style world models predict in latent space rather than pixels, which reduces pressure to model nuisance appearance but provides no guarantee that visual perturbations leave the action-conditioned predictions intact. The paper under review, "Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency" (2608.12939), introduces a diagnostic framework that evaluates this property on frozen, trained models. Its central object, Action-Conditioned Predictive Consistency (ACPC), measures how far a clean history and a visually perturbed view of it diverge after both are rolled forward under the same action sequence, and it is summarized at the checkpoint level by two complementary statistics: the Invariance Radius (IR) and the Separation Rate (SR).

The ACPC diagnostic

Given a clean history $h$ and its perturbed version $\tilde h$, a frozen encoder maps them to $z$ and $\tilde z$, and a frozen action-conditioned predictor rolls both representations forward for $H$ steps under an identical action sequence. The per-step predicted representations are projected into the planner's latent space and combined with uniform weights into a single trajectory-level distance:

$$\mathrm{ACPC}_H(h,\tilde h,\mathbf a)=\left(\sum_{k=1}^H\alpha_k\left\|\Pi(\hat z_k)-\Pi(\hat{\tilde z}_k)\right\|_2^2\right)^{1/2}.$$

Two design choices matter. Holding the action sequence fixed removes action differences from the comparison, so ACPC isolates the effect of the visual perturbation; and measuring after rollout, rather than at the encoder output, captures how the predictor amplifies or contracts the initial representation difference over multiple steps. ACPC measures rollout consistency, not prediction accuracy, and the paper is careful to distinguish encoder shift $\|z-\tilde z\|_2$, which is available without prediction, from the post-rollout divergence that ACPC measures.

The authors motivate the diagnostic with a local-geometry case study on PushT. Without noise augmentation, the median ratio of a history's perturbation radius to its nearest-clean-neighbor distance is 1.41 at the encoder and 1.86 after eight rollout steps, and none of the 128 anchor clouds is fully disjoint. After Gaussian-noise augmentation at $\sigma=0.08$, these ratios fall to 0.10 and 0.19, and fully disjoint fractions rise to 95.3% and 84.4%. Augmentation thus contracts perturbation-induced spread relative to nearby clean histories, and much of this contraction persists through prediction. The authors note plainly that this case study is descriptive: it does not establish task-relevant separation or planning robustness, which is precisely what ACPC and the IR–SR screen are designed to measure.

Theoretical guarantees

The paper proves two samplewise bounds that connect ACPC to quantities planners care about, both derived without distributional or smoothness assumptions and shown to be tight.

Prediction-error drift. If $e_h$ and $e_{\tilde h}$ are the errors of the two predicted rollouts against the same observed future, the reverse triangle inequality gives $\tilde h$0. The bound concerns the change in error, not absolute accuracy: both rollouts can be inaccurate even when ACPC is small.

Planning-cost stability. For a planner scoring candidate $\tilde h$1 by squared distance to a fixed goal embedding, the cost change satisfies $\tilde h$2 where $\tilde h$3 and $\tilde h$4 is the final-step displacement, bounded by candidate-specific ACPC divided by $\tilde h$5. Consequently, if the clean winner leads every competitor by more than the sum of their bounds, the perturbation cannot reverse the selection; analogous certificates preserve the top-$\tilde h$6 elite set. A corollary extends this to adaptive CEM: if the elite-set certificate holds at every iteration of paired clean and perturbed runs sharing random numbers, both runs return the same action. The guarantee is conditional—it applies only while the certificate is verified—and concerns the model-based choice, not environment return.

IR and SR at the checkpoint level

Pairwise ACPC does not by itself summarize a checkpoint. The Invariance Radius aggregates ACPC across 100 logged anchors, each with five perturbation draws, normalized by the anchor's typical clean one-step motion and summarized at the 90th percentile, so that highly sensitive anchors are not hidden by averaging. Lower IR indicates less perturbation sensitivity. Low IR alone, however, is compatible with representational collapse: a constant representation makes every paired rollout identical. The Separation Rate guards against this by checking whether rollouts of nearby histories with different state-coordinate labels remain farther apart than raw IR plus a margin of 0.10, when both are rolled under the anchor's recorded actions. Higher SR indicates better preservation of state distinctions.

A controlled ablation demonstrates that this complementarity is not hypothetical. Training LeWM on TwoRoom without the SIGReg regularizer collapses the representation (median latent distance drops from roughly 17.2–18.7 to 0.006, and clean success falls from 96.3–99.3% to 33.3%); the collapsed model attains the lowest raw IR (0.048) but its SR falls to 0.066, versus 0.967–0.984 for regularized checkpoints. SR thus exposes a failure mode that raw IR actively rewards.

Empirical findings

The evaluation uses LeWM (2608.12939) as the primary model family and PLDM (2608.12939) as a second architecture, across four tasks (TwoRoom, PushT, Reacher, OGBench-Cube), with a nine-point Gaussian-noise augmentation sweep per task and three independent training runs.

ACPC predicts error drift. In a 16-fold leave-one-trajectory-group-out ridge regression, a model augmented with eight-step ACPC under the recorded actions achieves the lowest cross-validated MAE in all 12 task–run cells, reducing MAE by $\tilde h$7 relative to a base model and $\tilde h$8 relative to a conservative control using the best of three destroyed-action eight-step ACPC variants (zeroed, swapped, or shuffled actions). The comparison against destroyed-action controls is the key design element: it shows the gain is attributable to the recorded action sequence, not merely to rolling out eight steps.

ACPC predicts the cost of CEM plan changes. Adding the 90th-percentile candidate-level ACPC at CEM's five-step horizon to a baseline regression (which already includes severity, training condition, one-step ACPC, and the clean best–second-best cost gap) reduces test MAE on $\tilde h$9 by $z$0, with improvement in all 12 task–run test cases. The paper is explicit that this measures whether ACPC reflects the model-cost effect of a plan change, not whether ACPC-guided planning improves task success.

Cross-task screening. Thresholds $z$1 selected on subsets of tasks transfer to held-out tasks: 13 of 14 splits select $z$2, and with two or three source tasks the first accepted checkpoint is within 0.5 augmentation-grid steps of recovery on average (balanced accuracy 0.900). Single-source selection is weaker (balanced accuracy 0.723–0.913), with Reacher alone choosing a tighter IR threshold that delays acceptance on other tasks by up to five levels. Notably, the selected IR threshold of 0.3 is the largest tested value, so the upper edge of the useful range remains unresolved.

Cross-stressor transfer. Under blur and resize, with Gaussian-noise-selected thresholds held fixed, the sign of the change in joint IR–SR score agrees with a prespecified five-point success criterion in 22 of 24 checkpoint pairs (balanced accuracy 0.889, Spearman $z$3). The two disagreements are boundary cases with a four-point success gain where the diagnostic and behavioral trends actually agree in direction.

Architecture transfer. PLDM, with a different architecture and training recipe, reproduces the qualitative pattern: augmented checkpoints generally have lower relative IR and higher SR, with generally improved noisy-evaluation success. The checkpoint-level findings are empirical rather than consequences of the pairwise bounds, and the paper labels them as such.

Limitations

The authors state the scope of their results directly. The planner analysis is post hoc and does not test ACPC-guided planning; the adaptive-CEM guarantee holds only while the elite-set certificate verifies; the full screen requires an unaugmented reference checkpoint from the same task, run, and family; IR normalization uses observed future frames and SR requires dataset state labels, so the screen applies only to the evaluated shifts and label definitions; and the checkpoint experiments cover Gaussian noise plus one severity each of blur and resize, with the IR threshold at the boundary of the tested grid. The local-sensitivity analysis relating ACPC to the Frobenius norm of the composed encoder–rollout Jacobian is likewise local and provides no global guarantee. Open questions left by the paper include whether ACPC can be used during planning to improve task success, whether the IR threshold range extends beyond 0.3, and how the diagnostics behave under more varied perturbations, state pairs, and JEPA architectures.

Conclusion

This paper provides a principled, frozen-model diagnostic for JEPA world models grounded in the bisimulation intuition that state equivalence is defined by action-conditioned consequences. Its theoretical contribution—samplewise, assumption-free bounds linking ACPC to prediction-error drift and planning-cost stability—is modest in proof complexity but exact and tight, and it is matched by consistent empirical evidence: multi-step ACPC carries predictive information about downstream error and planner-cost effects that encoder-only and one-step comparisons miss, the joint IR–SR screen transfers across tasks and stressors, and the pattern replicates on a second architecture. The framework's main restriction is that it diagnoses rather than repairs: it quantifies how perturbations propagate through a fixed predictor but does not itself modify training or planning, leaving the question of ACPC-guided intervention open.