MHWA: Multi-timescale Hierarchical World-Action Model

# MHWA: Multi-timescaleHierarchical World-Action Model
Source: [https://arxiv.org/html/2610.14727](https://arxiv.org/html/2610.14727)
## MHWA: Multi\-timescale Hierarchical World\-Action Model

Pengcheng Pan††thanks:These authors contributed equally\.Affiliation:Laboratory of Brain Atlas and Brain\-inspired Intelligence,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaAffiliation:School of Artificial Intelligence, University of Chinese Academy of Sciences, Beijing, ChinaAffiliation:State Key Laboratory of Brain Cognition and Brain\-inspired Intelligence Technology,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaEmail:[panpengcheng19@mails\.ucas\.ac\.cn](mailto:[email protected])Guoqing Ma11footnotemark:1Affiliation:Laboratory of Brain Atlas and Brain\-inspired Intelligence,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaAffiliation:School of Future Technology, University of Chinese Academy of Sciences, Beijing, ChinaAffiliation:State Key Laboratory of Brain Cognition and Brain\-inspired Intelligence Technology,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaEmail:[maguoqing22@mails\.ucas\.ac\.cn](mailto:[email protected])Yuhan ZhangAffiliation:Laboratory of Brain Atlas and Brain\-inspired Intelligence,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaAffiliation:School of Artificial Intelligence, University of Chinese Academy of Sciences, Beijing, ChinaAffiliation:State Key Laboratory of Brain Cognition and Brain\-inspired Intelligence Technology,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaEmail:[zhangyuhan2022@ia\.ac\.cn](mailto:[email protected])Yang ChenAffiliation:Laboratory of Brain Atlas and Brain\-inspired Intelligence,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaAffiliation:State Key Laboratory of Brain Cognition and Brain\-inspired Intelligence Technology,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaEmail:[yang\.chen@ia\.ac\.cn](mailto:[email protected])Yichen LiuAffiliation:Laboratory of Brain Atlas and Brain\-inspired Intelligence,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaAffiliation:School of Future Technology, University of Chinese Academy of Sciences, Beijing, ChinaAffiliation:State Key Laboratory of Brain Cognition and Brain\-inspired Intelligence Technology,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaEmail:[liuyichen2024@ia\.ac\.cn](mailto:[email protected])Ziheng LiAffiliation:Laboratory of Brain Atlas and Brain\-inspired Intelligence,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaAffiliation:School of Artificial Intelligence, University of Chinese Academy of Sciences, Beijing, ChinaAffiliation:State Key Laboratory of Brain Cognition and Brain\-inspired Intelligence Technology,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaEmail:[liziheng2024@ia\.ac\.cn](mailto:[email protected])Shan Yu††thanks:Corresponding author:shan\.yu@nlpr\.ia\.ac\.cnAffiliation:Laboratory of Brain Atlas and Brain\-inspired Intelligence,Institute of Automation, Chinese Academy of Sciences, Beijing, ChinaAffiliation:State Key Laboratory of Brain Cognition and Brain\-inspired Intelligence Technology,Institute of Automation, 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\.

## 1Introduction

Model\-Based Reinforcement Learning \(MBRL\) offers a promising avenue for sample efficient learning by enabling agents to plan within a learned World Model \(WM\)\[[1](https://arxiv.org/html/2610.14727#bib.bib4),[2](https://arxiv.org/html/2610.14727#bib.bib5)\]\. However, maintaining predictive consistency under multi\-step open loop rollouts remains a fundamental challenge\. State\-of\-the\-art approaches predominantly employ autoregressive Transformer architectures\[[3](https://arxiv.org/html/2610.14727#bib.bib7),[4](https://arxiv.org/html/2610.14727#bib.bib6),[5](https://arxiv.org/html/2610.14727#bib.bib26)\]\. While effective at short range prediction, autoregressive rollouts can exhibit increasing drift with rollout depth: once the model conditions on its own imperfect predictions, small local errors may compound under distribution shift, causing imagined trajectories to deviate from the real dynamics\[[6](https://arxiv.org/html/2610.14727#bib.bib9),[7](https://arxiv.org/html/2610.14727#bib.bib8)\]\.

A prevalent response to this challenge is to scale model capacity, since larger models can better absorb heterogeneity and noise\. A prominent example is the Jointly\-Optimized World\-Action \(JOWA\) model, which scales up to 150 million parameters to achieve strong performance in multi\-game offline reinforcement learning\[[5](https://arxiv.org/html/2610.14727#bib.bib26)\]\. While effective, such brute force scaling can be computationally costly, requiring substantial GPU resources and long training cycles, which raises the barrier to iteration\.

Hierarchical Reinforcement Learning \(HRL\) provides an alternative design axis by introducing temporal abstraction\[[8](https://arxiv.org/html/2610.14727#bib.bib15),[9](https://arxiv.org/html/2610.14727#bib.bib12)\]\. This motivates the emerging line of Hierarchical Model\-Based RL \(HMBRL\), or Hierarchical World Models \(HWM\)\[[10](https://arxiv.org/html/2610.14727#bib.bib1)\], which aim to separate slow varying context from fast local dynamics and thereby reduce the effective horizon over which global consistency must be maintained\. In principle, such multi\-timescale structure could improve longer range credit assignment and planning while retaining the sample efficiency benefits of world models\.

However, the HWM landscape remains relatively sparse compared to model free approaches; as noted by[Schiewer et al\. \[10\]](https://arxiv.org/html/2610.14727#bib.bib1), combining hierarchy and world modeling often yields systems that are “highly complex and difficult to control,” and the added hierarchy does not reliably outperform flat baselines\. A key reason is that existing designs face two coupled forms of interference\.

Conceptual \(semantic\) interference\.In many goal conditioned hierarchical formulations, higher\-level decisions \(e\.g\., subgoals\) are evaluated through the evolving competence of the lower\-level goal\-reaching policy\. As multiple levels learn simultaneously, the semantic meaning and value of a high\-level decision can drift over training, inducing a*moving target*learning signal and exacerbating non\-stationarity\[[11](https://arxiv.org/html/2610.14727#bib.bib10),[10](https://arxiv.org/html/2610.14727#bib.bib1)\]\. This instability is especially pronounced when the abstract interface is high\-dimensional or when the hierarchy relies on objectives whose targets implicitly change as the lower level improves\.

Structural interference\.A second issue is architectural: many hierarchical agents rely on a single universal low\-level network conditioned on goals, forcing one set of parameters to fit heterogeneous behaviors and dynamics regimes\. Such heavy parameter sharing can entangle incompatible modes and lead to gradient competition across subtasks\[[11](https://arxiv.org/html/2610.14727#bib.bib10)\]\. This motivates modular parameterizations that can specialize while remaining scalable\.

Together, these observations suggest that without explicit decoupling and flexible coordination, simply adding hierarchy can introduce complexity that becomes a burden rather than a benefit\.

To address these challenges, we propose theMulti\-timescaleHierarchicalWorldAction model \(MHWA\), a problem structured hierarchical world action model for reducing compounding errors\. MHWA treats hierarchy not as a standalone remedy, but as a design that must explicitly control the interference it introduces: slow context inference is decoupled from native resolution local dynamics, and low\-level computation is dynamically routed through context aware experts\.

Scope\.We treat multi\-timescale abstraction and context driven routing as an engineering hypothesis for reducing open loop drift under heterogeneous multi\-game dynamics, and provide a supporting error bound intuition in Appendix[C](https://arxiv.org/html/2610.14727#A3)\. Our claims are scoped to multi\-game offline Atari, which we use as a controlled benchmark for testing whether structural inductive bias can improve rollout stability without relying solely on brute force scaling\.

Our contributions are summarized as follows:

An interference aware hierarchy for reducing compounding errors\.We propose MHWA, a problem structured hierarchical world action model that treats hierarchy not as a standalone remedy, but as a mechanism that must be carefully decoupled\. Our central hypothesis is that temporal abstraction can mitigate compounding errors in autoregressive world models when conceptual interference between hierarchy levels and structural interference across dynamics modes are explicitly addressed\. MHWA implements this idea by coupling slow context inference with native resolution local dynamics through a strided High\-Level Manager and a context routed Low\-Level Worker\.

Isolated slow context learning for conceptual interference\.The High\-Level Manager learns a slow varying latent context using a skip step predictive objective \(ℒh​i​g​h\\mathcal\{L\}\_\{high\}\), while gradients from low\-level losses are blocked before reaching the manager\. This prevents the high\-level context from being driven by low\-level noise or moving target optimization signals\.

Context aware expert routing for structural interference\.The Low\-Level Worker replaces a monolithic feed forward parameterization with a sparsely activated MoE\. Its gate conditions jointly on the broadcast high\-level context and the instantaneous token state, enabling expert selection to be both mode aware and locally reactive under heterogeneous multi\-game dynamics\.

Structure versus scale under a controlled benchmark\.On a 14\-game Atari subset with 10% subsampled offline data, MHWA \(77M\) achieves an IQM HNS of 0\.884 across three training seeds, outperforming JOWA\-70M \(0\.529\) trained on the same data with the same batch setting\. Its performance is numerically comparable to the published JOWA\-150M result \(0\.833\)\[[5](https://arxiv.org/html/2610.14727#bib.bib26)\], with nearly half the parameters and 70\.88% fewer forward\+backward FLOPs per training sequence\.

## 2Related Work

Offline RL, world models, and compounding errors\.Offline RL learns policies from fixed datasets, where distribution shift between the behavior data and the learned policy is a central challenge\[[12](https://arxiv.org/html/2610.14727#bib.bib35),[13](https://arxiv.org/html/2610.14727#bib.bib36)\]\. Model\-free methods often mitigate this issue through policy constraints or conservative value estimation\[[14](https://arxiv.org/html/2610.14727#bib.bib37),[15](https://arxiv.org/html/2610.14727#bib.bib25),[16](https://arxiv.org/html/2610.14727#bib.bib42)\], while model\-based methods learn dynamics models for planning or policy improvement\[[17](https://arxiv.org/html/2610.14727#bib.bib38),[18](https://arxiv.org/html/2610.14727#bib.bib39)\]\. World models provide a natural interface for model\-based decision making\[[19](https://arxiv.org/html/2610.14727#bib.bib41),[20](https://arxiv.org/html/2610.14727#bib.bib14)\]\. Recent sequence modeling approaches, including Decision Transformer, IRIS, and JOWA, cast RL as autoregressive prediction over observation\-action histories\[[21](https://arxiv.org/html/2610.14727#bib.bib43),[4](https://arxiv.org/html/2610.14727#bib.bib6),[5](https://arxiv.org/html/2610.14727#bib.bib26)\]\. However, when such models roll out open loop and condition on their own predictions, local errors can compound over time\. MHWA targets this failure mode by introducing temporal abstraction into tokenized world action modeling\.

Temporal abstraction and conceptual interference\.Hierarchical RL decomposes behavior across temporal scales through options, managers, or subgoals\[[8](https://arxiv.org/html/2610.14727#bib.bib15),[22](https://arxiv.org/html/2610.14727#bib.bib40),[23](https://arxiv.org/html/2610.14727#bib.bib16),[24](https://arxiv.org/html/2610.14727#bib.bib13),[25](https://arxiv.org/html/2610.14727#bib.bib11),[26](https://arxiv.org/html/2610.14727#bib.bib22)\]\. Hierarchical world models similarly aim to separate slow varying context from fast local dynamics\[[10](https://arxiv.org/html/2610.14727#bib.bib1),[27](https://arxiv.org/html/2610.14727#bib.bib2),[28](https://arxiv.org/html/2610.14727#bib.bib3)\]\. This separation is appealing for reducing effective rollout horizon, but it can also introduce conceptual interference: high\-level decisions may be evaluated through a lower\-level policy whose competence and semantics evolve during training\[[11](https://arxiv.org/html/2610.14727#bib.bib10),[10](https://arxiv.org/html/2610.14727#bib.bib1)\]\. THICK\[[29](https://arxiv.org/html/2610.14727#bib.bib18)\]learns temporal abstraction from sparse latent context changes and predicts future latent scenarios associated with those changes\. In contrast, MHWA uses a fixed strided manager with skip step observation token supervision and stop\-gradient\-isolated context broadcast, avoiding direct dependence of the high\-level context on low\-level optimization signals\. Beyond temporal abstraction, perturbation\-based local learning has also been explored to improve training scalability and convergence efficiency in spiking neural networks\[[30](https://arxiv.org/html/2610.14727#bib.bib21)\]\.

Modularity and structural interference\.A second challenge is structural interference, where heterogeneous behaviors or dynamics regimes are forced through a heavily shared low\-level parameterization\[[11](https://arxiv.org/html/2610.14727#bib.bib10)\]\. Sparse Mixture\-of\-Experts models address a related problem by routing inputs to a subset of specialized parameters\[[31](https://arxiv.org/html/2610.14727#bib.bib24),[32](https://arxiv.org/html/2610.14727#bib.bib17),[33](https://arxiv.org/html/2610.14727#bib.bib44)\]\. MoE world models support modular dynamics modeling for multi\-task and multi\-agent learning and planning\[[34](https://arxiv.org/html/2610.14727#bib.bib19),[35](https://arxiv.org/html/2610.14727#bib.bib20)\]\. MHWA differs by coupling sparse routing with an explicit high\-level context channel: the worker gate conditions jointly on the broadcast slow context and the instantaneous token state, enabling expert selection to be both mode aware and locally reactive\. Thus, we do not claim skip step prediction, hierarchy, or MoE routing as standalone novelties; our contribution is an interference aware integration of these mechanisms for reducing drift in tokenized multi\-game offline world\-action modeling\.

Figure 1:The architecture of MHWA\.It is illustrated at a context\-update timestept=1\+m⋅kt=1\+m\\cdot k\. The model processes an interleaved sequence of observation \(oto\_\{t\}\) and action \(ata\_\{t\}\) tokens\.
\(Top\)At this specific strided step, the High\-Level Manager consumes the downsampled history to infer a slowly varying global contextzc​t​xz^\{ctx\}, supervised by a skip\-step prediction headpoh​i​g​hp^\{high\}\_\{o\}that forecasts the observation at the next coarse stepo^t\+k\\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 contextzc​t​xz^\{ctx\}to dynamically route computation to a sparse subset of experts \(MoE\)\. The worker’s output drives the world model heads \(po,pr,pdp\_\{o\},p\_\{r\},p\_\{d\}\) and the critic head \(hQh\_\{Q\}\)\.
## 3The Multi\-timescale Hierarchical World\-Action Model \(MHWA\)

We propose theMulti\-timescaleHierarchicalWorldAction model \(MHWA\), an asymmetric dual\-layer architecture for reducing compounding errors in tokenized autoregressive world models\. MHWA targets two interference modes: conceptual interference between hierarchy levels, and structural interference across heterogeneous dynamics modes\. It addresses the former by isolating a slow high\-level context through gradient detachment, and the latter by using context\-aware sparse routing in the low\-level dynamics model\.

### 3\.1Problem Formulation

We consider offline reinforcement learning as a Markov Decision Process \(MDP\)ℳ=⟨𝒮,𝒜,𝒫,ℛ,γ⟩\\mathcal\{M\}=\\langle\\mathcal\{S\},\\mathcal\{A\},\\mathcal\{P\},\\mathcal\{R\},\\gamma\\rangle\. In visual RL, the agent observes high\-dimensional observationsot∈𝒪o\_\{t\}\\in\\mathcal\{O\}rather than ground\-truth states, and we treat the stacked observation history as the state representation\. Given a static dataset𝒟=\{τi\}i=1N\\mathcal\{D\}=\\\{\\tau\_\{i\}\\\}\_\{i=1\}^\{N\}, whereτ=\(o1,a1,r1,…,oT,aT,rT\)\\tau=\(o\_\{1\},a\_\{1\},r\_\{1\},\\dots,o\_\{T\},a\_\{T\},r\_\{T\}\), the goal is to learn predictive transition and outcome models that support long\-term return maximization\.

Following tokenized world\-action models, each observation is encoded intoKKdiscrete visual tokens using a VQ\-VAE tokenizer, reusing the pretrained tokenizer from JOWA\. These visual tokens are interleaved with executed action tokens to form a single autoregressive sequence\. Token, positional, and task embeddings are combined before both Transformers, retaining JOWA’s explicit task conditioning\. An autoregressive predictor processes this sequence and produces embeddings consumed by four heads: \(i\) a next\-observation token headpop\_\{o\}, \(ii\) a reward headprp\_\{r\}, \(iii\) a termination headpdp\_\{d\}, and \(iv\) a Q\-value headhQh\_\{Q\}trained with offline TD learning\. In MHWA, this predictor is implemented by the Low\-Level Worker modulated by the High\-Level Manager, as described in Sec\.[3\.2](https://arxiv.org/html/2610.14727#S3.SS2)\. We refer to the components producing\(po,pr,pd\)\(p\_\{o\},p\_\{r\},p\_\{d\}\)as the*world\-part*and those producinghQh\_\{Q\}as the*action\-part*; both world\-modeling and TD/CQL critic losses back\-propagate through the autoregressive predictor\.

Standard autoregressive world models treat trajectories as uniform token sequences, where prediction errors may accumulate with rollout depth\. MHWA instead introduces temporal abstraction to infer global context less frequently, while preserving local dynamics prediction at the native token resolution\.

### 3\.2Asymmetric Dual\-Layer Architecture

MHWA is organized around the two interference modes above\. To mitigate*conceptual interference*, the high\-level representation should capture slow\-varying context without being dominated by low\-level reconstruction noise\. To mitigate*structural interference*, the low\-level model should avoid fitting heterogeneous dynamics with a single shared feed\-forward parameterization\. Accordingly, MHWA retains the tokenized world\-action interface and standard prediction heads, while replacing the monolithic autoregressive predictor with an asymmetric two\-level hierarchy \(Figure[1](https://arxiv.org/html/2610.14727#S2.F1)\): \(i\) a High\-Level Manager that infers slowly varying context at a reduced update frequency, and \(ii\) a Low\-Level Worker implemented as an autoregressive Transformer whose feed\-forward sublayers are replaced by sparse Mixture\-of\-Experts \(MoE\) layers\. The worker output drives the world\-model heads\(po,pr,pd\)\(p\_\{o\},p\_\{r\},p\_\{d\}\)and the critic head\(hQ\)\(h\_\{Q\}\)\. The manager provides an isolated slow context channel, while the worker combines this context with the instantaneous state for dynamic expert routing\. This couples temporal abstraction with modular local prediction, directly targeting conceptual and structural interference in heterogeneous multi\-game dynamics\.

### 3\.3High\-Level Manager: Context Inference

The High\-Level Manager targets conceptual interference by learning a slow context representation on a strided temporal scale\. Letxt∈ℝdx\_\{t\}\\in\\mathbb\{R\}^\{d\}denote the input embedding at atomic steptt, wherextx\_\{t\}represents either an observation \(oto\_\{t\}\) or an action \(ata\_\{t\}\) embedding because the input sequence interleaves observations and actions\. We introduce a high\-level indexmm, where each step corresponds to atomic timestampt=1\+m⋅kt=1\+m\\cdot kwith stride factorkk; we setk=4k=4\. At coarse stepmm, the manager consumes the strided token sequence𝐗s​l​o​w\(m\)=\(x1,x1\+k,…,x1\+m⋅k\)\\mathbf\{X\}^\{\(m\)\}\_\{slow\}=\(x\_\{1\},x\_\{1\+k\},\\dots,x\_\{1\+m\\cdot k\}\), akk\-fold temporal downsampling of the atomic sequence\. A causal Transformer encoder\[[36](https://arxiv.org/html/2610.14727#bib.bib23)\]computes:

hmh​i​g​h=TransformerH​i​g​h​\(𝐗s​l​o​w\(m\)\)h^\{high\}\_\{m\}=\\text\{Transformer\}\_\{High\}\(\\mathbf\{X\}^\{\(m\)\}\_\{slow\}\)\(1\)Fromhmh​i​g​hh^\{high\}\_\{m\}, a trainable Context Projector produces a continuous context vectorzmc​t​x=Projector⁡\(stopgrad⁡\(hmh​i​g​h\)\)z^\{ctx\}\_\{m\}=\\mathrm\{Projector\}\(\\mathrm\{stopgrad\}\(h^\{high\}\_\{m\}\)\), which parametrizes the slow\-varying context of the current dynamics regime rather than acting as a discrete subgoal\.

Context Broadcasting and Gradient Isolation\.We broadcastzmc​t​xz^\{ctx\}\_\{m\}to the Low\-Level Worker and hold it constant for the subsequent window ofkkatomic steps\. Stop\-gradient is placed between the manager and the Context Projector, leaving the projector, router, and worker trainable under the low\-level objectives\. Thus, the manager is trained through its own auxiliary objective instead of low\-level gradients, preventing short\-range prediction noise from overwriting the high\-level context\. This is the main mechanism for controlling conceptual interference between hierarchy levels\.

Implementation of Stride\.Temporal abstraction is implemented as a fixed token\-level stride on the interleaved sequence\. This matches the tokenized world\-action interface, where visual and action tokens are flattened into one stream\. The stride shortens the manager sequence for efficient slow\-context inference, while the Low\-Level Worker still handles fine\-grained native\-resolution prediction\. Although striding discards some tokens, VQ\-VAE visual codes contain sufficient redundancy for coarse context inference\.

Distinction from Static Execution\.The stridekksets the manager’s temporal sampling and prediction scale; the context\-refresh intervalrris set tokkby default\. Unlike static temporal abstractions\[[10](https://arxiv.org/html/2610.14727#bib.bib1)\]that lock low\-level policies or dynamics for fixed intervals, MHWA keeps the low\-level routing mechanism operating at the atomic timescalett\. The worker therefore remains locally reactive by conditioning on both the broadcast context and the instantaneous state\.

### 3\.4Low\-Level Worker: Context\-Aware Dynamic Routing

The Low\-Level Worker targets structural interference\. In heterogeneous multi\-game dynamics, a monolithic low\-level FFN must model diverse visual patterns, action effects, and transition modes with shared parameters, which can entangle incompatible behaviors\. MHWA replaces these FFN sublayers with sparse MoE layers\[[31](https://arxiv.org/html/2610.14727#bib.bib24),[32](https://arxiv.org/html/2610.14727#bib.bib17)\]\. The routing mechanism conditions on both the broadcast high\-level context and the instantaneous token state, allowing expert selection to reflect slow dynamics modes while remaining locally reactive\.

Letxtx\_\{t\}denote the intermediate token representation in the residual stream at token positiontt, and letzmc​t​xz^\{ctx\}\_\{m\}be the broadcast high\-level context form=⌊\(t−1\)/k⌋m=\\lfloor\(t\-1\)/k\\rfloor\. A gating projectionWg∈ℝ2​d×NW\_\{g\}\\in\\mathbb\{R\}^\{2d\\times N\}maps the concatenated input to routing logits:

lt=\(xt⊕zmc​t​x\)​Wg,l\_\{t\}=\(x\_\{t\}\\oplus z^\{ctx\}\_\{m\}\)W\_\{g\},\(2\)where⊕\\oplusdenotes concatenation andNNis the number of experts\{Ej\}j=1N\\\{E\_\{j\}\\\}\_\{j=1\}^\{N\}\. We apply Top\-KKsparse routing by selecting𝒯t=TopK⁡\(lt,Kroute\)\\mathcal\{T\}\_\{t\}=\\mathrm\{TopK\}\(l\_\{t\},K\_\{\\mathrm\{route\}\}\)and computing masked Softmax weights\. We setN=4N=4;Kroute=2K\_\{\\mathrm\{route\}\}=2:

wt,j=exp⁡\(lt,j\)∑r∈𝒯texp⁡\(lt,r\),j∈𝒯t;wt,j=0,j∉𝒯t\.w\_\{t,j\}=\\frac\{\\exp\(l\_\{t,j\}\)\}\{\\sum\_\{r\\in\\mathcal\{T\}\_\{t\}\}\\exp\(l\_\{t,r\}\)\},\\;j\\in\\mathcal\{T\}\_\{t\};\\qquad w\_\{t,j\}=0,\\;j\\notin\\mathcal\{T\}\_\{t\}\.\(3\)The MoE output is:

yt=xt\+∑j∈𝒯twt,j​Ej​\(xt\)\.y\_\{t\}=x\_\{t\}\+\\sum\_\{j\\in\\mathcal\{T\}\_\{t\}\}w\_\{t,j\}\\,E\_\{j\}\(x\_\{t\}\)\.\(4\)Conditioning the gate on bothzmc​t​xz^\{ctx\}\_\{m\}andxtx\_\{t\}makes routing both mode\-aware and locally reactive: the context term encourages specialization across slow\-varying dynamics regimes, while the state term preserves sensitivity to instantaneous visual and action\-dependent details\. This addresses structural interference by replacing a single shared low\-level transformation with sparse, context\-conditioned expert computation\. We use no capacity factor, token dropping, or truncation; all selected tokens are computed, and expert collapse is mitigated with an auxiliary load\-balancing regularizer\. We provide complete architectural and routing implementation details in Appendix[B](https://arxiv.org/html/2610.14727#A2)\.

### 3\.5Learning Objectives

We train the MHWA architecture jointly but with cross\-level stop\-gradient\. Accordingly, the total objective function comprisesfour distinct components\. The first two \(ℒb​a​s​e,ℒc​r​i​t​i​c\\mathcal\{L\}\_\{base\},\\mathcal\{L\}\_\{critic\}\) are adopted from baseline JOWA\[[5](https://arxiv.org/html/2610.14727#bib.bib26)\]\. The latter two \(ℒh​i​g​h,ℒa​u​x\\mathcal\{L\}\_\{high\},\\mathcal\{L\}\_\{aux\}\) are novel structural constraints introduced by MHWA, specifically designed to enforce hierarchical representation learning and dynamic expert specialization\.

1\. Base World Model Objectives \(ℒb​a​s​e\\mathcal\{L\}\_\{base\}\):Following standard autoregressive approaches, the Low\-Level Worker transforms the input sequence into a final hidden representationhth\_\{t\}\. This representation is then used to predict the next discrete observation tokenot\+1o\_\{t\+1\}, rewardrtr\_\{t\}, and termination signaldtd\_\{t\}via standard Cross\-Entropy loss\. This ensures the model accurately simulates the environment’s local dynamics:

ℒb​a​s​e=−∑t\[\\displaystyle\\mathcal\{L\}\_\{base\}=\-\\sum\_\{t\}\\Big\[log⁡po​\(ot\+1\|ht\)⏟Observation\+log⁡pr​\(rt\|ht\)⏟Reward\+log⁡pd​\(dt\|ht\)⏟Termination\]\\displaystyle\\underbrace\{\\log p\_\{o\}\(o\_\{t\+1\}\|h\_\{t\}\)\}\_\{\\text\{Observation\}\}\+\\underbrace\{\\log p\_\{r\}\(r\_\{t\}\|h\_\{t\}\)\}\_\{\\text\{Reward\}\}\+\\underbrace\{\\log p\_\{d\}\(d\_\{t\}\|h\_\{t\}\)\}\_\{\\text\{Termination\}\}\\Big\]\(5\)
2\. Offline Critic Objectives \(ℒc​r​i​t​i​c\\mathcal\{L\}\_\{critic\}\):For policy evaluation, we employ a distributional Q\-function learned via Temporal Difference \(TD\) learning, augmented with a Conservative Q\-Learning \(CQL\) penalty to mitigate value overestimation in offline regimes\[[15](https://arxiv.org/html/2610.14727#bib.bib25)\]\. The critic loss is defined as:

ℒc​r​i​t​i​c=ℒT​D​\(θQ\)\+αc​q​l​ℒC​Q​L​\(θQ\)\\mathcal\{L\}\_\{critic\}=\\mathcal\{L\}\_\{TD\}\(\\theta\_\{Q\}\)\+\\alpha\_\{cql\}\\mathcal\{L\}\_\{CQL\}\(\\theta\_\{Q\}\)\(6\)whereℒT​D\\mathcal\{L\}\_\{TD\}minimizes the distributional Bellman error \(using C51\) andℒC​Q​L\\mathcal\{L\}\_\{CQL\}penalizes Q\-values on out\-of\-distribution actions\.

3\. High\-Level Representation Learning \(ℒh​i​g​h\\mathcal\{L\}\_\{high\}\):A core challenge in hierarchical modeling is preventing the high\-level representation from collapsing into a trivial identity mapping of low\-level features or overfitting to transient, hard\-to\-predict details\. To encourage the High\-Level Manager to distill slow\-varying and semantically meaningful context, we introduce a skip\-step predictive objective\. Concretely, an MLP projection headph​i​g​hp\_\{high\}predicts the token at the next sampled position using only the corresponding high\-level hidden statehmh​i​g​hh^\{high\}\_\{m\}\. This skip\-step gap acts as an information bottleneck, discouraging reliance on ephemeral noise and promoting representations that remain predictive across multiple token steps:

ℒh​i​g​h=−∑mlogpoh​i​g​h\(o\(m\+1\)​k∣hmh​i​g​h\)\\mathcal\{L\}\_\{high\}=\-\\sum\_\{m\}\\log p\_\{o\}^\{high\}\(o\_\{\(m\+1\)k\}\\mid h^\{high\}\_\{m\}\)\(7\)
4\. Routing Regularization \(ℒa​u​x\\mathcal\{L\}\_\{aux\}\):To prevent expert collapse and ensure balanced resource utilization, we minimize the Kullback\-Leibler \(KL\) divergence between the batch\-wise expert selection frequencyp^\\hat\{p\}and the uniform distribution𝒰\\mathcal\{U\}:

ℒa​u​x=DK​L\(p^∥𝒰\)=∑i=1Ne​x​pp^ilogp^i\+const\\mathcal\{L\}\_\{aux\}=D\_\{KL\}\(\\hat\{p\}\\parallel\\mathcal\{U\}\)=\\sum\_\{i=1\}^\{N\_\{exp\}\}\\hat\{p\}\_\{i\}\\log\\hat\{p\}\_\{i\}\+\\text\{const\}\(8\)wherep^i=1B⋅T⋅Kroute​∑b,t𝕀⁡\(i∈𝒯b,t\)\\hat\{p\}\_\{i\}=\\frac\{1\}\{B\\cdot T\\cdot K\_\{\\mathrm\{route\}\}\}\\sum\_\{b,t\}\\mathbb\{I\}\(i\\in\\mathcal\{T\}\_\{b,t\}\)represents the normalized selection ratio of expertiiwithin a batch, hereBBdenotes the batch size,TTdenotes the number of routed token positions per sequence,KrouteK\_\{\\mathrm\{route\}\}is the number of active experts per token\. Minimizing this objective maximizes the entropy of the routing distribution, enforcing a balanced load across the expert bank\.

Total Objective:The final training objective aggregates the main generative tasks with these structural constraints:

ℒt​o​t​a​l=ℒb​a​s​e\+ℒc​r​i​t​i​c⏟Main Task\+λh​ℒh​i​g​h⏟Representation\+λa​ℒa​u​x⏟Regularization\\mathcal\{L\}\_\{total\}=\\underbrace\{\\mathcal\{L\}\_\{base\}\+\\mathcal\{L\}\_\{critic\}\}\_\{\\text\{Main Task\}\}\+\\underbrace\{\\lambda\_\{h\}\\mathcal\{L\}\_\{high\}\}\_\{\\text\{Representation\}\}\+\\underbrace\{\\lambda\_\{a\}\\mathcal\{L\}\_\{aux\}\}\_\{\\text\{Regularization\}\}\(9\)whereλh\\lambda\_\{h\}andλa\\lambda\_\{a\}are hyperparameters governing the strength of the hierarchical and sparsity constraints, respectively\. We setλh=1\\lambda\_\{h\}=1,λa=0\.1\\lambda\_\{a\}=0\.1\.

## 4Experiment

### 4\.1Experimental Setup

Dataset\.While previous multi\-game offline RL works such as JOWA utilized the DQN Replay Dataset\[[37](https://arxiv.org/html/2610.14727#bib.bib27)\], the original source is no longer publicly available\. Following the recommendation of the original authors, we utilize therlu\-atariversion of the dataset released as part of the RL Unplugged benchmark\[[38](https://arxiv.org/html/2610.14727#bib.bib28)\], which maintains the identical data distribution \(50M transitions per game per run\)\. However, since therlu\-ataribenchmark excludes the gameBerzerk, we select a subset of 14 games for training, maintaining the difficulty distribution \(easy, medium, hard\)\[[38](https://arxiv.org/html/2610.14727#bib.bib28)\], see Appendix[A](https://arxiv.org/html/2610.14727#A1)\. Consistent with\[[39](https://arxiv.org/html/2610.14727#bib.bib31),[40](https://arxiv.org/html/2610.14727#bib.bib32),[5](https://arxiv.org/html/2610.14727#bib.bib26)\], we use data from 2 out of 5 training runs\. Following the JOWA baseline\[[5](https://arxiv.org/html/2610.14727#bib.bib26)\], we uniformly draw 10% of transitions at random, resulting in approximately 10M transitions \(including frame stacking\) per game\.

Training\.MHWA has 77\.27M total parameters, with 60\.49M parameters activated per forward pass\. It usesK=36K=36visual tokens per frame and adopts the sameL=8L=8\-frame context window as JOWA\[[5](https://arxiv.org/html/2610.14727#bib.bib26)\]\. We train on 2 NVIDIA A40 GPUs with AdamW\[[41](https://arxiv.org/html/2610.14727#bib.bib30)\]and a learning rate of1×10−41\\times 10^\{\-4\}\. To match the baseline token exposure under constrained compute, we use batch size 1024 and 0\.875M gradient steps, compared with the published baseline batch size 512 and 1\.75M steps\. Thus, the total number of processed interaction tokens remains matched\. MHWA requires 21\.6 GPU\-days on 2 A40 GPUs\. We compare training computation using forward\+backward FLOPs with identical inputs: 8 frames, 296 tokens per sequence, and batch size 256 \(Table[2](https://arxiv.org/html/2610.14727#S4.T2)\)\. MHWA uses 42\.78% fewer FLOPs per sequence than JOWA\-70M and 70\.88% fewer than JOWA\-150M; the counting details are given in Appendix[B\.1](https://arxiv.org/html/2610.14727#A2.SS1)\.

Evaluation and Metrics\.We follow the Atari evaluation procedure of JOWA\[[5](https://arxiv.org/html/2610.14727#bib.bib26)\], reporting the average return over 16 episodes per game\. Performance is measured by Human Normalized Score \(HNS\),score−scorerandomscorehuman−scorerandom\\frac\{\\text\{score\}\-\\text\{score\}\_\{\\text\{random\}\}\}\{\\text\{score\}\_\{\\text\{human\}\}\-\\text\{score\}\_\{\\text\{random\}\}\}\. For aggregate comparison, we report IQM and Median HNS following RLiable\[[42](https://arxiv.org/html/2610.14727#bib.bib29)\]\.

Table 1:Performance comparison on 14 Atari games\.Baseline scores are from[Cheng et al\. \[5\]](https://arxiv.org/html/2610.14727#bib.bib26)\. 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 inbold\. “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±\\pmsample standard deviation over seed\-level aggregate metrics\. Full per\-seed results are provided in Appendix[D](https://arxiv.org/html/2610.14727#A4)\.GameRandomHumanMTBCMGDTEDTScaled\-QLJOWAOurs120M200M200M80M40M70M150M77MAssault222\.4742\.01203\.51741\.51915\.52292\.01578\.31402\.82302\.01068\.8Atlantis12850\.029028\.137812\.52565750\.02108166\.749100\.041662\.5228825\.02690387\.51690733\.3BeamRider363\.916926\.5786\.06011\.34149\.87023\.5880\.3854\.43498\.04202\.1Carnival0\.03800\.05560\.02610\.03678\.81940\.04685\.03812\.15316\.0881\.7Centipede2090\.912017\.04046\.64604\.03389\.93650\.35669\.34485\.04677\.04970\.9ChopperCommand811\.07387\.8256\.33300\.83443\.8853\.83118\.82568\.83812\.54547\.9DemonAttack152\.11971\.02611\.36549\.43455\.77936\.51233\.81409\.43547\.86593\.5NameThisGame2292\.38049\.06045\.06610\.57060\.06461\.77335\.67507\.111421\.013280\.0Phoenix761\.47242\.61446\.95120\.55320\.04961\.71608\.85070\.05348\.04476\.1Seaquest68\.442054\.7120\.02720\.03160\.4650\.01033\.11040\.02725\.01215\.0SpaceInvaders148\.01668\.7605\.0742\.5513\.8714\.4694\.7805\.3744\.72948\.3StarGunner664\.010250\.01493\.88625\.59550\.04728\.65737\.521700\.018150\.019100\.0TimePilot3568\.05229\.2762\.53866\.72812\.54072\.73268\.81100\.33669\.04700\.0Zaxxon32\.59173\.30\.0462\.5325\.50\.00\.0200\.02163\.01166\.7\#Superhuman0\-43333465\.7±\\pm0\.6Median HNS01\.0000\.1510\.5320\.5520\.4130\.3600\.5490\.5820\.676±\\pm0\.049IQM HNS01\.0000\.2920\.5370\.5540\.4430\.3970\.5310\.8330\.867±\\pm0\.047

### 4\.2Overall Performance Comparison

We compare our method against strong multi\-task offline RL baselines, including MTBC, MGDT\[[39](https://arxiv.org/html/2610.14727#bib.bib31)\], EDT\[[43](https://arxiv.org/html/2610.14727#bib.bib34)\], Scaled\-QL\[[40](https://arxiv.org/html/2610.14727#bib.bib32)\], and the SOTA model\-based approach, JOWA\[[5](https://arxiv.org/html/2610.14727#bib.bib26)\]\. The baseline results in Table[1](https://arxiv.org/html/2610.14727#S4.T1)are sourced from the original JOWA paper\[[5](https://arxiv.org/html/2610.14727#bib.bib26)\]under the same 10% low\-data protocol on the 14 shared games\. For MHWA, Table[1](https://arxiv.org/html/2610.14727#S4.T1)reports the mean and sample standard deviation of seed\-level IQMs, while Table[2](https://arxiv.org/html/2610.14727#S4.T2)reports IQM aggregated across the seed\-by\-game score matrix\. We additionally train JOWA\-70M with three independent seeds on the same 14\-gamerlu\-ataridataset and with the same batch setting as MHWA\. Using RLiable’s seed\-by\-game aggregation\[[42](https://arxiv.org/html/2610.14727#bib.bib29)\], MHWA\-77M achieves IQM HNS0\.884with a 95% stratified\-bootstrap confidence interval of \[0\.736, 1\.031\], compared with 0\.529 \[0\.459, 0\.602\] for the matched JOWA\-70M baseline \(Table[2](https://arxiv.org/html/2610.14727#S4.T2)\)\. This comparison demonstrates the benefit of the hierarchical routing architecture at comparable model capacity\. MHWA is also numerically comparable to the published JOWA\-150M result \(0\.833\)\.

Under the JOWA held\-out fine\-tuning protocol, MHWA also obtains 0\.643±\\pm0\.007 IQM DNS on five unseen Atari games, close to JOWA\-150M \(0\.647\) and above JOWA\-70M \(0\.603\); see Appendix[E](https://arxiv.org/html/2610.14727#A5)\.

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\.
### 4\.3Open\-loop Dynamics Modeling

To evaluate error accumulation under autoregressive rollout, we measure open\-loop reconstruction quality over stepst=5​…​8t=5\\ldots 8within the context window\. For each game, we sample held\-out trajectory segments disjoint from the training data, condition each model on the first four ground\-truth frames and corresponding actions, and then roll out fromt=5t=5without injecting ground\-truth observations\. Actions are kept fixed to the dataset actions, so differences in MSE/PSNR/SSIM against ground\-truth frames reflect dynamics\-modeling errors rather than policy differences\.

As shown in Figure[2](https://arxiv.org/html/2610.14727#S4.F2), the monolithic JOWA\-70M baseline shows a marked deterioration as rollout depth increases, indicating that reducing capacity without architectural changes amplifies error accumulation even within the context window\. In contrast, MHWA\-77M curbs error propagation and achieves predictive stability comparable to JOWA\-150M\. A longer\-horizon diagnostic fromt=5t=5tot=54t=54further shows that MHWA outperforms the smaller monolithic baseline and remains comparable to the larger baseline; see Appendix[G](https://arxiv.org/html/2610.14727#A7)\.

Overall, these results suggest that MHWA improves open\-loop stability through structural inductive bias rather than relying solely on brute\-force scaling\.

Figure 2:Quantitative comparison of open\-loop prediction within the context window \(T=5T=5to88\)\.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\.
### 4\.4Ablation Studies

We evaluate design choices on a fixed 6\-game subset using a 44M MHWA variant \(Appendix[A](https://arxiv.org/html/2610.14727#A1)\)\. This controlled ablation setting follows the subset\-ablation protocol used by JOWA\[[5](https://arxiv.org/html/2610.14727#bib.bib26)\]and keeps the study computationally feasible under our available budget\. For each variant, we first average the 16 evaluation episodes within each game to obtain one HNS per game, and then compute IQM and Median HNS over the six per\-game scores\. Table[4](https://arxiv.org/html/2610.14727#S4.T4)reports the resulting aggregate scores, together with the relative IQM drop compared to the full model\. We provide detailed per\-game raw scores in Appendix[F](https://arxiv.org/html/2610.14727#A6)\. To isolate architectural effects from capacity,w/o MoEwidens the dense FFN andMoE\-only \(w/o Manager\)widens the experts to keep parameters at≈\\approx44M\.

Temporal abstraction\.The stride ablations show that the high\-level context update frequency is a key factor\. Collapsing the hierarchy tok=1k=1causes the largest degradation among all variants, with a 62\.1% relative IQM drop, indicating that temporal abstraction is central to MHWA’s benefit\. The defaultk=4k=4performs best in this 44M/6\-game setting, whilek=2k=2andk=8k=8remain lower but competitive, with relative IQM drops of 20\.9% and 8\.3%, respectively\. This suggests that insufficient abstraction is especially harmful, whereas overly coarse abstraction has a milder effect in this setting\.

Routing mechanism\.TheMoE\-only \(w/o Manager\)variant removes the manager, Context Projector, skip\-step head, andℒh​i​g​h\\mathcal\{L\}\_\{high\}, retaining the task embedding and low\-level MoE; its IQM is 0\.445\. Removing the state input to the router yields IQM 0\.359\. The larger drop from removing state routing \(36\.3%\) suggests that local state information is especially important for reactive fine\-grained prediction\. The 21\.1% drop from removing the manager pathway further indicates that trajectory\-dependent context helps expert selection across slow\-varying dynamics modes\.

MoE modularity\.Replacing the sparse MoE with a dense FFN at matched parameter count \(w/o MoE\) reduces IQM to 0\.450, corresponding to a relative IQM drop of 20\.2%\. This is consistent with modular specialization being beneficial under heterogeneous multi\-game dynamics\.

High\-level supervision and routing regularization\.Removing the skip\-step objectiveℒh​i​g​h\\mathcal\{L\}\_\{high\}yields a moderate drop \(IQM:0\.564→0\.5030\.564\\rightarrow 0\.503, 10\.8%\), suggesting that explicit high\-level supervision helps the manager learn stride\-level predictive context rather than short\-range noisy features\. Removing the load\-balancing regularizerℒa​u​x\\mathcal\{L\}\_\{aux\}causes a larger drop \(IQM:0\.564→0\.4440\.564\\rightarrow 0\.444, 21\.3%\), supporting its role in preventing expert collapse and maintaining stable specialization\.

Gradient isolation\.The completed factorial ablation \(Table[3](https://arxiv.org/html/2610.14727#S4.T3)\) shows that stop\-gradient improves IQM both withℒh​i​g​h\\mathcal\{L\}\_\{high\}\(0\.533→0\.5640\.533\\rightarrow 0\.564\) and without it \(0\.387→0\.5030\.387\\rightarrow 0\.503\)\. Withoutℒh​i​g​h\\mathcal\{L\}\_\{high\}, the manager has fixed weights but supplies trajectory\-dependent hidden states to the trainable Context Projector\. Replacing those states with zeros or Gaussian noise reduces IQM from 0\.503 to 0\.406 or 0\.381 \(Appendix[F\.1](https://arxiv.org/html/2610.14727#A6.SS1)\)\.

Expert configuration and context refresh\.The default four\-expert, Top\-2 configuration performs best among the tested alternatives; four\-expert Top\-1 obtains 0\.542 and eight\-expert Top\-2 obtains 0\.511 \(Appendix[F\.2](https://arxiv.org/html/2610.14727#A6.SS2)\)\. Separating manager stride from context refresh, retraining withk=4,r=2k=4,r=2gives IQM 0\.558 versus 0\.564 for the defaultr=4r=4, with Median HNS 0\.484 versus 0\.463 \(Appendix[F\.3](https://arxiv.org/html/2610.14727#A6.SS3)\)\.

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\.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\.ΔIQM\(%\)\\Delta\_\{\\mathrm\{IQM\}\}\(\\%\)denotes the relative IQM drop compared to the full model, i\.e\.,\(IQMFull−IQMvariant\)/IQMFull×100\(\\mathrm\{IQM\}\_\{\\mathrm\{Full\}\}\-\\mathrm\{IQM\}\_\{\\mathrm\{variant\}\}\)/\\mathrm\{IQM\}\_\{\\mathrm\{Full\}\}\\times 100\. Bold and underline mark the largest and second\-largestΔIQM\(%\)\\Delta\_\{\\mathrm\{IQM\}\}\(\\%\), respectively\. All variants are controlled to have≈\\approx44M parameters for fair comparison\. Detailed per\-game raw scores are provided in Appendix[F](https://arxiv.org/html/2610.14727#A6)\.\(a\)Full MHWA\(b\)w/o High\-Level Loss\(c\)w/o Stride \(k=1k=1\)
Figure 3:PCA visualization of expert activations in the last MoE layer\.\(a\) Full MHWA mapsAssaultandPhoenixclosely 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\.
### 4\.5Interpretability of Representation Learning

To probe whether MHWA captures semantically meaningful dynamics, we visualize the PCA\-projected expert activation frequencies of the last MoE layer\.AssaultandPhoenixare both vertical space shooters sharing identical mechanics, and the projection shows that this similarity is reflected in the full MHWA model \(Fig\.[3](https://arxiv.org/html/2610.14727#S4.F3)a\) via highly similar expert\-usage patterns\. In contrast, this correspondence is not preserved when removing high\-level supervision \(Fig\.[3](https://arxiv.org/html/2610.14727#S4.F3)b\) or the temporal stride \(Fig\.[3](https://arxiv.org/html/2610.14727#S4.F3)c\), suggesting that theℒh​i​g​h\\mathcal\{L\}\_\{high\}and time abstraction design is necessary for the model to recover such task\-level commonalities from expert activations\.

## 5Discussion and Conclusion

We introduced MHWA, a multi\-timescale hierarchical world\-action model that combines a strided, gradient\-isolated high\-level manager with a context\-aware MoE low\-level worker\. The manager reduces how often global context must be inferred, while the worker remains reactive at the native token resolution and routes computation through specialized experts\. On the 10% offline Atari benchmark, MHWA outperforms a matched three\-seed JOWA\-70M baseline, achieves numerically comparable performance to JOWA\-150M with 77M parameters and 70\.88% fewer training FLOPs per sequence, and shows competitive adaptation to held\-out Atari games\.

The current evidence is intentionally scoped\. We evaluate primarily in discrete\-action offline Atari with a pretrained tokenizer and fixed stride, and extending to adaptive temporal abstraction\[[44](https://arxiv.org/html/2610.14727#bib.bib33)\], continuous control, and broader domains remains future work\.

## References

- \[1\]\(2018\)World models\.arXiv preprint arXiv:1803\.101222\(3\)\.Cited by:[§1](https://arxiv.org/html/2610.14727#S1.p1.1)\.
- \[2\]D\. Hafner, J\. Pasukonis, J\. Ba, and T\. Lillicrap\(2023\)Mastering diverse domains through world models\.arXiv preprint arXiv:2301\.04104\.Cited by:[§1](https://arxiv.org/html/2610.14727#S1.p1.1)\.
- \[3\]J\. Robine, M\. Höftmann, T\. Uelwer, and S\. Harmeling\(2023\)Transformer\-based world models are happy with 100k interactions\.arXiv preprint arXiv:2303\.07109\.Cited by:[§1](https://arxiv.org/html/2610.14727#S1.p1.1)\.
- \[4\]V\. Micheli, E\. Alonso, and F\. Fleuret\(2022\)Transformers are sample\-efficient world models\.arXiv preprint arXiv:2209\.00588\.Cited by:[§1](https://arxiv.org/html/2610.14727#S1.p1.1),[§2](https://arxiv.org/html/2610.14727#S2.p1.1)\.
- \[5\]J\. Cheng, R\. Qiao, Y\. Ma, B\. Li, G\. Xiong, Q\. Miao, Y\. Li, and Y\. Lv\(2024\)Scaling offline model\-based rl via jointly\-optimized world\-action model pretraining\.arXiv preprint arXiv:2410\.00564\.Cited by:[§1](https://arxiv.org/html/2610.14727#S1.p1.1),[§1](https://arxiv.org/html/2610.14727#S1.p14.1),[§1](https://arxiv.org/html/2610.14727#S1.p2.1),[§2](https://arxiv.org/html/2610.14727#S2.p1.1),[§3\.5](https://arxiv.org/html/2610.14727#S3.SS5.p1.1),[§4\.1](https://arxiv.org/html/2610.14727#S4.SS1.p1.1),[§4\.1](https://arxiv.org/html/2610.14727#S4.SS1.p2.1),[§4\.1](https://arxiv.org/html/2610.14727#S4.SS1.p3.1),[§4\.2](https://arxiv.org/html/2610.14727#S4.SS2.p1.1),[§4\.4](https://arxiv.org/html/2610.14727#S4.SS4.p1.1),[Table 1](https://arxiv.org/html/2610.14727#S4.T1),[Table 1](https://arxiv.org/html/2610.14727#S4.T1.6.1),[Table 2](https://arxiv.org/html/2610.14727#S4.T2.6.4.1.1)\.
- \[6\]E\. Talvitie\(2017\)Self\-correcting models for model\-based reinforcement learning\.InProceedings of the AAAI conference on artificial intelligence,Vol\.31\.Cited by:[§1](https://arxiv.org/html/2610.14727#S1.p1.1)\.
- \[7\]K\. Asadi, D\. Misra, S\. Kim, and M\. L\. Littman\(2019\)Combating the compounding\-error problem with a multi\-step model\.arXiv preprint arXiv:1905\.13320\.Cited by:[§1](https://arxiv.org/html/2610.14727#S1.p1.1)\.
- \[8\]R\. S\. Sutton, D\. Precup, and S\. Singh\(1999\)Between mdps and semi\-mdps: a framework for temporal abstraction in reinforcement learning\.Artificial intelligence112\(1\-2\),pp\. 181–211\.Cited by:[§1](https://arxiv.org/html/2610.14727#S1.p3.1),[§2](https://arxiv.org/html/2610.14727#S2.p2.1)\.
- \[9\]D\. Hafner, K\. Lee, I\. Fischer, and P\. Abbeel\(2022\)Deep hierarchical planning from pixels\.Advances in Neural Information Processing Systems35,pp\. 26091–26104\.Cited by:[§1](https://arxiv.org/html/2610.14727#S1.p3.1)\.
- \[10\]R\. Schiewer, A\. Subramoney, and L\. Wiskott\(2024\)Exploring the limits of hierarchical world models in reinforcement learning\.Scientific Reports14\(1\),pp\. 26856\.Cited by:[§1](https://arxiv.org/html/2610.14727#S1.p3.1),[§1](https://arxiv.org/html/2610.14727#S1.p4.1),[§1](https://arxiv.org/html/2610.14727#S1.p5.1),[§2](https://arxiv.org/html/2610.14727#S2.p2.1),[§3\.3](https://arxiv.org/html/2610.14727#S3.SS3.p4.1)\.
- \[11\]S\. Pateria, B\. Subagdja, A\. Tan, and C\. Quek\(2021\)Hierarchical reinforcement learning: a comprehensive survey\.ACM Computing Surveys \(CSUR\)54\(5\),pp\. 1–35\.Cited by:[§1](https://arxiv.org/html/2610.14727#S1.p5.1),[§1](https://arxiv.org/html/2610.14727#S1.p6.1),[§2](https://arxiv.org/html/2610.14727#S2.p2.1),[§2](https://arxiv.org/html/2610.14727#S2.p3.1)\.
- \[12\]S\. Levine, A\. Kumar, G\. Tucker, and J\. Fu\(2020\)Offline reinforcement learning: tutorial, review, and perspectives on open problems\.arXiv preprint arXiv:2005\.01643\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p1.1)\.
- \[13\]S\. Fujimoto, D\. Meger, and D\. Precup\(2019\)Off\-policy deep reinforcement learning without exploration\.InInternational conference on machine learning,pp\. 2052–2062\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p1.1)\.
- \[14\]S\. Fujimoto and S\. S\. Gu\(2021\)A minimalist approach to offline reinforcement learning\.Advances in neural information processing systems34,pp\. 20132–20145\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p1.1)\.
- \[15\]A\. Kumar, A\. Zhou, G\. Tucker, and S\. Levine\(2020\)Conservative q\-learning for offline reinforcement learning\.Advances in neural information processing systems33,pp\. 1179–1191\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p1.1),[§3\.5](https://arxiv.org/html/2610.14727#S3.SS5.p3.1)\.
- \[16\]I\. Kostrikov, A\. Nair, and S\. Levine\(2021\)Offline reinforcement learning with implicit q\-learning\.arXiv preprint arXiv:2110\.06169\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p1.1)\.
- \[17\]T\. Yu, G\. Thomas, L\. Yu, S\. Ermon, J\. Y\. Zou, S\. Levine, C\. Finn, and T\. Ma\(2020\)Mopo: model\-based offline policy optimization\.Advances in Neural Information Processing Systems33,pp\. 14129–14142\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p1.1)\.
- \[18\]R\. Kidambi, A\. Rajeswaran, P\. Netrapalli, and T\. Joachims\(2020\)Morel: model\-based offline reinforcement learning\.Advances in neural information processing systems33,pp\. 21810–21823\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p1.1)\.
- \[19\]D\. Ha and J\. Schmidhuber\(2018\)Recurrent world models facilitate policy evolution\.Advances in neural information processing systems31\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p1.1)\.
- \[20\]D\. Hafner, T\. Lillicrap, J\. Ba, and M\. Norouzi\(2019\)Dream to control: learning behaviors by latent imagination\.arXiv preprint arXiv:1912\.01603\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p1.1)\.
- \[21\]L\. Chen, K\. Lu, A\. Rajeswaran, K\. Lee, A\. Grover, M\. Laskin, P\. Abbeel, A\. Srinivas, and I\. Mordatch\(2021\)Decision transformer: reinforcement learning via sequence modeling\.Advances in neural information processing systems34,pp\. 15084–15097\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p1.1)\.
- \[22\]T\. D\. Kulkarni, K\. Narasimhan, A\. Saeedi, and J\. Tenenbaum\(2016\)Hierarchical deep reinforcement learning: integrating temporal abstraction and intrinsic motivation\.Advances in neural information processing systems29\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p2.1)\.
- \[23\]P\. Bacon, J\. Harb, and D\. Precup\(2017\)The option\-critic architecture\.InProceedings of the AAAI conference on artificial intelligence,Vol\.31\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p2.1)\.
- \[24\]A\. S\. Vezhnevets, S\. Osindero, T\. Schaul, N\. Heess, M\. Jaderberg, D\. Silver, and K\. Kavukcuoglu\(2017\)Feudal networks for hierarchical reinforcement learning\.InInternational conference on machine learning,pp\. 3540–3549\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p2.1)\.
- \[25\]O\. Nachum, S\. S\. Gu, H\. Lee, and S\. Levine\(2018\)Data\-efficient hierarchical reinforcement learning\.Advances in neural information processing systems31\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p2.1)\.
- \[26\]H\. Zhang, Y\. Cheng, C\. L\. P\. Chen, and X\. Wang\(2026\)Diffusion subgoal planning for long\-horizon offline goal\-conditioned reinforcement learning\.arXiv preprint arXiv:2609\.34575\.External Links:[Link](https://arxiv.org/abs/2609.34575)Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p2.1)\.
- \[27\]V\. Shaj Kumar, S\. Gholam Zadeh, O\. Demir, L\. Douat, and G\. Neumann\(2023\)Multi time scale world models\.Advances in Neural Information Processing Systems36,pp\. 26764–26775\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p2.1)\.
- \[28\]V\. Saxena, J\. Ba, and D\. Hafner\(2021\)Clockwork variational autoencoders\.Advances in neural information processing systems34,pp\. 29246–29257\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p2.1)\.
- \[29\]C\. Gumbsch, N\. Sajid, G\. Martius, and M\. V\. Butz\(2023\)Learning hierarchical world models with adaptive temporal abstractions from discrete latent dynamics\.InThe Twelfth International Conference on Learning Representations,Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p2.1)\.
- \[30\]G\. Ma, Y\. Chen, Y\. Zhang, L\. Guo, Y\. Wu, and S\. Yu\(2026\)Orthogonal space constraint enhances learning scalability and convergence efficiency without gradient backpropagation\.Neurocomputing679,pp\. 133139\.External Links:ISSN 0925\-2312,[Document](https://dx.doi.org/10.1016/j.neucom.2026.133139),[Link](https://www.sciencedirect.com/science/article/pii/S0925231226005369)Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p2.1)\.
- \[31\]N\. Shazeer, A\. Mirhoseini, K\. Maziarz, A\. Davis, Q\. Le, G\. Hinton, and J\. Dean\(2017\)Outrageously large neural networks: the sparsely\-gated mixture\-of\-experts layer\.arXiv preprint arXiv:1701\.06538\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p3.1),[§3\.4](https://arxiv.org/html/2610.14727#S3.SS4.p1.1)\.
- \[32\]W\. Fedus, B\. Zoph, and N\. Shazeer\(2022\)Switch transformers: scaling to trillion parameter models with simple and efficient sparsity\.Journal of Machine Learning Research23\(120\),pp\. 1–39\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p3.1),[§3\.4](https://arxiv.org/html/2610.14727#S3.SS4.p1.1)\.
- \[33\]C\. Riquelmeet al\.\(2021\)Scaling vision with sparse mixture of experts\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p3.1)\.
- \[34\]C\. Tang, Y\. Liu, Y\. Wu, W\. Han, Q\. Yin, X\. Zheng, W\. Zeng, and Q\. Zhang\(2025\)MoE\-world: a mixture\-of\-experts architecture for multi\-task world models\.Electronics14\(24\),pp\. 4884\.Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p3.1)\.
- \[35\]Z\. Zhao, Z\. Zhao, K\. Xu, Y\. Fu, J\. Chai, Y\. Zhu, and D\. Zhao\(2025\)Learning and planning multi\-agent tasks via an MoE\-based world model\.InAdvances in Neural Information Processing Systems,Vol\.38,pp\. 45724–45757\.External Links:[Document](https://dx.doi.org/10.52202/085713-1364)Cited by:[§2](https://arxiv.org/html/2610.14727#S2.p3.1)\.
- \[36\]A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, Ł\. Kaiser, and I\. Polosukhin\(2017\)Attention is all you need\.Advances in neural information processing systems30\.Cited by:[§3\.3](https://arxiv.org/html/2610.14727#S3.SS3.p1.1)\.
- \[37\]R\. Agarwal, D\. Schuurmans, and M\. Norouzi\(2020\)An optimistic perspective on offline reinforcement learning\.InInternational conference on machine learning,pp\. 104–114\.Cited by:[§4\.1](https://arxiv.org/html/2610.14727#S4.SS1.p1.1)\.
- \[38\]C\. Gulcehre, Z\. Wang, A\. Novikov, T\. Paine, S\. Gómez, K\. Zolna, R\. Agarwal, J\. S\. Merel, D\. J\. Mankowitz, C\. Paduraru,et al\.\(2020\)Rl unplugged: a suite of benchmarks for offline reinforcement learning\.Advances in neural information processing systems33,pp\. 7248–7259\.Cited by:[Table 5](https://arxiv.org/html/2610.14727#A1.T5),[Table 5](https://arxiv.org/html/2610.14727#A1.T5.4),[Appendix A](https://arxiv.org/html/2610.14727#A1.p1.1),[§4\.1](https://arxiv.org/html/2610.14727#S4.SS1.p1.1)\.
- \[39\]K\. Lee, O\. Nachum, M\. S\. Yang, L\. Lee, D\. Freeman, S\. Guadarrama, I\. Fischer, W\. Xu, E\. Jang, H\. Michalewski,et al\.\(2022\)Multi\-game decision transformers\.Advances in neural information processing systems35,pp\. 27921–27936\.Cited by:[§4\.1](https://arxiv.org/html/2610.14727#S4.SS1.p1.1),[§4\.2](https://arxiv.org/html/2610.14727#S4.SS2.p1.1)\.
- \[40\]A\. Kumar, R\. Agarwal, X\. Geng, G\. Tucker, and S\. Levine\(2022\)Offline q\-learning on diverse multi\-task data both scales and generalizes\.arXiv preprint arXiv:2211\.15144\.Cited by:[§4\.1](https://arxiv.org/html/2610.14727#S4.SS1.p1.1),[§4\.2](https://arxiv.org/html/2610.14727#S4.SS2.p1.1)\.
- \[41\]I\. Loshchilov and F\. Hutter\(2017\)Decoupled weight decay regularization\.arXiv preprint arXiv:1711\.05101\.Cited by:[§4\.1](https://arxiv.org/html/2610.14727#S4.SS1.p2.1)\.
- \[42\]R\. Agarwal, M\. Schwarzer, P\. S\. Castro, A\. C\. Courville, and M\. Bellemare\(2021\)Deep reinforcement learning at the edge of the statistical precipice\.Advances in neural information processing systems34,pp\. 29304–29320\.Cited by:[§4\.1](https://arxiv.org/html/2610.14727#S4.SS1.p3.1),[§4\.2](https://arxiv.org/html/2610.14727#S4.SS2.p1.1)\.
- \[43\]Y\. Wu, X\. Wang, and M\. Hamaya\(2023\)Elastic decision transformer\.Advances in neural information processing systems36,pp\. 18532–18550\.Cited by:[§4\.2](https://arxiv.org/html/2610.14727#S4.SS2.p1.1)\.
- \[44\]A\. Neitz, G\. Parascandolo, S\. Bauer, and B\. Schölkopf\(2018\)Adaptive skip intervals: temporal abstraction for recurrent dynamical models\.Advances in Neural Information Processing Systems31\.Cited by:[§5](https://arxiv.org/html/2610.14727#S5.p2.1)\.

## Appendix AAtari Game Subset

We follow the game difficulty taxonomy introduced in the RL Unplugged Atari benchmark\[[38](https://arxiv.org/html/2610.14727#bib.bib28)\]and report the properties of the games used in our multi\-game training set\. Compared to the 20\-game suite reported in prior work, we do*not*include games used for fine\-tuning/transfer evaluation, and we additionally excludeBerzerkbecause it is not available in therlu\-atarirelease\. The resulting training set contains 14 Atari games, listed in Table[5](https://arxiv.org/html/2610.14727#A1.T5)along with their difficulty labels and action\-space sizes\.

Table 5:Atari games used in our multi\-game training set, with difficulty labels from\[[38](https://arxiv.org/html/2610.14727#bib.bib28)\]and discrete action\-space sizes\.#### 6\-game subset for ablations\.

Following the subset protocol used in JOWA, we conduct ablations on a fixed 6\-game subset:Assault,Carnival,Centipede,NameThisGame,Phoenix, andSpaceInvaders\. This subset is chosen to keep a reasonably balanced mix of difficulty levels and discrete action\-space sizes\.

## Appendix BArchitecture and Parameterization Details

The High\-Level Manager is also a 4\-layer, 8\-head, 512\-dimensional Transformer, but uses a dense MLP with expansion ratio 4 instead of MoE\. Its output is detached before a trainable512→512512\\rightarrow 512Context Projector and broadcast across the following stride window\. Low\-level losses train the Context Projector, router, and worker, whileℒh​i​g​h\\mathcal\{L\}\_\{high\}trains the manager\. Both Transformers receive the shared token, positional, and task embeddings\. The tokenized interface uses 37 tokens per block, consisting of 36 visual tokens and one action token, with 8 blocks per sequence\. The critic head uses one Q\-head, MLP dimension 768, and dropout 0\.01\.

The Low\-Level Worker is a 4\-layer causal Transformer with 8 attention heads and embedding dimension 512\. Each feed\-forward sublayer is replaced by an MoE layer withN=4N=4experts and Top\-K=2K=2routing\. Each expert is a512→2048→512512\\rightarrow 2048\\rightarrow 512GELU MLP\. The routing gate is a bias\-free linear projection over the concatenation of the current token state and the broadcast high\-level context\. We do not use token dropping or capacity truncation; all selected tokens are computed, and exactly two experts are active per token\.

The full MHWA model contains 77\.27M parameters, with approximately 60\.49M parameters activated per forward pass including the tokenizer\.

### B\.1Training Computation

We count forward\+backward FLOPs for the main experimental configurations using 8 frames and 296 tokens per sequence, with batch size 256\. The count includes the complete forward graph, world\-model, critic, and auxiliary losses, and backward computation\. For MHWA, it includes both hierarchy levels, all four MoE layers, four stored experts with Top\-2 routing, load balancing, and high\-level observation prediction\. MoE computation is counted for the activated experts\.

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

## Appendix CA Simple Error\-Propagation View

We provide a simple formal view of when a strided context model can have a smaller one\-step open\-loop rollout\-error bound than a flat model\. This analysis is intended to isolate the core mechanism of MHWA, rather than to provide a general performance guarantee\. We focus on open\-loop state prediction under a fixed action sequence\.

Let the latent state bext=\(st,ct\)x\_\{t\}=\(s\_\{t\},c\_\{t\}\), wherest∈𝒮s\_\{t\}\\in\\mathcal\{S\}denotes the fast state andct∈𝒞c\_\{t\}\\in\\mathcal\{C\}denotes the slow context\. Let𝒜\\mathcal\{A\}be any metric action space, either discrete or continuous\. The true dynamics are

st\+1=f⁡\(st,ct,at\)\.s\_\{t\+1\}=f\(s\_\{t\},c\_\{t\},a\_\{t\}\)\.\(10\)In MHWA, at macro\-stepmm, the manager predicts a contextc^m\\hat\{c\}\_\{m\}, which is held fixed for the nextkklocal steps, while the worker rolls out

s^t\+1=f^​\(s^t,c^m,a~t\),\\hat\{s\}\_\{t\+1\}=\\hat\{f\}\(\\hat\{s\}\_\{t\},\\hat\{c\}\_\{m\},\\tilde\{a\}\_\{t\}\),\(11\)wherea~t\\tilde\{a\}\_\{t\}denotes the action representation used by the model\.

#### Assumptions\.

We use the following standard assumptions for open\-loop dynamics prediction\.

Lipschitz continuity\.There exist constantsLs,Lc,La≥0L\_\{s\},L\_\{c\},L\_\{a\}\\geq 0such that

d𝒮​\(f⁡\(s,c,a\),f⁡\(s′,c′,a′\)\)≤Ls​d𝒮​\(s,s′\)\+Lc​d𝒞​\(c,c′\)\+La​d𝒜​\(a,a′\),d\_\{\\mathcal\{S\}\}\\\!\\left\(f\(s,c,a\),f\(s^\{\\prime\},c^\{\\prime\},a^\{\\prime\}\)\\right\)\\leq L\_\{s\}d\_\{\\mathcal\{S\}\}\(s,s^\{\\prime\}\)\+L\_\{c\}d\_\{\\mathcal\{C\}\}\(c,c^\{\\prime\}\)\+L\_\{a\}d\_\{\\mathcal\{A\}\}\(a,a^\{\\prime\}\),\(12\)whered𝒮d\_\{\\mathcal\{S\}\},d𝒞d\_\{\\mathcal\{C\}\}, andd𝒜d\_\{\\mathcal\{A\}\}are metrics on the fast\-state, context, and action spaces\.

Slow context drift\.There existsρ≥0\\rho\\geq 0such that

d𝒞​\(ct,ct\+1\)≤ρ\.d\_\{\\mathcal\{C\}\}\(c\_\{t\},c\_\{t\+1\}\)\\leq\\rho\.\(13\)Therefore, for a local stept=m​k\+jt=mk\+jwithj∈\{0,…,k−1\}j\\in\\\{0,\\ldots,k\-1\\\},

d𝒞​\(cm​k,ct\)≤j​ρ\.d\_\{\\mathcal\{C\}\}\(c\_\{mk\},c\_\{t\}\)\\leq j\\rho\.\(14\)
Manager context error\.At each macro\-stepmm, the manager’s context prediction satisfies

d𝒞​\(cm​k,c^m\)≤βH\.d\_\{\\mathcal\{C\}\}\(c\_\{mk\},\\hat\{c\}\_\{m\}\)\\leq\\beta\_\{H\}\.\(15\)
Local one\-step modeling and interface error\.Assume the local dynamics model has one\-step approximation error

d𝒮​\(f⁡\(s,c,a\),f^​\(s,c,a\)\)≤εf\.d\_\{\\mathcal\{S\}\}\\\!\\left\(f\(s,c,a\),\\hat\{f\}\(s,c,a\)\\right\)\\leq\\varepsilon\_\{f\}\.\(16\)We further collect action\-representation and latent/interface approximation errors into a single quantityη\\eta:

d𝒮​\(f⁡\(s,c,a\),f^​\(s,c,a~\)\)≤η,η:=εf\+La​δa\+δx,d\_\{\\mathcal\{S\}\}\\\!\\left\(f\(s,c,a\),\\hat\{f\}\(s,c,\\tilde\{a\}\)\\right\)\\leq\\eta,\\qquad\\eta:=\\varepsilon\_\{f\}\+L\_\{a\}\\delta\_\{a\}\+\\delta\_\{x\},\(17\)whereLa​δaL\_\{a\}\\delta\_\{a\}bounds the error induced by action approximation andδx\\delta\_\{x\}collects additional latent or interface approximation error\. In continuous\-action or continuous\-latent settings, these terms need not arise from discretization and may be smaller when no discretization interface is used\.

#### Derivation\.

Define the open\-loop rollout error

et:=d𝒮​\(st,s^t\)\.e\_\{t\}:=d\_\{\\mathcal\{S\}\}\(s\_\{t\},\\hat\{s\}\_\{t\}\)\.\(18\)For a local stept=m​k\+jt=mk\+j, we have

et\+1\\displaystyle e\_\{t\+1\}=d𝒮​\(st\+1,s^t\+1\)\\displaystyle=d\_\{\\mathcal\{S\}\}\(s\_\{t\+1\},\\hat\{s\}\_\{t\+1\}\)=d𝒮​\(f⁡\(st,ct,at\),f^​\(s^t,c^m,a~t\)\)\.\\displaystyle=d\_\{\\mathcal\{S\}\}\\\!\\left\(f\(s\_\{t\},c\_\{t\},a\_\{t\}\),\\hat\{f\}\(\\hat\{s\}\_\{t\},\\hat\{c\}\_\{m\},\\tilde\{a\}\_\{t\}\)\\right\)\.\(19\)By the triangle inequality,

et\+1\\displaystyle e\_\{t\+1\}≤d𝒮​\(f⁡\(st,ct,at\),f⁡\(s^t,c^m,at\)\)\+d𝒮​\(f⁡\(s^t,c^m,at\),f^​\(s^t,c^m,a~t\)\)\.\\displaystyle\\leq d\_\{\\mathcal\{S\}\}\\\!\\left\(f\(s\_\{t\},c\_\{t\},a\_\{t\}\),f\(\\hat\{s\}\_\{t\},\\hat\{c\}\_\{m\},a\_\{t\}\)\\right\)\+d\_\{\\mathcal\{S\}\}\\\!\\left\(f\(\\hat\{s\}\_\{t\},\\hat\{c\}\_\{m\},a\_\{t\}\),\\hat\{f\}\(\\hat\{s\}\_\{t\},\\hat\{c\}\_\{m\},\\tilde\{a\}\_\{t\}\)\\right\)\.\(20\)Applying the Lipschitz condition in Eq\. \([12](https://arxiv.org/html/2610.14727#A3.E12)\) to the first term gives

d𝒮​\(f⁡\(st,ct,at\),f⁡\(s^t,c^m,at\)\)\\displaystyle d\_\{\\mathcal\{S\}\}\\\!\\left\(f\(s\_\{t\},c\_\{t\},a\_\{t\}\),f\(\\hat\{s\}\_\{t\},\\hat\{c\}\_\{m\},a\_\{t\}\)\\right\)≤Ls​d𝒮​\(st,s^t\)\+Lc​d𝒞​\(ct,c^m\)\.\\displaystyle\\leq L\_\{s\}d\_\{\\mathcal\{S\}\}\(s\_\{t\},\\hat\{s\}\_\{t\}\)\+L\_\{c\}d\_\{\\mathcal\{C\}\}\(c\_\{t\},\\hat\{c\}\_\{m\}\)\.\(21\)Using the local error definition in Eq\. \([17](https://arxiv.org/html/2610.14727#A3.E17)\), Eq\. \([20](https://arxiv.org/html/2610.14727#A3.E20)\) yields

et\+1≤Ls​et\+Lc​d𝒞​\(ct,c^m\)\+η\.e\_\{t\+1\}\\leq L\_\{s\}e\_\{t\}\+L\_\{c\}d\_\{\\mathcal\{C\}\}\(c\_\{t\},\\hat\{c\}\_\{m\}\)\+\\eta\.\(22\)It remains to bound the context mismatch\. By the triangle inequality, slow context drift, and manager context error,

d𝒞​\(ct,c^m\)\\displaystyle d\_\{\\mathcal\{C\}\}\(c\_\{t\},\\hat\{c\}\_\{m\}\)≤d𝒞​\(ct,cm​k\)\+d𝒞​\(cm​k,c^m\)\\displaystyle\\leq d\_\{\\mathcal\{C\}\}\(c\_\{t\},c\_\{mk\}\)\+d\_\{\\mathcal\{C\}\}\(c\_\{mk\},\\hat\{c\}\_\{m\}\)≤j​ρ\+βH\.\\displaystyle\\leq j\\rho\+\\beta\_\{H\}\.\(23\)Substituting Eq\. \([23](https://arxiv.org/html/2610.14727#A3.E23)\) into Eq\. \([22](https://arxiv.org/html/2610.14727#A3.E22)\), we obtain the MHWA one\-step open\-loop error bound

et\+1=em​k\+j\+1≤Ls​et\+η\+Lc​\(βH\+j​ρ\)\.e\_\{t\+1\}=e\_\{mk\+j\+1\}\\leq L\_\{s\}e\_\{t\}\+\\eta\+L\_\{c\}\(\\beta\_\{H\}\+j\\rho\)\.\(24\)

#### Comparison to a flat model\.

For a controlled comparison, we match the local one\-step errorη\\etaand isolate the effect of context inference\. Let a flat model re\-infer context at every step, with per\-step context error bounded byβF\\beta\_\{F\}\. Its analogous error recursion is

et\+1F≤Ls​etF\+η\+Lc​βF\.e^\{F\}\_\{t\+1\}\\leq L\_\{s\}e^\{F\}\_\{t\}\+\\eta\+L\_\{c\}\\beta\_\{F\}\.\(25\)For matched current rollout erroret=etFe\_\{t\}=e\_\{t\}^\{F\}, MHWA has a one\-step upper bound no larger than the flat model’s whenever

βH\+j​ρ≤βF\.\\beta\_\{H\}\+j\\rho\\leq\\beta\_\{F\}\.\(26\)Since the worst case within a stride window isj=k−1j=k\-1, a sufficient condition over the whole window is

βH\+\(k−1\)​ρ≤βF\.\\beta\_\{H\}\+\(k\-1\)\\rho\\leq\\beta\_\{F\}\.\(27\)This condition identifies when the strided context yields a one\-step error bound no larger than the flat model’s\.

#### Implications for transfer across settings\.

The comparison above does not rely on discrete actions or tokenized latent interfaces: discrete\-action Atari is one special case, while continuous actions and continuous latent interfaces are covered by the same metric\-space formulation, with interface effects absorbed intoη\\eta\. Likewise, in this narrow architectural sense, the comparison does not depend on whether the model is trained offline or online; it depends on the achieved error terms\(βH,βF,η,ρ\)\(\\beta\_\{H\},\\beta\_\{F\},\\eta,\\rho\)rather than on how they are learned\.

## Appendix DPer\-seed Results on the 14\-game Benchmark

For each seed and game, we average 16 evaluation episodes before computing HNS\. Tables[7](https://arxiv.org/html/2610.14727#A4.T7)and[8](https://arxiv.org/html/2610.14727#A4.T8)report per\-seed raw returns and aggregate metrics\. Mean and Sample SD summarize the three seed\-level values in each row\. The three MHWA seed\-level IQMs are 0\.837, 0\.922, and 0\.844\. RLiable aggregation of the3×143\\times 14seed\-by\-game score matrix yields IQM 0\.884 with a 95% stratified\-bootstrap CI of \[0\.736, 1\.031\]\.

Table 7:Per\-seed results of MHWA\-77M on the 14\-game Atari benchmark\.### D\.1Matched JOWA\-70M Results

JOWA\-70M is trained with three independent seeds on the same 14\-gamerlu\-ataridataset and with the same batch setting as MHWA\. Its seed\-by\-game IQM is 0\.529 with a 95% stratified\-bootstrap CI of \[0\.459, 0\.602\]\. Table[8](https://arxiv.org/html/2610.14727#A4.T8)reports the complete per\-seed results and game\-level cross\-seed variation, complementing the aggregate comparison in Table[2](https://arxiv.org/html/2610.14727#S4.T2)\.

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

## Appendix EHeld\-out Game Fine\-tuning

We evaluate transfer to unseen games using the held\-out Atari fine\-tuning protocol used by JOWA\. The held\-out games areGravitar,MsPacman,Pong,Robotank, andYarsRevenge\. Each game is fine\-tuned separately using 5k expert\-level transitions drawn from the last 20% of DQN\-Replay\. To match JOWA’s data exposure, which uses batch size 32 for 50k updates, we use batch size 1024 for 1562 updates, yielding 1,599,488 samples versus 1\.60M samples for JOWA\.

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±\\pmsample standard deviation over seed\-level aggregate DNS for the bottom rows\.GameRandomDQNJOWA\-70MJOWA\-150MMHWA Seed 0MHWA Seed 1MHWA Seed 2MHWA MeanGravitar173\.0473\.0387\.4273\.3375\.2271\.8284\.4310\.5MsPacman307\.33085\.6908\.02016\.71437\.52006\.21971\.01804\.9Pong\-20\.719\.514\.617\.717\.018\.118\.117\.7Robotank2\.263\.913\.825\.012\.411\.39\.311\.0YarsRevenge3092\.918089\.916339\.217506\.216056\.017760\.419280\.617699\.0Mean DNS––0\.5760\.6470\.6100\.6060\.6260\.614±\\pm0\.010Median DNS––0\.7150\.6150\.6740\.6110\.5990\.628±\\pm0\.040IQM DNS––0\.6030\.6470\.6480\.6350\.6450\.643±\\pm0\.007

Across the three MHWA seeds, the IQM DNS values are 0\.648, 0\.635, and 0\.645, yielding 0\.643±\\pm0\.007\. MHWA is competitive onPongandYarsRevenge, remains reasonably close onMsPacmanandGravitar, and is weaker onRobotank, suggesting that the transfer benefit is meaningful but not uniform across all held\-out games\. As a sanity check, we also ran one more expensive fine\-tuning run with batch size 32, which reached IQM DNS 0\.668; since this result is available for only one seed, we do not use it as the main comparison\.

## Appendix FPer\-game Ablation Results

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](https://arxiv.org/html/2610.14727#S4.T4)\.GameRandomHumanFullw/oℒh​i​g​h\\mathcal\{L\}\_\{high\}w/oℒa​u​x\\mathcal\{L\}\_\{aux\}w/o MoEMoE\-onlyw/o Statek=1k=1k=2k=2k=8k=8Assault222\.4742\.01086\.0989\.9930\.1917\.3902\.7745\.81003\.6843\.21043\.6Carnival0\.03800\.0626\.7220\.0380\.0340\.0680\.0420\.0463\.8300\.0240\.0Centipede2090\.912017\.05387\.74527\.04199\.04530\.02982\.12801\.33902\.03228\.03964\.0NameThisGame2292\.38049\.09186\.58480\.04160\.06950\.08246\.65690\.63350\.06770\.08270\.0Phoenix761\.47242\.63924\.04214\.44350\.02628\.14610\.04040\.91503\.13120\.04300\.0SpaceInvaders148\.01668\.7813\.9630\.91205\.0961\.9110\.0485\.0796\.9950\.0745\.0Median HNS––0\.4630\.4250\.4390\.4120\.3860\.3640\.1830\.4460\.469IQM HNS––0\.5640\.5030\.4440\.4500\.4450\.3590\.2140\.4460\.517

### F\.1Hierarchical Context Controls

The factorial ablation in Table[3](https://arxiv.org/html/2610.14727#S4.T3)retains the manager, Context Projector, context\-conditioned router, and stride\-kkbroadcasting in all four configurations\. Withℒh​i​g​h\\mathcal\{L\}\_\{high\}on, the manager is trained by high\-level supervision and, when stop\-gradient is off, also by low\-level objectives\. Withℒh​i​g​h\\mathcal\{L\}\_\{high\}off, the manager is fixed\-weight when stop\-gradient is on and trained only by low\-level objectives when it is off\. The Context Projector, router, and worker remain trainable throughout\. MoE\-only removes the complete manager pathway and widens the low\-level experts to match model size\.

We train two additional information controls under theℒh​i​g​h\\mathcal\{L\}\_\{high\}\-off, stop\-gradient\-on setting\. Both preserve the fixed\-weight manager, trainable Context Projector, worker, MoE, parameter count, andk=4k=4broadcasting\. They replace the manager hidden state after stop\-gradient with zeros or Gaussian noise during both training and evaluation\. Zero input produces a learned constant context through the projector bias; Gaussian noise is scaled using batch RMS and varies across context blocks\.

Table 11:Information controls for the fixed\-weight manager pathway\.All variants are separately trained in the 44M/6\-game setting withℒh​i​g​h\\mathcal\{L\}\_\{high\}off and stop\-gradient on\.The trajectory\-conditioned hidden state improves IQM over both controls\. Its Median HNS is close to that of the zero\-input variant, while IQM captures the larger separation across the six\-game score distribution\.

### F\.2Expert Number and Routing Sparsity

Table[12](https://arxiv.org/html/2610.14727#A6.T12)varies expert number, Top\-KK, and expert\-MLP expansion in the 44M/6\-game setting\. Letr0r\_\{0\}denote the default expansion ratio\. TheN=2N=2, Top\-1 configuration doubles the expansion, whileN=8N=8, Top\-4 halves it to keep total expert parameters comparable to the default\. TheN=4N=4, Top\-1 variant keeps the default expert bank; theN=8N=8, Top\-2 variant keeps the default per\-expert expansion\.

Table 12:Expert configuration study\.IQM HNS on the six\-game ablation benchmark\.Four\-expert Top\-1 retains 96% of the default IQM, and the eight\-expert alternatives retain 87–91%\. The two\-expert configuration shows the largest drop, reflecting the effect of coarser capacity partitioning\.

For each configuration, we aggregate final\-layer expert activation frequencies for the six games, fit PCA independently, and rank the 15 pairwise distances in the three\-dimensional projection\. Table[13](https://arxiv.org/html/2610.14727#A6.T13)reports the rank of the Assault–Phoenix pair\. This neighborhood is closest for both Top\-2 configurations, while Top\-1 reorganizes cross\-task sharing\. Under four\-expert Top\-1, Centipede–SpaceInvaders moves from the default’s 15th\-ranked pair to the 6th\-ranked pair; their raw returns are 8950\.4 and 917\.8, respectively, and the six\-game Median HNS is 0\.516\. For two experts, the normalized usage vector has one effective degree of freedom\.

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\.ExpertsNNTop\-KKAssault–Phoenix distance rank4214114821844
### F\.3Manager Stride and Context\-Refresh Interval

The manager sampling stridekkdetermines the temporal spacing of its inputs and the high\-level prediction scale\. The context\-refresh intervalrrdetermines how often the projected context is refreshed for low\-level routing\. The default usesk=r=4k=r=4\.

We first vary onlyrrat inference using the same checkpoint trained withk=r=4k=r=4\. The manager retains itsk=4k=4input spacing and maximum context length, with sliding\-window refreshes atr=2r=2orr=1r=1\. We then retrain withk=4,r=2k=4,r=2throughout training and evaluation, allowing the router and worker to adapt to the alternative schedule\. Table[14](https://arxiv.org/html/2610.14727#A6.T14)separates these results\.

Table 14:Context\-refresh schedules at fixed manager stridek=4k=4\.The inference\-only variants use the default checkpoint; the retrained variant usesr=2r=2during both training and evaluation\.Table 15:Per\-game results after retraining withk=4,r=2k=4,r=2\.Raw returns under the same 44M/6\-game protocol\.After retraining,r=2r=2improves Carnival and Centipede and lowers returns on the other four games\. The resulting aggregate performance is comparable to the default schedule, with a higher Median and a slightly lower IQM\. Changing the refresh schedule redistributes game\-level performance while preserving the manager’s learned temporal scale\.

## Appendix GLong\-horizon Open\-loop Diagnostics

We further evaluate open\-loop rollouts beyond the context window\. Starting from four ground\-truth context frames, each model autoregressively predicts frames fromt=5t=5tot=54t=54while conditioning on the same action sequence\. As shown in Figure[4](https://arxiv.org/html/2610.14727#A7.F4), the result shows that MHWA exhibits smoother degradation than the smaller monolithic baseline and remains broadly comparable to the 150M baseline\. Importantly, the most visible anomaly in JOWA\-150M aroundt=17t=17is not a uniform collapse across games; it is mainly driven by a sharp fluctuation on DemonAttack att=17t=17–2020, while the same game/timestep does not show a comparable spike for the other models\. This diagnostic is intended to test the smoothness of error accumulation rather than to claim indefinitely stable long\-horizon generation\.

Figure 4:Long\-horizon open\-loop reconstruction diagnostic fromt=5t=5tot=54t=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\.
## Appendix HResponsible Research and Existing Assets

#### Broader impacts\.

This work studies offline model\-based reinforcement learning and world modeling in Atari benchmark environments\. A potential positive impact is improved compute efficiency for training world models, which may lower the resource barrier for research on sample\-efficient decision making\. The main risks are indirect: more capable model\-based RL systems could eventually be used in domains where inaccurate rollouts, distribution shift, or unsafe planning decisions may cause harm\. Our experiments are limited to benchmark video games and do not involve deployment, human\-subject data, personal data, or safety\-critical decision making\. We therefore view the immediate societal risk of this work as limited, while noting that future applications of model\-based RL in real\-world control should require domain\-specific safety evaluation, uncertainty estimation, and deployment safeguards\.

#### Safeguards\.

The paper does not release a high\-risk foundation model, image/text generator, scraped dataset, personal dataset, or deployed decision\-making system\. The code release is intended for reproducing Atari benchmark experiments\. The datasets used are public benchmark datasets, and the trained models operate in Atari\-like simulated environments\. As a result, we do not identify a need for special access restrictions beyond standard research\-use documentation and citation of the underlying assets\.

#### Existing assets and licenses\.

We use the publicrlu\-atariversion of the Atari replay data from RL Unplugged, which follows the DQN Replay data distribution used in prior offline Atari work\. We cite the original DQN Replay Dataset, RL Unplugged, JOWA, and other baseline methods used for comparison\. We also reuse the pretrained tokenizer interface from JOWA for a controlled comparison within the same tokenized world\-action framework\. Our supplementary code release will include instructions for obtaining the public datasets and pretrained components from their original sources, and users should follow the corresponding licenses and terms of use of those upstream assets\. We do not introduce a new dataset\.