Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency

2026-08-13

AI Summary

Although JEPA-based world models perform prediction in latent space, they remain susceptible to visual perturbations, leading to distorted state representations and biased action-conditioned predictions. This work proposes Action-Conditioned Prediction Consistency (ACPC), a diagnostic framework that, for the first time, operationalizes the twin-simulation concept into a computable metric by quantifying the divergence between multi-step forward rollouts from clean and perturbed observations under identical action sequences. To enable robustness evaluation across tasks and architectures, we introduce two complementary metrics: Invariance Radius (IR) and Separation Rate (SR). Empirical results demonstrate that ACPC effectively predicts perturbation-induced errors in prediction and planning costs, and the IR–SR pair exhibits strong generalization and diagnostic capability across models such as LeWM and PLDM.

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.

Problem

• JEPA
• world models
• visual perturbations
• action-conditioned prediction
• predictive consistency

Innovation

• Action-Conditioned Predictive Consistency
• Joint-Embedding Predictive Architecture
• Bisimulation
• Invariance Radius
• Separation Rate