PaperlayerView exact revision on arXiv

Source-linked AI summary

Preserving Unstable Modes Through Inverse Dynamics in JEPA World Models

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

arXiv:2610.07540v1cs.LGcs.ROcs.SYeess.SYmath.OC

TL;DR

The standard JEPA recipe can minimize prediction and anti-collapse losses while discarding unstable modes needed for feedback stabilization. The paper adds endpoint action reconstruction and proves reachable-direction preservation for linear systems. EP-IDM improves nonlinear visual control results across CartPole, Walker2D, and PointMaze.

  • Problem

    Next-step prediction with anti-collapse regularization can reach its minimum while discarding unstable modes needed for stabilization from the learned representation.

    Source: §1 · p. 2

  • Method

    EP-IDM reconstructs the action sequence from initial and final latent states to encourage preservation of action-reachable directions.

    Source: §1 · p. 2, §5.1 · p. 8

  • Results

    EP-IDM improves nonlinear visual control across CartPole, Walker2D, and PointMaze, including latent-LQR success from 0% to 100% on CartPole.

    Source: §1 · p. 3, §7 · p. 15

  • Takeaways & Limitations

    For linear systems, exact action reconstruction preserves the finite-horizon reachable subspace, while its controllability-Gramian dominant eigenspace converges to the controllable unstable subspace as H grows.

    Source: §1 · p. 2, §7 · p. 14

    Exact recovery of arbitrary action sequences is not feasible, and the result requires sufficiently rich action excitation.

    Source: §5.2 · pp. 8–9

Abstract

from arXiv · show

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.

1 Introduction

Robust feedback from visual observations requires representations that retain unstable modes, but prediction with anti-collapse regularization can discard them. The paper proposes endpoint inverse dynamics and studies its theoretical and empirical benefits.

1 Introduction

  • Next-step prediction with anti-collapse regularization can reach its minimum while discarding every unstable mode, preventing stabilization from the learned latent state.The paper identifies preserving unstable modes as necessary for control from latent representations.Source: §1 · p. 2
  • EP-IDM reconstructs the action sequence from initial and final latent states to encourage representations that preserve action-reachable directions.Exact reconstruction with sufficient action excitation preserves every direction reachable within H steps.Source: §1 · p. 2
  • As H grows, the dominant eigenspace of the finite-horizon controllability Gramian converges to the controllable unstable subspace.Source: §1 · p. 3
  • EP-IDM preserves unstable directions in synthetic linear systems and linearized CartPole, with further validation on nonlinear visual CartPole, Walker2D, and PointMaze.Source: §1 · p. 3

2 Related Work

Related work studies latent predictive world models and feedback stabilization, but the paper focuses on whether standard JEPA objectives preserve information needed to stabilize open-loop unstable systems.

Joint embedding predictive world models

  • Prior JEPA work develops action-conditioned prediction, planning, anti-collapse methods, and objectives for shaping latent dynamics, while this work addresses selective loss of stabilization-relevant directions.Source: §2 · p. 3
  • Ivashkov et al. (2026) use one-step inverse dynamics to prevent collapse; this paper instead reconstructs the full action sequence and studies feedback stabilization of open-loop unstable systems.The paper connects exact action reconstruction to finite-horizon reachability and its convergence to the controllable unstable subspace.Source: §2 · p. 3

Latent feedback stabilization

  • Prior stabilization work studies unknown observation maps, sample complexity, and low-dimensional unstable representations; this paper asks whether standard JEPA training learns representations sufficient for stabilization.Source: §2 · p. 3

3 Setup

The setup defines a physical system observed through potentially multimodal measurements and a latent world model trained to predict future representations. It also specifies a simple latent-space planning objective and the stabilization requirement for unstable systems.

3 Setup

  • The physical system has state x_t, action a_t, and observation y_t, which can include images, proprioception, both, or other modalities.Source: §3 · pp. 3–4
  • The model learns an observation encoder, action encoder, and action-conditioned predictor that maps observations into a lower-dimensional latent space and forecasts latent evolution.In experiments, each encoder input uses two consecutive images, their difference, and the current proprioceptive measurement.Source: §3 · p. 4
  • Standard training combines next-step prediction with SIGReg to prevent total representation collapse, although dispersed latent variability does not ensure preservation of any particular physical direction.Source: §3 · p. 4
  • Latent planning optimizes an action sequence toward a goal, but the stated goal-reaching cost may be insufficient for complex tasks requiring task-specific costs or decoded-state costs.The implementation uses receding-horizon planning.Source: §3 · pp. 4–5
  • For an open-loop unstable linear system, latent feedback stabilization requires the encoder to preserve every unstable or marginally stable direction that feedback must correct.Source: §3 · p. 5

