Preserving Unstable Modes Through Inverse Dynamics in JEPA World Models

Leonardo F. Toso1 · Yann LeCun3 · James Anderson1 · Oumayma Bounou2

1 Columbia University   2 New York University   3 New York University, AMI Labs

PaperCodeCheckpoints

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 representations and their dynamics from visual data. However, we demonstrate that next step prediction combined with anti-collapse regularization provides no such guarantee: the training loss can be minimized while controllable unstable modes are collapsed, making stabilization from the learned representation impossible. To circumvent this, we augment predictive 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; hence, no state direction produced by an action sequence within H steps can be discarded. Moreover, as H grows, the dominant eigenspace of the finite-horizon controllability Gramian reveals the controllable unstable subspace. Although our theory considers linear systems, experiments on nonlinear visual control tasks (CartPole, Walker2D, and PointMaze) validate our findings and show the benefits of control-aware representation learning.

Problem

The collapse of unstable modes

Prediction does not guarantee stabilization. In the linear setting, a controller using only the learned latent observations can stabilize the physical system if and only if every unstable or marginally stable mode are observable through the learned encoder, assuming the physical system is stabilizable and its original observations are observable.

Prediction loss can be minimized and still discard unstable modes.

There exist open-loop unstable linear systems and data distributions for which prediction loss and SIGReg are simultaneously minimized while the encoder discards the entire unstable subspace. The representation can remain non-constant and predictable, yet make stabilization from latent observations impossible.

Synthetic unstable linear system and linearized CartPole: eigenvector comparisons and closed-loop trajectories under SIGReg and EP-IDM latent LQR controllers
Linear-system examples. Top: a synthetic open-loop unstable system. Bottom: linearized CartPole. Left: true unstable and stable eigenvectors versus the learned predictor’s dominant eigenvector. Center: latent LQR trajectories with SIGReg. Right: latent LQR trajectories with EP-IDM. In both examples, SIGReg learns a misaligned unstable direction and the resulting controller fails to stabilize the physical dynamics.

Method

World-model architecture

JEPA architecture with an encoder, action-conditioned latent predictor, and endpoint inverse dynamics loss
Forward prediction learns latent dynamics. EP-IDM reconstructs the full action sequence from the first and last encoded observations.

At each step, an encoder maps the observation to a latent state, and an action-conditioned predictor rolls out future latent representations. Single- or multi-step prediction losses align these predictions with the encoded ground-truth observations. Endpoint inverse dynamics (EP-IDM) estimates the intermediate action sequence from the initial and final latent states.

Preserving unstable modes through inverse dynamics

EP-IDM reconstructs the complete action sequence over the prediction horizon from the initial and final latent representations. This objective encourages the encoder to retain action-induced directions that are crucial for control.

Theoretical Guarantee

Under mild assumptions, including sufficiently rich action excitation and exact reconstruction by an affine inverse model, the encoder is injective on the H-step reachable subspace.

ker(F) ∩ ℛH = {0}

No reachable state direction is discarded.

F = WC is the linear state-to-latent map: C maps the physical state to an observation, and W is the linear part of the observation encoder.

ℛH = range([AH−1B, …, AB, B]) is the H-step reachable subspace: the state directions generated by action sequences over H steps. A and B are the system’s state-transition and input matrices.

In the linear-system examples above, endpoint inverse dynamics preserves the unstable direction in latent space, allowing the latent LQR controller to stabilize the physical system.

Numerical experiments

0%

Prediction + SIGReg

100%

Prediction + EP-IDM

CartPole state norm under latent LQR over 300 steps
CartPole state norm during a 300-step rollout.
CartPole stabilization · encoders trained from scratch · 10 trials
Training objectiveLQR successLQR MFSCEM successGBP success
Ground truth100%1.000100%100%
1SP + SIG0%0.202100%90%
MSP + SIG0%0.247100%100%
1SP + EP-IDM100%0.999100%100%
MSP + EP-IDM + SIG100%0.996100%100%

Success: final state norm ≤ 0.7 after 300 steps. MFS: mean fraction of steps within that threshold. 1SP / MSP: one-step / multi-step prediction. SIG: SIGReg. CEM: cross-entropy planning. GBP: gradient-based planning.

Walker2D: preserving a walking gait

EP-IDM preserves the cyclic structure and range of joint angles and angular velocities, whereas the SIGReg models fail to capture the limit cycle induced by the locomotion policy.

Walker2D right-hip angle versus angular velocity phase portraits comparing ground truth, EP-IDM, and SIGReg models
Walker2D · right-hip phase portraits across 600-step closed-loop rollouts. EP-IDM retains the cyclic structure observed in ground-truth dynamics.

CartPole: local dynamics and stability

CartPole ground-truth and learned vector fields, three-step prediction errors, and zero-action planning costs
CartPole diagnostics · vector fields, three-step prediction error, and zero-action planning costs.
CartPole empirical regions of attraction and Lyapunov differences under learned latent LQR controllers
Local stability · empirical regions of attraction and Lyapunov decrease under latent LQR. These diagnostics evaluate the controllers in the ground-truth system.

PointMaze: goal reaching with stable dynamics

All action-reconstruction variants trained from scratch reach at least 90% CEM success and at least 70% LQR success in the PointMaze experiments.

Videos

WALKER2D

Preserving a walking gait

Locomotion requires continuous feedback. EP-IDM preserves the cyclic structure of the walking gait.

CARTPOLE

Stabilizing the upright equilibrium

Small angular disturbances reveal whether the learned representation retains the directions needed for feedback.

POINTMAZE

Navigating with an open-loop stable system

Action-reconstruction variants trained from scratch achieve at least 90% CEM success.