Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency

The paper addresses a critical gap in evaluating Joint-Embedding Predictive Architectures (JEPAs) for world models. JEPAs learn world models that predict in a compact latent space rather than in pixels, reducing pressure to model nuisance appearance. However, this design provides no guarantee against visual perturbations: they can still alter the encoded representation and affect subsequent action-conditioned predictions. The authors note that encoder distances alone do not show how a perturbation propagates through prediction. Over a multi-step rollout, the predictor may amplify or contract the initial representation difference.

The central research question is: when a clean history and its perturbed view are rolled forward under the same action sequence, how far apart are their predicted trajectories?

The paper introduces Action-Conditioned Predictive Consistency (ACPC), which rolls a clean history and its perturbed view forward under the same action sequence and measures the distance between their predicted trajectories. Formally, given a clean history h and its perturbed version h̃, with encoder Eθ and action-conditioned predictor Fθ, the predicted representations after k actions are:

ẑk = Fθk(z, a0:k−1), z̃ˆk = Fθk(z̃, a0:k−1)

ACPC is defined as the distance between weighted predicted rollouts:

ACPCH(h, h̃, a) = Ḡa(Eθ(h)) − Ḡa(Eθ(h̃))2

where Ḡa combines all H predicted steps with weights αk. The paper uses uniform weights αk = 1/H and H = 8 for the main diagnostic, with H = 5 for planner analysis to match the CEM planning horizon.

Proposition 1 (Prediction-error change bound): The paper proves that for every paired sample, eh̃ − eh ≤ ACPCH(h, h̃, a), where eh and eh̃ are prediction errors against the same observed future. This follows directly from the reverse triangle inequality. The bound concerns the change in error, not absolute prediction accuracy.

Proposition 2 (Planning-cost bounds): For a planner scoring candidates by squared distance to a goal embedding, the paper proves C̃j − Cj ≤ bj, where bj = rj(xj − g2 + x̃j − g2). If the winner changes, the perturbed winner's excess cost satisfies 0 ≤ Cw̃ − Cw ≤ bw̃ + bw. The paper also provides conditions (Proposition 3) for preserving the selected candidate or elite set when clean cost gaps exceed combined bounds.

Corollary 1 (Conditional CEM stability): If the elite-set preservation condition holds at every iteration of CEM, the clean and perturbed runs remain aligned and select identical actions.

Invariance Radius (IR): Summarizes perturbation-induced rollout spread across histories. For each anchor history, ACPC is computed with M perturbation draws, normalized by the anchor's typical one-step clean motion si, averaged across draws, then summarized by the q90 quantile across anchors. Lower IR indicates less sensitivity to visual perturbations.

Separation Rate (SR): Checks whether different states remain distinguishable after rollout. For each anchor, a nearby history with a different state-coordinate label is selected, both are rolled forward under the anchor's recorded actions, and SR is the fraction whose different-state distance exceeds raw IR plus a margin δ = 0.10. Higher SR indicates better separation.

The paper emphasizes that Low ACPC alone, however, does not rule out representational collapse: a constant representation would make every paired rollout identical. Therefore, IR and SR must be used together: IR checks whether perturbed views remain close to their clean counterparts. SR checks whether this invariance also preserves differences between states.

The paper evaluates on four visual control tasks: TwoRoom navigation, PushT planar manipulation, Reacher arm control, and OGBench-Cube (Cube) 3D manipulation. The primary model family is LeWM, with PLDM as a second architecture. The main experiments add Gaussian noise with σ = 0.08 to history observations while keeping the goal image clean. Additional experiments use Gaussian blur (k = 15) and resize (scale 0.25).

The Gaussian-noise sweep contains nine training conditions: one without noise augmentation and eight with full-sequence Gaussian-noise augmentation at σmax ∈ 0.01,..., 0.08. Each task–condition pair has three independent training runs (seeds 3072/3073/3074), evaluated with seeds 42/43/44 using 100 episodes per evaluation seed.

Using a PushT case study comparing unaugmented and noise-augmented checkpoints, the paper finds that without augmentation, the median r/NN ratio is 1.41 at the encoder and 1.86 after eight rollout steps, with none of the 128 anchor clouds fully disjoint. After Gaussian-noise augmentation, ratios fall to 0.10 and 0.19, with fully disjoint fractions rising to 95.3% at the encoder and 84.4% after rollout.

At evaluation noise σ = 0.08, unaugmented LeWM checkpoints lose 25.9 ± 2.4, 74.6 ± 5.6, 41.0 ± 1.4, and 22.1 ± 2.8 percentage points in success rate on TwoRoom, PushT, Reacher, and Cube respectively. Every augmentation level that meets the success-rate criterion has lower IR and higher SR than its unaugmented reference. Across 108 LeWM checkpoint rows, 77 pass the IR threshold tIR = 0.3, and all 77 also pass tSR = 0.95.

Training LeWM on TwoRoom with four SIGReg weights shows that without SIGReg, the representation collapses: the median latent distance shrinks from 17.2–18.7 to 0.006, and clean success falls from 96.3–99.3% to 33.3%. Although this model has the lowest raw IR (0.048), its SR falls to 0.066, compared with 0.967–0.984 for nonzero SIGReg. This demonstrates that SR exposes a loss of state separation that raw IR does not capture.

