PaperlayerView exact revision on arXiv

Source-linked AI summary

Hamiltonian JEPA: Action-Conditioned World Models with an Inherited Control State

Tamim Zoabi, Ameen Ali, Lior Wolf

arXiv:2609.33497v1cs.LG

TL;DR

Planning from pixels requires a representation that is not only stable but also organized around how actions move the system. H-JEPA separates a perceptual code from an inherited-covariance control state and uses structured latent dynamics with tied port-inverse consistency. It matches or exceeds evaluated baselines across four benchmarks, with its largest reported gain on OGB-Cube.

  • Problem

    Action-conditioned JEPAs use one embedding for perception and control, while stable marginal embeddings alone do not specify which directions actions move or which coordinates should guide planning.

    Source: §1 · p. 1

  • Method

    H-JEPA regularizes a wide perceptual code toward isotropy, projects it through a fixed orthonormal slice into a control state, and uses phase-conditioned port-Hamiltonian dynamics with tied port-inverse consistency.

    Source: §1 · p. 2

  • Results

    H-JEPA matches or exceeds reconstruction, full-space-prior, and action-decoding baselines on four pixel-based control benchmarks, with its largest gain on OGB-Cube (91.9 percent against 79.3 for the strongest baseline).

    Source: §1 · p. 2

  • Takeaways & Limitations

    Ablations support structured dynamics and, on PushT, regularizing perception and tying the inverse readout to the action port as components of the approach.

    Source: §6 · p. 10

    The experiments are confined to simulation; physical deployment would require application-specific validation of safety, robustness, and failure behavior.

    Source: ETHICS STATEMENT · p. 10

Abstract

from arXiv · show

Planning from pixels needs more than a latent space that is stable and predictable. The state the planner scores must also be organized by how actions move the system. Joint-embedding predictive architectures (JEPAs) avoid pixel reconstruction by predicting future representations, but existing action-conditioned JEPAs ask one embedding to serve both perception and control. We introduce H-JEPA, which separates the two. A wide perceptual code is regularized toward a well-scaled isotropic geometry with a Bures-Wasserstein prior, and a fixed orthonormal slice of that code is the control state, which inherits the code's covariance without any objective of its own. The state evolves under phase-conditioned dissipative port-Hamiltonian dynamics whose input port has orthonormal columns. Port-inverse consistency (PIC) reads the executed action back through the transpose of that port. We show that this readout is exactly the rollout error projected onto the port directions, so PIC is a parameter-free reweighting of prediction error and not an auxiliary action decoder. Untying the readout from the port breaks this identity and loses half of the gain. H-JEPA matches or exceeds reconstruction-free baselines, including the action-decoding Delta-JEPA, on four pixel-based control benchmarks after at most $10$ training epochs, and its largest gain is on OGB-Cube ($91.9$ against $79.3$ percent). Ablations on PushT and OGB-Cube separate the contributions of the structured predictor, PIC, the prediction horizon, the state rank, and the anti-collapse prior.

1 INTRODUCTION

Planning from pixels requires action-sensitive predictions: a stable embedding alone does not ensure that candidate action sequences are ranked differently. H-JEPA separates perception from the planner-facing control state and reports performance at least matching evaluated baselines across four benchmarks.

1 INTRODUCTION

  • Action sensitivity is central to planning because a model that blurs differences between action sequences can rank candidate plans alike.Reconstruction-based models may spend capacity on visual detail irrelevant to control, while JEPAs predict future representations instead of pixels.Source: §1 · p. 1
  • Stable marginal embeddings do not specify which directions actions move, which coordinates should define the planning metric, or how much capacity control needs.The cited discussion also identifies shared-coordinate regularization and prediction losses, latent-MSE scale degeneracy, and Sub-JEPA’s regularizer subspaces as limitations of prior designs.Source: §1 · pp. 1–2
  • H-JEPA uses a wide, isotropy-regularized perceptual code and a fixed orthonormal slice as the control state, then applies phase-conditioned dissipative port-Hamiltonian dynamics.Port-inverse consistency reads actions through the port transpose; its residual equals rollout error projected onto port directions, reweighting prediction error without a learned decoder.Source: §1 · p. 2
  • H-JEPA matches or exceeds reconstruction, full-space-prior, and action-decoding baselines on four pixel-based control benchmarks, with its largest gain on OGB-Cube (91.9 percent against 79.3 for the strongest baseline).The benchmarks are Two-Room, Reacher, PushT, and OGB-Cube, evaluated under the same planning protocol.Source: §1 · p. 2

