跳转至

MHWA: Multi-Timescale Hierarchical World-Action Model

会议: NeurIPS 2026
arXiv: 2610.14727
领域: 强化学习
关键词: 世界模型、离线强化学习、分层强化学习、混合专家模型、时间抽象

一句话总结

MHWA 针对自回归世界模型在开环外推中误差随步数累积漂移的顽疾,提出了时间步长跨越的全局慢时标 Manager 与局部上下文感知门控稀疏 MoE Worker 的非对称分层架构,在减少 70.88% 训练 FLOPs 的同时比肩两倍参数量的 JOWA 表现。

研究背景与动机

基于自回归 Transformer 的世界模型在强化学习与决策智能中展现出强劲的样本效率,使智能体能够在想象的轨迹中进行策略评估与搜索规划。然而,自回归序列生成的本质决定了模型在开环预测时必然面临复合漂移问题:当自回归展开逐步依赖自身之前生成的不完美预测结果时,细微的局部观测扰动与分布偏移会迅速恶化并沿时域扩散,导致生成轨迹在短规划时域内便迅速脱离真实物理环境规律。以往主流的应对思路通常依赖于暴力缩放参数量,例如拥有 1.5 亿参数的 JOWA 模型通过堆叠模型容量来强行吸收异构模式与噪声,但这显著增加了计算开销与训练周期。

分层时域抽象原本是抑制长期累积漂移的优雅途径,它通过分离慢速全局上下文与快速局部动态来缩短宏观时序视界。但在世界模型中引入分层机制常常伴随着两类难以调和的耦合干扰。其一是概念语义干扰(Conceptual Interference):在传统目标条件分层中,高层子目标必须由低层执行策略的当下能力来评价,两者同时训练导致高层奖励信号持续漂移,产生严重的非平稳性目标。其二是结构干扰(Structural Interference):若使用单一且高度参数共享的紧耦合低层网络来拟合跨游戏、跨模态的异构动态,容易引发梯度竞争与模式混淆,使得分层系统不仅极其复杂难以收敛,甚至无法超越平铺基线。

针对这两大核心症结,本文摒弃了单纯将分层视作独立灵丹妙药的思路,而是提出了带显式解耦控制的层级世界-动作模型。核心 idea:通过步长下采样与停止梯度将慢速全局上下文表征与低层损失隔离以根除概念干扰,同时利用结合即时状态与广播上下文的双重门控动态稀疏路由低层 MoE 专家以消解异构结构干扰。

方法详解

整体框架

MHWA 继承了离散视觉 Token 与动作交错排列的序列决策接口(每帧图像经预训练 VQ-VAE 编码为 36 个视觉离散 Token,与 1 个动作 Token 组成时序块,共包含 8 帧交互历史),但在预测器架构上打破了原有的扁平自回归设计,构建了非对称双层世界-动作模型体系。输入序列一方面经过步长下采样送入高层管理者(High-Level Manager)提炼宏观动态上下文;另一方面全量输入送入低层执行者(Low-Level Worker),并结合管理者广播的上下文向量,通过稀疏混合专家(MoE)结构进行细粒度自回归预测。低层输出最终同时驱动世界预测头(下一观测 Token、即时奖励、终止标志)与离线价值评判头(分布 Q 估计)。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
    A["交错序列输入<br/>观测 Token + 动作 Token"] --> B["时间步长下采样<br/>步长因子 k=4 抽取慢速序列"]
    A --> E["低层原子序列<br/>全量密集时序 Token 嵌入"]

    subgraph S1["高层管理者:上下文推断"]
        direction TB
        B --> C["因果 Transformer 编码器<br/>跨步聚合慢速隐状态"]
        C --> D["跨步预测头 + 梯度截断投影<br/>隔步预测下一宏观观测"]
    end

    D -->|"广播固定上下文向量<br/>stop-gradient 阻止回传"| F["上下文感知动态路由<br/>拼接即时状态与宏观向量"]

    subgraph S2["低层执行者:动态稀疏专家网络"]
        direction TB
        E --> F
        F --> G["Top-2 门控稀疏分配<br/>4 专家中激活最优子集"]
        G --> H["MoE 前馈网络计算<br/>无截断无丢弃混合更新"]
    end

    H --> I["双功能预测输出<br/>世界模型头 + 分布式 CQL Critic"]

关键设计

