🤖 AI

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

HaM-World introduces a structured world model that integrates Mamba-based selective memory with a decomposed Soft-Hamiltonian latent space to stabilize long-horizon planning and significantly improve robustness against out-of-distribution dynamics shifts.

Original authors: Haoyun Tang, Haodong Cui, Keyao Xu, Kun Wang, Zhandong Mei

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haoyun Tang, Haodong Cui, Keyao Xu, Kun Wang, Zhandong Mei

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to play a complex game, like juggling or balancing a pole. To do this well, the robot needs a "world model"—an internal simulation that lets it imagine what will happen if it takes a certain action before actually doing it.

The paper introduces HaM-World, a new type of world model designed to make these mental simulations more stable and accurate, especially over long periods or when the rules of the game change slightly.

Here is how it works, broken down into simple concepts and analogies:

1. The Problem: The "Daydream" Gets Foggy

Current world models are like a student trying to memorize a story by reading it one sentence at a time. If the story is short, they remember it perfectly. But if they try to imagine the story 100 steps into the future, the details get blurry, and the story falls apart.

The Issue: As the robot plans further ahead, its predictions become unstable. It also struggles if the environment changes (e.g., the floor becomes slippery or the robot's arm gets heavier).
The Cause: The paper argues that existing models mix up too many things in their "brain." They confuse the robot's position (where it is), its momentum (how fast it's moving), and the context (what the goal is or if the floor is wet). This mix-up causes errors to pile up quickly.

2. The Solution: A Specialized "Brain" with Three Rooms

HaM-World solves this by organizing the robot's internal state into three distinct "rooms" or compartments, rather than one big messy room:

Room A: The Physics Engine (q and p)
This room handles the pure physics: position ( $q$ ) and momentum ( $p$ ). The authors use a mathematical concept called Hamiltonian dynamics here.
- The Analogy: Think of this as a billiard table. On a billiard table, energy is conserved; balls bounce off each other in predictable ways. HaM-World tries to make the robot's position and momentum behave like billiard balls. This creates a stable "backbone" for the simulation, preventing the robot from imagining impossible physics (like a ball suddenly speeding up for no reason).
- The "Soft" Twist: Real life isn't a perfect billiard table; there is friction, and the robot has motors. So, HaM-World uses a "Soft-Hamiltonian" approach. It keeps the stable billiard-table structure but allows for a "residual" layer that learns the messy, real-world stuff like friction and motor control.
Room B: The Context Room (c)
This room holds information that doesn't follow the billiard-table rules.
- The Analogy: This is the robot's notebook. It writes down things like "the goal is to spin the finger," "the floor is slippery," or "I need to be careful." These are semantic details that don't fit into the pure physics of position and speed, but are crucial for knowing what to do.
Room C: The Memory Room (Mamba)
This is the robot's short-term memory.
- The Analogy: Sometimes the robot can't see everything (partial observability), or there is a delay between when it decides to move and when it actually moves. A standard model might forget what happened 5 seconds ago. HaM-World uses a special memory system called Mamba that acts like a selective filter. It remembers the important parts of the past needed to make a decision right now, ensuring the robot doesn't get "lost" in its own daydream.

3. How It Works Together

When the robot plans a move, it doesn't just guess. It runs a simulation (a "rollout") inside its head:

It looks at the current state.
It uses the Memory Room to fill in any missing details from the past.
It updates the Physics Room using the stable billiard-table rules, tweaked by the motor controls.
It updates the Context Room with new information about the goal or environment.
It repeats this for many steps into the future to see which action leads to the best result.

Because the physics part is so stable (like the billiard table) and the memory part is so good at filling in gaps, the robot's "daydream" stays clear and accurate even for long planning horizons.

4. The Results: Better at Everything

The authors tested HaM-World on four standard robot tasks (like a cheetah running or a finger spinning an object).

Accuracy: It made far fewer mistakes when imagining the future compared to other top models. In fact, its long-term prediction errors were less than half of the next-best model.
Robustness: When the researchers changed the rules of the game (e.g., made the robot heavier, slowed down its actions, or hid parts of its vision), HaM-World kept performing well. It was the only model that maintained the highest score in every single difficult test condition.
Efficiency: It learned to control the robots faster and with fewer training attempts than other methods.

Summary

HaM-World is like giving a robot a better internal map. Instead of a blurry, mixed-up sketch, it gives the robot a structured map with a stable physical foundation (the billiard table), a place for important notes (the context), and a reliable memory (the filter). This allows the robot to plan complex, long-term actions without getting confused or losing its way, even when the world around it changes.

Technical Summary: HaM-World

1. Problem Statement

World models enable model-based planning by learning latent dynamics to perform imagined rollouts. However, as the planning horizon extends or the environment dynamics shift, these imagined rollouts often become unstable. The authors argue this instability stems from two missing structural components in planner-facing latent representations:

Approximate Markov Completeness: Partial observability and action delays render single-frame latents insufficient as Markov states, yet existing models often lack history-conditioned memory to address this.
Geometric Organization: Monolithic latent representations entangle configuration, momentum, and task semantics. This entanglement causes multi-step prediction errors to grow without structural constraints, leading to poor physical-structure extrapolation under perturbed control regimes.

Current approaches primarily capture statistical correlations rather than the underlying generative structure (e.g., invariances and conservation laws), resulting in planners that struggle with long-horizon consistency and out-of-distribution (OOD) generalization.

2. Methodology: HaM-World

The authors propose HaM-World (HMW), a structured world model that decomposes the latent state into a unified interface shared across dynamics prediction, reward/value estimation, imagined rollouts, and Cross-Entropy Method (CEM) action search. The architecture consists of three core components:

