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

Tamim Zoabi; Ameen Ali; Lior Wolf

cs.LGSep 27, 2026

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

Authors: Tamim Zoabi, Ameen Ali, Lior Wolf

Organizations: The Blavatnik School of Computer Science, Tel Aviv University

Abstract

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.

Figures & tables

Relaxed choice	Observed failure	Adopted constraint
BatchNorm in the projector, or a standardizer on $s$	Hides collapse: $s$ keeps unit marginals while $h$ degenerates to rank one	Per-sample affine-free $\mathrm{LN}_{0}$ , no normalizer on $s$ (§ 3.1 , § 3.2 )
LSUV-style whitening of the projector activations at initialization ( Mishkin and Matas, 2016 )	Amplifies projector weights about $10^{4}$ -fold, point-mass collapse within tens of steps	Isotropy as a loss with bounded per-step influence (§ 3.3 )
Learnable projection $U$	Co-adapts with the losses and collapses the rank of the state	Fixed seeded orthonormal $U$ (§ 3.2 )
Learned action embedding	The PIC regression target collapses to effective rank about $1.3$	Raw normalized actions (§ 3.1 )
SIGReg, or the Euclidean penalty $\\|\widehat{\Sigma}-I\\|_{F}^{2}$	The restoring force fades as variance shrinks, collapse within about $50$ to $100$ steps	Bures-Wasserstein prior (§ 3.3 )
Unconstrained port gain	The gain trades off against the encoder scale, so optimization moves scale and not content	Orthonormal port via differentiable QR (§ 3.4 )

Table 1: Design choices we relaxed during development, the failure each one produced in this architecture, and the constraint adopted instead. These observations motivate the constraints of § 3 and complement the controlled ablations of § 4.1 .

Method	Two-Room	Reacher	PushT	OGB-Cube
PLDM	$93.73\pm 1.03$	$64.33\pm 2.14$	$76.13\pm 1.70$	$57.27\pm 1.53$
LeWM	$74.93\pm 0.42$	$79.87\pm 0.90$	$84.53\pm 1.50$	$64.13\pm 1.89$
Sub-JEPA	$90.60\pm 0.53$	$81.00\pm 2.40$	$63.73\pm 0.12$	$62.67\pm 1.45$
Delta-JEPA	$\mathbf{100.00\pm 0.00}$	$81.33\pm 0.50$	$89.07\pm 1.90$	$79.27\pm 1.81$
H-JEPA (ours)	$\mathbf{100.00\pm 0.00}$	$\mathbf{86.13\pm 0.23}$	$\mathbf{90.40\pm 0.60}$	$\mathbf{91.93\pm 1.30}$

Table 2: Planning success rate (%) on the four LeWM benchmarks over

500

test episodes per task, mean

\pm

std over three training seeds. Baselines are evaluated under the Delta-JEPA protocol, in which Delta-JEPA trains for

50

epochs. H-JEPA trains for

10

epochs and is evaluated at the per-task checkpoint epoch selected on a held-out seed (§ 4 ).

Predictor	Params total / dyn.	Success	Representation
LeWM, native embedding	$18.0$ M / $10.8$ M	$20\%$	SIGReg
AdaLN transformer on H-JEPA, $r{=}64$	$9.6$ M / $3.3$ M	$84\%$	H-JEPA
Port-Hamiltonian $+$ PIC, $r{=}64$	$7.2$ M / $0.86$ M	$92\%$	H-JEPA
Port-Hamiltonian $+$ PIC, $r{=}192$	$11.6$ M / $5.4$ M	$\mathbf{94\%}$	H-JEPA

Table 3: Predictor class on PushT at a matched representation and horizon (

K_{p}{=}5

in all rows). Parameters are total model followed by dynamics module. H-JEPA rows share the ViT-T/14 encoder and

0.79

M projector.

