HaM-World: Soft-Hamiltonian World Models with Selective Memory for Planning
Haoyun Tang ⋅ Haodong Cui ⋅ Keyao Xu ⋅ Zhan-Dong Mei ⋅ Kun Wang
Abstract
World models enable model-based planning through learned latent dynamics, but imagined rollouts often become unstable as the planning horizon grows or the dynamics distribution shifts. We argue that this instability arises from two missing structures in planner-facing latents: history-conditioned memory for approximate Markov completeness, and geometric organization that separates configuration, momentum, and task semantics. We propose HaM-World, a structured world model that decomposes the latent state into a canonical $(q, p)$ subspace and a context subspace $c$, while incorporating Mamba selective state-space memory as a history-conditioned input to the same latent dynamics. Within this unified interface, $(q, p)$ evolves under a Soft-Hamiltonian dynamics composed of an energy-derived Hamiltonian vector field and learnable residual and control dynamics, while $c$ captures semantic, dissipative, and non-conservative factors. This design provides the planner with a single latent representation shared across dynamics prediction, reward and value estimation, imagined rollouts, and cross-entropy method (CEM) planning. On four DeepMind Control Suite tasks, HaM-World achieves the highest average AUC (+9.5% over strong baselines), reduces long-horizon rollout error to 45% of a competitive model, and wins 11 out of 12 $k \in \{3,5,7\}$ rollout MSE metrics. Under 12 out-of-distribution perturbations, HaM-World consistently attains the highest returns, with average gains of 10.2% on Finger Spin and 13.6% on Reacher Easy. Mechanism diagnostics further demonstrate bounded energy drift under action-free rollouts, structured energy variation under policy control, and coherent control-induced energy transfer, supporting the effectiveness of the proposed Soft-Hamiltonian latent dynamics. Code: https://anonymous.4open.science/r/HaM_World-47CD
Chat is not available.
Successful Page Load