4 Prediction Does Not Guarantee Stabilization

The linear analysis establishes that accurate prediction and anti-collapse regularization do not ensure stabilization-relevant information is retained. It identifies an observability condition for latent stabilization and constructs a case where the standard objective discards unstable modes.

4 Prediction Does Not Guarantee Stabilization

  • A latent controller can converge while the physical state diverges if the encoder hides an unstable mode, despite accurate latent prediction.The analysis considers open-loop unstable linear systems and an affine encoder with a linear predictor.Source: §4 · pp. 5–6

4.1 Necessary and Sufficient Conditions for Stabilization of LDS

  • Latent stabilization is possible if and only if the state-to-latent map makes the system detectable, meaning every unstable or marginally stable mode is observable through that map.This condition assumes the physical system is stabilizable and observable.Source: §4.1 · p. 6, §4 · p. 6

4.2 The Collapse of Unstable Modes

  • There are systems and data distributions where prediction loss and SIGReg are simultaneously minimized even though the encoder discards the entire unstable subspace.The construction retains stable coordinates while omitting unstable ones; the regularizer can still be minimized when the retained latent distribution is standard Gaussian.Source: §4.2 · p. 6
  • In the constructed case, the standard objective is minimized for every regularization weight, yet the physical system cannot be stabilized from the resulting latent state.The linear analysis shows this failure even with a linear observation-to-latent map and little compression.Source: §4.2 · p. 7
  • In synthetic unstable dynamics and linearized CartPole, SIG learns a misaligned unstable direction and its latent-LQR controller fails to stabilize the physical system.Figure 2 compares learned and true eigenvectors alongside closed-loop trajectories.Source: §4.2 · p. 7, Figure 2 · p. 7

5 Preserving Unstable Directions Through Inverse Dynamics

EP-IDM augments JEPA training with action-sequence reconstruction from endpoint latents. Under stated excitation and exact-reconstruction conditions, this preserves finite-horizon reachable directions; the section also relates long-horizon controllability-Gramian eigenspaces to unstable modes.

5.1 Endpoint Inverse Dynamics

  • EP-IDM reconstructs the action sequence from initial and final latent states, augmenting world-model training with an inverse-dynamics objective.The inverse dynamics model is jointly optimized with the encoders and predictor, using the EP-IDM loss in place of the regularizer.Source: §5 · p. 8, §5.1 · p. 8

5.2 Preservation of Reachable Directions

  • With sufficiently rich action excitation and zero affine inverse-dynamics loss, the encoder is injective on the H-step reachable subspace.Thus, reachable unstable or marginally unstable modes are retained when they lie in that subspace; exact recovery is an idealized regime, while experiments use the loss as a soft inductive bias.Source: §5.2 · p. 8
  • Exact recovery of an arbitrary action sequence is not feasible, but minimizing EP-IDM can encourage preservation of action-induced directions required for control.The exact-reconstruction theorem therefore describes a stronger regime than the soft objective used empirically.Source: §5.2 · pp. 8–9

5.3 Reachability and Unstable Dynamics

  • In the illustrated linear systems, EP-IDM preserves the unstable direction and the learned latent dynamics support stabilization.The evaluation also depicts CartPole, Walker2D, and PointMaze visual-control tasks.Source: §5.2 · p. 9, Figure 3 · p. 10
  • As H grows, the dominant eigenspace of the finite-horizon controllability Gramian converges to the unstable subspace.The relative weighting of reachable directions continues to change with H, and unstable-mode effects grow exponentially; the theorem assumes sufficient growth in unstable controllability energy and excludes unit-circle eigenvalues.Source: §5.3 · p. 9

6 Experiments

The experiments test whether action reconstruction preserves information needed for stabilization and control across CartPole, Walker2D, and PointMaze. Across the unstable tasks, EP-IDM supports latent stabilization and coherent walking, while standard SIGReg models fail under local latent LQR or to preserve the Walker2D gait.

