LeHami: Symplectic Latent Dynamics for World Models

Romina Yalovetzky; Randall Balestriero; Yann LeCun

cs.LGOct 11, 2026

LeHami: Symplectic Latent Dynamics for World Models

Authors: Romina Yalovetzky, Randall Balestriero, Yann LeCun

Organizations: Courant Institute, New York University · Brown University · AMI Labs

Abstract

Joint-embedding predictive architectures (JEPAs) learn world models directly in latent space, but unconstrained autoregressive predictors may contract the latent distribution over long rollouts. We introduce Symplectic Hamiltonian JEPA (LeHami), which replaces the standard predictor with an action-conditioned Hamiltonian flow over a structured latent state (q, p) integrated using an exactly symplectic leapfrog step. This construction preserves phase-space volume at every rollout step and rules out uniform contraction, asymptotically stable fixed points, and merging of distinct latent states, independently of the learned parameters or action sequence. We obtain the momentum-like component from temporal derivatives of the learned position-like representation, providing an explicit second-order state for prediction and control. On conservative and actuated visual pendulum environments, LeHami substantially improves long-horizon velocity prediction over an unconstrained JEPA world model and empirically maintains constant latent volume where the baseline contracts. The learned flow also supports accurate integration at unseen sub-frame step sizes and numerical time reversal. These results show how geometric structure can provide inference-time guarantees that are complementary to representation regularization in latent world models.

Figures & tables

Figure 1: LeHami: Learning symplectic world models end-to-end. Bottom left: Actuated pendulum frames and applied torque sequence $u_{t}$ (green). Upper left: Latent phase space $(\mathbf{q}_{t},\mathbf{p}_{t})$ , where transitions are generated by a learned scalar $H_{\phi}(\mathbf{q}_{t},\mathbf{p}_{t};u_{t})=K_{\phi}(\mathbf{p}_{t})+V_{\phi}(\mathbf{q}_{t},u_{t})$ via a symplectic step. Grey lines indicate level sets; the blue curve shows the rollout. Right: Three properties of the Hamiltonian flow: (i) Volume preservation: Latent volume is exactly preserved under any action sequence. (ii) Continuous step size: $H_{\phi}$ integrates at arbitrary time intervals, including between observed frames. (iii) Time-reversibility: Integrates backward by inverting time ( $t\rightarrow-t$ ) and reversing actions.

Figure 2: LeHami multi-step training. Left: at the anchor, the encoder block (Enc + SG) maps a window of

W+1

frames, drawn here as the causal window

\mathbf{x}_{t-W},\dots,\mathbf{x}_{t}

, to the latent state

(\mathbf{q}_{t},\mathbf{p}_{t})

. The Hamiltonian flow (Ham flow) advances it by one leapfrog step of the separable Hamiltonian

H_{\phi}=K_{\phi_{1}}(\mathbf{p})+V_{\phi_{2}}(\mathbf{q},\mathbf{u})

, with the action entering the potential, and is applied

K

times to its own outputs. Middle: each prediction

(\tilde{\mathbf{q}}_{t+k},\tilde{\mathbf{p}}_{t+k})

is scored by the prediction loss (C) against the encoding of the corresponding frames, shown here for the last step

k=K

. The regularizer (R) acts on the encoder’s embeddings of

\mathbf{q}

(blue) and

\mathbf{p}

(orange). Right: the encoder block as instantiated in this work: a pixel encoder (Enc) outputs

\mathbf{q}_{t}

, and a fixed Savitzky–Golay stencil (SG) derives

\mathbf{p}_{t}

from

\mathbf{q}_{t-3},\dots,\mathbf{q}_{t}

Figure 3: The two pendulum environments. Top: the conservative pendulum, with no actuation. Bottom: the actuated pendulum. Beneath each row is its respective torque $u_{t}$ , drawn as the piecewise-constant signal, held constant for three frames.

Figure 4: Prediction performance during conservative pendulum rollouts. Wrapped MAE of $\theta$ (left) and $\omega$ (middle) across the test rollout steps ( $K$ ). Metrics are computed using a $k$ -NN probe ( $k=25$ ) fitted at $K=1$ and kept frozen. Shaded areas indicate variance across three training seeds, averaged over test episodes. Cumulative log-volume of the one-step latent map, averaged over test episodes, for one training seed (right) .

