HaM-World: Soft-Hamiltonian World Models with Selective Memory for Planning

2026-05-07

AI Summary

Existing world models exhibit instability in long-horizon planning or under distributional shift due to latent variables lacking historical memory and geometric structure. This work proposes a structured world model that disentangles the latent state into a Hamiltonian phase space (q, p) and a contextual subspace c, and for the first time integrates soft Hamiltonian dynamics with the Mamba selective state space mechanism to enable history-conditioned modeling while achieving approximate Markov completeness. The resulting framework unifies dynamics prediction, reward estimation, and action planning. Evaluated on the DeepMind Control Suite, it achieves an average AUC of 117.9 (+9.5%), reduces long-horizon prediction error to 45% of the baseline, and attains state-of-the-art performance across all 12 out-of-distribution perturbations, with OOD returns improved by 10.2% on Finger Spin and 13.6% on Reacher Easy.

Abstract

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.

Problem

• world models
• planning instability
• latent dynamics
• Markov completeness
• OOD generalization

Innovation

• Soft-Hamiltonian dynamics
• Selective memory
• Structured world models
• Latent state decomposition
• Model-based planning