DreamSAC: Learning Hamiltonian World Models via Symmetry Exploration

Jinzhou Tang; Fan Feng; Minghao Fu; Wenjun Lin; Biwei Huang; Keze Wang

cs.CVMar 8, 2026

DreamSAC: Learning Hamiltonian World Models via Symmetry Exploration

Authors: Jinzhou Tang, Fan Feng, Minghao Fu, Wenjun Lin, Biwei Huang, Keze Wang

Organizations: UC San Diego · Sun Yat-sen University

Abstract

Learned world models excel at interpolative generalization but fail at extrapolative generalization to novel physical properties. This limitation arises because they learn statistical correlations rather than the environment's underlying generative rules, such as physical invariances and conservation laws. We argue that learning these invariances is key to robust extrapolation. To achieve this, we first introduce Symmetry Exploration, an unsupervised exploration strategy where an agent is intrinsically motivated by a Hamiltonian-based curiosity bonus to actively probe and challenge its understanding of conservation laws, thereby collecting physically informative data. Second, we design a Hamiltonian-based world model that learns from the collected data, using a novel self-supervised contrastive objective to identify the invariant physical state from raw, view-dependent pixel observations. Our framework, DreamSAC, trained on this actively curated data, significantly outperforms state-of-the-art baselines in 3D physics simulations on tasks requiring extrapolation.

Figures & tables

Figure 2: Overview of DreamSAC. (Right) Our world model maps observations $x_{t}$ to object-centric slots $Z_{t}$ via SAVi [ 24 ] . We structure each slot $z_{t}^{i}$ into generalized coordinates ( $q_{t}^{i}$ ) and canonical momenta ( $p_{t}^{i}$ ). The dynamics are twofold: the stochastic state $Z_{t+1}$ is computed by integrating our $G$ -invariant Hamiltonian $H_{\phi}$ , while the deterministic state $h_{t+1}$ is updated by a GRU. (Left) Symmetry Exploration: To efficiently learn $H_{\phi}$ , a policy $\pi_{\theta}$ is trained entirely in imagination to maximize our Symmetry-Aware Curiosity reward $r_{sym}$ . This incentivizes the policy to work to break symmetry. The imagined policy is then executed in the real environment to collect challenging data, which refines the world model.

	DeepMind Control Suite											GymFetch
	Cheetah		Acrobot		Hopper		Walker		Humanoid			FetchPush		FetchReach
Method	H=16	H=100	H=16	H=100	H=16	H=100	H=16	H=100	H=5	H=10	H=16	H=8	H=16	H=8	H=16
DreamerV3+Policy	0.7981	0.7507	0.7723	0.9392	1.0355	—	4.3769	—	—	—	—	1.275	2.030	0.855	1.492
DreamerV3+Random	0.8747	0.4048	0.8423	1.7547	0.6406	0.9239	2.2527	4.3760	5.2171	6.0500	6.5078	1.048	1.932	0.962	1.670
DreamerV3+RND	0.6364	0.4578	0.2109	0.5628	1.0643	1.2764	2.8976	3.2160	4.2077	5.7294	6.9389	0.976	1.708	0.574	0.682
DreamSAC+Random	0.1565	0.3367	0.2532	0.9347	0.5227	0.9762	2.4667	3.6391	5.6391	6.1402	6.8589	0.675	0.790	0.498	0.652

