Preserving Unstable Modes Through Inverse Dynamics in JEPA World Models

Leonardo F. Toso; Yann LeCun; James Anderson; Oumayma Bounou

cs.LGOct 6, 2026

Preserving Unstable Modes Through Inverse Dynamics in JEPA World Models

Authors: Leonardo F. Toso, Yann LeCun, James Anderson, Oumayma Bounou

Organizations: Columbia University · New York University, AMI Labs · New York University

Abstract

Robotic systems often exhibit unstable modes, along which small perturbations and disturbances can cause unbounded growth unless corrected through feedback. Controlling such systems from high-dimensional visual observations requires representations that preserve these modes. Joint-embedding predictive architectures (JEPAs) provide a natural framework for learning such representations and their dynamics from visual data. However, we demonstrate that next step prediction combined with anti-collapse regularization does not guarantee that controllable unstable modes are preserved: the training loss can be minimized while these modes are collapsed, making stabilization from the learned representation impossible. To address this, we augment world-model training with an action reconstruction objective (i.e., an inverse dynamics loss) that encourages control-aware representations, namely, visual representations that preserve crucial features for control. We prove that exact action reconstruction makes the encoder injective on the finite-horizon reachable subspace. Thus, the encoder cannot discard any state direction reachable by an action sequence within $H$ steps. Moreover, we show that, as $H$ grows, the dominant eigenspace of the finite-horizon controllability Gramian converges to the controllable unstable subspace. We establish our theoretical results for linear systems and demonstrate empirically that our findings extend to nonlinear visual control tasks (CartPole, Walker2D, and PointMaze), highlighting the benefits of control-aware representation learning.

Figures & tables

Figure 1 : Architecture of our proposed world model. At each step, an encoder ( $E$ ) maps observation $y_{t}$ to a latent state $z_{t}$ , and a predictor ( $P$ ) rolls out latent predictions $\hat{z}_{t+1},\ldots,\hat{z}_{t+H}$ conditioned on actions. Training losses are applied at each predicted latent: 1SP/MSP (single- or multi-step prediction) aligns $\hat{z}_{t+k}$ with the encoded ground-truth $z_{t+k}$ . The EP-IDM estimates the intermediate actions $\hat{\mathbf{a}}_{t,H}=(\hat{a}_{t},\hat{a}_{t+1},\ldots,\hat{a}_{t+H-1})$ from $z_{t}$ and $z_{t+H}$ .

Figure 2 : SIG vs. EP-IDM across two examples. Each row compares a representation trained with SIGReg (SIG) against one trained with the endpoint inverse-dynamics (action-reconstruction) loss (EP-IDM). Top row: synthetic open-loop unstable LDS. Bottom row: linearized CartPole. Left: true unstable (

v_{u}

, green) and stable (

v_{s}

, grey) eigenvectors versus the learned predictor’s dominant eigenvector, for SIG (blue) and EP-IDM (dashed orange). Center: closed-loop trajectories under the latent LQR controller from four initial states (circles) for the model trained with SIGReg. Right: closed-loop trajectories under the latent LQR controller for the model trained with the inverse dynamics loss (EP-IDM). For the bottom row

x_{3}

, and

x_{4}

corresponds to the pole angle and angular velocity, respectively. The latent dimension is one for the example in the top row and three for the example in the bottom row.

Figure 3 : CartPole , Walker2D , and PointMaze visual-control tasks used in our evaluation. Each row depicts ten frames from a successful trial (initial state: blue border, final state: green border). CartPole: a continuous-action inverted-pendulum balancing task controlled via LQR. Walker2D: a bipedal robot locomotion task controlled via iCEM ( Pinneri et al., 2021 ) . PointMaze: a 2-D point-mass robot navigation task planned via CEM.

Figure 4 : CartPole state norm

\|x_{t}\|_{2}

under the designed latent LQR controller over

300

steps.

Figure 5 : Walker2D. Phase portraits of right-hip angle vs. angular velocity for GT dynamics, 1SP+EP-IDM, 1SP+SIG, and MSP+SIG. Each trajectory is a

500

-step closed-loop rollout: GT uses real MuJoCo physics, model panels roll out in latent space under an MLP-decoded state, with a SAC ( Huang et al., 2022 ; Haarnoja et al., 2018 ) policy acting on the decoded observation.

