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

2026-10-02

AI Summary

This study addresses the causal dynamics information collapse problem in Joint-Embedding Predictive Architecture (JEPA) world models, wherein visual features are preserved while the physical consequences of actions are lost. We formally define this failure mode for the first time and propose Action-anchored Visual Learning (AVL), a method that constructs a visual invariance path using actions as auxiliary anchors to enforce the encoding and utilization of causal dynamics within the latent space. Experimental results demonstrate that AVL significantly improves planning success rates under visual perturbations across four robotic control tasks while maintaining performance in standard environments. These findings validate the effectiveness of the proposed framework for robust causal dynamics modeling.

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.

Problem

• Joint Embedding Predictive Architecture
• World Models
• Causal Dynamics Information Collapse
• Robotic Control

Innovation

• Joint Embedding Predictive Architecture
• Causal Dynamics Information Collapse
• Action-Grounded Vision-Invariance Latent
• World Models
• Robotic Control