6.1 Setup

  • The experiments cover unstable-equilibrium stabilization, walking-gait preservation, and goal reaching, using latent LQR, CEM, GBP, and Walker2D iCEM.CartPole success is defined by final state norm after 300 steps; Walker2D evaluates gait preservation and forward locomotion, and PointMaze measures final distance to goal.Source: §6 · p. 10, §6.1 · p. 11

6.2 CartPole: Stabilization Around an Unstable Equilibrium

  • CEM and GBP succeed with SIGReg models, although their local latent linearizations do not yield stabilizing LQR controllers.They optimize finite action sequences through nonlinear predictor rollouts, unlike LQR’s reliance on local linearization.Source: §6.2 · p. 11
  • 0% latent-LQR success with SIGReg contrasts with 100% after adding EP-IDM on CartPole.The result is for encoder and predictor models trained from scratch; Figure 4 shows the SIGReg state norm growing without bound while EP-IDM drives it toward equilibrium.Source: Figure 4 · p. 10, §6.2 · pp. 11–12, Table 1 · p. 12

6.3 Walker2D: Preserving a Walking Gait

  • EP-IDM preserves Walker2D’s right-hip limit cycle, while 1SP+SIG and MSP+SIG fail to capture it.Decoded trajectories also show coherent walking with EP-IDM but erratic, physically implausible motion with the SIGReg models.Source: §6.3 · p. 12
  • 2.91 m/s and 9.15 m of forward displacement with 1SP+EP-IDM are closer to ground truth than the SIGReg models’ near-zero displacement.Ground truth achieves 3.61 m/s and 14.43 m; the learned-model comparison uses latent iCEM across ten trials, and Figure 7 shows coherent gait only with EP-IDM.Source: §6.3 · pp. 12–13, Table 2 · p. 13, Figure 7 · p. 13

6.4 Diagnostics

  • 1SP+EP-IDM reproduces CartPole’s local vector field with low three-step prediction error and stabilizes almost every evaluated initial condition.Its latent LQR controller also yields ∆V < 0 around equilibrium; MSP+SIG captures only parts of the vector field and stabilizes a small neighborhood.Source: §6.4 · p. 13, Figure 8 · p. 14, Figure 9 · p. 14
  • On PointMaze, all action-reconstruction variants trained from scratch achieve at least 90% CEM success and at least 70% LQR success.PointMaze tests goal reaching with open-loop stable point-mass dynamics.Source: §6.1 · p. 11, §6.4 · pp. 13–14

7 Conclusion and Future Work

The paper argues that next-step prediction with anti-collapse regularization can discard unstable modes needed for latent feedback stabilization. It proposes EP-IDM and reports theoretical guarantees for linear systems alongside empirical results on nonlinear visual-control tasks.

7 Conclusion and Future Work

  • Next-step prediction with standard anti-collapse regularization can be minimized even when the encoder discards controllable unstable modes.A latent representation can avoid collapse on the training distribution yet remain insufficient for a stabilizing feedback controller.Source: §7 · p. 14
  • EP-IDM reconstructs the action sequence from initial and final latent representations; exact reconstruction preserves the finite-horizon reachable subspace in linear systems.As the horizon grows, the dominant eigenspace of the finite-horizon controllability Gramian converges to the controllable unstable subspace.Source: §7 · pp. 14–15
  • On nonlinear CartPole, EP-IDM raises latent-LQR success from 0% to 100%, preserves the Walker2D limit cycle, and supports 9.15 m average forward displacement over 500 time steps.The paper also reports strong planning performance on PointMaze.Source: §7 · p. 15
  • Future work could combine control-aware representations with latent-geometry objectives, including bisimulation and temporal straightening.The passage presents these as possible directions for efficient and robust planning.Source: §7 · p. 15

A Appendix

The appendix outlines additional experiments, implementation details, ablations, linear examples, and theoretical proofs.

A Appendix

  • Appendix A.1 extends experiments, A.2 tests frozen pretrained encoders, A.3 details implementations, A.4 gives ablations, and A.5–A.6 present linear examples and proofs.Source: A Appendix · p. 19

A.1 Additional Experimental Validation (Extending Section 6)

Additional validation examines representative CartPole and PointMaze trajectories and learned local dynamics, then tests long-horizon prediction along feedback trajectories.

