AVL-JEPA: Preventing Causal Dynamics Information Collapse In Joint Embedding Predictive Architecture World Models

Yikang Qiao; Ling Zhang; Ziying Song; Duan Huang

cs.ROOct 2, 2026

AVL-JEPA: Preventing Causal Dynamics Information Collapse In Joint Embedding Predictive Architecture World Models

Authors: Yikang Qiao, Ling Zhang, Ziying Song, Duan Huang

Organizations: Central South University · Yanshan University

Abstract

Joint embedding predictive architectures (JEPAs) predict future latent representations without reconstructing observations, enabling world models to focus on high-level semantic dynamics. However, a JEPA can preserve high dimensional visual information while discarding information about the physical consequences of actions. We call this failure mode causal dynamics information collapse and propose action-grounded vision-invariance latent (AVL) to prevent this collapse. We first use the executed action as an auxiliary dynamics anchor that encourages the model to preserve dynamics information, and then use a vision-invariance pathway which aligns perturbed and clean latent predictions without discarding dynamics information, forcing the model to fully understand and utilize causal dynamics information. We validate AVL on four robotic control tasks (TwoRoom, PushT, OGBench Cube, and Reacher), showing that it substantially improves success rates under visual perturbations while preserving clean-environment performance. We further evaluate physical consequence alignment, clean-noisy dynamics consistency, and the causal effect of targeted transition subspace erasure. Collectively, these results indicate that dynamic information causally relevant to planning is preserved from collapse under AVL.

Figures & tables

Figure 1: Measured PCA projection of predicted

\Delta z

transition responses for LeWM and AVL (ours). Both panels use one shared transform and identical action candidates; different colors denote different candidate action categories.

Figure 2: Overview of the proposed AVL method. The action-grounded pathway anchors clean latent transitions to executed actions through a transition-attributed action head that subtracts a state-only zero-transition baseline. The vision-invariance pathway aligns noisy and clean latent predictions while freezing the action head. The alignment creates a potential collapse shortcut, but the frozen action loss requires the latent to retain action-relevant dynamics, thereby targeting causal dynamics information collapse.

Figure 3: Clean–noisy candidate-transition geometry on TwoRoom for LeWM and AVL (ours). Red triangles indicate states satisfying

D_{\mathrm{nuis}}/D_{\mathrm{act}}>0.5

together with a clean–noisy Top-1 candidate flip. The transparent planes show the same ratio threshold in the

(D_{\mathrm{act}},D_{\mathrm{nuis}})

and

\log_{10}(D_{\mathrm{nuis}}/D_{\mathrm{act}})

views.

Task	Model	Clean–noisy Spearman $\uparrow$	Top-1 agreement (%) $\uparrow$	Action-to nuisance ratio $\uparrow$	Nuisance drift $\downarrow$	Noisy regret $\downarrow$
TwoRoom	LeWM	0.89	58	2.2	8.9	0.162
	Action-grounded	0.90	60	2.0	8.7	0.088
	AVL (ours)	0.98	76	2.9	6.8	0.069
PushT	LeWM	0.809	45.3	2.35	6.15	0.387
	Action-grounded	0.970	76.0	7.21	3.60	0.313
	AVL (ours)	0.998	93.0	19.23	1.07	0.272

Table 1: Clean–noisy dynamics consistency on TwoRoom and PushT (100 matched states; three noisy repeats; Gaussian

\sigma=0.05

). Higher clean–noisy Spearman, Top-1 agreement, and action-to-nuisance ratio are better; lower nuisance drift and noisy regret are better. Best values within each task and metric are bolded.

Figure 4: Counterfactual physical consequence alignment, selection quality, and TwoRoom selection accuracy. The first panel reports predicted-to-physical Spearman correlation, the second reports normalized regret, and the third reports Top-1 and Top-3 selection accuracy. Higher Spearman and selection accuracy are better, whereas lower regret is better.

Figure 5: Targeted erasure sweep on TwoRoom and PushT. Each panel reports the absolute difference between the erased and corresponding non-erased baseline values across

k\in\{1,2,3,5,8\}

for LeWM and AVL.

Model	TwoRoom	PushT	Reacher	Cube
PLDM	91.7 $\pm$ 1.3	78.1 $\pm$ 1.7	78.3 $\pm$ 2.4	65.2 $\pm$ 1.5
SD-JEPA	86.3 $\pm$ 6.2	84.6 $\pm$ 4.7	85.0 $\pm$ 5.2	69.6 $\pm$ 5.3
LeWM	88.4 $\pm$ 2.3	96.0 $\pm$ 3.5	86.4 $\pm$ 4.4	73.2 $\pm$ 6.8
Action-grounded	92.8 $\pm$ 2.4	95.5 $\pm$ 1.9	82.4 $\pm$ 3.9	69.1 $\pm$ 5.6
Vision-invariance	85.2 $\pm$ 3.2	91.6 $\pm$ 3.7	84.0 $\pm$ 2.8	72.4 $\pm$ 7.8
AVL (ours)	93.2 $\pm$ 1.6	96.0 $\pm$ 1.3	85.2 $\pm$ 4.4	78.2 $\pm$ 5.3