2 RELATED WORK

Related work spans reconstruction-based latent dynamics, representation-predictive JEPAs, and methods that use inverse dynamics or structured mechanics for control. H-JEPA differs in how its planner state and action readout are tied to the learned model.

2 RELATED WORK

  • LeWM applies a Gaussian regularizer to action-conditioned JEPA embeddings, Sub-JEPA applies it in benchmark-selected random subspaces, and Delta-JEPA adds a learned decoder on latent differences.H-JEPA instead uses a single fixed subspace as the control state, retains a full-space prior whose geometry reaches that state by inheritance, and uses the forward input port for inverse readout.Source: §2 · p. 2
  • Reconstruction-based latent dynamics models often predict pixels and rewards, whereas H-JEPA trains an end-to-end pixel JEPA without pixels or rewards as targets, teachers, or stop-gradients.Its collapse defense replaces the Euclidean covariance metric in Barlow Twins- and VICReg-family moment matching with Bures-Wasserstein distance on Gaussians.Source: §2 · p. 2
  • Prior structured-dynamics work includes energy-gradient models for specified physical states or reconstruction; H-JEPA uses port-Hamiltonian structure as an inductive bias for reconstruction-free planning representations.The orthonormal port removes a scale freedom that would otherwise interact with the representation objective.Source: §2 · p. 3
  • Inverse-dynamics approaches recover control-relevant features, while H-JEPA ties action readout to the forward input map’s transpose, which is an exact pseudoinverse when the port is orthonormal.The related work also notes stated assumptions for multi-step inverse prediction and known limits on when it suffices to identify a control-endogenous state.Source: §2 · p. 3

3 METHOD

H-JEPA separates perception from control: a distributional prior shapes a wide perceptual code, while a fixed slice supplies the planner’s state. Phase-conditioned port-Hamiltonian dynamics predict that state, and port-inverse consistency (PIC) reweights rollout error along the action-port directions.

3 METHOD

  • The planner rolls out candidate action sequences from offline RGB trajectories and ranks them by distance to an encoded goal state, without rewards or simulator state.The fixed slice is the only quantity the planner scores.Source: §3 · p. 3

3.1 PERCEPTION ON A SPHERE

  • Per-frame codes are centered and fixed to a common radius, removing scale freedom that could change transition magnitudes without improving geometry.Actions are averaged over the frameskip window; the model uses no learned action embedding.Source: §3.1 · p. 4

3.2 A CONTROL STATE WITH INHERITED COVARIANCE

  • A fixed orthonormal projection defines the control state, which inherits the code’s covariance without a separate state-space objective or normalization.At full rank, the target covariance has one exact null direction; for lower rank, one direction receives shrinkage tied to the relative state dimension.Source: §3.2 · p. 4

3.3 A BURES-WASSERSTEIN ISOTROPY PRIOR

  • The Bures-Wasserstein prior matches the code’s first two moments to the most isotropic covariance available on the centered hyperplane.Its penalty maintains a constant pull on a direction as its standard deviation approaches zero, unlike the fading pull described for Euclidean and SIGReg penalties.Source: §3.3 · pp. 4–5

3.4 PHASE-CONDITIONED PORT-HAMILTONIAN DYNAMICS

  • Phase-conditioned port-Hamiltonian dynamics combine energy-driven motion, positive-semidefinite dissipation, and action input through an orthonormal port.The continuous-time field is discretized with one explicit Euler step; its energy properties need not hold exactly after discretization.Source: §3.4 · p. 5

3.5 TRAINING OBJECTIVE

  • Training matches multi-step predicted states to encoded futures over the same rollout horizon used by the planner.The horizon is five steps, and predicted states feed back into subsequent rollout steps.Source: §3.5 · p. 6
  • PIC reads actions through the transpose of the same port that injects them, making it a parameter-free reweighting of rollout error along port directions rather than an action-decoding task.The combined loss weights error along those directions more heavily than error in other directions.Source: §3.5 · pp. 6–7

4 EXPERIMENTS

H-JEPA is best or tied on all four planning benchmarks, and ablations examine its predictor, action readout, prior, horizon, and state rank. The experiments also probe what its learned state represents and how its dynamics behave.

4 EXPERIMENTS

  • H-JEPA is best or tied on every task, leads by 4.8 points on Reacher and 12.7 points on OGB-Cube, and trains for 10 epochs versus Delta-JEPA’s 50.Two-Room performance saturates together with Delta-JEPA, while PushT is within one standard deviation of it; Sub-JEPA’s reported results use a different evaluation protocol.Source: §4 · p. 7
  • Physical probes show task-dependent state and flow patterns: H-JEPA is strongest on PushT agent location, while control flow is 0.13 times passive flow on PushT and 1.53 times on OGB-Cube.H-JEPA is stronger on OGB-Cube joint velocity and end-effector position; planning quality is not monotone in object-pose decodability.Source: §4 · pp. 7–8

4.1 ABLATIONS AND ANALYSIS

  • 94 percent success with the selected r=192 model compares with 20 percent for LeWM’s transformer and 84 percent for the same transformer on H-JEPA’s representation.At r=64, the port-Hamiltonian predictor with PIC reaches 92 percent while reducing the dynamics module from 3.3M to 0.86M parameters.Source: §4.1 · p. 8
  • Replacing the port-Hamiltonian step with an MLP costs 8.3 PushT points and 7.6 OGB-Cube points when PIC is retained.Removing PIC costs 18.3 points on PushT and 2.9 on OGB-Cube; the latter difference is within evaluation noise at N=300.Source: §4.1 · p. 8, Table 4 · p. 8
  • 93.3 percent with the tied readout compares with 83.7 percent for an untied fixed-map readout and 75.0 percent without PIC.The untied readout breaks the identity that makes PIC a pure reweighting of rollout error, accounting for the other half of the gain.Source: §4.1 · p. 9
  • Moving the Bures-Wasserstein prior from code h to scored state s lowers PushT success from 93.3 to 86.3, while SIGReg on h reaches 36.7 percent.The SIGReg result applies inside H-JEPA, where the prior alone must preserve code directions outside the control-state slice.Source: §4.1 · p. 9
  • On PushT, matching prediction and PIC horizons at three or five steps reaches 94 percent, versus 12 or 34 percent with one-step horizons.Keeping Kp=5 but reducing PIC to one step reaches 90 percent; the authors set both horizons to the planning horizon.Source: Table 5 · p. 8, §4.1 · p. 9
  • The useful state rank is task dependent: Two-Room saturates across tested ranks, Reacher and OGB-Cube degrade at full rank, and PushT improves up to r=D.On PushT, the candidate-ranking remainder beyond probed physical variables has Spearman correlation 0.997 and recovers 89 percent of CEM elites, versus 0.66 and 38 percent for the physical part.Source: Table 6 · p. 8, §4.1 · p. 9

5 LIMITATIONS

H-JEPA’s evidence is limited to offline pixel control in simulation with short-horizon CEM planning; longer horizons, real observations, and online data collection remain untested.

5 LIMITATIONS

  • The evidence covers offline pixel control in simulation with short-horizon CEM; longer horizons, real observations, and online data collection remain untested.The orthonormal port removes scale ambiguity but gives up state-dependent action gain, and state rank is selected per task rather than learned.Source: §5 · p. 9
  • H-JEPA uses mechanics as a structural bias, not as a claim of latent system identification.The learned coordinates may rotate around the planning objective, and component-level physical alignment is task dependent.Source: §5 · p. 9
  • At r=D, the state has one null direction, and inherited covariance is only as isotropic as the trained code.Measured covariance spectra remain spread around one.Source: §5 · p. 9

6 CONCLUSION

H-JEPA separates perceptual representation learning from the planner-facing state and combines covariance inheritance with structured, action-sensitive prediction. It matches or exceeds evaluated baselines across four pixel-based benchmarks, while the work remains confined to simulated control.

