LeHami: Symplectic Latent Dynamics for World Models
Abstract Overview
LeHami is a JEPA-based world model that replaces the standard unconstrained autoregressive predictor with an action-conditioned Hamiltonian flow over a structured latent state split into position-like and momentum-like components. The latent dynamics are integrated using an exactly symplectic leapfrog step, providing phase-space volume preservation at every rollout step and ruling out rollout collapse modes such as uniform contraction, asymptotically stable fixed points, and merging of distinct latent states. The momentum-like latent is obtained from temporal derivatives of the learned position-like representation via a fixed Savitzky–Golay stencil, explicitly representing second-order physical state for prediction and control. Experiments on conservative and actuated visual pendulum environments evaluate long-horizon prediction, latent volume behavior, continuous-time re-stepping, time reversal, and goal-conditioned planning under matched compute budgets.
Novelty
The paper's core novelty is embedding an exact symplectic Hamiltonian flow inside a pixel-based JEPA world model to preserve phase-space volume at inference time independently of learned parameters or action sequences. It also introduces a zero-parameter Savitzky–Golay stencil over encoder trajectories to extract canonical momentum-like representations without requiring a decoder or supervised state coordinates.
Results
On conservative and actuated visual pendulum benchmarks, LeHami improves long-horizon velocity prediction over the LeWM baseline and empirically maintains constant latent volume while the baseline steadily contracts. The learned continuous vector field supports integration at unseen sub-frame step sizes and achieves numerical time reversal with round-trip error below 10^-4%. In goal-conditioned planning under a 15 TFLOP budget, LeHami provides a compute-efficient structural prior for signed velocity, though an untied-window LeWM baseline attains higher success rates on certain clauses.
Key Points
- LeHami replaces the unconstrained JEPA predictor with a separable, action-conditioned Hamiltonian flow over a canonical (q, p) latent split integrated via an exact symplectic leapfrog step.
- The symplectic formulation guarantees per-step phase-space volume preservation and prevents specific rollout collapse modes, with empirical validation showing constant latent volume over extended horizons.
- LeHami enables sub-frame time re-stepping and near-exact numerical time reversal while providing compute-efficient representations of signed velocity for goal-conditioned planning.