Can inverse dynamics preserve unstable modes in JEPA world models?

Prediction loss can hide a collapsed unstable mode

The paper asks whether a JEPA world model trained the usual way keeps the physical directions a feedback controller must see. For an open-loop unstable system, the authors state a controller acting on latent states can stabilize the system if and only if the encoder keeps every unstable mode: if the controller cannot see a mode, it cannot correct it. Their first result is negative. With SIGReg as the anti-collapse regularizer, the standard objective of next-step prediction plus regularization can be driven to its minimum while the encoder discards all unstable modes.

The distinction matters because it separates two things readers often treat as one. Total representation collapse means the encoder maps almost all observations to one constant, and SIGReg prevents that by keeping the aggregate latent distribution dispersed and roughly isotropic. The paper's point is narrower and more damaging: a representation can stay dispersed on the training distribution and still drop the specific directions that feedback must observe. A low prediction loss therefore does not certify that the latent space is usable for stabilization.

Reconstructing actions from two endpoints

The proposed fix is an endpoint inverse-dynamics loss. At each step an encoder maps observations to a latent state, and a predictor rolls latent states forward conditioned on actions; single- or multi-step prediction aligns predicted latents with encoded ground-truth latents. EP-IDM adds a second requirement: from the latent at the start of a window and the latent H steps later, a model must estimate the intermediate action sequence. If two different states were mapped to the same latent, the actions leading away from them would generally differ, so the reconstruction task penalizes that merging.

The authors prove, for linear systems, that exact action reconstruction makes the encoder injective on the finite-horizon reachable subspace. In plain terms, the encoder cannot throw away any state direction that some action sequence can reach within H steps. They also show that as H grows, the dominant eigenspace of the finite-horizon controllability Gramian converges to the controllable unstable subspace, which explains why the unstable directions become the most identifiable ones under this objective. These results are established for linear systems; the nonlinear experiments are presented as evidence that the findings extend, not as a proof.

CartPole stabilization separates the two objectives

The clearest comparison is nonlinear CartPole, a continuous-action inverted-pendulum balancing task controlled with LQR. The paper reports that EP-IDM increases latent-LQR success from 0% to 100%. The same section describes the contrast in behavior: models trained with EP-IDM, both from scratch and with a frozen iBOT encoder, keep the pole upright throughout a sampled rollout, while the prediction-plus-SIGReg model, frozen iBOT without EP-IDM, and DINO-WM fail to stabilize the system.

The diagnostics explain what the success rate is measuring. The EP-IDM model recovers the local vector field around the upright equilibrium, shows low prediction error near that equilibrium, and produces a broad region of attraction with the expected decrease in the Lyapunov certificate. The comparison models have less accurate local dynamics even when their terminal latent cost is minimized near the goal. That is the paper's central warning in concrete form: reaching the goal in latent cost is not the same as learning dynamics from which a stabilizing controller can be designed.

Walker2D and PointMaze test different claims

Walker2D and PointMaze probe cases beyond a single unstable equilibrium, and their numbers should not be read as one combined result. On Walker2D, a bipedal locomotion task planned with iCEM, the paper reports that EP-IDM preserves the limit cycle and supports successful planning with 9.15 m of average forward displacement over 500 time steps. On PointMaze, a 2-D point-mass navigation task, the authors note that all action-reconstruction variants trained from scratch reach at least 90% CEM success and at least 70% LQR success.

PointMaze also serves as a check on systems with no unstable modes. The reported table shows ground truth at 100% CEM, 80% LQR, and 90% GBP success; from-scratch variants such as 1SP+IDM at 100% CEM and 70% LQR, MSP+EP-IDM at 90% CEM and 80% LQR, and MSP+SIG at 90% CEM and 60% LQR. The authors' reading is that action reconstruction remains effective there rather than being useful only when unstable modes exist. Because these are success rates under different planners, they compare planner-specific outcomes, not a single unified score.

Where the guarantee stops

The theoretical claims are stated for linear systems and depend on exact action reconstruction with sufficient action excitation. The convergence result for the controllability Gramian excludes eigenvalues on the unit circle, so its instantiation to the linearized CartPole example concerns only the strictly unstable mode shown in the left panel. The nonlinear tasks are empirical demonstrations, not extensions of the proof.

A second boundary is the planning cost. The paper notes that the goal-reaching cost used in latent planning is deliberately simple and may be insufficient for more complex control tasks, which can require task-specific stage costs, learned latent reward models, or costs defined through a state decoder, as done for Walker2D. The CartPole evidence also comes from a specific controller, latent LQR, so the 0% to 100% shift describes that controller and that task rather than every downstream use of the representation.

Judge a world model by what its latent controller can do

The paper's practical priority is a change in evaluation. A low prediction loss and a non-collapsed latent distribution are not sufficient evidence that a JEPA world model is ready for feedback control of an unstable system. The check that matters is whether a controller designed on the learned latent dynamics can actually stabilize the plant, and the CartPole comparison shows those two assessments can diverge completely: 0% latent-LQR success for prediction plus SIGReg versus 100% with EP-IDM.

The scenario split follows from the evidence. When the system has unstable modes that feedback must correct, the paper's argument and its CartPole result make action reconstruction the relevant training addition, with the linear theory explaining why reachable directions survive. When the system has no unstable modes, as in PointMaze, the reported from-scratch action-reconstruction variants still reach at least 90% CEM and at least 70% LQR success, so the objective is not presented as harmful there. The deciding question is not whether the latent predicts well, but whether the directions control depends on are still present.

Sources used in this answer