6 CONCLUSION

  • H-JEPA separates perceptual representation learning from the state used for control, without reconstructing pixels.A Bures-Wasserstein prior regularizes the perceptual code, and a fixed orthonormal projection transfers its covariance geometry to a task-dependent control state.Source: §6 · p. 9
  • A phase-conditioned port-Hamiltonian predictor structures latent dynamics, while tied PIC reweights rollout error along action-input directions without an additional decoder.Source: §6 · p. 9
  • Ablations support structured dynamics and, on PushT, regularizing perception and tying the inverse readout to the action port; H-JEPA matches or exceeds baselines across four benchmarks.The authors describe covariance inheritance with structured, action-sensitive prediction as a promising combination for effective planning.Source: §6 · p. 10
  • The experiments are confined to standard simulated control benchmarks; physical deployment would require application-specific safety, robustness, and failure-behavior validation.Source: ETHICS STATEMENT · p. 10

ACKNOWLEGMENTS

The work acknowledges institutional and foundation support, including support for Tamim Zoabi’s PhD thesis research at Tel Aviv University.

ACKNOWLEGMENTS

  • Support came from Tel Aviv University, the Blavatnik Family Foundation, Israel’s Ministry of Innovation, Science & Technology, and the Michael J. Fox Foundation.Tamim Zoabi’s contribution formed part of PhD thesis research at Tel Aviv University.Source: ACKNOWLEGMENTS · p. 10

A ARCHITECTURE OVERVIEW

H-JEPA separates perception from planning: a prior shapes the spherical code, while a fixed orthonormal slice supplies the planner-facing state. The predictor is trained open-loop, with PIC reading actions through the port transpose.

A ARCHITECTURE OVERVIEW

  • The Bures-Wasserstein prior shapes the spherical code, while a fixed orthonormal slice provides the control state scored by the planner.The predictor is supervised open-loop, and PIC reads actions back through the port transpose.Source: Figure 1 · p. 14, A ARCHITECTURE OVERVIEW · p. 14

B DERIVATION OF THE PIC IDENTITY

PIC’s action readout is tied to the predictor’s input port, making its residual a projection of open-loop rollout error. Untying the readout adds a port-mismatch term and reduces the reported gain.

B DERIVATION OF THE PIC IDENTITY

  • The PIC residual equals the open-loop state error projected onto the port subspace and scaled by 1/Δt.Summed across steps, PIC weights error along the d_a port directions by 1 + r/(d_aΔt^2) relative to the other r − d_a directions.Source: B DERIVATION OF THE PIC IDENTITY · p. 15
  • On exact transitions the transpose readout returns the executed action; otherwise it projects residual displacement without amplifying noise.Its magnitude is bounded by transition energy, and LPIC is order one at initialization in normalized units.Source: B DERIVATION OF THE PIC IDENTITY · p. 15
  • The untied readout adds a term that vanishes only when the learned port equals the fixed frame at every visited phase.The added term pulls the port toward a fixed frame while prediction asks it to follow the state; the tied readout leaves this term identically zero.Source: B DERIVATION OF THE PIC IDENTITY · p. 15
  • The untied control recovers 75.0 to 83.7, versus 93.3 with the tied readout.Weighting error along action-entry directions is qualitatively related to controllability, while predictive accuracy plays an analogous role for observability.Source: B DERIVATION OF THE PIC IDENTITY · p. 15

C DEFAULT CONFIGURATION AND TRAINING DETAILS

The appendix specifies default configuration and training details, then describes physical-state probing on three tasks. Probe results are reported as mean ± standard deviation across three MLP-probe seeds.

C DEFAULT CONFIGURATION AND TRAINING DETAILS

D FULL PHYSICAL PROBING TABLES

  • Physical probes use 20,000 frames, trajectory-level train/test splits, train-split z-scoring, and OLS and 2×256 ReLU MLP probes.The MLP probe runs for 100 epochs with three seeds; angular and quaternion targets are evaluated in raw coordinates, so wrap-around and sign ambiguity depress correlations for all methods alike.Source: D FULL PHYSICAL PROBING TABLES · p. 16
  • The selected H-JEPA checkpoints probe the control state at task-specific ranks: 192 for PushT, 64 for Reacher, and 32 for OGB-Cube.LeWM and Sub-JEPA use their official projected 192-dimensional embeddings; Tables 8–10 report the three task results.Source: D FULL PHYSICAL PROBING TABLES · p. 16, Table 8 · p. 16, Table 9 · p. 16

