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

Morgan Byrd; Robert Wright; Sehoon Ha

cs.ROSep 30, 2026

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

Authors: Morgan Byrd, Robert Wright, Sehoon Ha

Organizations: Georgia Institute of Technology, Atlanta, GA, 30308, USA · Georgia Tech Research Institute, Atlanta, GA, 30308, USA

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.

Figures & tables

Panel of Figure 1 — Fig. 1 : Primary tasks for evaluation. We test on the DM Control Suite Reacher, PushT, OGBench Cube, and TwoRoom tasks.

Fig. 3 : Tasks with example distractors (N=2).

	LeWM	SMWM	CF-JEPA
Model components
Forward loss	O	O	O
SIGReg	O	X	O
Inverse dynamics	X	O (full $z$ )	O ( $z^{c}$ )
Adversarial dynamics	X	X	O ( $z^{u}$ )
Latent factorization	X	X	O

TABLE I : Model components and the loss weight that is tuned per task,

\lambda

for the inverse dynamics loss of SMWM and

\alpha

for our inverse loss. LeWM has neither, so no weight applies.

Task	LeWM	SMWM	CF-JEPA
Reacher	0.86	0.66	0.61 (±0.03)
PushT	0.96	0.83	0.78 (±0.04)
OGBench Cube	0.74	0.84	0.81 (±0.03)
TwoRoom	0.87	0.99	0.95 (±0.04)

TABLE II : Nominal task performance. We report the values for LeWM and SMWM directly from the source papers and list our values as mean

\pm

standard deviation across three training seeds.

Task	LeWM	SMWM	CF-JEPA
Reacher	0.07 (±0.05)	0.1 (±0.03)	0.48 (±0.12)
PushT	0.01 (±0.01)	0.86 (±0.02)	0.87 (±0.03)
OGBench Cube	0.42 (±0.02)	0.84 (±0.01)	0.85 (±0.02)
TwoRoom	0.33 (±0.03)	0.98 (±0.05)	0.95 (±0.03)

TABLE III : Task performance under standard distractor settings (N=2) For all tasks, we report mean

\pm

one standard deviation.

Task	LeWM	SMWM	CF-JEPA
Reacher ( $N\!=\!2$ )	0.08	0.1 (±0.03)	0.48 (±0.12)
PushT ( $N\!=\!2$ )	0.01	0.86 (±0.02)	0.87 (±0.03)
PushT ( $N\!=\!16$ )	–	0.36 (±0.36)	0.76 (±0.)
PushT ( $N\!=\!32$ )	–	0.57 (±0.01)	0.68 (±0.06)
PushT ( $N\!=\!64$ )	–	0.27 (±0.25)	0.57 (±0.01)
OGBench Cube ( $N\!=\!2$ )	0.41	0.84 (±0.01)	0.85 (±0.02)

TABLE IV : Task performance as the number of distractors

N

increases. For each task, we report values up to the collapse of SMWM, except for OGBench Cube, where SMWM does not collapse.

Fig. 4 : Simplified ManiSkill Reach task used for further distractor evaluation.

Success Threshold (m)	LeWM	SMWM	CF-JEPA
0.06 (easiest)	0.55 (±0.04)	0.79 (±0.05)	0.83 (±0.06)
0.05	0.51 (±0.05)	0.62 (±0.04)	0.72 (±0.08)
0.04	0.44 (±0.03)	0.58 (±0.02)	0.64 (±0.1)
0.03	0.36 (±0.03)	0.38 (±0.02)	0.53 (±0.14)
0.025	0.32 (±0.06)	0.30 (±0.04)	0.48 (±0.17)
0.02 (hardest)	0.21 (±0.05)	0.22 (±0.04)	0.35 (±0.12)

TABLE V : Task performance under N=16 distractor settings for ManiSkill Reach task. We report mean

\pm

standard deviation across three seeds for each model.

Panel of Figure 10 — Fig. 5 : Only the controllable latent of CF-JEPA retains the agent. Reconstructions are decoded from frozen encoders trained with distractors and shown on Reacher test frames without distractors. Compared with the target (1st row), the controllable latent $z^{c}$ of CF-JEPA reconstructs the arm (2nd row), whereas the uncontrollable latent $z^{u}$ only blurs it (3rd row). SMWM and LeWM lose the arm entirely (4th and 5th rows).

	LeWM	SMWM	CF-JEPA
Nominal	48.2 (±0.8)	3.71 (±0.07)	42.3 (±0.1)
Distractor (N=2)	2.24 (±0.23)	1. (±0.01)	12.9 (±0.35)

TABLE VI : Latent vector participation ratio under nominal and distracted conditions for the Reacher task. LeWM and SMWM both collapse much further than CF-JEPA, correlating with the worse decoder and control performance.

Explore similar work

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models
Aug 6, 2026Xi Zeng, Haojie Ren, Ziying Song +2Latent World Models Joint-Embedding Predictive Architectures
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
D-JEPA: A Decision-Aligned Latent World Model
Sep 21, 2026Shuaijun Liu, Chengyu Wu, Qifu Wen +5Latent World Models World Models
The Hong Kong University of Science and Technology (Guangzhou) · Boston University · Shanghai Jiao Tong University
Subspace-Decomposed JEPAs: Disentangling Progression and Content in Latent World Models
May 29, 2026Lucas Thil, Jesse Read, Rim Kaddah +1Joint-Embedding Predictive Architectures Latent Prediction
LIX, École Polytechnique IRT SystemX Palaiseau, France · IRT SystemX Palaiseau, France · Safran Tech Chateaufort, France

Aug 6, 2026, cs.CVJ/K move · Enter open · S save

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

Xi Zeng, Haojie Ren, Ziying Song, Yuanbo Nie+1

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

Latent World Models Joint-Embedding Predictive Architectures World Models Inference Dynamics Latent Variable

Open paper AskPDF arXiv

We propose PhyLatent, a dynamics-relevant training objective for Joint-Embedding Predictive Architecture (JEPA) world models. Our key observation is that preventing global latent collapse does not necessarily ensure that the learned representation preserves physically meaningful state and action relationships. We identify three failure modes: sensitivity to appearance changes that leave the physical state unchanged, insufficient separation of distinct physical states, and insufficient separation of different action-conditioned futures. We refer to these as Physical Invariance Collapse, Physical Distinguishability Collapse, and Counterfactual Dynamics Collapse, respectively. PhyLatent targets these failures through three coordinated training pathways, implemented with static visual invariance, physical state grounding, future representation alignment, counterfactual branch separation, and latent denoising. On OGBench-Cube, PhyLatent reduces the three collapse rates by 43.9%, 27.9%, and 47.1%, respectively, while improving model predictive control (MPC) success by 12.0 percentage points (17.2% relative). Across four visual-control tasks, average planning success increases by 6.62 percentage points (8.3% relative). These results show that global non-collapse alone is insufficient for learning a reliable JEPA world-model state space, and that explicitly preserving dynamics-relevant structure can improve closed-loop planning.

Abstract

Explore similar work

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

D-JEPA: A Decision-Aligned Latent World Model

Subspace-Decomposed JEPAs: Disentangling Progression and Content in Latent World Models