A. Latent Factorization ( $q, p, c$ )

The latent state $z_t$ is explicitly split into:

Canonical Subspace $(q, p)$ : Represents configuration and momentum. This subspace is governed by a Soft-Hamiltonian prior.
Context Subspace $c$ : Captures semantic, dissipative, and non-conservative factors (e.g., contact switches, friction, task semantics) that do not strictly obey energy conservation.

B. Soft-Hamiltonian Dynamics

Instead of enforcing strict energy conservation (which is too restrictive for controlled systems with friction and contact), HaM-World employs a "Soft-Hamiltonian" update for the $(q, p)$ pair:
$\Delta q_t = (1 - \alpha) \Delta q_{net} + \alpha \frac{\partial H}{\partial p}$
$\Delta p_t = (1 - \alpha) \Delta p_{net} - \alpha \frac{\partial H}{\partial q} + G_t a_t$

$H_\phi(q, p)$ : A learnable energy field.
$\alpha$ : A mixing coefficient that balances data-driven network updates with the Hamiltonian vector field.
$G_t a_t$ : An explicit learned control drive added to the momentum update.
Mechanism: This design allows the model to maintain an energy-organized geometric backbone for stability while retaining the flexibility to fit contact, friction, and task objectives via residual dynamics and control drives.

C. Selective Memory (Mamba)

To address partial observability and action delays, HaM-World integrates Mamba selective state-space models as a history-conditioned input ( $h_t$ ) to the latent dynamics.

Unlike parallel recurrent rollout models, Mamba provides a single history-conditioned input to the same latent dynamics used by the planner.
This ensures the planner interface remains unified across training, imagined rollouts, and CEM search without requiring a separate recurrent state for the planner.

D. Training Objectives

The model is trained with a composite loss function that includes:

Multi-step Rollout Consistency ( $L_{roll}$ ): Constrains error accumulation over the planning horizon.
Hamiltonian Alignment ( $L_{ham}$ ): A soft directional bias aligning the network branch with the Hamiltonian vector field.
Action-Free Energy Regularization ( $L_{energy}$ ): Suppresses energy drift in the absence of actions.
Standard Objectives: JEPA-style representation alignment, one-step dynamics consistency, and reward/value regression (using symlog two-hot targets).

3. Key Contributions

Structured Planner-Facing Dynamics: The introduction of a $q/p/c$ latent factorization that couples Mamba selective memory with Soft-Hamiltonian $(q, p)$ dynamics. This provides a single latent interface for CEM planning, reward estimation, and dynamics prediction.
Mechanism-Aware Evaluation: A comprehensive evaluation protocol including standard control performance, long-horizon rollout consistency (MSE at $k \in \{3, 5, 7\}$ ), and 12 OOD perturbations (dynamics shifts, delays, masking).
Empirical Gains: Demonstrated superior performance across multiple metrics, specifically in long-horizon stability and OOD robustness, supported by mechanism diagnostics showing bounded energy drift and coherent control-induced energy transfer.

4. Experimental Results

Evaluated on four DeepMind Control Suite (DMC) tasks (Reacher Easy, Finger Spin, Cheetah Run, Cartpole Swingup) with a 100k-step sample budget:

Control Performance: HaM-World achieved the highest Average Area Under the Curve (Avg. AUC) of 117.9, outperforming TD-MPC2 (+9.5%), SAC (+66.3%), and DreamerV3 (+101.2%). It achieved the highest final return on Finger Spin and Reacher Easy.
Long-Horizon Stability: HaM-World reduced long-horizon rollout error to 45% of the strong baseline (TD-MPC2) and 14.3% of DreamerV3. It won 11 out of 12 MSE cells for $k \in \{3, 5, 7\}$ .
OOD Generalization: Under 12 OOD perturbations (including mass changes, damping, friction, action delays, and observation masking), HaM-World achieved the highest return in every condition.
- On Finger Spin, it showed a 10.2% average OOD-return gain over TD-MPC2.
- On Reacher Easy, it showed a 13.6% average OOD-return gain over TD-MPC2.
Mechanism Diagnostics:
- Energy Drift: In action-free rollouts, the learned Hamiltonian energy remained nearly stable (approx. 1% drift), confirming the Soft-Hamiltonian structure.
- Control Coupling: Analysis showed that action "push" (projection of control onto $\nabla_p H$ ) consistently correlated with energy layer crossing, validating the controlled-Hamiltonian design.
- Ablation: Removing the geometric structure (A1) or memory (A2) significantly degraded performance. Notably, memory was the dominant prerequisite for usability under partial observability, while the Soft-Hamiltonian structure lifted the ceiling for rollout robustness.

5. Significance and Claims

The paper claims that HaM-World demonstrates that treating Markov completeness (via selective memory) and geometric consistency (via Soft-Hamiltonian structure) as coupled stages of unified latent dynamics significantly improves planner performance.

Structural vs. Statistical: The gains are attributed to the model's ability to incorporate physical regularities (energy organization) rather than relying solely on statistical correlation fitting.
Unified Interface: The work shows that a single latent state, shared by prediction, value estimation, and planning, can be made stable and effective for long-horizon tasks if it respects the separation of conservative dynamics ( $q, p$ ) and non-conservative context ( $c$ ).
Robustness: The method provides a structural solution to the instability of imagined rollouts, enabling better extrapolation under distribution shifts and perturbations without requiring separate recurrent planner states.

Limitations: The authors note the current work is limited to state-based observations on four DMC tasks. Scaling to pixel inputs, broader morphologies, and discontinuous contact scenarios remains an open challenge. The diagnostics support the learned interface but do not prove the dynamics are globally Hamiltonian.