model	probe	MAE $\theta$ [ $10^{-2}$ rad]	MAE $\omega$ [rad/s]	$R^{2}(\omega)$
LeHami	ridge	$23.0\pm 2.0$	$\mathbf{1.168\pm 0.001}$	$\mathbf{+0.491\pm 0.005}$
	25-NN	$6.8\pm 1.8$	$\mathbf{0.419\pm 0.045}$	$\mathbf{+0.889\pm 0.024}$
	MLP	$2.9\pm 0.4$	$\mathbf{0.184\pm 0.019}$	$\mathbf{+0.987\pm 0.004}$
LeWM	ridge	$\mathbf{1.4\pm 0.1}$	$2.422\pm 0.001$	$-0.014\pm 0.001$
	25-NN	$\mathbf{1.1\pm 0.0}$	$2.508\pm 0.002$	$-0.103\pm 0.001$
	MLP	$\mathbf{1.9\pm 0.1}$	$2.419\pm 0.001$	$-0.011\pm 0.001$

Table 1: Performance in physics prediction of the actuated pendulum at the anchor K=0 . Mean

\pm

s.d.over three training seeds, test split. Bold marks the better model within a probe tier. A negative

R^{2}

is below the predict-the-mean floor.

Figure 5: The learned flow behaves like a continuous-time vector field, and is approximately reversible, showcased with the conservative pendulum. Left: MAE of $\theta$ integrated with step sizes never seen in training (to the ground-truth from the integration of the ODE). The model is trained only at $\Delta t=0.1$ s; halving and quartering the step (2 and 4 sub-steps per frame, respectively) traces the same trajectory, while doubling it breaks. Right: Round-trip error after rolling $K$ steps out and back by sign flip $\Delta t\to-\Delta t$ in the $H_{\phi}$ .

Clause	Model	Success	TFLOP	Replanning interval
Position	LeWM $W{=}1$	$\mathbf{0.800\pm 0.019}$	12.87	$R=9$ , 4 solves
	LeWM $W{=}3$ (tied)	$0.633\pm 0.020$	13.65	$R=9$ , 4 solves
	LeWM $W{=}3$ (untied)	$0.696\pm 0.016$	13.65	$R=9$ , 4 solves
	LeHami	$0.480\pm 0.084$	8.05	$R=9$ , 4 solves
Velocity	LeWM $W{=}1$	$0.184\pm 0.027$	8.78	$R=15$ , 2 solves
	LeWM $W{=}3$ (tied)	$0.473\pm 0.014$	13.65	$R=9$ , 4 solves

Table 2: Goal-reaching success on the actuated pendulum, by clause. Best success probability per model at a matched budget of at most

15

TFLOP/episode. Mean

\pm

s.d. over three model seeds. The floor probabilities (corresponding to

u=0

) are

0.060

0.240

and

0.040

on each of the three clauses respectively. Bold marks the better model within a clause. The planner re-solves every

R

steps towards the same goal at horizon 30, so it makes

\lceil 30/R\rceil

solves, the last of which covers only the remaining steps.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 6 — Figure 6: The temporal ConvNet encoder inside the LeHami framework , drawn at $W=3$ . A window of frames goes through a stem applied once per frame, then a single Conv1d over time with a symmetric tied kernel $(a,b,a)$ , then the baseline’s spatial trunk and projector, unchanged, to give $\mathbf{q}_{t}$ . The same encoder is used by the $W{=}3$ LeWM baselines; the main LeHami model uses $W{=}1$ , i.e. no temporal mix. The rollout is initialised with the causal Savitzky–Golay momentum $\mathbf{p}_{t}$ , computed from the current and three previous positions only, and advanced one frame at a time by the leapfrog step (orange) of $H_{\phi}=K_{\phi}(\mathbf{p}/s)+V_{\phi}(\mathbf{q},\mathbf{u})$ , where the learnable per-dimension scale $s$ (blue) rescales the momentum and therefore appears only in the position update.

Figure 7: Frames from the three rod lengths of ( $L\in\{0.7,1.0,1.4\}\,\mathrm{m}$ ), rendered to scale at $64\times 64$ . The rod length is therefore readable from a single frame, which is what makes it a positive control for the probes rather than a prediction target.

split, in episodes (train / val / test):
	conservative pendulum	actuated pendulum
episodes used	300	900
frames per episode	80	80
$\Delta t$	$0.10\,$ s	$0.10\,$ s
episodes	$\mathbf{210/45/45}$	$\mathbf{630/135/135}$
design cells	3 rod lengths	9 ( $3L\times 3m$ )

Appendix

Table 3: The two environments and their episode-level splits. Both use the same protocol, described in Appendix D.3 .