Table 4

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 1: H-JEPA. The Bures-Wasserstein prior $\mathcal{L}_{\mathrm{iso}}$ acts only on the spherical code $h_{t}$ . A fixed orthonormal slice $U^{\top}$ yields the control state $s_{t}$ . The phase-conditioned port-Hamiltonian predictor $\mathcal{F}_{\theta}$ is supervised open-loop ( $\mathcal{L}_{\mathrm{pred}}$ ), PIC reads actions back through the port transpose ( $\mathcal{L}_{\mathrm{PIC}}$ ), and planning scores $s$ alone.

Quantity	Value
Perceptual dimension $D$	$192$
Encoder	ViT-T/14 from scratch, MLP projector $192\to 2048\to 192$
Output normalization	affine-free LayerNorm, $\\|h\\|=\sqrt{D}$ , per sample
Projection $U$	fixed random orthonormal, seeded QR
History and phase	$2$ frames, $c=[s,v]$ with $v=s-s_{\mathrm{prev}}$ , left-pad $v_{0}=0$
Net action dimension $d_{a}$	$2$ on Two-Room, Reacher, PushT, $5$ on OGB-Cube

Appendix

Table 7: Default H-JEPA configuration. Task-specific entries are shown explicitly for rank, action dimension, and checkpoint epoch.

Property	Method	Linear MSE $\downarrow$	Linear $\rho$ $\uparrow$	MLP MSE $\downarrow$	MLP $\rho$ $\uparrow$
Agent location	H-JEPA	$0.028\pm 0.003$	$0.986$	$0.008\pm 0.003$	$0.996$
	LeWM	$0.062\pm 0.012$	$0.968$	$0.017\pm 0.006$	$0.992$
	Sub-JEPA	$0.065\pm 0.009$	$0.966$	$0.016\pm 0.004$	$0.992$
Block location	H-JEPA	$0.152\pm 0.025$	$0.926$	$0.049\pm 0.024$	$0.979$
	LeWM	$0.038\pm 0.011$	$0.982$	$0.024\pm 0.014$	$0.990$
	Sub-JEPA	$0.033\pm 0.009$	$0.984$	$0.017\pm 0.007$	$0.992$

Appendix

Table 8: Physical latent probing on PushT (

\rho

is the Pearson correlation).

Property	Method	Linear MSE $\downarrow$	Linear $\rho$ $\uparrow$	MLP MSE $\downarrow$	MLP $\rho$ $\uparrow$
Finger position	H-JEPA	$0.002\pm 0.000$	$0.999$	$0.000\pm 0.000$	$1.000$
	LeWM	$0.002\pm 0.000$	$0.999$	$0.000\pm 0.000$	$1.000$
	Sub-JEPA	$0.002\pm 0.000$	$0.999$	$0.000\pm 0.000$	$1.000$
Joint position	H-JEPA	$0.601\pm 0.125$	$0.588$	$0.742\pm 0.111$	$0.562$
	LeWM	$0.636\pm 0.120$	$0.573$	$0.738\pm 0.116$	$0.550$
	Sub-JEPA	$0.639\pm 0.125$	$0.575$	$0.726\pm 0.129$	$0.559$

Appendix

Table 9: Physical latent probing on Reacher (

\rho

is the Pearson correlation).

Property	Method	Linear MSE $\downarrow$	Linear $\rho$ $\uparrow$	MLP MSE $\downarrow$	MLP $\rho$ $\uparrow$
Joint position	H-JEPA	$0.445\pm 0.044$	$0.670$	$0.554\pm 0.056$	$0.685$
	LeWM	$0.579\pm 0.060$	$0.645$	$0.775\pm 0.109$	$0.647$
	Sub-JEPA	$0.492\pm 0.067$	$0.676$	$0.557\pm 0.059$	$0.675$
Joint velocity	H-JEPA	$0.740\pm 0.027$	$0.432$	$0.713\pm 0.039$	$0.541$
	LeWM	$1.069\pm 0.038$	$0.088$	$2.674\pm 0.236$	$0.050$
	Sub-JEPA	$0.988\pm 0.032$	$0.203$	$1.595\pm 0.226$	$0.101$

Appendix

Table 10: Physical latent probing on OGB-Cube (

\rho

is the Pearson correlation).