Model	LQR (SR)	LQR (MFS)	CEM (SR)	GBP (SR)
Ground truth (GT)	100%	1.000	100%	100%
1SP + SIG	0%	0.202	100%	90%
MSP + SIG	0%	0.247	100%	100%
1SP + EP-IDM	100%	0.999	100%	100%
MSP + EP-IDM + SIG	100%	0.996	100%	100%

Table 1 : CartPole stabilization success rate (%).

Figure 6 : Trajectory frames for Walker2D . Each row shows ten frames uniformly sampled from a $500$ -step closed-loop rollout under the SAC policy ( Haarnoja et al., 2018 ) , decoded from the trained model’s latent state via an MLP decoder.

Model	Avg Velocity (m/s)	Displacement (m)	Final Height (m)
Ground truth (GT)	3.61	14.43	1.21
1SP + SIG	0.08	0.06	0.98
MSP + SIG	0.46	0.66	0.85
1SP + EP-IDM	2.91	9.15	0.92

Table 2 : Walker2D locomotion under latent iCEM ( Pinneri et al., 2021 ) . All metrics are averaged over ten trials. Final height is the torso height (m) at the end of the episode.

Figure 7 : Planning trajectory frames for Walker2D . Each row shows ten frames uniformly sampled from a $500$ -step closed-loop rollout under latent iCEM planning for the trained models. We select the best trial, with respect to the forward displacement, across the ten trials for the trained models.

Figure 8 : CartPole diagnostics. Left: Phase portraits of ground-truth and learned vector fields. Center: $H=3$ open-loop prediction error across $(\theta,\dot{\theta})$ . Right: Zero-action planning cost.

Encoder	Objective	CEM (SR)	LQR (SR)	GBP (SR)
GT		100%	80%	90.0% *
From scratch	1SP + IDM	100%	70%	100%
	MSP + MS-IDM	100%	80%	90.0%
	MSP + EP-IDM	90.0%	80%	90.0%
	1SP + EP-IDM	100%	80%	80.0%
	MSP + SIG	90.0%	60%	100%

Table 3 : PointMaze navigation success rate (%).

Figure 9 : Local stability on CartPole . First two panels: Empirical region of attraction of the learned LQR controller when deployed in the ground-truth system. The dashed line indicates the ground-truth LQR boundary. Last two panels: Lyapunov decrease $\Delta V=V(x^{\prime})-V(x)$ under the latent LQR controller.

Appendix figures & tables27 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 10 : CartPole trajectory visualization for the trained models under the latent LQR controller.

Figure 11 : MSP+EP-IDM+SIG model on CartPole . Left: Phase portrait of ground-truth vs. learned vector fields. Center: $H=3$ open-loop prediction error. Right: Zero-action planning cost.

Figure 12 : Local stability of MSP+EP-IDM+SIG on CartPole . Left: Empirical region of attraction (ROA) of the learned LQR policy in the GT environment. Dashed line: GT LQR boundary. Right: Lyapunov decrease $\Delta V=V(x^{\prime})-V(x)$ under the learned policy using the GT DARE solution.

Figure 13 : MSP+IBOT+PR-EP-IDM model on CartPole . Panel descriptions are in Fig. 11 .

Figure 14 : Local stability of MSP+IBOT+PR-EP-IDM on CartPole . Panel descriptions are the same as those in Fig. 12 .

Figure 15 : MSP+IBOT+PR-SIG model on CartPole . Panel descriptions are the same as those in Fig. 11 .

Figure 16 : MSP+IBOT model on CartPole . Panel descriptions are the same as those in Fig. 11 .

Figure 17 : Local stability of MSP+IBOT on CartPole . Left: Empirical region of attraction (ROA) of the learned LQR policy in the GT environment. Panel descriptions are as those in Fig. 12 .

Figure 18 : DINO-WM model on CartPole . Panel descriptions are the same as those in Fig. 11 .

Encoder	Objective	CEM ( $1/L^{2}$ )	CEM (FE)
GT		0.0003	0.30202
From scratch	1SP + SIG	0.0003	0.19956
	1SP + IDM + MS-IDM	0.0003	0.43608
	1SP + EP-IDM	0.0003	0.13683
	MSP + SIG	0.0003	0.07667
	MSP + EP-IDM + SIG	0.0003	0.11985

Appendix

Table 4 : CartPole CEM diagnostic metrics (action context 1).

1/L^{2}