1. 步长下采样上下文推断与梯度隔离:根除高低层概念干扰

为解决分层学习中高层语义随低层能力动态波动而带来的非平稳性漂移,高层管理者(High-Level Manager)在时间维度上对输入的交错 Token 序列进行 \(k\) 倍跨步下采样(默认 \(k=4\))。管理者仅摄取粗粒度子序列 \(\mathbf{X}^{(m)}_{slow} = (x_1, x_{1+k}, \dots, x_{1+m \cdot k})\),经因果 Transformer 编码后提取宏观慢变隐状态 \(h^{high}_m\)。关键机制在于,模型在管理者输出后设计了一个可训练的上下文投影器(Context Projector)以生成连续上下文向量 \(z^{ctx}_m = \mathrm{Projector}(\mathrm{stopgrad}(h^{high}_m))\)。在此处施加停止梯度操作(stop-gradient),使得低层的世界模型预测误差及强化学习 TD/CQL 损失完全被阻断在管理者外部。高层管理者完全由独立的跨步预测辅助任务驱动更新,必须凭借自身的隐状态直接预测未来下采样步长处的观测 Token \(o_{(m+1)k}\):

\[ \mathcal{L}_{high} = -\sum_{m} \log p_{o}^{high}(o_{(m+1)k} \mid h^{high}_m) \]

这种瓶颈机制强迫管理者从时序流中过滤掉局部的瞬态高频噪声,提取出跨时段稳定的慢时标宏观环境表征,并且彻底解耦了优化目标,避免高层被低层的短期拟合误差带偏。

2. 状态与全局上下文双重调节的动态专家路由:抑制异构动态结构干扰

针对多游戏和多动态模式下低层共享参数相互争夺容量引起的结构干扰,低层执行者(Low-Level Worker)将每个 Transformer 模块中的传统前馈网络(FFN)替换为稀疏门控混合专家网络(MoE)。与传统仅依赖输入 Token 自主路由的 MoE 不同,MHWA 采用双模态门控机制,将当前中间特征 \(x_t\) 与所在窗口内固定的广播上下文向量 \(z^{ctx}_m\) 进行拼接映射,计算 \(N\) 个专家的路由得分:

\[ l_t = (x_t \oplus z^{ctx}_m) W_g \]

模型选定激活得分前 \(K_{route}=2\) 名的专家(总专家数 \(N=4\)),并对其进行归一化 Softmax 加权合成输出。该设计的巧妙之处在于:全局上下文项 \(z^{ctx}_m\) 驱动专家在不同游戏机制或慢速动态状态之间形成宏观模式特化(如区分垂直飞行与横向规避等宏观阶段),而即时状态项 \(x_t\) 赋予系统对当前原子动作与瞬时图像细微变化的敏锐反应能力。全流程不施加 Capacity Factor 丢弃或截断,确保所有选定 Token 均得到完整计算。

3. 联合平衡正则化与分布值函数训练:兼顾专家特化与离线保守性

为防止 MoE 路由出现退化(专家坍缩为单模态占用),MHWA 引入了批次维度的专家使用率 KL 散度约束作为负载均衡正则项 \(\mathcal{L}_{aux} = D_{KL}(\hat{p} \parallel \mathcal{U})\),驱动门控网络在整个专家库上逼近均匀分布,其中 \(\hat{p}_i = \frac{1}{B \cdot T \cdot K_{route}} \sum_{b,t} \mathbb{I}(i \in \mathcal{T}_{b,t})\)。在决策层面,模型联合训练世界预测头与基于 C51 的分布式 Q 评判头,并在离线价值更新中融合保守 Q 学习(CQL)惩罚项 \(\alpha_{cql} \mathcal{L}_{CQL}\),抑制分布外不可信动作的 Q 值虚高。总目标函数在无互相梯度污染的前提下实现多目标联合优化:

\[ \mathcal{L}_{total} = \mathcal{L}_{base} + \mathcal{L}_{critic} + \lambda_h \mathcal{L}_{high} + \lambda_a \mathcal{L}_{aux} \]

其中超参数配置为 \(\lambda_h=1\) 与 \(\lambda_a=0.1\)。

损失函数 / 训练策略