Table 1: World model image prediction accuracy. We compare the Mean Squared Error (MSE) (lower is better) at 1M steps of our full model against the DreamerV3 [ 17 ] baselines (random policy, sample from policy replay buffers) and ablations of our exploration strategy (random policy, RND [ 4 ] . Models are evaluated across various DMCS and GymFetch tasks at varying rollout horizons (H).

	Structural & Interaction Generalization				Parametric Generalization
Model	Reacher-hard		FetchReach		Walker-walk		Cheetah-run
	Unseen View	Unseen Goal	Unseen Object	Unseen Goal	Unseen Gravity	Unseen Dist.	Unseen Friction	Unseen Dist.
DreamerV3+Policy	265.33 $\pm$ 10.33	919.73 $\pm$ 9.12	0.65 $\pm$ 0.12	0.76 $\pm$ 0.11	189.76 $\pm$ 23.27	125.44 $\pm$ 35.70	118.79 $\pm$ 62.23	80.38 $\pm$ 28.43
DreamerV3+RND 0-shot	79.31 $\pm$ 52.98	892.90 $\pm$ 10.37	—	—	86.75 $\pm$ 39.72	34.28 $\pm$ 6.29	9.42 $\pm$ 4.27	7.71 $\pm$ 1.41
DreamerV3+RND	313.97 $\pm$ 27.31	927.36 $\pm$ 18.73	0.70 $\pm$ 0.07	0.72 $\pm$ 0.13	167.52 $\pm$ 21.57	113.40 $\pm$ 32.99	97.43 $\pm$ 67.82	103.42 $\pm$ 27.91
DreamSAC 0-shot	149.64 $\pm$ 37.62	934.21 $\pm$ 7.98	—	—	124.78 $\pm$ 21.77	67.22 $\pm$ 30.13	27.53 $\pm$ 10.98	107.31 $\pm$ 45.41

Table 2: Comprehensive comparison of extrapolative generalization on all OOD tasks. We report final performance (Success Rate for FetchReach tasks, Mean Reward for all others) after 2M unsupervised pre-training steps and 500K task adaptation steps (mean

\pm

std. over 5 seeds). The tasks evaluate Structural OOD (Unseen View: new viewpoints, Unseen Object: new object counts, Unseen Goal) and Parametric OOD (Unseen Gravity: 1.5x gravity, Unseen Friction: 2.0x friction, Unseen Dist.: physical properties domain shift).

Figure 3: Qualitative analysis of DreamSAC’s internal mechanisms. (a) t-SNE projections show our full model (with $\mathcal{L}_{\text{vr}}$ ) learns viewpoint-invariant representations, unlike an ablation without it. (b) The learned Hamiltonian $H_{\phi}$ (red dashed line) is conserved during a zero-action rollout, confirming the model learned a physical invariant (energy conservation). (c) Reward and Loss curves comparing DreamSAC with different baselines. (d) The latent states $(\boldsymbol{q},\boldsymbol{p})$ demonstrate physics-awareness: representations for pre-train (yellow) and fine-tune (blue) mix for familiar In-Distribution properties, but clearly separate to learn novel Out-of-Distribution properties.

	Hopper		Quadruped		Walker
Method	Hop	Stand	Run	Escape	Stand	Walk
DreamerV3+Policy	354.264	929.851	901.963	207.691	903.556	965.217
DreamerV3+RND	389.142	937.651	867.324	179.169	942.684	979.340
DreamSAC (Ours)	366.117	967.865	911.488	236.538	963.251	996.502

Table 3: In-Distribution (ID) Task Generalization. DreamSAC is compared against baselines on standard downstream tasks after pre-training and fine-tuning except for DreamerV3+Policy which is trained from scratch. We report the final Mean Reward (avg. over 5 seeds), showing highly competitive performance.

	Reacher †	Cheetah	Walker	Avg. OOD
Model	Unseen View	Seen Env.	1.5x Gravity	MSE
Ours w/o $\mathcal{L}_{vr}$	212.37 $\pm$ 32.47	0.5182	1.0681	0.7932
Ours w/o $H_{\phi}$	159.63 $\pm$ 27.12	0.8313	4.9673	2.8993
Ours w/o SAVI	279.68 $\pm$ 19.24	0.6179	1.1882	0.9031
DreamSAC (Full)	321.90 $\pm$ 13.28	0.4052	1.0044	0.7048

Table 4: Ablation study of DreamSAC components. We report final performance on representative OOD tasks. Removing any component causes a significant performance drop on the task it is designed to solve. † Metric is Reward; all other metrics are MSE.

Temperature $\tau$	0.05	0.07 (Ours)	0.1	0.2	0.5
Mean Reward	312.4	321.9	308.7	245.3	112.8

Table 5: Sensitivity analysis of the temperature parameter

\tau

on Reacher Hard (Unseen View) mean reward (mean over 5 seeds).

Figure 4: Visual illustration of the linear annealing schedule for an environment with

T_{anneal}=10^{6}

Environment	$T_{anneal}$ (Steps)
GymFetch Manipulation
FetchReach	$2\times 10^{5}$
FetchPush	$5\times 10^{5}$
FetchSlide	$1\times 10^{6}$
DeepMind Control Suite
Acrobot Swingup	$2\times 10^{5}$

Table 6: Annealing duration

T_{anneal}

for different environments. We set longer annealing phases for environments with complex dynamics or high-dimensional state spaces ( e.g. , Humanoid, Quadruped) to allow sufficient exploration via RND before switching to symmetry probing.

Parameter	Value
Training Config
Batch Size	50
Sequence Length ( $T$ )	64
Total Training Steps	$2\times 10^{6}$ (Pretrain) + $5\times 10^{5}$ (Adapt)
Optimizer	AdamW
Learning Rate (World Model)	$1\times 10^{-4}$

Table 7: Hyperparameters used for DreamSAC training.

Method	Mean Reward
DreamerV3+ICM	469.72
DreamerV3+Plan2Explore	379.28
DreamSAC (Ours)	499.91

Table 8: Mean Reward on the DMCS Walker-walk Unseen Gravity OOD task (Gravity

1\times\to 1.5\times

). Comparison ensures identical backbone architectures (with Hamiltonian).

$\lambda_{s}$ Value	Mean Reward	Behavior
0.0 (No Reg.)	967.43	Jitter
0.01 (Ours)	996.50	Coherent
0.1 (High Reg.)	921.66	Over-smooth

Table 9: Sensitivity analysis of Action Smoothness

\lambda_{s}

on Walker-walk . Note that even with

\lambda_{s}=0

, the model maintains decent performance, indicating that the Hamiltonian prior is robust. The smoothing parameter primarily serves to filter out high-frequency jitter that creates false energy deltas.

Metric	Euler Inference	Leapfrog Inference (Ours)
Prediction MSE ( $\downarrow$ )	0.215	0.198
Energy Drift ( $\sigma_{H}$ , $\downarrow$ )	0.128	0.015

Table 10: Comparison of Inference Integrators on Acrobot (

H=100

). While MSE remains similar, the Symplectic Leapfrog integrator (Ours) drastically reduces Energy Drift, confirming its ability to enforce physical conservation laws that explicit Euler fails to maintain over long horizons.

Method	Training Time / Step	GPU Memory Usage
DreamerV3	1.00 $\times$	1.00 $\times$
DreamSAC (Ours)	1.35 $\times$	1.12 $\times$

Table 11: Computational cost comparison (normalized relative to DreamerV3).

World Model Learning

World model learning builds predictive models of how environments change under actions to support planning, simulation, and decision-making.

See all 414 papers

Correcting a learned physical invariant improves world-model rollouts
Aug 24, 2026Richard BaoWorld Models World Model Learning
Independent Researcher
Physically Native World Models: A Hamiltonian Perspective on Generative World Modeling
May 1, 2026Sen Cui, Jingheng MaLatent Dynamics Modeling World Models
Tsinghua University
PhysiCalWorld: In-Context Physics Calibration for World Models
Oct 21, 2026Yunhao Luo, Wentao He, Yinpei Dai +5World Models Action-Conditioned World Models
University of Michigan

Aug 24, 2026J/K move · Enter open · S save

Correcting a learned physical invariant improves world-model rollouts

Richard Bao

Independent Researcher

World Models World Model Learning Latent World Models World Model Evaluation Scientific Discovery

Open paper AskPDF arXiv

World models can predict video without learning dynamics that they reliably preserve. We test whether a frozen DreamerV3 trained only on pendulum video learns a scalar that its own latent transition treats as approximately conserved. A label-free search recovers the same energy-like invariant across independently trained conservative models, while the same procedure finds no comparable invariant in matched damped models. During autonomous rollouts, this quantity drifts. Projecting the latent state back toward its initial level set reduces rollout error in all three conservative models, whereas matched random constraints usually increase it. These results distinguish a dynamically meaningful invariant from a merely decodable correlate and reveal a concrete failure mode: a world model can learn a physical constraint from pixels yet violate that constraint when it imagines forward.

Abstract

World Model Learning

Explore similar work

Correcting a learned physical invariant improves world-model rollouts

Physically Native World Models: A Hamiltonian Perspective on Generative World Modeling

PhysiCalWorld: In-Context Physics Calibration for World Models