Figure 8: The LeHami embeddings on the actuated pendulum are canonical. $R^{2}$ of the pendulum angle $\theta$ (left) and angular velocity $\omega$ (right) decoded from the position half $q$ (solid) and momentum half $p$ (hatched) of the frozen encoder output. Probes are 25-NN under five episode-grouped folds fitted on embedded anchor frames ( $K=0$ ) pooled across test episodes. $R^{2}$ of $\theta$ is evaluated over $(\sin\theta,\cos\theta)$ . Bars indicate the mean over three training seeds; error bars denote standard deviation. LeWM has no partitioned momentum block and is reported as a single joint latent $z$ .

Metric	Model	$0.75$ PFLOP	$1.5$ PFLOP	$3$ PFLOP	$6$ PFLOP
MAE $\theta$ [rad]	LeHami $W{=}1$	$\mathbf{0.283}$ (ep 9)	$\mathbf{0.207}$ (ep 19)	$\mathbf{0.252}$ (ep 39)	$\mathbf{0.222}$ (ep 79)
MAE $\theta$ [rad]	LeWM $W{=}1$	0.590 (ep 9)	0.871 (ep 19)	0.487 (ep 39)	0.321 (ep 79)
MAE $\omega$ [rad/s]	LeHami $W{=}1$	$\mathbf{1.002}$ (ep 9)	$\mathbf{0.929}$ (ep 19)	$\mathbf{0.934}$ (ep 39)	$\mathbf{0.902}$ (ep 79)
MAE $\omega$ [rad/s]	LeWM $W{=}1$	1.625 (ep 9)	2.146 (ep 19)	1.620 (ep 39)	1.270 (ep 79)

Appendix

Table 4: MAE at a fixed training budget for the conservative pendulum. Test error at the checkpoint with the lowest validation prediction loss within each compute budget (epoch in parentheses), decoded through a frozen

k

-NN probe at

k=20

to ensure a shared scale. Checkpoints were saved every 10 epochs; evaluated on one seed per model.

Model	Total	Encoder	Predictor	Projector	Pred. proj.
LeHami $W{=}1$	3.03	0.315	2.234	0.465	—
LeWM $W{=}1$	4.44	0.315	3.310	0.530	0.268
LeWM $W{=}3$	4.45	0.318	3.310	0.530	0.268

Appendix

Table 5: Parameter counts of the models trained on the conservative pendulum environment. The values are shown in millions, measured from the checkpoints. The predictor is where the two families differ: LeHami’s separable Hamiltonian flow (

K(p)+V(q)

MLPs, hidden_dim

=1024

) against LeWM’s 6-layer autoregressive transformer. pred. proj. is LeWM’s predictor projection head, which LeHami has no counterpart for; the encoder and the latent projector are matched across families, and the

W{=}3

encoder adds only

\sim

3k parameters over

W{=}1

Figure 9: Conservative pendulum: rollout latents projected onto pairs of principal components. Each panel is one pair of the leading PCs of the encoder’s $(\mathbf{q},\mathbf{p})$ over the test episodes; colour is the episode. The learned $H_{\phi}$ is conserved along each rollout, so an episode traces a closed curve in every projection rather than drifting across them.

Figure 10: Latent representation space projected into the first two principal components (PCs). The embeddings correspond to LeHami (left) and unstuctured LeWM (right) , shown as a function (z-axis) of the Hamiltonian learned (left) and the physical energy corresponding to the episodes. The embeddings of four episodes are shown in colors. The planes with the same color correspond to the learned Hamiltonian and energy, respectively. While the rollout episode’s embeddings in LeHami are restrictive in each of the planes corresponding to the (constant) respective learned Hamiltonian, in LeWM, the rollout episode’s representations move across planes of energy, mixing up.

Figure 11: The learned energy landscape against the physical one. (left) level sets of the true energy of the conservative pendulum; (right) level sets of the learned $H_{\phi}$ over the same $60$ test rollouts, each started from its own anchor. Distinct energies give disjoint closed curves in both, so the learned generator stratifies the latent space the way the physical energy stratifies phase space.

rollout step $k$	0	1	5	10	20	25	50
$R^{2}_{q}$ (position block)	$1.000$	$0.997$	$0.985$	$0.950$	$0.929$	$0.646$	$0.373$
$R^{2}_{p}$ (momentum block)	$1.000$	$0.935$	$0.867$	$0.613$	$0.586$	$-0.324$	$-1.165$

Appendix

Table 6: Where the LeHami’s two blocks lose the reference, conservative pendulum.