模型训练采用两块 NVIDIA A40 GPU,优化器选用 AdamW,学习率设为 \(1 \times 10^{-4}\)。为在受限算力下严格对齐基线 JOWA 的样本暴露量,将批大小扩大至 1024、梯度步数设为 0.875M(等效于 JOWA 的批大小 512、步数 1.75M)。总共 77.27M 参数中,每次前向传播仅激活约 60.49M 参数。每序列的前向+反向 FLOPs 仅需 48.199 GFLOPs,较 JOWA-70M 减少 42.78%,较 JOWA-150M 大幅降低 70.88%。

实验关键数据

主实验

在 14 款跨越简易、中等、困难三档的 Atari 游戏测试中(使用 RL Unplugged 的 10% 稀疏离线数据协议),MHWA-77M 与多任务决策 Transformer(MGDT)、弹性决策 Transformer(EDT)、高扩展 Q 学习(Scaled-QL)以及最强世界模型基线 JOWA 进行了严密对比。

模型 参数量 激活参数量 GFLOPs/序列 14游戏 Median HNS 14游戏 IQM HNS (RLiable)
MGDT 200M 200M - 0.532 0.537
EDT 200M 200M - 0.552 0.554
Scaled-QL 80M 80M - 0.413 0.443
JOWA-70M (对齐复现) 70M 70M 84.238 0.549 0.529 [0.459, 0.602]
JOWA-150M (官方公布) 150M 150M 165.524 0.582 0.833
MHWA (本文) 77M 60.49M 48.199 0.676 ± 0.049 0.884 [0.736, 1.031]

