AI Research 技术调研

Planning in Stochastic Environments with a Learned Model (Stochastic MuZero)

18分钟阅读

一句话定位

给 muzero 的价值等价隐模型引入 afterstate（动作生效后、环境随机转移前的中间状态），把”状态转移”拆成”确定性动作效果 → 随机性机会事件”两段，用一个不带 decoder、不做重建、端到端训练的定制 VQ-VAE 学随机的 chance code，并把 MCTS 扩展出 chance node，使 MuZero 首次能在 2048、西洋双陆棋这类本质随机的环境里工作，同时在确定性的围棋上不掉性能；ICLR 2022 Spotlight，发表于 OpenReview，无 arXiv 版本。

背景与定位

muzero 用价值等价（value-equivalent）的隐模型 + MCTS 在围棋、国际象棋、将棋、Atari 上做到了不依赖规则先验的超人表现，sampled-muzero 把它扩展到高维/连续动作空间，muzero-unplugged 把它扩展到离线强化学习，但这一整条谱系的模型都是确定性函数——隐式假设环境动力学也是确定性的。真实世界的很多环境本身是随机的、部分可观测的，或者复杂到有限容量的智能体只能把它当作随机的来对待，这类环境上纯确定性模型的规划效果会明显变差。

论文把此前处理这个问题的路线分成三类并逐一指出局限：observation models（Oh et al., 2015；Chiappa et al., 2017；Kaiser et al., 2020 的 SimPLe）直接拟合观测和奖励的动力学，再配合 Dyna 式的 model-free 更新，但对高维图像观测建模计算代价高、多步展开容易累积误差；latent models（muzero、value-prediction-network、Henaff et al., 2017）在隐空间里做多步展开预测 reward/value/policy，MuZero 证明了这条路能做到 SOTA，但用的是确定性动力学函数；stochastic latent models 中，dreamer-v2（Hafner et al., 2021）的 RSSM 用确定性 recurrent state 加多维 multinomial 隐变量建模随机性，但学到的模型只是用来生成合成数据训练 model-free actor-critic，效果不及 MuZero 这类”模型直接参与规划”的方案；vq-models-planning（Ozair et al., 2021）用 VQ-VAE 学随机 transition model 并接 MCTS，在国际象棋上匹配了 MuZero，但局限于监督设定（用专家数据而非在线 RL 交互学习），且用显式 decoder 加 reconstruction loss、两阶段训练（先学观测表示，再学 transition model），难以直接搬到强化学习场景。

本文自称是第一个在纯在线强化学习环境中”empirically effective”地学习随机价值等价模型并用于规划的方法，把它具体实例化为 Stochastic MuZero。

模型架构

Afterstate 分解：把一次状态转移 s_t —(动作 a_t)⇒ as_t —(机会事件 c_t)⇒ s_{t+1} 拆成两段确定性映射：afterstate as_t 是”动作已生效、但环境还未完成随机转移”的假设中间状态（例如双陆棋里一方走完棋但对方还没掷骰子时的局面）；给定 afterstate，一个有限个数 M 的机会结果 c 决定了具体转移到哪个下一状态。这样 V(as_t) = Q(s_t, a_t)，Pr(s_{t+1}|as_t) = Pr(c|as_t)，学随机模型的问题被规约为学习 afterstate 和 chance outcome。

模型由 5 个函数组成（对应 MuZero 三件套 h/g/f 之外新增 φ、ψ）：

  • 表示函数 h：o≤t → 初始隐状态 s⁰_t（同 MuZero）。
  • afterstate 动力学函数 φ：(s^k_t, a_{t+k}) → afterstate as^k_t。
  • 动力学函数 g：(as^k_t, c_{t+k+1}) → 下一隐状态 s^{k+1}_t 和奖励 r^{k+1}_t。
  • 预测函数 f：s^k_t → 策略 p^k_t、价值 v^k_t（同 MuZero）。
  • afterstate 预测函数 ψ：as^k_t → 一个 Q 值 Q^k_t，以及未来机会结果的分布 σ^k_t = Pr(c_{t+k+1}|as^k_t)。

