CF-JEPA: Improving Robustness of JEPA World Models via Controllability Factorization

2026-09-30

AI Summary

This study addresses the vulnerability of JEPA-based world models to distractor signals in visual control, which often leads to latent space collapse. To this end, we propose CF-JEPA, a framework built upon a JEPA-style latent world model with a pixel-reconstruction-free self-supervised architecture. By introducing a controllability factorization mechanism, CF-JEPA explicitly decomposes the latent space into controllable and uncontrollable subspaces, effectively isolating distracting information while extracting task-relevant features for control. Experimental results demonstrate that CF-JEPA achieves performance comparable to conventional methods under nominal conditions while significantly outperforming them in the presence of distractors. Notably, it is the only approach capable of preventing latent space collapse. Furthermore, its practical utility is validated through simulated robotic tasks.

Abstract

Controlling an agent with vision requires being able to separate useful information from irrelevant background information. JEPA-style latent world models seem like a natural approach for this, as they do not perform pixel-level reconstruction; however, they are still sensitive to these distractor signals and experience latent collapse. In this work, we introduce Controllability Factorized JEPA (CF-JEPA), a JEPA-style world model which splits the latent space into controllable and uncontrollable subspaces. This factorization allows us to capture all the distractor information into the uncontrollable region, while we use the control-relevant latent information for our task. With this, we show comparable performance across 2D and 3D control tasks under nominal conditions and improved performance under distracted conditions, where CF-JEPA is the only model that does not experience latent collapse. We also validate our model under distracted conditions for a simulated robot task, highlighting the practical application of such a scheme.

Problem

• JEPA
• world models
• latent collapse
• distractor signals
• robustness

Innovation

• Controllability Factorization
• JEPA
• Latent World Models
• Robustness
• Latent Collapse