AVL-JEPA: Preventing Causal Dynamics Information Collapse In Joint Embedding Predictive Architecture World Models
This arXiv preprint introduces AVL-JEPA, a training method that stops joint-embedding world models from discarding action-conditioned physical dynamics while preserving visual representations. Across four simulated robot control tasks, it reaches 81.5% success on PushT and 88.7% on TwoRoom under Gaussian input noise, versus 15.3% and 65.7% for the baseline, while maintaining competitive clean-environment performance. The work matters for robot planners that must stay reliable when camera or state observations are noisy.
Can an auxiliary action-anchoring objective plus clean-versus-perturbed latent alignment prevent action-conditioned dynamics information collapse in JEPA world models, and does that improve planning under visual perturbations?
Joint-embedding predictive architectures learn future latents without reconstructing pixels, but common collapse-prevention regularizers only keep the latent non-degenerate; they can still discard the information needed to distinguish how different actions change the physical state. As a result, two actions with different physical outcomes can map to nearly the same predicted latent transition, and closed-loop model-predictive control degrades severely when observations are perturbed.
Earlier approaches included pixel-reconstructing latent world models such as the Dreamer family and reconstruction-free JEPA or regularized predictors such as VICReg, I-JEPA, V-JEPA, LeWM with SIGReg, PLDM, and SD-JEPA. These mainly constrain global latent statistics such as variance, covariance, rank, or distribution, preventing full representation collapse but not explicitly preserving action-conditioned physical consequences.
AVL trains in two phases on a LeWM-style JEPA. First, an action-grounded head predicts the executed multi-step action from a transition-attributed latent response, subtracting a zero-transition baseline so the prediction captures change caused by actions rather than static state. Second, a vision-invariance phase randomly perturbs the current observation with Gaussian noise, aligns the noisy predicted future latent toward the clean target for the same state, and keeps the action head frozen with its loss still active so alignment cannot be achieved by discarding dynamics information. The planning checkpoint removes the auxiliary head and uses an unchanged LeWM interface.
Across 10 seeds and 50 episodes per seed, AVL achieved clean success rates of 93.2% on TwoRoom, 96.0% on PushT, 85.2% on Reacher, and 78.2% on Cube, compared with LeWM at 88.4%, 96.0%, 86.4%, and 73.2%; on PushT the means tie but AVL reduces seed standard deviation from 3.5 to 1.3 points. The paper reports that under Gaussian noise sigma 0.10, AVL success reached 81.5% on PushT and 88.7% on TwoRoom, versus 15.3% and 65.7% for the baseline; under saturation shift +0.35, AVL reached 62.0% and 86.7%, versus 14.0% and 54.8%. Mechanism checks on TwoRoom showed predicted-to-physical Spearman correlation rising from 0.35 to 0.53 and normalized regret falling from 0.152 to 0.061, with clean-noisy ranking consistency rising from 0.89 to 0.98. On PushT, clean-noisy Spearman was 0.998 for AVL versus 0.809 for LeWM, Top-1 agreement 93.0% versus 45.3%, and action-to-nuisance ratio 19.23 versus 2.35.
The authors state the method is simulation-only and still depends on the SIGReg regularizer from LeWM rather than replacing it; real-robot experiments are future work. Evident limitations include evaluation on four standard simulated tasks without hardware transfer, no reported code or model release, and added two-phase training overhead, with up to about 62 minutes per epoch on Cube using an RTX 5090. The auxiliary action head is removed at deployment, so test-time benefit relies on the learned latent preserving dynamics without the head's direct signal.
Model-based control and robot learning teams can use this kind of regularized JEPA world model to make manipulation planners robust to camera noise, brightness, and saturation changes without pixel reconstruction. Warehouse and logistics robot providers, domestic service robot developers, and general-purpose robot foundation model teams are plausible adopters, but real-robot validation and scaling beyond four benchmark tasks are needed first. If real-robot results hold, integration into MPC stacks for vision-based manipulation could occur in 1-3 years; broad deployment likely takes 3+ years because safety-critical perception, compute, and latency constraints must still be addressed.
The full text is not republished here because the paper's license does not allow it. Read the original on arXiv.