HaM-World: Soft-Hamiltonian World Models with Selective Memory for Planning
About
World models enable model-based planning through learned latent dynamics, but imagined rollouts become unstable as the planning horizon grows or the dynamics distribution shifts. We argue that this instability reflects 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 (HMW), a structured world model that decomposes the latent state into a canonical (q, p) subspace and a context subspace c, while using Mamba selective state-space memory as the history-conditioned input to the same latent dynamics. Within this interface, (q, p) evolves through an energy-derived Hamiltonian vector field plus learnable residual/control dynamics, while c captures semantic, dissipative, and non-conservative factors. This gives the planner a single latent state shared by dynamics prediction, reward/value estimation, imagined rollouts, and CEM action search. On four DeepMind Control Suite tasks, HaM-World reaches the highest Avg. AUC (117.9, +9.5%), reduces long-horizon rollout error to 45% of a strong baseline model, and wins 11/12 k in {3,5,7} MSE cells. Under 12 OOD perturbations spanning dynamics shifts, action delay, and observation masking, HaM-World achieves the highest return in every condition, with average OOD-return gains of 10.2% on Finger Spin and 13.6% on Reacher Easy. Mechanism diagnostics further show bounded action-free Hamiltonian-energy drift, structured energy variation under policy rollouts, and coherent control-induced energy transfer, supporting the intended Soft-Hamiltonian dynamics design.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Finger Spin | DeepMind Control Suite Finger Spin | Final Return (100k steps)254 | 5 | |
| Multi-task Control Performance and Consistency | DeepMind Control Suite Average | Average AUC117.9 | 5 | |
| Reacher Easy | DeepMind Control Suite Reacher Easy | Final Return150.6 | 5 | |
| Reinforcement Learning | Finger Spin OOD | Score (Friction 0.5)242.8 | 5 | |
| Reinforcement Learning | Reacher Easy OOD | Score (mass×0.7)158.5 | 5 | |
| Cartpole Swingup | DeepMind Control Suite Cartpole Swingup | Final Return (100k steps)58.9 | 5 | |
| Cheetah Run | DeepMind Control Suite Cheetah Run | Final Return (100k steps)184.4 | 5 |