: inverse squared mean stabilization length. FE: final error

\|x_{T}\|

. PR denotes a learned projector on top of the frozen encoder.

Figure 19 : PointMaze visualization for our trained models under CEM planning.

Figure 20 : Latent goal-distance fields $\|z-z_{\mathrm{g}}\|$ across the PointMaze U-maze for four trained world models.

Figure 21 :

k

-step closed-loop prediction error along trajectories generated by the ground-truth LQR controller. Curves show the physical-state error between the ground-truth trajectory and the trajectory decoded from each latent rollout.

Encoder	Objective	LQR (SR)	LQR (MFS)	CEM (SR)	GBP (SR)
DINOv2	1SP	0%	0.301	0%	0%
	MSP	0%	0.232	0%	0%
	1SP + PR-EP-IDM	0%	0.025	20%	0%
	MSP + PR-SIG	0%	0.460	100%	0%
	MSP + PR-IDM	0%	0.455	0%	0%
	MSP + PR-EP-IDM	0%	0.513	40%	10%

Appendix

Table 5 : CartPole control performance with frozen pretrained encoders. SR: success rate. MFS: mean fraction of steps within the success threshold. PR denotes a learned projector. DINO-WM ( Zhou et al., 2024 ) is an external baseline.

Figure 22 : CartPole state norm

\|x_{t}\|_{2}

under the learned latent LQR controller over

300

steps. Left: DINOv2. Right: iBOT.

Task	Image size	Proprioception	Action	Frame skip	Training trajectories
CartPole	$128\times 128$	$4$	$1$	$5$	$787$
Walker2D	$64\times 64$	$17$	$6$	$5$	$3,200$
PointMaze	$64\times 64$	$4$	$2$	$5$	$1,600$

Appendix

Table 6 : Task-specific observation and action dimensions. The frame skip is the number of simulator steps represented by one model transition.

Hyperparameter	Value	Hyperparameter	Value
Optimizer	AdamW	Epochs	$200$
Batch size	$64$	Prediction horizon $H$	$3$
Initial learning rate	$10^{-4}$	Minimum learning rate	$10^{-6}$
Weight decay	$10^{-3}$	Warm-up epochs	$5$
Gradient-norm threshold	$1$	Latent dimension	$192$
Random seed	$42$	Active loss weights	$1$

Appendix

Table 7 : World-model training hyperparameters.

Task	Planner	Horizon	Execute	Population	Elites	Iterations
CartPole	CEM	$10$	$1$	$300$	$30$	$10$
Walker2D	iCEM	$3$	$1$	$200$	$20$	$5$
PointMaze	CEM	$25$	$25$	$300$	$30$	$10$

Appendix

Table 8 : CEM and improved-CEM planning hyperparameters. Here, “execute” denotes the number of actions applied before replanning.

Objective	LQR (SR)	LQR (MFS)	CEM (SR)	GBP (SR)
1SP + IDM + MS-IDM	100%	0.999	100%	100%
1SP + EP-IDM	100%	0.999	100%	100%
MSP + EP-IDM + SIG	100%	0.996	100%	100%

Appendix

Table 9 : CartPole ablation over inverse-dynamics formulations for an encoder trained from scratch.

		LQR (SR)		CEM (SR)		GBP (SR)
Encoder	Objective	Context 1	Context 5	Context 1	Context 5	Context 1	Context 5
From scratch	1SP + SIG	0	0	100	100	90	90
	1SP + IDM + MS-IDM	100	100	100	100	100	100
	1SP + EP-IDM	100	100	100	100	100	100
	MSP + SIG	0	40	100	100	100	100
DINOv2	MSP	0	0	0	0	0	0

Appendix

Table 10 : CartPole success rates (%) with action contexts one and five. PR denotes a learned projector on top of the frozen encoder.

		LQR (SR)		CEM (SR)		GBP (SR)
Encoder	Objective	Context 1	Context 5	Context 1	Context 5	Context 1	Context 5
From scratch	1SP + IDM	70	80	100	90	100	80
	1SP + EP-IDM	80	80	100	100	80	80
	1SP + IDM + MS-IDM	80	70	90	100	90	90
	MSP + SIG	60	80	90	70	100	60
	MSP + MS-IDM	80	80	100	100	90	90

Appendix

Table 11 : PointMaze success rates (%) with action contexts one and five.

