Hamiltonian JEPA: Action-Conditioned World Models with an Inherited Control State
cs.LG
Submitted: 2026-09-27
Updated: 2026-09-27
License: http://creativecommons.org/licenses/by/4.0/
The gist: Planning from pixels needs more than a latent space that is stable and predictable.
Terminology
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.
Sources
- LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics
- Rethinking Weight Tying: Pseudo-Inverse Tying for LM Stable Training and Updates
- Dream to Control: Learning Behaviors by Latent Imagination
- Sensorimotor World Models: Perception for Action via Inverse Dynamics
- LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels
- DeepMind Control Suite
- The information bottleneck method
- Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding
- Sub-JEPA: Subspace Gaussian Regularization for Stable End-to-End World Models
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks