PaperlayerView exact revision on arXiv

Source-linked AI summary

MHWA: Multi-timescale Hierarchical World-Action Model

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

arXiv:2610.14727v1cs.AI

TL;DR

Open-loop autoregressive world models can accumulate drift, while hierarchical designs may add complexity and fail to reliably outperform flat baselines. MHWA combines a strided high-level context channel with context-aware routing to low-level experts. On a 14-game offline Atari benchmark, it outperforms matched JOWA-70M and is numerically comparable to JOWA-150M with fewer parameters and training FLOPs.

  • Problem

    Autoregressive world models can compound prediction errors during open-loop rollouts, while hierarchical designs may be difficult to control and do not reliably outperform flat baselines.

    Source: §1 · pp. 1–2

  • Method

    MHWA combines a strided high-level manager for slow context inference with a context-aware low-level mixture of experts routed using both context and current token state.

    Source: §1 · p. 2, §3 · p. 4, §3.3 · p. 5

  • Results

    0.884 IQM HNS for MHWA (77M) versus 0.529 for matched JOWA-70M, and numerically comparable to JOWA-150M (0.833) [5], with 70.88% fewer training FLOPs per sequence.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The paper treats multi-timescale abstraction and context-driven routing as an engineering hypothesis for testing structural inductive bias against brute-force scaling on multi-game offline Atari.

    Source: §1 · p. 2

    The evidence is scoped primarily to discrete-action offline Atari with a pretrained tokenizer and fixed stride.

    Source: §1 · p. 2

Abstract

from arXiv · show

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.

1 Introduction

Open-loop autoregressive world models can accumulate prediction drift, while hierarchical designs may add complexity and fail to outperform flat baselines. MHWA decouples slow context inference from local dynamics and routes low-level computation through context-aware experts.

1 Introduction

  • Open-loop autoregressive rollouts can compound small prediction errors under distribution shift, causing imagined trajectories to diverge from real dynamics.Scaling model capacity is a prevalent response, but can require substantial GPU resources and long training cycles.Source: §1 · pp. 1–2
  • Hierarchical world models aim to separate slow context from fast local dynamics, but can be complex to control and do not reliably outperform flat baselines.The introduction identifies conceptual interference from moving high-level learning targets and structural interference from shared low-level parameters fitting heterogeneous behaviors.Source: §1 · p. 2
  • MHWA treats hierarchy as a design that must control interference, decoupling slow context inference from native-resolution dynamics and routing low-level computation through context-aware experts.The high-level manager uses a skip-step predictive objective, while its gradients are isolated from low-level losses; the worker gate uses both broadcast context and the current token state.Source: §1 · pp. 2–3
  • The paper frames multi-timescale abstraction and context-driven routing as an engineering hypothesis, with claims scoped to multi-game offline Atari.The benchmark is used to test whether structural inductive bias can improve rollout stability without relying solely on brute-force scaling.Source: §1 · p. 2
  • 0.884 IQM HNS for MHWA (77M) exceeds matched JOWA-70M (0.529) and is numerically comparable to JOWA-150M (0.833) [5], with 70.88% fewer training FLOPs per sequence.Evaluated across three seeds on a 14-game Atari subset using 10% subsampled offline data.Source: §1 · p. 3

2 Related Work

Related work situates MHWA among autoregressive world models, hierarchical temporal abstraction, and modular low-level parameterizations. Its distinction is an interference-aware integration of these mechanisms for tokenized multi-game offline world-action modeling.

Offline RL, world models, and compounding errors

  • Autoregressive sequence models can accumulate local errors during open-loop rollouts, and MHWA targets this failure mode with temporal abstraction in tokenized world-action modeling.The discussion places this work in offline RL, world models, and compounding errors.Source: §2 · p. 3

Temporal abstraction and conceptual interference

  • Hierarchical world models separate slow context from fast local dynamics, while MHWA uses a fixed strided manager, skip-step observation supervision, and stop-gradient-isolated context broadcast.This design avoids direct dependence of high-level context on low-level optimization signals.Source: §2 · p. 3

Modularity and structural interference

  • MHWA couples sparse expert routing to a high-level context channel, with the worker gate using both slow context and instantaneous token state.The paper frames its contribution as an interference-aware integration, not as standalone novelty in skip-step prediction, hierarchy, or MoE routing.Source: §2 · p. 3

3 The Multi-timescale Hierarchical World-Action Model (MHWA)

MHWA pairs a slowly updated high-level context channel with an atomic-resolution worker, using context-aware routing to address conceptual and structural interference. Its formulation retains tokenized world-action prediction and offline critic learning while adding objectives for high-level context and balanced expert use.

3.1 Problem Formulation

  • The offline formulation learns predictive transition and outcome models from a static trajectory dataset to support long-term return maximization.Observations and actions are interleaved as autoregressive tokens; observations are encoded with a VQ-VAE tokenizer reused from JOWA.Source: §3.1 · p. 4

3.2 Asymmetric Dual-Layer Architecture

  • The asymmetric hierarchy combines a strided High-Level Manager with a native-resolution Low-Level Worker whose sparse experts route using context and current state.The worker drives world-model and critic heads.Source: §3 · p. 4, §3.2 · p. 5

3.3 High-Level Manager: Context Inference

  • At stride k = 4, the manager infers slow-varying context from a downsampled interleaved sequence and broadcasts it for k atomic steps.A skip-step prediction objective trains the manager separately; low-level gradients are blocked from reaching it.Source: §3.3 · p. 5, §3.5 · p. 6

3.4 Low-Level Worker: Context-Aware Dynamic Routing

  • The worker replaces shared feed-forward layers with sparse experts, routing from both broadcast context and the instantaneous token state to address heterogeneous dynamics.Top-K routing selects 2 of 4 experts; an auxiliary load-balancing regularizer mitigates expert collapse.Source: §3.4 · pp. 5–6

3.5 Learning Objectives

  • MHWA retains autoregressive world-model and offline critic objectives while adding high-level representation and routing-regularization losses as structural constraints.The routing regularizer minimizes divergence between batch-wise expert selection frequency and a uniform distribution.Source: §3.5 · pp. 6–7
  • The base world-model loss predicts next observation tokens, rewards, and termination, while the offline critic uses distributional TD learning with a CQL penalty.The CQL term penalizes Q-values on out-of-distribution actions.Source: §3.5 · p. 6

4 Experiment

Across a 14-game offline Atari benchmark, MHWA outperforms a parameter-matched JOWA-70M baseline and is numerically comparable to JOWA-150M. Open-loop tests and ablations examine predictive stability and the contributions of temporal abstraction, routing, supervision, and expert representations.

4.1 Experimental Setup

  • Experiments use a 14-game rlu-atari subset, 10% randomly sampled transitions, and IQM and Median HNS aggregated following RLiable.MHWA has 77.27M total parameters, with 60.49M activated per forward pass; its training FLOPs per sequence are 70.88% lower than JOWA-150M.Source: §4.1 · p. 7

4.2 Overall Performance Comparison

  • 0.884 IQM HNS for MHWA-77M exceeds matched JOWA-70M (0.529) and is numerically comparable to published JOWA-150M (0.833).Table 2 compares seed-by-game scores; the training-computation comparison uses identical input shapes.Source: §4.2 · pp. 7–8, Table 2 · p. 8

4.3 Open-loop Dynamics Modeling

  • MHWA curbs open-loop error propagation over t=5–8 and remains stable, matching JOWA-150M while JOWA-70M deteriorates with rollout depth.The comparison measures MSE, PSNR, and SSIM; actions are fixed to dataset actions.Source: §4.3 · p. 8, Figure 2 · p. 9

4.4 Ablation Studies

  • 62.1% relative IQM drop occurs at k=1, while k=4 performs best and k=2 and k=8 remain competitive in the 44M, six-game setting.The ablation variants are controlled to approximately 44M parameters.Source: §4.4 · p. 9, Table 4 · p. 10
  • Removing state routing yields IQM 0.359 (36.3% drop), while dense FFN replacement drops IQM 20.2% at matched parameter count.Removing the manager pathway gives IQM 0.445 (21.1% drop), less than removing state routing.Source: §4.4 · p. 9
  • 0.564→0.503 (10.8%) without Lhigh and 0.564→0.444 (21.3%) without Laux; stop-gradient improves IQM with or without Lhigh.Four-expert Top-2 performs best among tested alternatives; k=4, r=2 yields IQM 0.558 versus 0.564 at default r=4.Source: §4.4 · p. 9, Table 3 · p. 10

4.5 Interpretability of Representation Learning

  • Assault and Phoenix show similar last-layer expert usage in full MHWA, but not without high-level supervision or temporal stride.The games share vertical shooting mechanics; the visualization uses PCA-projected expert activation frequencies.Source: §4.5 · p. 10, Figure 3 · p. 10

5 Discussion and Conclusion

MHWA combines a strided high-level manager with a context-aware MoE worker. On the 10% offline Atari benchmark, it outperforms matched JOWA-70M, is numerically comparable to JOWA-150M with fewer parameters and training FLOPs, and shows competitive adaptation to held-out games.

5 Discussion and Conclusion

  • MHWA outperforms a matched three-seed JOWA-70M baseline and is numerically comparable to JOWA-150M on the 10% offline Atari benchmark.The model has 77M parameters and uses 70.88% fewer training FLOPs per sequence.Source: §5 · p. 10
  • The model combines a strided, gradient-isolated high-level manager with a context-aware MoE low-level worker.The manager reduces how often global context is inferred, while the worker routes computation through specialized experts at native token resolution.Source: §5 · p. 10
  • MHWA shows competitive adaptation to held-out Atari games.The evaluation is primarily scoped to discrete-action offline Atari with a pretrained tokenizer and fixed stride; adaptive temporal abstraction, continuous control, and broader domains remain future work.Source: §5 · p. 10

A Atari Game Subset

The training set contains 14 Atari games, while ablations use a fixed six-game subset. The architecture pairs a strided high-level manager with a context-conditioned MoE worker.

A Atari Game Subset

  • The 14-game training set excludes Berzerk and games used for fine-tuning or transfer evaluation, and Table 5 lists difficulty labels and action-space sizes.The subset follows the RL Unplugged difficulty taxonomy and retains a mix of difficulty levels.Source: A Atari Game Subset · p. 14, Table 5 · p. 14
  • Ablations use a fixed six-game subset selected to balance difficulty levels and discrete action-space sizes.The games are Assault, Carnival, Centipede, NameThisGame, Phoenix, and SpaceInvaders.Source: A Atari Game Subset · p. 14

B Architecture and Parameterization Details

  • The high-level manager is a four-layer, eight-head, 512-dimensional Transformer whose output is detached, projected, and broadcast across the next stride window.Low-level losses train the Context Projector, router, and worker, while Lhigh trains the manager.Source: B Architecture and Parameterization… · p. 14
  • The low-level worker uses four MoE layers with four experts and Top-2 routing based on the current token state and broadcast high-level context.All selected experts are computed, without token dropping or capacity truncation.Source: B Architecture and Parameterization… · p. 14
  • The full MHWA model contains 77.27M parameters, with approximately 60.49M activated per forward pass including the tokenizer.Source: B Architecture and Parameterization… · p. 14

B.1 Training Computation

Training computation is counted as forward-plus-backward FLOPs at a matched input shape, including the model, critic, auxiliary losses, and backward computation.

B.1 Training Computation

  • Forward-plus-backward FLOPs are counted at a matched input shape of eight frames and 296 tokens per sequence, with batch size 256.Source: §B.1 · p. 14, Table 6 · p. 15
  • The computation count includes the complete forward graph, world-model, critic, auxiliary losses, and backward computation.Source: §B.1 · p. 14
  • For MHWA, the count includes both hierarchy levels, all four MoE layers, load balancing, and high-level observation prediction.MoE computation is counted for activated experts.Source: §B.1 · p. 15

C A Simple Error-Propagation View

The analysis gives a conditional one-step open-loop error comparison between a strided-context model and a flat model under fixed actions. It identifies when MHWA’s bound is no larger and states that the comparison extends beyond discrete-action Atari under its error-term formulation.

C A Simple Error-Propagation View

  • The comparison is framed for open-loop state prediction under a fixed action sequence and is intended to isolate MHWA’s core mechanism, not provide a general performance guarantee.Source: C A Simple Error-Propagation View · p. 15

Assumptions

Derivation

Local one-step modeling and interface error

  • The error accounting includes local dynamics approximation error and combines action-representation and latent/interface approximation errors into η.The text notes that interface terms may be smaller when continuous actions or latents require no discretization interface.Source: C A Simple Error-Propagation View · p. 15

Comparison to a flat model

  • MHWA’s one-step open-loop error bound is no larger than the flat model’s when βH + jρ ≤ βF, with j = k − 1 giving the worst case within a stride window.The comparison matches local one-step error to isolate the effect of context inference.Source: C A Simple Error-Propagation View · p. 16

Implications for transfer across settings

  • The comparison does not depend on discrete actions, tokenized latents, or offline versus online training in this narrow architectural sense.It depends on the achieved error terms βH, βF, η, and ρ.Source: C A Simple Error-Propagation View · p. 16

D Per-seed Results on the 14-game Benchmark

Across three seeds, MHWA's 14-game seed-by-game IQM HNS is 0.884; the matched JOWA-70M result is 0.529. The section reports per-seed results for both models.

D Per-seed Results on the 14-game Benchmark

D.1 Matched JOWA-70M Results

  • 0.529 is the matched JOWA-70M seed-by-game IQM, with a 95% stratified-bootstrap CI of [0.459, 0.602].Table 8 reports per-seed results and game-level cross-seed variation.Source: Table 8 · p. 17, §D.1 · p. 17

E Held-out Game Fine-tuning

The held-out-game fine-tuning evaluation reports a three-seed aggregate and uneven game-level transfer results. The following ablation material describes per-game score reporting and controls for the hierarchical context pathway.