Chance outcome 建模：论文提出了 VQ-VAE 的一个定制变体——codebook 大小固定为 M（全部实验统一用 M=32），每个 code 是固定的 one-hot 向量；encoder 输出的 embedding 直接被建模为 categorical 变量，取最近 code 等价于 one_hot(argmax(embedding))，相当于零温度前向 + straight-through 反向估计的 Gumbel-softmax 重参数化。与 Ozair et al.（2021）不同，这里没有显式 decoder，也不用 reconstruction loss，整个模型和 MuZero 一样端到端训练。

Stochastic MCTS：搜索树里 decision node 和 chance node 交替出现（根节点总是 decision node）。chance node 对应一个 afterstate，展开时查询模型得到 Q 值和先验分布 Pr(c|as)，选择阶段按该先验采样 c（实际用 Ozair et al., 2021 附录 A.3 的拟随机采样公式 argmax_c Pr(c|as)/(N(c)+1)）；decision node 的展开和边选择沿用标准 MuZero 的 pUCT 公式不变。

2048 网络配置：representation / afterstate dynamics / dynamics / encoder 四个函数都用 10 个 block 的 ResNet v2 风格 pre-activation 残差塔，配 Layer Normalisation 和 ReLU，每个线性层输出宽度 256；输入是 4×4 棋盘展平成 16 维向量，每个数字用 31-bit 二进制编码，总输入维度 496；value/reward 沿用 MuZero 的 categorical 表示，用 601 个 bin，覆盖 [0, 600]，并用可逆变换 h(x)=sign(x)(√(x+1)−1+εx)（ε=0.001）先缩放目标。

Backgammon 网络配置：board 表示为 28 维向量（24 个点位的双方筹码数 + 2×2 个吃子/成功棋子数）；一个 action 由最多 4 个 micro-action 组成（每个 micro-action = 起始位置×6 + 骰子点数，起始位置 26 种可能），采用 sampled-muzero（Hubert et al., 2021）提出的自回归策略头展开 4 步生成完整 action；value 用标量表示（非 categorical），用 Monte Carlo 回报，discount=1；网络同样是 10-block ResNet v2 + LayerNorm + ReLU，宽度 256，codebook size 32。

Go 网络配置：沿用 sampled-muzero（Hubert et al., 2021）的设置，用 n-step bootstrapping + target network 提升数据效率。为了让 Stochastic MuZero 和 MuZero 计算量相当（Stochastic MuZero 每次 MCTS 展开要过 afterstate dynamics + dynamics 两次网络，相当于 2 倍 network expansions），论文把 Stochastic MuZero 的 chance 和 dynamics 网络深度设为 MuZero dynamics 网络深度的一半。

数据

无预采集数据集，全部通过自对弈/自交互在线生成轨迹：

  • 2048：replay buffer 保留最近 125,000 条序列（每条最长 200 步），用 (Schaul et al., 2016) 的 prioritized replay，priority exponent α=1，重要性采样权重 β=1；训练 20M 步，batch size 1024；value target 用 n-step TD(λ) bootstrapping，n=10，discount=0.999，λ=0.5；根节点注入 Dirichlet 噪声 alpha=0.25、fraction=0.1；温度调度 [1.0, 0.5, 0.1] 分别对应训练步 [1e5, 2e5, 3e5]，之后贪婪选择；MCTS 搜索预算 100 次模拟/步。
  • Backgammon：replay buffer 保留 100,000 局完整对局，均匀采样（非 prioritized）；训练 8M 步，batch size 1024；根节点用自适应 Dirichlet 噪声，alpha=1/√(合法动作数)；因动作空间维度高，沿用 sampled-muzero 提出的 sample-based search；MCTS 搜索预算 1600 次模拟/步。
  • Go：9×9 和 19×19 两种棋盘。9×9 下 MuZero 基线训练用 200 次模拟/评测用 800 次，Stochastic MuZero 训练用 400 次/评测用 1600 次；19×19 下 MuZero 训练用 400 次/评测用 800 次，Stochastic MuZero 训练用 800 次/评测用 1600 次；两条 Elo 曲线各自锚定”最终 MuZero 基线 = 2000 Elo”。
  • 论文未涉及仿真 vs 真实、跨具身或多模态数据混合——三个环境都是纯自对弈/自交互生成的同分布轨迹数据。