The paper tests whether ACPC helps predict how much a visual perturbation changes multi-step prediction error. Using three nested regressions (Base with one-step ACPC, Base+Control8 with destroyed-action controls, and Base+ACPC8 with eight-step ACPC under recorded actions), the results show that Base+ACPC8 attains the lowest cross-validated MAE in all 12 task–run cells. The reduction is 55.9 ± 4.7% relative to Base and 51.3 ± 3.5% relative to Base+Control8.

Adding planner-horizon ACPC (five-step, matching CEM's prediction horizon) lowers cross-task test MAE by 15.2 ± 2.0% for predicting the extra clean-model cost of perturbation-induced plan changes. The error decreases in all 12 task–run test cases.

Across 14 source/test splits, 13 choose (tIR, tSR) = (0.3, 0.95); only the Reacher-only split chooses (0.1, 0.95). With two or three source tasks, the first accepted checkpoint is within 0.5 grid steps of recovery on average and never differs by more than one step (balanced accuracy 0.900, precision/recall 0.913/0.953). Single-source selection is less reliable, with balanced accuracy averaging 0.855.

"Across all four PLDM tasks, the augmented checkpoints generally have lower relative IR and higher SR than the unaugmented checkpoint. Planning success under evaluation noise also generally improves from the unaugmented checkpoint. These results reproduce the qualitative low-IR, high-SR pattern under a second world-model architecture."

In 22 of 24 comparisons, the sign of ∆S agrees with whether the augmented checkpoint meets the prespecified success criterion (balanced accuracy 0.889). Across all pairs, larger ∆S is associated with larger gains in planning success (Spearman ρ = 0.835). The two remaining cases are boundary cases with four-point success gains (one point below the five-point criterion) and positive ∆S.

The paper emphasizes that ACPC diagnoses how visual perturbations affect predicted rollouts and planning costs; it does not modify CEM. The full IR–SR screen requires an unaugmented reference checkpoint and uses observed future frames for normalization and dataset state labels for SR pairs. The chosen IR threshold is at the edge of the tested range. Future work should test whether using ACPC during planning improves task success and evaluate broader perturbation severities, state pairs, robustness-training methods, and other JEPA world models.

Based on the paper, here are specific improvements to AI systems and what the improved systems can do:

1. Robustness-Aware World Model Training

• Improvement: Add ACPC as a training-time regularizer that penalizes divergence between clean and perturbed predicted trajectories, in addition to or instead of standard noise augmentation.

• What the improved system can do: Learn world models that are explicitly invariant to visual perturbations across multi-step rollouts, not just at the encoder level. This reduces compounding prediction errors in downstream planning, leading to higher task success under sensor noise, blur, or resolution changes without requiring exhaustive augmentation schedules.

2. Perturbation-Aware Model Selection and Checkpoint Recovery

• Improvement: Use the IR–SR screen (thresholds tIR=0.3, tSR=0.95) as an automated early-stopping or checkpoint-selection criterion during training, rather than relying solely on validation success rate.

• What the improved system can do: Automatically identify the first training checkpoint that is both invariant to perturbations (low IR) and preserves state separation (high SR). This reduces training time by up to one grid step (as shown in cross-task screening) and ensures the deployed model is robust before any evaluation, avoiding late-stage catastrophic failures under noise.

3. Planner Stability Certification

• Improvement: Integrate the planning-cost bound (Proposition 2) and conditional CEM stability (Corollary 1) into the planner as a runtime safety check. Before executing an action, compute ACPC-based bounds on cost changes for the current history and reject plans where the clean-vs-perturbed cost gap exceeds a threshold.

• What the improved system can do: Guarantee that the selected action sequence remains near-optimal even if the current observation is perturbed. This prevents erratic plan switching in real-time control (e.g., robotic manipulation) under sensor noise, improving execution reliability without retraining.

4. Perturbation-Impact Prediction for Online Adaptation

• Improvement: Train a meta-model that takes ACPC values (computed over a short history) as input to predict the expected change in multi-step prediction error and planning cost. Use this prediction to dynamically adjust the planner's horizon or exploration noise.

• What the improved system can do: In deployment, the system can detect when a perturbation is likely to cause large rollout divergence (high ACPC) and respond by shortening the planning horizon, increasing action smoothing, or switching to a more conservative policy—reducing the risk of task failure before it occurs.

5. Representation Collapse Prevention via SR-Aware Regularization

• Improvement: Add a separation-rate-based loss term during training that penalizes models where SR drops below a threshold, complementing existing invariance losses (e.g., SIGReg).

• What the improved system can do: Avoid the failure mode where a model becomes trivially invariant (e.g., constant representation) while still being robust. This ensures the latent space retains meaningful state distinctions, enabling both perturbation robustness and accurate long-horizon predictions—critical for tasks requiring fine-grained control (e.g., precise manipulation).

6. Cross-Task Robustness Transfer

• Improvement: Use the cross-task screening result (13/14 splits selecting the same IR–SR thresholds) to pre-train a shared robustness criterion across tasks. Apply this criterion to select a single universal checkpoint for a family of related tasks.

• What the improved system can do: Deploy one world model across multiple environments (e.g., navigation, manipulation, arm control) with a single robustness threshold, reducing the need for task-specific tuning. This enables faster adaptation to new tasks by starting from a checkpoint known to be both invariant and separable, improving sample efficiency and success rates under noise.

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.