Task	$\dfrac{\\|R\nabla H\\|}{\\|J\nabla H\\|}$	$\dfrac{\\|G\bar{a}\\|}{\\|\mathrm{passive}\\|}$	$\dfrac{\operatorname{tr}(G^{\top}RG)}{\operatorname{tr}(R)}$	$\operatorname{corr}(\operatorname{tr}R,\\|v\\|)$
PushT	4.35	0.13	0.001	0.48
Reacher	1.16	0.56	0.008	0.24
OGB-Cube	1.11	1.53	0.184	0.13

Appendix

Table 11: Learned port-Hamiltonian structure on held-out transitions. The relative contribution of dissipation, passive dynamics, and direct control changes markedly across tasks.

Task	Passive $\rightarrow\Delta y$	$G\bar{a}\rightarrow\Delta y$	$\nabla H\rightarrow\Delta y$	$s\rightarrow y$
PushT	0.881	0.875	0.899	0.953
Reacher	0.351	0.033	0.291	0.514
OGB-Cube	0.268	0.323	0.312	0.561

Appendix

Table 12: Linear physical alignment of the learned predictor.

R^{2}

is measured on held-out transitions using ridge regression.

Model	$\\|U^{\top}e\\|_{2}^{2}$	$r/D$	$1-\\|U^{\top}e\\|_{2}^{2}$	min	median	max	MP interval
			exceptional	measured eigenvalues of $\mathrm{Cov}(s)$
PushT, $r=64$	$0.356$	$0.333$	$0.644$	$0.236$	$0.713$	$2.48$	$[0.42,\,1.83]$
PushT, $r=192$	$1.000$	$1.000$	$0.000$	$1.3\times 10^{-6}$	$0.877$	$3.31$	$[0.15,\,2.60]$
OGB-Cube, $r=32$	$0.116$	$0.167$	$0.884$	$0.015$	$0.625$	$1.46$	$[0.56,\,1.56]$

Appendix

Table 13: Measured geometry of the control state on

n=512

encoded frames from the selected checkpoints. The Marchenko-Pastur (MP) interval is the range of sample eigenvalues that an exactly isotropic state would show at this

n

and

r

( Marčenko and Pastur, 1967 ) .

Figure 2: Eigenvalues of

\mathrm{Cov}(s)

for the selected checkpoints. The dashed line marks the exceptional eigenvalue

1-\|U^{\top}e\|_{2}^{2}

of the target covariance and the dotted line marks one. At

r=D=192

the predicted null direction appears as the last eigenvalue.

Figure 3: Eigenvalue derivatives of the three penalties on synthetic data. The plot shows

|\partial\ell/\partial\lambda|

against

\lambda

. Bures-Wasserstein rises as

\lambda^{-1/2}

when

\lambda\to 0

, while the Euclidean penalty and SIGReg flatten to constants. By equation 20 , a flat curve means a force on the standard deviation that vanishes like

\sqrt{\lambda}

, and the

\lambda^{-1/2}

rise means a constant force. Vertical offsets reflect the native normalization of each statistic and are not comparable across curves. The comparison is between slopes.

Figure 4: H-JEPA training on PushT. Training and validation losses across the

10

-epoch run for (a) open-loop latent prediction, (b) port-inverse consistency (PIC), and (c) the Bures-Wasserstein isotropy objective. All objectives decrease smoothly, with closely tracking training and validation curves.

Figure 5: H-JEPA training on OGB-Cube. Training and validation losses for (a) open-loop latent prediction, (b) PIC, and (c) isotropy. All three objectives fall rapidly during the first epochs and subsequently improve smoothly, with little separation between training and validation curves. The benchmark model uses the epoch-

3

OGB-Cube checkpoint selected in § 4 . The extended trace shows that lower representation-level losses after convergence are not required for the reported planning result.

Explore similar work

Aug 13, 2026cs.LG

Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency

Joint-embedding predictive architectures (JEPAs) learn world models that predict in a compact latent space rather than in pixels, reducing the pressure to model nuisance appearance. Yet this provides no guarantee against visual perturbations: they can still alter the encoded representation and affect subsequent action-conditioned predictions. Bisimulation captures this requirement precisely: two observations should be treated as the same state only when their action-conditioned consequences agree. Guided by this criterion, we introduce Action-Conditioned Predictive Consistency (ACPC), a diagnostic that measures how far a clean history and a visually perturbed view of it diverge after being rolled forward under the same action sequence. We prove that this divergence bounds the perturbation-induced change in multi-step prediction error and planner cost. Building on pairwise ACPC, we define two complementary measures: the Invariance Radius (IR) summarizes clean-perturbed rollout spread, while the Separation Rate (SR) checks whether different states remain distinguishable after rollout. Experiments on four visual control tasks show that pairwise ACPC predicts perturbation-induced prediction and cost changes. On LeWM, the IR-SR screen transfers across tasks, and the joint diagnostic remains informative under blur and resize. PLDM exhibits similar diagnostic trends under a different architecture.

Guo An, Zijing Wu, Honghua Dong +7

Huawei · University of Science and Technology of China · Zhejiang University +3

Action Prediction World Models World

Ask PDF

Aug 20, 2026cs.LG

Orthogonal JEPA: Factorized Predictive States for Latent World Models

World models construct latent states that support prediction, planning, and reasoning about an underlying system. Joint-embedding predictive architectures (JEPAs) offer a direct way to learn such states by predicting targets in representation space instead of reconstructing every detail of the observation. Standard JEPAs, however, organize all predictable content through one target embedding and one prediction pathway. In complex systems, this monolithic state can allocate redundant capacity to dominant signals while providing weak or conflicting gradients to less dominant predictive structure. We introduce \method, a latent world-modeling framework based on orthogonal predictive factorization. Learned basis matrices analyze each target state into multiple components, and a dedicated prediction branch estimates each component from a shared context representation. Predictive regression preserves the factor magnitudes required for state synthesis, an orthogonality objective discourages repeated directions, factor-activity regularization maintains variation in projected targets, and online variance regularization discourages coordinate-wise encoder collapse. Predicted components are synthesized into a complete latent state that can be used by a readout, decoder, planner, or autoregressive rollout. The same predictive-state mechanism applies when the target is temporally future, spatially hidden, or another partial observation of the same system. Experiments on controlled vision, single-cell transcriptomics, longitudinal health records, continuous control, and molecular dynamics evaluate representation quality, forecasting, planning, and long-horizon stability.

Taoyong Cui, Pheng Ann Heng, Wanli Ouyang

The Chinese University of Hong Kong (CUHK)

Joint-Embedding Predictive Architectures Next-State Prediction Latent World Models World Models Regularization Single-Cell Rna

Ask PDF

Aug 6, 2026cs.CV

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

We propose PhyLatent, a dynamics-relevant training objective for Joint-Embedding Predictive Architecture (JEPA) world models. Our key observation is that preventing global latent collapse does not necessarily ensure that the learned representation preserves physically meaningful state and action relationships. We identify three failure modes: sensitivity to appearance changes that leave the physical state unchanged, insufficient separation of distinct physical states, and insufficient separation of different action-conditioned futures. We refer to these as Physical Invariance Collapse, Physical Distinguishability Collapse, and Counterfactual Dynamics Collapse, respectively. PhyLatent targets these failures through three coordinated training pathways, implemented with static visual invariance, physical state grounding, future representation alignment, counterfactual branch separation, and latent denoising. On OGBench-Cube, PhyLatent reduces the three collapse rates by 43.9%, 27.9%, and 47.1%, respectively, while improving model predictive control (MPC) success by 12.0 percentage points (17.2% relative). Across four visual-control tasks, average planning success increases by 6.62 percentage points (8.3% relative). These results show that global non-collapse alone is insufficient for learning a reliable JEPA world-model state space, and that explicitly preserving dynamics-relevant structure can improve closed-loop planning.

Xi Zeng, Haojie Ren, Ziying Song +2

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

Latent World Models Joint-Embedding Predictive Architectures World Models Inference Dynamics Latent Variable

Ask PDF