Table 2: Clean closed-loop results.

Figure 6: Success-rate degradation under Gaussian policy input noise. Panels (a)–(c) report TwoRoom and panels (d)–(f) report PushT.

Figure 7: Closed-loop OGBench Cube rollouts. The top row shows a successful AVL (ours) rollout, whereas the bottom row shows a failed LeWM rollout under the same task. Each row contains three temporally ordered agent observations and the corresponding goal observation.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Notation	Description
Observations and latents
$o_{t}$	Observation at time step $t$
$\tilde{o}_{t}$	Visually perturbed observation at time step $t$
$o_{t+H}$	Observation after $H$ low level action steps
$z_{t}=f_{\theta}(o_{\leq t})$	Current latent produced by the encoder
$\tilde{z}_{t}$	Latent of the perturbed observation

Appendix

Table 3: Summary of the notation used throughout the paper.

Figure 8: Input-level visual perturbation examples on an archived real PushT RGB frame. The clean image is shown alongside the Gaussian, brightness, and saturation settings used in the evaluation protocol. These examples illustrate the appearance transformation only; they are not episode-level performance measurements.

Figure 9: Candidate-transition subspace erasure sweep on TwoRoom for the directly computed dimensions

k\in\{1,2,3,5,8,12,16,20,30,40\}

. Targeted PCA removal is compared with same-dimensional random-subspace removal and the unmodified base.

Figure 10: Additional representative closed-loop rollouts for Reacher, PushT, and TwoRoom. Each row shows three temporally ordered agent observations and the corresponding goal observation.

Model	TwoRoom	PushT	Reacher	Cube
LeWM	65.7 $\pm$ 3.9	15.3 $\pm$ 5.8	42.6 $\pm$ 3.4	52.6 $\pm$ 7.2
Action-grounded	76.0 $\pm$ 1.6	10.0 $\pm$ 2.8	82.6 $\pm$ 3.3	64.2 $\pm$ 6.5
Vision-invariance	82.6 $\pm$ 2.9	31.3 $\pm$ 4.1	84.0 $\pm$ 2.8	68.0 $\pm$ 9.2
AVL (ours)	88.7 $\pm$ 1.9	81.5 $\pm$ 4.1	78.2 $\pm$ 4.4	72.7 $\pm$ 7.4

Appendix

Table 4: Gaussian policy input noise,

\sigma=0.10

Model	TwoRoom	PushT	Reacher	Cube
LeWM	78.0 $\pm$ 4.8	39.6 $\pm$ 3.3	11.6 $\pm$ 6.9	45.4 $\pm$ 3.7
Action-grounded	86.0 $\pm$ 1.4	30.5 $\pm$ 2.6	7.3 $\pm$ 0.9	42.0 $\pm$ 1.6
Vision-invariance	80.7 $\pm$ 3.7	62.6 $\pm$ 2.4	14.1 $\pm$ 6.6	48.5 $\pm$ 4.9
AVL (ours)	88.9 $\pm$ 2.4	72.0 $\pm$ 0.7	14.7 $\pm$ 8.1	48.7 $\pm$ 5.7

Appendix

Table 5: Brightness shift,

+0.20

Model	TwoRoom	PushT	Reacher	Cube
LeWM	54.8 $\pm$ 4.5	14.0 $\pm$ 3.9	24.4 $\pm$ 2.4	70.8 $\pm$ 5.2
Action-grounded	36.0 $\pm$ 1.1	31.3 $\pm$ 1.8	54.0 $\pm$ 4.3	69.3 $\pm$ 7.4
Vision-invariance	85.3 $\pm$ 3.4	22.4 $\pm$ 1.6	49.5 $\pm$ 2.4	71.4 $\pm$ 3.8
AVL (ours)	86.7 $\pm$ 3.4	62.0 $\pm$ 4.2	65.3 $\pm$ 0.9	74.0 $\pm$ 1.6

Appendix

Table 6: Saturation shift,

+0.35

Model	TwoRoom	PushT	Reacher	Cube
LeWM	83.5 $\pm$ 2.6	85.6 $\pm$ 1.7	80.2 $\pm$ 4.3	72.5 $\pm$ 8.2
Action-grounded	90.7 $\pm$ 2.9	92.0 $\pm$ 1.6	83.1 $\pm$ 5.2	71.3 $\pm$ 8.9
Vision-invariance	81.4 $\pm$ 3.4	92.6 $\pm$ 3.4	82.8 $\pm$ 4.1	70.9 $\pm$ 9.8
AVL (ours)	92.0 $\pm$ 1.6	96.7 $\pm$ 0.9	80.5 $\pm$ 5.8	74.6 $\pm$ 7.3

Appendix

Table 7: Gaussian policy input noise,

\sigma=0.03