E Held-out Game Fine-tuning

  • Held-out games were fine-tuned separately with 5k expert-level transitions, using matched exposure of about 1.60M samples for MHWA and JOWA.The protocol uses batch size 1024 for 1562 MHWA updates to match JOWA's batch-size-32, 50k-update exposure.Source: E Held-out Game Fine-tuning · p. 18
  • 0.643±0.007 is MHWA's three-seed IQM DNS, with competitive results on Pong and YarsRevenge, closer results on MsPacman and Gravitar, and weaker results on Robotank.A single additional batch-size-32 run reached IQM DNS 0.668 but is not the main comparison because it has only one seed.Source: E Held-out Game Fine-tuning · p. 18

F Per-game Ablation Results

  • Table 10 reports 16-episode average raw scores per game and variant, plus Median and IQM HNS over the six-game subset.All variants use the 44M controlled setting.Source: Table 10 · p. 18

F.1 Hierarchical Context Controls

  • The factorial controls retain the manager, Context Projector, context-conditioned router, and stride-k broadcasting, while varying high-level supervision and stop-gradient.MoE-only removes the full manager pathway and widens the low-level experts to match model size.Source: §F.1 · p. 18
  • Trajectory-conditioned hidden state improves IQM over zero-input and Gaussian-noise controls, although its Median HNS is close to the zero-input variant.The IQM reflects a larger separation across the six-game score distribution.Source: §F.1 · pp. 18–19

F.2 Expert Number and Routing Sparsity

The expert study reports how expert count and routing sparsity relate to aggregate performance and routing neighborhoods. The context-refresh study compares schedules while holding the manager stride at four.

F.2 Expert Number and Routing Sparsity

  • Four-expert Top-1 retains 96% of default IQM, eight-expert alternatives retain 87–91%, and the two-expert configuration has the largest drop.Table 12 reports IQM HNS in the six-game ablation benchmark.Source: Table 12 · p. 19, §F.2 · p. 19
  • Four-expert Top-1 shifts Centipede–SpaceInvaders from the default's 15th-ranked PCA-distance pair to sixth, while the Assault–Phoenix pair is closest under both Top-2 configurations.The two-expert normalized usage vector has one effective degree of freedom.Source: Table 13 · p. 19, §F.2 · p. 19

F.3 Manager Stride and Context-Refresh Interval

  • The manager stride k sets input spacing and high-level prediction scale, while refresh interval r sets how often projected context is refreshed for routing; the default is k = r = 4.Table 14 compares refresh schedules at fixed manager stride k = 4.Source: §F.3 · p. 19, Table 14 · p. 20
  • The refresh study tests inference-only r changes with the default checkpoint and a retrained k = 4, r = 2 schedule.The inference-only variants keep the manager's k = 4 input spacing and maximum context length; retraining lets the router and worker adapt.Source: §F.3 · p. 20
  • After retraining with r = 2, aggregate performance is comparable to default, with higher Median and slightly lower IQM; returns improve on Carnival and Centipede but fall on four games.The changed schedule redistributes game-level performance while preserving the manager's learned temporal scale.Source: Table 15 · p. 20, §F.3 · p. 20

G Long-horizon Open-loop Diagnostics

Long-horizon open-loop diagnostics find smoother degradation for MHWA than for the smaller monolithic baseline, while noting eventual manifold departure. The responsible-research discussion limits experiments to benchmark games and describes asset use and safeguards.

G Long-horizon Open-loop Diagnostics

  • From t = 5 to t = 54, MHWA shows smoother degradation than the smaller monolithic baseline and remains broadly comparable to the 150M baseline, but does not eliminate eventual manifold departure.The diagnostic assesses smoothness of error accumulation, not indefinitely stable generation; the visible JOWA-150M fluctuation is mainly driven by DemonAttack at t = 17–20.Source: G Long-horizon Open-loop Diagnostics · p. 20, Figure 4 · p. 21

H Responsible Research and Existing Assets

  • Experiments are limited to Atari benchmark games, with no deployment, human-subject or personal data, or safety-critical decision making; future real-world control applications warrant domain-specific safeguards.The paper characterizes immediate societal risk as limited and identifies inaccurate rollouts, distribution shift, and unsafe planning as potential indirect risks.Source: H Responsible Research and Existing… · p. 21
  • The work uses public RL Unplugged Atari replay data and a pretrained JOWA tokenizer interface, introduces no new dataset, and directs users to follow upstream asset licenses and terms.The code release is intended to reproduce Atari benchmark experiments, and the models operate in Atari-like simulated environments.Source: H Responsible Research and Existing… · p. 21
Loading 2610.14727v1…