Predictor	Objective	LQR (SR)	LQR (MFS)	CEM (SR)	CEM (held)	GBP (SR/held)
MLP	MSP + SIG	20%	0.571	20%	20%	0% / 0%
	1SP + EP-IDM	0%	0.204	100%	100%	0% / 0%
Transformer	MSP + SIG	20%	0.461	0%	0%	0% / 0%
	1SP + EP-IDM	80%	0.869	100%	100%	100% / 100%

Appendix

Table 12 : CartPole ablation: Action context one, only proprioceptive information (state).

Objective	Proprioception	LQR (SR)	CEM (SR)	CEM (held)	GBP (SR)	GBP (held)
MSP	Yes	0%	100%	100%	60%	40%
	No	0%	0%	0%	0%	0%
MSP + PR-SIG	Yes	40%	90%	90%	60%	60%
	No	0%	0%	0%	0%	0%
MSP + PR-IDM	Yes	100%	80%	80%	70%	70%
	No	0%	0%	0%	0%	0%

Appendix

Table 13 : CartPole ablation with a frozen iBOT encoder and action context five, with and without proprioceptive inputs. PR denotes the learned projector. SR: success rate, held: fraction of trials in which stabilization is preserved.

Encoder	$\\|x_{0}\\|{=}0.005$	$\\|x_{0}\\|{=}0.01$	$\\|x_{0}\\|{=}0.02$	$\\|x_{0}\\|{=}0.05$	$\\|x_{0}\\|{=}0.10$	$\\|x_{0}\\|{=}0.20$
iBOT	1.07 $\times$	1.11 $\times$	1.07 $\times$	1.07 $\times$	1.02 $\times$	1.03 $\times$
DINOv2	25.45 $\times$	26.07 $\times$	26.38 $\times$	25.67 $\times$	21.83 $\times$	25.12 $\times$

Appendix

Table 14 : Ablation on the encoder gain

\|\Delta z_{t}\|/\|\Delta x_{t}\|

as a function of initial-state perturbation magnitude

\|x_{0}\|

. We start from a perturbed upright equilibrium, the cart pole is stepped with zero action for 60 steps. We measure the mean latent displacement

\|\Delta z_{t}\|=\|z_{t}-z_{t-1}\|

divided by the mean physical-state displacement

\|\Delta x_{t}\|=\|x_{t}-x_{t-1}\|

Encoder	Objective	$\\|P(z_{\mathrm{eq}},0)-z_{\mathrm{eq}}\\|$
DINOv2	1SP	5.569
	MSP	7.215
	1SP + PR-EP-IDM	0.4026
	MSP + PR-SIG	2.515
	MSP + PR-EP-IDM	0.343
iBOT	1SP	0.059

Appendix

Table 15 : Fixed-point residual of the learned predictor at the encoded upright CartPole equilibrium under zero action. The residual measures the one-step drift from this encoded equilibrium. PR denotes a learned projector on top of the frozen encoder.

Figure 23 : Synthetic open-loop stable LDS. This example considers a two-dimensional open-loop stable system. Panels descriptions are the same as those in Fig. 2 .

Figure 24 : Sign of $\Delta V=V(x^{\prime})-V(x)$ after one closed-loop step under the learned latent LQR controller, where $V(x)=x^{\top}Sx$ and $S\succ 0$ is the ground-truth Lyapunov function matrix. We also depict the $\Delta V=0$ contour (separatrix), i.e., the boundary separating states where the learned controller decreases $V$ (blue) from those where it does not (red).

Explore similar work

CF-JEPA: Improving Robustness of JEPA World Models via Controllability Factorization
Sep 30, 2026Morgan Byrd, Robert Wright, Sehoon HaLatent World Models Controllability
Georgia Institute of Technology, Atlanta, GA, 30308, USA · Georgia Tech Research Institute, Atlanta, GA, 30308, USA
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
Toward Physically Grounded JEPA World Models for Goal-Conditioned Robotic Planning
Sep 3, 2026Muyuan Liu, Yue Huang, Zheng Liang +1Inverse Dynamics Robot Planning
GENISOM AI, Beijing, China

Sep 30, 2026, cs.ROJ/K move · Enter open · S save

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

Morgan Byrd, Robert Wright, Sehoon Ha

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

Latent World Models Controllability World Models Low-Latency Latent Planning Distractor

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

Toward Physically Grounded JEPA World Models for Goal-Conditioned Robotic Planning