训练方法

总损失 L_total = L_MuZero + L_chance。L_MuZero 与原始 MuZero 完全一致（策略用 MCTS 访问分布的交叉熵，价值用 n-step return 回归，奖励用真实即时奖励回归，围棋/双陆棋用 MSE 型损失、2048 用 categorical 损失，具体形式见附录 A）。L_chance 是新增项：

L_chance = w · [ Σ_k l_Q(z_{t+k}, Q^k_t) + Σ_k l_σ(c_{t+k+1}, σ^k_t) ] + β · Σ_k ‖c_{t+k+1} − c̃_{t+k+1}‖²

其中 Q^k_t 用 MSE 朝 value target z_{t+k} 训练，σ^k_t 用 categorical 损失朝 encoder 产生的 one-hot chance code 训练，最后一项是标准 VQ-VAE commitment cost（约束 encoder 输出的连续 embedding 靠近其量化后的 code）。

优化器：2048 用 Adam，学习率 0.0003；Backgammon 用带权重衰减的 Adam（Loshchilov & Hutter, 2017），学习率 0.0003，权重衰减 0.0001；两者 batch size 均为 1024。

Reproducibility 实验：在 2048、Backgammon、Go 三个环境各跑 9 个不同随机种子（为控制算力用了比正式实验更少的训练步），观察到不同种子间性能方差很小，证明方法对随机初始化鲁棒。

论文没有单独的消融实验章节（不同于 sampled-muzero），而是在附录 F 做了一次”chance analysis”：统计每个 chance node 的 Pr(c|as) 分布并对若干局取平均，发现确定性环境（Go）下该分布几乎坍缩到单一 code，随机环境下会用到多个 code；在 Backgammon 中，non-negligible 概率的 code 数量恰好等于 21——正好是两个骰子的所有组合数，说明学到的 chance code 确实对应了物理上正确的随机性结构。

Infra（训练 / 推理工程）

  • 网络用 JAX（Bradbury et al., 2018）+ Haiku（Hennigan et al., 2020）实现。
  • 全部实验用 Google Cloud 第二代 TPU（Google, 2018）。
  • Backgammon：1 个 TPU 用于训练 + 16 个 TPU 用于 acting（自对弈生成数据），耗时约 27 小时，论文换算为约等效单张 V100 GPU 10 天。
  • 2048：1 个 TPU 用于训练 + 4 个 TPU 用于 acting，耗时 80 小时/实验，等效单张 V100 约 8 天。
  • Go：论文只说”沿用 MuZero（Schrittwieser et al., 2020）中相同的 setup”，未给出针对本文实验的新 TPU 数字或耗时——未披露。
  • 推理侧只披露了 MCTS 模拟次数（如上”数据”节所列各环境训练/评测预算），未换算为墙钟延迟、控制频率（Hz）或边缘硬件指标——未披露。
  • 代码未开源：Reproducibility Statement 明确写道”We did not release the full code as it relies on a lot of proprietary internal infrastructure, limiting its usefulness”，转而在附录提供了详细伪代码（Network/MCTS/Self-play/Replay 各模块的类和函数签名）与全部环境专属超参数表，作为可复现性的替代方案。GitHub 上存在若干第三方非官方复现（如 DHDev0/Stochastic-muzero），但均非 DeepMind 官方发布，本文未提供任何官方代码/项目页链接。