A.1 Additional Experimental Validation (Extending Section 6)

  • The additional experiments relate CartPole and PointMaze closed-loop performance to learned geometry and local dynamics, and assess long-horizon prediction along feedback trajectories.Source: §A.1 · p. 19

A.1.1 CartPole

CartPole evaluations show that EP-IDM models stabilize the pole in sampled LQR rollouts, while several models without EP-IDM fail. Local-dynamics and CEM diagnostics distinguish model accuracy from the ability of planners to compensate for imperfect dynamics.

A.1.1 CartPole

  • EP-IDM models, trained from scratch or with frozen iBOT, keep the pole upright in the sampled LQR rollout, unlike SIGReg-only, frozen-iBOT without EP-IDM, and DINO-WM models.The comparison is for the randomly selected trajectory shown in Fig. 10.Source: §A.1.1 · p. 19
  • The scratch-trained MSP+EP-IDM model recovers local dynamics and a broad attraction region; iBOT with EP-IDM through a learned projector shows similar performance.The diagnostics include low prediction error near equilibrium and the expected decrease in the Lyapunov certificate; iBOT and DINO-WM without EP-IDM have less accurate local dynamics even when terminal latent costs are minimized near the goal.Source: §A.1.1 · p. 19
  • CEM can compensate for imperfect local dynamics through repeated nonlinear replanning, whereas LQR depends directly on the local model around equilibrium.Table 4 reports inverse squared stabilization length and final error as CEM diagnostics.Source: §A.1.1 · p. 19

A.1.2 PointMaze

PointMaze is open-loop stable, so the evaluation tests task-relevant representations and nonlinear planning rather than recovery of an unstable subspace. EP-IDM models and MSP+SIG reach the goal under CEM in the displayed comparison, while the DINO-WM rollout fails.

A.1.2 PointMaze

  • PointMaze is open-loop stable, and the evaluation checks whether IDM retains task-relevant directions without degrading CEM planning.Both EP-IDM models and MSP+SIG reach the goal in the displayed rollout, whereas DINO-WM fails.Source: §A.1.2 · p. 19
  • Scratch-trained models produce smooth latent goal-distance fields that decrease toward the goal and broadly respect maze geometry, giving CEM an informative terminal cost.The DINO-WM field is less aligned with the navigation geometry and can direct planning toward latent-near states without a feasible path.Source: §A.1.2 · pp. 19–20
  • For this open-loop stable task, the representation objective is smooth planning geometry rather than recovery of unstable modes.This task therefore checks whether task-relevant directions are retained without degrading nonlinear planning.Source: §A.1.2 · pp. 19–20

A.1.3 Closed-Loop Prediction Error

Closed-loop rollout comparisons show that prediction-error behavior depends on the training objective: 1SP+EP-IDM stays close to the ground-truth trajectory, while multi-step prediction alone still drifts. Bounded rollout error is complementary to, but does not establish, preservation of unstable directions.

A.1.3 Closed-Loop Prediction Error

  • 1SP+EP-IDM stays close to the ground-truth trajectory over the full horizon, while MSP+EP-IDM error grows steadily and multi-step prediction with EP-IDM and SIGReg keeps error bounded.Multi-step training alone does not prevent rollout drift in an unstable system, where small model errors can be amplified by the dynamics.Source: §A.1.3 · p. 20
  • Figure 21 measures physical-state error between the ground-truth trajectory and decoded latent rollouts over k-step closed-loop prediction under the ground-truth LQR controller.The comparison is over trajectories generated by that controller.Source: Figure 21 · p. 20
  • SIGReg improves representation conditioning, but low rollout error alone does not guarantee preservation of unstable directions required by LQR.The passage describes this as complementary to the stability results.Source: §A.1.3 · p. 21

A.2 Pre-trained encoders

This section examines how representation properties relate to control through experiments and ablations on frozen-backbone performance, alongside proofs concerning reconstruction constraints, latent feedback stabilization, and finite-horizon controllability Gramian eigenspaces.

A.2 Pre-trained encoders

  • With frozen iBOT, MSP+PR-EP-IDM achieves 100% LQR and CEM success, whereas no DINOv2 configuration stabilizes under LQR, showing that inverse dynamics cannot recover information discarded by the backbone.In closed loop, MSP+iBOT+PR-EP-IDM keeps CartPole near the upright equilibrium for the full 300-step rollout; DINOv2-based models cross the stabilization threshold, except MSP, whose state norm remains at most 10.Source: §A.2 · p. 21