R^{2}

of the predicted

\mathbf{q}_{k}

and

\mathbf{p}_{k}

against the encoded ground-truth

\mathbf{q}

and

\mathbf{p}

at the same frame, test split. These are latent-space quantities in the model’s own coordinates, not physical ones: they measure how far the flow drifts from the encoder trajectory it was trained to reproduce, and being self-referential they cannot by themselves establish that a physical quantity survives. What survives decoding is in Table 7 ..

anchor ( $K=0$ ), $R^{2}(\omega)$ from the frozen encoder features: LeHami $+0.975$ LeWM $-0.667$
	MAE $\theta$ [rad, wrapped]			MAE $\omega$ [rad/s]
band	LeHami	LeWM	ratio	LeHami	LeWM	ratio
$k=1$ – $10$	0.107	0.100	0.93	0.364	1.020	2.80
$k=11$ – $20$	0.225	0.287	1.27	0.669	1.169	1.75
$k=21$ – $70$	0.596	0.981	1.65	1.534	2.331	1.52
same anchor, LeHami read by the linear tier: ridge $+0.448$ , against the 25-NN’s $+0.975$ above

Appendix

Table 7: Performance in physics prediction of the conservative pendulum. Wrapped MAE of angle

\theta

and

\omega

predicted by the

k

-NN probe (with

k=25

) fitted on each model’s own

K=1

output and applied frozen along the rollout, averaged over the rollout band. Ratio is LeWM

\div

LeHami, so

>1

favours the LeHami.

Figure 12: DINO-LeHami vs. DINO-WM on the conservative pendulum. Both models use a frozen DINOv2 encoder followed by a learned projector; Wrapped MAE of $\theta$ (left) and $\omega$ (right) along the test rollout ( $K$ steps, $\Delta t=0.05$ s), read by a $k$ -NN probe ( $k=25$ ) fitted at $K=1$ and kept frozen. One training seed per model, both at epoch $60$ .

Figure 13: Log-volume factor of the one-step latent map on the conservative pendulum. One Jacobian per test episode at the anchor ( $45$ states; bars mark the mean). Left: the full map, $n=64$ for LeHami on $(\mathbf{q},\mathbf{p})$ and $n=192$ for LeWM on its three-embedding context. Right: the newest block $\partial\mathbf{z}_{t+1}/\partial\mathbf{z}_{t}$ , $n=64$ for both. The dashed line is the true flow. One training seed per model.

Figure 14: The learned energy landscape under actuation. The same projection as Fig. 11 , for the actuated pendulum. The torque does work, so $H_{\phi}$ is not constant along a rollout and the trajectories cross the level sets instead of nesting into disjoint shells. Symplecticity is unaffected: it constrains the map, not the value of $H_{\phi}$ .

window	SIGReg	within-episode %	$\mathrm{std}(q)$	$\mathrm{std}(p)$	$\mathrm{std}(p)/\mathrm{std}(q)$
$W=1$	joint	$0.01\pm 0.00$	1.356	0.006	0.0043
	q	$15.32\pm 0.29$	0.980	0.101	0.1031
	split	$\mathbf{48.89\pm 16.39}$	1.174	0.236	0.2014
$W=3$	joint	$4.90\pm 0.30$	1.405	0.051	0.0362
	q	$8.87\pm 0.54$	0.979	0.051	0.0519
	split	$\mathbf{47.06\pm 2.56}$	1.124	0.201	0.1784

Appendix

Table 8: Within- versus between-episode variance of the encoder’s latent coordinates, by SIGReg placement. For each latent dimension,

V_{\text{within}}

is the mean squared deviation of an episode’s trajectory about its own mean and

V_{\text{between}}

is the variance of those per-episode means; the column reports

V_{\text{within}}/(V_{\text{within}}+V_{\text{between}})

, averaged over dimensions. A value near zero means the latent is an episode identifier rather than a state. Measured on the encoder’s own latents ( gt_q , gt_p ) over the test episodes at a common horizon

K=0\ldots 74

; mean

\pm

standard deviation over three training seeds. Bold marks the largest within-episode share per window. The unstructured baseline attains

42.8\%

on this measure but cannot be tabulated here, as its evaluation path stores latents only at

K=0

Model	$n$	$R^{2}(\omega)$	$R^{2}(\|{}\omega\|{})$	gap
LeWM $W{=}1$	3	$-0.048$	$+0.032$	$+0.080$
LeWM $W{=}3$	3	$-0.115$	$+0.921$	$+1.036$