评测 benchmark

  • 2048（Figure 2）：Stochastic MuZero（100 次模拟规划、学到的随机模型）性能匹配 AlphaZero（100 次模拟、完美随机模拟器），并且只用了四分之一的训练数据就超过了此前 SOTA——Jaśkowski（2016）基于完美模拟器的 expectimax 树搜索加大量手工领域知识；MuZero（确定性学到的模型）表现明显更差。评测阶段随着搜索预算提升到中等水平（大致相当于 3-ply lookahead）时，Stochastic MuZero 超过 Jaśkowski（2016）的强度，但模拟次数进一步增加后出现收益递减，论文将其归因于学到的模型本身不完美。
  • Backgammon（Figure 3/5）：Stochastic MuZero（学到的随机模型，1600 次模拟/步）达到与 AlphaZero（完美随机模拟器，同样 1600 次模拟）以及超人级开源引擎 GNUbg Grandmaster（3-ply lookahead，平均约 20 种合法着法 × 21 种骰子结果的分支因子）相当的水平；MuZero（确定性学到的模型）表现差。评测阶段模拟预算 scale 到超过 10³ 次时，Stochastic MuZero 的强度超过 GNUbg Grandmaster。
  • Go（Figure 4，9×9 与 19×19）：Stochastic MuZero 在两种棋盘尺寸下都能匹配 MuZero 的 Elo 曲线（Elo 锚定到 MuZero 最终基线 = 2000），验证了引入随机性建模没有牺牲确定性环境下的表现；代价是相同表现下需要约 2 倍的 network expansions（因 chance/decision node 交替展开），论文用把 chance/dynamics 网络深度减半来抵消这部分额外计算，使两者计算量大致等价。
  • Reproducibility（Figure 5.4，9 个随机种子/环境）：各环境下不同种子间性能方差很小，证明方法鲁棒。
  • 论文没有报告 success rate 之类的统一指标，而是用各自领域惯用的评测方式：2048 用累计得分/Elo 曲线，Backgammon 用对战胜率/Elo，Go 用 Elo。

创新点与影响

  • 核心贡献：提出 afterstate 分解，把随机环境的状态转移拆成”确定性动作效果”和”随机机会事件”两段，将学习随机价值等价模型的问题规约为学习 afterstate 和 chance outcome；给出一个不带 decoder、不用 reconstruction loss、端到端训练的定制 VQ-VAE 变体来学 chance code；把 MCTS 扩展为交替的 decision/chance node 结构，使树搜索能在随机环境里做前瞻规划。
  • 改变了什么：这是第一个在纯在线强化学习设置下（无专家数据、无两阶段训练）证明”学到的随机模型 + MCTS”能在实践中有效工作的方法；在 2048（单人随机）和 Backgammon（双人零和随机）两个经典随机博弈上匹配或超过了依赖完美模拟器的 AlphaZero 或专门针对该问题设计的手工 SOTA 方法，同时在确定性的围棋上不牺牲 MuZero 的性能——即 Stochastic MuZero 是 MuZero 的严格泛化，使这一支 value-equivalent model + MCTS 的方法论第一次覆盖了确定性和随机环境两大类问题。
  • 作者自陈局限：(1) 相同表现下比 MuZero 需要约 2 倍的 network expansions（源于 chance/decision node 交替展开），只能靠减半网络深度部分抵消这部分额外计算；(2) 完整代码未开源，理由是依赖大量专有内部基础设施，可复现性依赖论文提供的详细伪代码和超参数而非可运行代码；(3) 2048 上搜索预算继续增加时因学到模型本身的不完美出现收益递减，说明随机环境下模型质量仍是规划效果的瓶颈。

原始链接

一手源存档（sources/）

关系图谱

反向链接