E INTERPRETING THE PORT-HAMILTONIAN PREDICTOR

The appendix examines how the structured predictor is used across tasks, whether its fields align with physical changes, how history is encoded, and how the control state inherits covariance. It also contrasts anti-collapse penalties near vanishing variance.

E INTERPRETING THE PORT-HAMILTONIAN PREDICTOR

  • PushT is dominated by passive dissipative flow, while OGB-Cube is control-dominated and has greater overlap between actuated and dissipative subspaces.On PushT, dissipation increases with latent speed (correlation 0.48); these measurements are structural signatures in learned coordinates, not physical coefficients in SI units.Source: E INTERPRETING THE PORT-HAMILTONIAN… · p. 17, Table 11 · p. 17
  • On PushT, the passive field, control field, and Hamiltonian gradient each explain about 88–90% of physical-increment variance, while the control state predicts physical state with R2 = 0.95.The analysis fits linear ridge probes from each field to physical state increments.Source: E INTERPRETING THE PORT-HAMILTONIAN… · p. 18, Table 12 · p. 18
  • The structured predictor is an inductive bias for organizing latent dynamics; physical interpretability emerges when the task supports it, but is not required for planning.Reacher and OGB-Cube retain control-relevant physical information in the planning state even when predictor components are not linearly aligned with privileged coordinates.Source: E INTERPRETING THE PORT-HAMILTONIAN… · p. 18

F HISTORY MECHANISM ABLATION

  • Replacing the learned history mixer with the parameter-free phase c = [s, v] removes a mixer that largely ignored history while leaving the planner’s metric untouched.About 90% of the learned mixer’s Frobenius energy concentrated on the latest state, and zeroing difference channels barely changed one-step predictions (cosine similarity about 0.99).Source: F HISTORY MECHANISM ABLATION · p. 18

G GEOMETRY OF THE INHERITED CONTROL STATE

H ANTI-COLLAPSE PENALTIES NEAR VANISHING VARIANCE

  • Near vanishing variance, Bures-Wasserstein maintains a constant force on standard deviation, whereas penalties with bounded eigenvalue derivatives exert a force that vanishes linearly.The figure compares slopes, not absolute vertical offsets, which reflect each statistic’s native normalization.Source: H ANTI-COLLAPSE PENALTIES NEAR… · p. 19, Figure 3 · p. 20

I TRAINING BEHAVIOR

H-JEPA’s PushT losses improve across the 10-epoch run, while OGB-Cube reaches its useful regime within three epochs. Training and validation curves remain closely matched, and the three objectives can be optimized jointly without warm-up schedules or alternating optimization.

Training behavior

  • H-JEPA models train end-to-end for 10 epochs, with evaluation checkpoints fixed per task after a held-out pilot seed.The chosen checkpoints were epochs 8, 6, 10, and 3 for Two-Room, Reacher, PushT, and OGB-Cube, respectively; Delta-JEPA trains for 50 epochs under its protocol.Source: I TRAINING BEHAVIOR · p. 20
  • PushT improves steadily across all 10 epochs, with closely tracking training and validation curves.Figure 4 plots open-loop latent prediction, PIC, and Bures-Wasserstein isotropy losses.Source: I TRAINING BEHAVIOR · p. 20, Figure 4 · p. 21
  • All three objectives can be optimized jointly without loss-specific warm-up schedules or alternating optimization.Source: I TRAINING BEHAVIOR · p. 20
  • OGB-Cube reaches its useful regime within three epochs, although representation-level losses continue to decrease afterward.The benchmark uses the epoch-3 checkpoint; the extended trace indicates that lower post-convergence losses are not required for the reported planning result.Source: I TRAINING BEHAVIOR · p. 20, Figure 5 · p. 21

Related papers

Loading 2609.33497v1…