Appendix

Table 9: The comparison of the

R^{2}

due to the time window in the encoder. They correspond to the

k

-NN (

k{=}25

) probe applied on encoded anchor latents over held-out episodes, across 3 seeds. The

n

refers to the history size of the predictor, which is the same in both. Note that

W=3

refers here to the tied variant, which exacerbates the sign-blindness gap.

Figure 15: Performance in physics predictions in rollout of the actuated pendulum. Wrapped MAE of $\theta$ (left) and $\omega$ (right) along the test rollout, 25-NN probe fitted at $K=1$ and applied frozen. Same encoder, same schedule, same regularizer weight. Shaded area corresponds to the standard deviation across the three architecture seeds.

Figure 16: The $50$ planning goals. Left: the goal velocity $\omega^{\color[rgb]{0,0.5,0.5}*}$ ; dashed lines mark $\pm\tau_{\omega}/2$ , beyond which a sign error alone fails the velocity clause. $92\%$ of the goals lie outside them. Right: the goal angle $\theta^{\color[rgb]{0,0.5,0.5}*}$ , wrapped; the green lines mark upright, which no goal is within $\tau_{\theta}$ of. Goals are recorded states of held-out test episodes at $K=30$ frames after the start, shared identically by every model, rung and seed.

Clause	Model	Success	TFLOP	Replanning Interval ( $R$ )
Velocity, magnitude	LeWM $W{=}1$	$0.251\pm 0.051$	8.78	$R=15$ , 2 solves
	LeWM $W{=}3$ (tied)	$0.638\pm 0.008$	13.65	$R=9$ , 4 solves
	LeWM $W{=}3$ (untied)	$\mathbf{0.813\pm 0.052}$	13.65	$R=9$ , 4 solves
	LeHami	$0.547\pm 0.022$	10.98	$R=6$ , 5 solves
	$u=0$ coast	0.300	—	—
Velocity, sign	LeWM $W{=}1$	$0.656\pm 0.017$	5.85	$R=30$ , 1 solve

Appendix

Table 10: Velocity clause decomposition into magnitude and directional sign components. Evaluated under the same computational budget (

\leq 15

TFLOP/episode) and protocol as Table 2 .

		rod length $L$		bob mass $m$
probe	model	$K=0$	$K=76$	$K=0$	$K=76$
ridge	LeHami	0.999	0.900 [0.871, 0.951]	0.999	0.977 [0.958, 0.993]
	LeWM	0.999	0.983 [0.976, 0.990]	0.999	0.993
$25$ -NN	LeHami	1.000	0.997 [0.992, 1.000]	1.000	1.000
	LeWM	1.000	1.000	1.000	1.000

Appendix

Table 11: Performance in prediction of rod’s length and mass in the actuated pendulum. Both episode-constant parameters are recovered from the latent by probes fitted at anchor (

K=0

) on training episodes and applied unchanged along each model’s own rollout over

135

held-out episodes produced by LeHami and LeWM.

R^{2}

is reported at the anchor and at the end of the horizon (

K=76

). Three seeds for LeHami (min–max in brackets) for both ridge and

25

-NN probes. Both models carry both parameters. While the

k

-NN (k=25) probe is exact at every

K

rollout, it is observed some drift for the ridge, affecting more LeHami. Note that The

R^{2}

columns therefore measure how far the frozen readout drifts, not whether the parameter is still identifiable.

Explore similar work

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models
Aug 6, 2026Xi Zeng, Haojie Ren, Ziying Song +2Representation Learning World Model Learning
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
Orthogonal JEPA: Factorized Predictive States for Latent World Models
Aug 20, 2026Taoyong Cui, Pheng Ann Heng, Wanli OuyangJoint-Embedding Predictive Architecture Orthogonal Representation Learning
The Chinese University of Hong Kong (CUHK)
Is Forward Prediction Enough? Physical State Grounding for JEPA World Models
Aug 7, 2026Haodong Yan, Jiaguan Zhu, Mingyuan Jia +12Physical Reasoning World Models for Robotics
1The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · 2COCO Matrix

Aug 6, 2026J/K move · Enter open · S save

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

Xi Zeng, Haojie Ren, Ziying Song, Yuanbo Nie+1

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

Representation Learning World Model Learning World Model-Based Planning Joint-Embedding Predictive Architecture State Representation Learning

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

Orthogonal JEPA: Factorized Predictive States for Latent World Models

Is Forward Prediction Enough? Physical State Grounding for JEPA World Models