在 14 款游戏中,MHWA 在平均达到超人水平表现(#Superhuman)的游戏数量上取得 5.7 款(最高单种子 6 款),显著超越了 JOWA-70M 的 4.0 款。在计算能耗几乎只有 JOWA-150M 三成(48.2 vs 165.5 GFLOPs)的前提下,MHWA 的 IQM 达到了 0.884,在统计置信区间上完全比肩甚至略微超出 150M 的大模型。

消融实验

在固定的 6 款代表性游戏子集(Assault, Carnival, Centipede, NameThisGame, Phoenix, SpaceInvaders)上,研究团队对 44M 参数量的等效受控变体进行了全面的结构消融分析。

模型配置 核心变动说明 Median HNS IQM HNS 相对降幅 \(\Delta_{\mathrm{IQM}}\) (%)
Full Model 完整模型 (\(k=4\), Top-2 MoE) 0.463 0.564 -
\(k=1\) (无步长抽象) 压缩分层至逐步更新,丧失时间抽象 0.183 0.214 62.1%
w/o State Routing 门控仅依赖全局上下文,缺失即时状态 0.364 0.359 36.3%
w/o \(\mathcal{L}_{aux}\) 移除负载均衡正则,导致专家坍缩 0.439 0.444 21.3%
MoE-only (w/o Manager) 移除高层管理者与全局通道,仅留低层 MoE 0.386 0.445 21.1%
\(k=2\) (过细步长) 上下文更新步长设为 2 0.446 0.446 20.9%
w/o MoE (Dense FFN) 稠密前馈网络替代 MoE,参数等量缩放 0.412 0.450 20.2%
w/o \(\mathcal{L}_{high}\) 移除高层跨步自监督预测损失 0.425 0.503 10.8%
\(k=8\) (过粗步长) 上下文更新步长设为 8 0.469 0.517 8.3%

在梯度隔离方面,消融实验证实:当移除管理者与低层之间的停止梯度操作(使低层损失反向传播进入管理者)时,在启用 \(\mathcal{L}_{high}\) 下 IQM 从 0.564 降至 0.533;而在未启用 \(\mathcal{L}_{high}\) 时,缺失停止梯度更是导致 IQM 从 0.503 暴跌至 0.387,充分验证了防止下层噪点干扰宏观上下文的重要性。

关键发现

  • 时间抽象步长的双向敏感性:将步长退化为 \(k=1\) 会引发高达 62.1% 的性能崩溃,是全架构中最脆弱的核心支柱;但步长过大(\(k=8\))仅轻微回落 8.3%,表明世界模型对粗粒度背景的容忍度明显高于无抽象的高频扰动。
  • 即时反应对动态拟合的决定性作用:移除路由中的即时状态输入会导致 36.3% 的断崖式掉点,说明专家网络绝不能退化为纯静态的任务分类器,必须在宏观语境下实时追踪局部运动与动作反馈。
  • 开环动力学漂移的有效遏制:在预测未来第 5 至 8 帧的严格开环图像重构评测中,单体架构 JOWA-70M 的 MSE 与 SSIM 随深度加深迅速恶化,而 MHWA 在保持低参状态下始终维持平缓的误差曲线,在图像保真度上超越了参数多出一倍的 JOWA-150M。
  • 表征可解释性与跨任务对齐:对最后一层 MoE 专家激活频率进行 PCA 降维投影发现,垂直射击机制完全相同的 Assault 与 Phoenix 投影距离在全模型下最近(在 15 组配对中排名第 1),而一旦剥离高层自监督损失或时间步长抽象,该机制相似性立即溃散。

亮点与洞察

  • 结构归纳偏置击溃暴力参数堆叠:论文有力地证明了,自回归世界模型面临的长期发散与模式竞争并不一定必须依靠数百兆参数的模型硬抗;引入精巧的非对称多时间尺度分层与 MoE 路由,可以在节省 70% 训练算力的同时反超两倍体量的单体密集网络。
  • 停止梯度在分层结构中的认知解耦价值:很多分层模型盲目进行端到端全量微调,反而让高层表征成为低层高频残差拟合的奴隶。MHWA 通过显式梯度截断配合跨步自监督预测,构筑了高层表征的信息瓶颈,这一技巧为各类层级架构设计提供了通用的防坍缩范式。
  • 从离线预训练到零样本迁移的高效泛化:在未经预训练的 5 款全新留出游戏(如 Pong, MsPacman 等)微调测试中,MHWA 取得了 0.643 的 IQM 表现,逼近 JOWA-150M(0.647),证明其学到的时空抽象具备扎实的域外迁移能力。

局限与展望

  • 固定启发式步长的理论边界:当前时域步长 \(k=4\) 采用固定的离散超参数,对不同动态速率(如慢变棋盘与瞬变弹幕)缺乏自适应能力;未来若能引入基于信息熵或事件驱动的动态自适应跳步机制将更具扩展性。
  • 实验域局限于离散图像控制:验证主要集中在 14 款离散动作的 Atari 复古游戏环境中,尚未扩展至连续控制机械臂、物理仿真具身导航或自动驾驶世界模型等更为复杂的高维真实场景。
  • 对初始视觉分词器质量的依赖:模型沿用了 JOWA 预训练好的离线 VQ-VAE 码本,世界模型的推断能力上限受限于静态分词对高速动态物体的重建保真度,未来可探索端到端协同更新的分层分词体系。

相关工作与启发

  • vs JOWA (Jointly-Optimized World-Action Model):JOWA 采用单一密集的自回归 Transformer 统一建模观测与动作,依靠将参数扩大至 150M 来吸收多任务扰动;MHWA 针对其开环滚出易发散的缺陷,设计了慢速管理者与 MoE 路由架构,以 77M 参数和三成计算量取得了同等乃至更优的稳定控制表现。
  • vs 传统分层强化学习与分层世界模型 (HRL / HWM):传统分层模型(如 Feudal Networks 或目标条件 HWM)由于子目标持续演变导致严重的概念干扰;MHWA 摒弃了不稳定的低层目标达成闭环,转而通过停止梯度与慢速自监督预测构造鲁棒的全局上下文向量,在保持局部极速反应的同时规避了学习目标非平稳性。
  • vs 传统 MoE 序列决策模型 (MoE-World):以往 MoE 主要基于全局任务 ID 或当前 Token 进行静态孤立路由;MHWA 提出了全局广播上下文与即时局部特征拼接的双模态动态门控机制,兼顾了宏观模式专精与瞬时动作响应。

评分

  • 新颖性: ⭐⭐⭐⭐ [成功解决了分层世界模型中长期存在的概念干扰与结构干扰,设计了清晰优雅的步长截断与上下文双模态路由机制]
  • 实验充分度: ⭐⭐⭐⭐⭐ [提供了严谨的 3 种子 RLiable 置信区间统计、完备的消融对比、开环图像生成诊断及流形可视化证据]
  • 写作质量: ⭐⭐⭐⭐⭐ [逻辑层层递进,对先前分层失败机理剖析深刻,结构图与公式推导清晰自洽]
  • 价值: ⭐⭐⭐⭐⭐ [为受限算力下低成本构建稳定抗漂移的高效世界模型提供了突破性的结构范式]