A.3.1 Model Architecture

  • The scratch-trained visual encoder uses a ViT-Tiny and a 192-dimensional latent, while frozen DINOv2 or iBOT features are mapped to that latent space by a trainable projector.The latent predictor is an action-conditioned six-block transformer; action decoders are MLPs, and modules are jointly trained with AdamW using prediction, regularization, and action-reconstruction objective coefficients of one.Source: §A.3.1 · p. 22

A.3.2 Task and Dataset Details; A.3.3 Planning and Control Details

  • The experiments specify separate datasets and evaluation protocols for CartPole, Walker2D, and PointMaze, and use terminal latent-space matching for planning with CEM or gradient-based optimization.Latent LQR is obtained by linearizing the learned dynamics at the desired equilibrium and solving the discrete algebraic Riccati equation.Source: §A.3.2 · p. 23, §A.3.3 · p. 23

A.3.5 Additional Action Reconstruction Losses

  • The empirical inverse-dynamics variants differ in decoder inputs: multi-step reconstruction uses the full latent trajectory, while one-step reconstruction predicts each action from adjacent latent states.SIGReg matches random one-dimensional projections of the latent distribution to a standard isotropic Gaussian; the theoretical guarantees use endpoint reconstruction, while multi-step and one-step objectives are empirical comparisons.Source: §A.3.5 · pp. 24–25, §A.3.4 · p. 24

A.4 Ablations

  • EP-IDM and the joint one-step/multi-step objective each achieve perfect CartPole success with all three controllers, and adding SIGReg to multi-step prediction with endpoint reconstruction does not degrade performance.Scratch-trained action-reconstruction models do not benefit from longer action context, while a longer context can improve planning with frozen iBOT encoders.Source: §A.4 · p. 25
  • Removing pixels leaves EP-IDM successful under LQR with a transformer predictor but not an MLP, while removing proprioception from frozen-iBOT models reduces success to zero across planners.The latter result indicates that the image encoder and short temporal context do not reliably recover velocity from pixels alone.Source: §A.4 · pp. 25–26
  • iBOT has a latent-to-physical displacement gain near one, whereas DINOv2 amplifies small perturbations by approximately 22–26 times, making the upright equilibrium difficult to represent as a stationary latent point.The text attributes DINOv2’s sensitivity to fine visual perturbations and describes iBOT as mapping nearby images to nearby latents.Source: §A.4 · pp. 26–27, Table 14 · p. 27, Table 15 · p. 28

A.5 Additional Details on Examples of Section 4

  • For both unstable linear systems, SIGReg-trained controllers have regions with increasing Lyapunov value, whereas EP-IDM yields a decrease throughout the evaluated region except at equilibrium.In the open-loop stable example, the standard SIGReg world-model objective can work, and EP-IDM recovers the dominant stable direction.Source: §A.5 · pp. 29–30

A.6.1 Supporting Lemmas

  • If an affine decoder reconstructs a random vector from a linear map plus an offset almost surely, and its support contains an open set, that linear map has full column rank.Source: §A.6.1 · p. 30

A.6.2 Proof of Lemma 1

  • A latent controller can stabilize all initial states only if (A,F) is detectable; with stabilizable dynamics, detectability also yields an observer-based stabilizing controller.Source: §A.6.2 · p. 31

A.6.3 Proof of Theorem 1

  • With exact action reconstruction on an open conditional support, F is injective on the H-step reachable subspace and any unstable subspace contained within it.The proof also gives the dimension constraint mH ≤ min{d, n}.Source: §A.6.3 · p. 31

A.6.4 Proof of Theorem 2

  • Whenever γβ < α², the Gramian’s unstable block grows exponentially, its stable block remains bounded, and its cross block grows strictly more slowly than the unstable block.Source: §A.6.4 · p. 33

The spectral-norm distance between the two subspaces is then

  • The top r-dimensional eigenspace converges to the unstable subspace as the spectral-norm distance between the compared subspaces tends to zero with H.The perturbation-to-eigengap ratio vanishes for sufficiently large H.Source: The spectral-norm distance between the… · pp. 33–34
Loading 2610.07540v1…