MHWA: Multi-timescale Hierarchical World-Action Model

Pengcheng Pan; Guoqing Ma; Yuhan Zhang; Yang Chen; Yichen Liu; Ziheng Li; Shan Yu

cs.AIOct 10, 2026

MHWA: Multi-timescale Hierarchical World-Action Model

Authors: Pengcheng Pan, Guoqing Ma, Yuhan Zhang, Yang Chen, Yichen Liu, Ziheng Li, Shan Yu

Organizations: Laboratory of Brain Atlas and Brain-inspired Intelligence, Institute of Automation, Chinese Academy of Sciences, Beijing, China · School of Artificial Intelligence, University of Chinese Academy of Sciences, Beijing, China · State Key Laboratory of Brain Cognition and Brain-inspired Intelligence Technology, Institute of Automation, Chinese Academy of Sciences, Beijing, China · School of Future Technology, University of Chinese Academy of Sciences, Beijing, China

Abstract

Autoregressive world models have emerged as a powerful paradigm for sample efficient decision making, yet their open loop rollouts often exhibit increasing drift that leads to error accumulation even within a relatively small planning horizon. Hierarchical temporal abstraction is a natural candidate to mitigate such drift by reducing the effective prediction horizon; however, existing hierarchical designs can introduce added structural and conceptual complexity, sometimes making them difficult to control and limiting gains over flat baselines. To this end, we introduce MHWA (Multi-timescale Hierarchical World-Action model), which employs a strided high level context channel to reduce global context update frequency, effectively curbing the propagation of compounding drift. Crucially, via context aware routing, the gating network dynamically reconfigures a low level mixture of experts by using both global context and current observations, thereby enabling flexible high-low coordination, helping alleviate interference commonly observed in static hierarchies. Empirically, using RLiable's seed-by-game aggregation, MHWA (77M parameters) achieves an IQM human normalized score (HNS) of 0.884 across three training seeds on a 14-game Atari subset using 10% subsampled offline data, outperforming a matched JOWA-70M baseline (0.529) and achieving numerically comparable performance to JOWA-150M while using nearly half the parameters and 70.88% fewer training FLOPs per sequence.

Figures & tables

Figure 1: The architecture of MHWA. It is illustrated at a context-update timestep

t=1+m\cdot k

. The model processes an interleaved sequence of observation (

o_{t}

) and action (

a_{t}

) tokens. (Top) At this specific strided step, the High-Level Manager consumes the downsampled history to infer a slowly varying global context

z^{ctx}

, supervised by a skip-step prediction head

p^{high}_{o}

that forecasts the observation at the next coarse step

\hat{o}_{t+k}

. We refer to the original, complete sequence as the dense trajectory, in contrast to the downsampled input used by the manager. (Bottom) The Low-Level Worker operates at the atomic resolution. It employs a Context-Aware Dynamic Routing mechanism where a Gating Network utilizes both the current dense input and the refreshed high-level context

z^{ctx}