Model	TwoRoom	PushT	Reacher	Cube
LeWM	83.3 $\pm$ 1.8	65.3 $\pm$ 4.8	73.9 $\pm$ 2.7	71.3 $\pm$ 9.2
Action-grounded	90.8 $\pm$ 1.6	74.7 $\pm$ 5.5	82.0 $\pm$ 2.4	68.7 $\pm$ 8.9
Vision-invariance	83.1 $\pm$ 1.8	91.2 $\pm$ 0.9	87.4 $\pm$ 2.0	69.3 $\pm$ 8.2
AVL (ours)	92.0 $\pm$ 1.6	95.3 $\pm$ 0.9	82.5 $\pm$ 3.9	74.4 $\pm$ 9.4

Appendix

Table 8: Gaussian policy input noise,

\sigma=0.05

Model	TwoRoom	PushT	Reacher	Cube
LeWM	83.6 $\pm$ 4.6	87.6 $\pm$ 3.8	17.6 $\pm$ 3.5	44.0 $\pm$ 7.4
Action-grounded	92.0 $\pm$ 4.3	93.0 $\pm$ 2.9	31.3 $\pm$ 0.9	44.7 $\pm$ 6.8
Vision-invariance	81.3 $\pm$ 2.9	92.0 $\pm$ 2.8	21.4 $\pm$ 4.2	44.1 $\pm$ 5.5
AVL (ours)	92.7 $\pm$ 2.4	96.6 $\pm$ 1.7	30.0 $\pm$ 5.7	46.0 $\pm$ 7.1

Appendix

Table 9: Brightness shift,

-0.10

Model	TwoRoom	PushT	Reacher	Cube
LeWM	82.8 $\pm$ 1.3	91.2 $\pm$ 3.8	15.6 $\pm$ 4.7	56.8 $\pm$ 4.3
Action-grounded	90.6 $\pm$ 0.9	92.0 $\pm$ 3.4	20.0 $\pm$ 4.3	63.6 $\pm$ 6.8
Vision-invariance	81.3 $\pm$ 4.7	91.3 $\pm$ 2.4	18.1 $\pm$ 2.8	55.2 $\pm$ 5.2
AVL (ours)	92.7 $\pm$ 0.9	92.0 $\pm$ 2.2	24.7 $\pm$ 1.9	54.1 $\pm$ 1.6

Appendix

Table 10: Brightness shift,

+0.10

Model	TwoRoom	PushT	Reacher	Cube
LeWM	54.0 $\pm$ 4.8	12.0 $\pm$ 3.6	14.8 $\pm$ 8.2	37.6 $\pm$ 5.6
Action-grounded	43.2 $\pm$ 4.9	30.5 $\pm$ 3.5	20.3 $\pm$ 9.0	44.7 $\pm$ 4.1
Vision-invariance	82.0 $\pm$ 3.2	22.6 $\pm$ 0.9	22.5 $\pm$ 6.2	42.9 $\pm$ 4.3
AVL (ours)	91.3 $\pm$ 3.4	40.7 $\pm$ 1.3	24.7 $\pm$ 5.3	47.3 $\pm$ 4.7

Appendix

Table 11: Brightness shift,

-0.20

Model	TwoRoom	PushT	Reacher	Cube
LeWM	83.2 $\pm$ 1.8	91.6 $\pm$ 3.2	56.4 $\pm$ 6.5	70.2 $\pm$ 6.1
Action-grounded	92.0 $\pm$ 1.6	78.0 $\pm$ 1.8	85.3 $\pm$ 5.3	72.0 $\pm$ 5.9
Vision-invariance	83.1 $\pm$ 1.9	85.3 $\pm$ 5.2	81.6 $\pm$ 5.7	69.6 $\pm$ 5.4
AVL (ours)	94.0 $\pm$ 1.6	93.6 $\pm$ 3.7	82.7 $\pm$ 3.8	73.7 $\pm$ 5.7

Appendix

Table 12: Saturation shift,

+0.15

Explore similar work

Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency
Aug 13, 2026Guo An, Zijing Wu, Honghua Dong +7World Models Joint-Embedding Predictive Architecture
Huawei · University of Science and Technology of China · Zhejiang University +3
Preserving Unstable Modes Through Inverse Dynamics in JEPA World Models
Oct 6, 2026Leonardo F. Toso, Yann LeCun, James Anderson +1Representation Learning World Model Learning
Columbia University · New York University, AMI Labs · New York University
PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models
Aug 6, 2026Xi Zeng, Haojie Ren, Ziying Song +2Representation Learning World Model Learning
School of Mechanical and Aerospace Engineering, Nanyang Technological University, Singapore · The University of Sheﬃeld, Sheﬃeld, United Kingdom · School of Artiﬁcial Intelligence (School of Software), Yanshan University, Qinhuangdao, China

Aug 13, 2026J/K move · Enter open · S save

Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency

Guo An, Zijing Wu, Honghua Dong, Yuhao Yan+6

Huawei · University of Science and Technology of China · Zhejiang University+3

World Models Joint-Embedding Predictive Architecture Action-Conditioned World Models World Model Evaluation

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency

Preserving Unstable Modes Through Inverse Dynamics in JEPA World Models

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models