to dynamically route computation to a sparse subset of experts (MoE). The worker’s output drives the world model heads (

p_{o},p_{r},p_{d}

) and the critic head (

h_{Q}

Game	Random	Human	MTBC	MGDT	EDT	Scaled-QL	JOWA			Ours
			120M	200M	200M	80M	40M	70M	150M	77M
Assault	222.4	742.0	1203.5	1741.5	1915.5	2292.0	1578.3	1402.8	2302.0	1068.8
Atlantis	12850.0	29028.1	37812.5	2565750.0	2108166.7	49100.0	41662.5	228825.0	2690387.5	1690733.3
BeamRider	363.9	16926.5	786.0	6011.3	4149.8	7023.5	880.3	854.4	3498.0	4202.1
Carnival	0.0	3800.0	5560.0	2610.0	3678.8	1940.0	4685.0	3812.1	5316.0	881.7
Centipede	2090.9	12017.0	4046.6	4604.0	3389.9	3650.3	5669.3	4485.0	4677.0	4970.9

Table 1: Performance comparison on 14 Atari games. Baseline scores are from Cheng et al. [5] . We report raw scores for each game and aggregate metrics (Median and IQM Human-Normalized Scores) at the bottom. The best raw score for each game, excluding reference scores, is highlighted in bold . “Ours” denotes our 77M-parameter MHWA model. For Ours, per-game raw scores are averaged over three independently trained seeds; bottom aggregate rows report mean

\pm

sample standard deviation over seed-level aggregate metrics. Full per-seed results are provided in Appendix D .

Model	IQM HNS	95% CI	GFLOPs/sequence
JOWA-70M	0.529	[0.459, 0.602]	84.238
MHWA-77M	0.884	[0.736, 1.031]	48.199
JOWA-150M [ 5 ]	0.833	–	165.524

Table 2: Matched performance and training computation. JOWA-70M and MHWA use three training seeds on the same 14-game dataset and batch setting. We report seed-by-game IQM HNS and 95% stratified-bootstrap CIs. JOWA-150M performance is the published reference. Forward+backward FLOPs use identical input shapes.

Figure 2: Quantitative comparison of open-loop prediction within the context window (

T=5

8

). We analyze reconstruction quality (MSE, PSNR, SSIM) of MHWA-77M against JOWA-150M and a parameter-matched JOWA-70M. While the reduced-capacity JOWA-70M degrades rapidly as rollout depth increases, MHWA remains stable and, in terms of MSE and SSIM, matches or exceeds the larger JOWA-150M. The gap between MHWA and JOWA-70M is consistent with the benefit of structural inductive bias under constrained parameters.

	Stop-gradient on	Stop-gradient off
$\mathcal{L}_{high}$ on	0.564	0.533
$\mathcal{L}_{high}$ off	0.503	0.387

Table 3: High-level supervision and gradient isolation. IQM HNS under the 44M/6-game ablation protocol. The Context Projector and worker remain trainable in all four settings.

Metric	Full	w/o $\mathcal{L}_{high}$	w/o $\mathcal{L}_{aux}$	w/o MoE	MoE-only	w/o State	$k=1$	$k=2$	$k=8$
Median HNS	0.463	0.425	0.439	0.412	0.386	0.364	0.183	0.446	0.469
IQM HNS	0.564	0.503	0.444	0.450	0.445	0.359	0.214	0.446	0.517
$\Delta_{\mathrm{IQM}}(\%)$	–	10.8	21.3	20.2	21.1	36.3	62.1	20.9	8.3

Table 4: Ablation analysis on the 6-game subset (44M). For each variant, we average 16 evaluation episodes within each game and report Median and IQM human-normalized score (HNS) over the six resulting per-game scores.

\Delta_{\mathrm{IQM}}(\%)

denotes the relative IQM drop compared to the full model, i.e.,

(\mathrm{IQM}_{\mathrm{Full}}-\mathrm{IQM}_{\mathrm{variant}})/\mathrm{IQM}_{\mathrm{Full}}\times 100

. Bold and underline mark the largest and second-largest

\Delta_{\mathrm{IQM}}(\%)

, respectively. All variants are controlled to have

\approx

44M parameters for fair comparison. Detailed per-game raw scores are provided in Appendix F .

Panel of Figure 3 — Figure 3: PCA visualization of expert activations in the last MoE layer. (a) Full MHWA maps Assault and Phoenix closely together, indicating they utilize similar expert combinations, which aligns with their shared vertical shooting mechanics. (b, c) This structural similarity vanishes without high-level supervision or temporal stride, indicating a failure to capture shared cross-game dynamics.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Game	Difficulty	$\|\mathcal{A}\|$	Split
Assault	Medium	7	Train
Atlantis	Hard	18	Train
BeamRider	Medium	9	Train
Carnival	Medium	6	Train
Centipede	Medium	18	Train
ChopperCommand	Easy	18	Train

Appendix

Table 5: Atari games used in our multi-game training set, with difficulty labels from [ 38 ] and discrete action-space sizes.

Model	GFLOPs/sequence	TFLOPs/training step
MHWA-77M	48.199	12.339
JOWA-70M	84.238	21.565
JOWA-150M	165.524	42.374

Appendix

Table 6: Forward+backward training computation at matched input shape.

Game	Seed 0	Seed 1	Seed 2	Mean	Sample SD
Assault	1,209.3	1,109.1	888.1	1,068.8	164.3
Atlantis	3,320,456.3	1,078,587.5	673,156.2	1,690,733.3	1,425,865.1
BeamRider	4,077.6	4,920.0	3,608.8	4,202.1	664.4
Carnival	1,300.0	305.0	1,040.0	881.7	516.1
Centipede	5,689.3	3,705.4	5,518.1	4,970.9	1,099.3
ChopperCommand	3,300.0	5,793.8	4,550.0	4,547.9	1,246.9

Appendix

Table 7: Per-seed results of MHWA-77M on the 14-game Atari benchmark.

Game	Seed 0	Seed 1	Seed 2	Mean	Sample SD
Assault	1,408.3	858.1	1,208.2	1,158.2	278.4
Atlantis	3,167,275.0	125,787.5	2,885,637.5	2,059,566.7	1,680,611.9
BeamRider	853.5	608.5	915.0	792.3	162.1
Carnival	3,243.8	1,293.8	690.0	1,742.5	1,334.7
Centipede	8,937.4	4,619.2	4,102.1	5,886.2	2,655.0
ChopperCommand	675.0	681.3	737.5	697.9	34.4

Appendix

Table 8: Per-seed results of matched JOWA-70M on the 14-game Atari benchmark.

Game	Random	DQN	JOWA-70M	JOWA-150M	MHWA Seed 0	MHWA Seed 1	MHWA Seed 2	MHWA Mean
Gravitar	173.0	473.0	387.4	273.3	375.2	271.8	284.4	310.5
MsPacman	307.3	3085.6	908.0	2016.7	1437.5	2006.2	1971.0	1804.9
Pong	-20.7	19.5	14.6	17.7	17.0	18.1	18.1	17.7
Robotank	2.2	63.9	13.8	25.0	12.4	11.3	9.3	11.0
YarsRevenge	3092.9	18089.9	16339.2	17506.2	16056.0	17760.4	19280.6	17699.0
Mean DNS	–	–	0.576	0.647	0.610	0.606	0.626	0.614 $\pm$ 0.010

Appendix

Table 9: Held-out Atari fine-tuning results. We report per-game raw scores and aggregate DQN-normalized scores (DNS). MHWA is evaluated over three seeds; the MHWA Mean column averages raw scores across the three seeds for per-game rows and reports mean

\pm

sample standard deviation over seed-level aggregate DNS for the bottom rows.

Game	Random	Human	Full	w/o $\mathcal{L}_{high}$	w/o $\mathcal{L}_{aux}$	w/o MoE	MoE-only	w/o State	$k=1$	$k=2$	$k=8$
Assault	222.4	742.0	1086.0	989.9	930.1	917.3	902.7	745.8	1003.6	843.2	1043.6
Carnival	0.0	3800.0	626.7	220.0	380.0	340.0	680.0	420.0	463.8	300.0	240.0
Centipede	2090.9	12017.0	5387.7	4527.0	4199.0	4530.0	2982.1	2801.3	3902.0	3228.0	3964.0
NameThisGame	2292.3	8049.0	9186.5	8480.0	4160.0	6950.0	8246.6	5690.6	3350.0	6770.0	8270.0
Phoenix	761.4	7242.6	3924.0	4214.4	4350.0	2628.1	4610.0	4040.9	1503.1	3120.0	4300.0
SpaceInvaders	148.0	1668.7	813.9	630.9	1205.0	961.9	110.0	485.0	796.9	950.0	745.0

Appendix

Table 10: Per-game raw scores for ablation variants on the 6-game subset. All variants use the 44M controlled setting. For each game and variant, we report the raw score averaged over 16 evaluation episodes. The bottom rows report Median and IQM HNS over the six per-game scores, matching Table 4 .

Hidden state supplied to Context Projector	Median HNS	IQM HNS
Trajectory-conditioned manager hidden state	0.425	0.503
Zero hidden state	0.424	0.406
Gaussian noise	0.410	0.381

Appendix

Table 11: Information controls for the fixed-weight manager pathway. All variants are separately trained in the 44M/6-game setting with

\mathcal{L}_{high}

off and stop-gradient on.

Experts $N$	Top- $K$	MLP expansion	IQM HNS
2	1	$2r_{0}$	0.429
8	4	$r_{0}/2$	0.488
4	1	$r_{0}$	0.542
8	2	$r_{0}$	0.511
4	2	$r_{0}$	0.564

Appendix

Table 12: Expert configuration study. IQM HNS on the six-game ablation benchmark.

Experts $N$	Top- $K$	Assault–Phoenix distance rank
4	2	1
4	1	14
8	2	1
8	4	4

Appendix

Table 13: Routing neighborhoods across expert configurations. Rank of the Assault–Phoenix distance among 15 game pairs in each configuration’s PCA projection; rank 1 is the closest pair.

Setting	Training $r$	Evaluation $r$	Median HNS	IQM HNS
Default	4	4	0.463	0.564
Inference-only	4	2	0.441	0.568
Inference-only	4	1	0.441	0.533
Retrained	2	2	0.484	0.558

Appendix

Table 14: Context-refresh schedules at fixed manager stride

k=4

. The inference-only variants use the default checkpoint; the retrained variant uses

r=2

during both training and evaluation.

Game	Default ( $k=4,r=4$ )	Retrained ( $k=4,r=2$ )
Assault	1086.0	1030.3
Carnival	626.7	758.8
Centipede	5387.7	7129.1
NameThisGame	9186.5	8205.0
Phoenix	3924.0	3751.3
SpaceInvaders	813.9	728.4

Appendix

Table 15: Per-game results after retraining with

k=4,r=2

. Raw returns under the same 44M/6-game protocol.

Figure 4: Long-horizon open-loop reconstruction diagnostic from

t=5

t=54

. MHWA shows smoother long-horizon degradation than the smaller monolithic baseline and remains broadly comparable to, while being less affected by sharp fluctuations than, the 150M baseline; however, it does not eliminate eventual manifold departure.

Hierarchical World Models

See all 11 papers

Learning Goal-Conditioned Subgoals through Frozen Hierarchical World Models
Oct 19, 2026Xiang Kaibo, Liu QingfengHierarchical World Models World Model-Based Planning
Department of Industrial and Systems Engineering, Hosei University
Hierarchical Planning with Latent World Models
Apr 3, 2026Wancong Zhang, Basile Terver, Artem Zholus +7Hierarchical Planning Hierarchical Robot Control
FAIR at Meta · New York University · Mila - Québec AI Institute +1
WAM-RL: World-Action Model Reinforcement Learning with Reconstruction Rewards and Online Video SFT
Jun 16, 2026Zezhong Qian, Xiaowei Chi, Yu Qi +3World Model Learning Action-Conditioned World Models
State Key Laboratory of Multimedia Information Processing, School of Computer Science, Peking University · Northeastern University · Tsinghua University

Oct 19, 2026J/K move · Enter open · S save

Learning Goal-Conditioned Subgoals through Frozen Hierarchical World Models

Xiang Kaibo, Liu Qingfeng

Department of Industrial and Systems Engineering, Hosei University

Hierarchical World Models World Model-Based Planning World Models Goal-Conditioned Planning Latent-Space Planning Efficient Inference

Open paper AskPDF arXiv

Hierarchical world models plan over long horizons by passing predicted states from a coarse model to a lower-level controller. Choosing these states can be expensive: the upper-level planner repeatedly differentiates through candidate macro-action trajectories. We replace this search with a lightweight goal-conditioned network trained through the frozen high-level predictor. The network supplies a conditioning vector, while the predictor generates the next latent waypoint. Training uses hindsight goals and state-prediction losses on offline trajectories, without search-solution targets or macro-action labels. Inference retains an autoregressive high-level rollout, but requires no high-level candidate optimization. We evaluate this approach on the shared-latent HWM checkpoints distributed with H-JEPA, using three backbone seeds and three head initializations per backbone. With a common evaluation protocol, mean success is 70.4% on Cube, 48.7% on Push-T, 100.0% on FourRoom, and 71.1% on Visual AntMaze. The method improves over the corresponding HWM search baseline on three environments and has the same mean on Push-T. Dedicated measurements on an NVIDIA L40S give 1.28-4.39x end-to-end speedups for batched evaluation. Ablations identify state-adaptive conditioning as a consistent contributor, while predictor-mediated subgoals and unrestricted conditions have environment-dependent benefits.

Abstract

Hierarchical World Models

Explore similar work

Learning Goal-Conditioned Subgoals through Frozen Hierarchical World Models

Hierarchical Planning with Latent World Models

WAM-RL: World-Action Model Reinforcement Learning with Reconstruction Rewards and Online Video SFT