AI Research 技术调研

Policy Improvement by Planning with Gumbel (Gumbel MuZero)

21分钟阅读

一句话定位

用 Gumbel-Top-k 无放回采样替换 AlphaZero/MuZero 根节点的 Dirichlet 噪声 + PUCB 探索机制、用 Sequential Halving（而非累积遗憾意义下的 PUCB）做根节点的简单遗憾（simple regret）优化，证明只要动作价值评估正确就严格保证策略提升（policy improvement）——哪怕每步只用 2 次模拟也能可靠学习；同时用”补全后的 Q 值”（completed Q-values）改写策略网络蒸馏目标和非根节点动作选择，在围棋、国际象棋、Atari 上匹配当时最优结果，且模拟预算越小、相对 MuZero/AlphaZero 的提升越明显。ICLR 2022 Spotlight，发表于 OpenReview，无 arXiv 版本；配套 JAX 库 mctx 随论文开源，后来成为 MCTS+深度 RL 社区事实上的参考实现。

背景与定位

2018 年 AlphaZero（Silver et al., 2018）用同一套算法在围棋、国际象棋、将棋上做到当时最优，随后被 MuZero（Schrittwieser et al., 2020，见 muzero）推广为”连环境规则都是学出来的”价值等价隐模型 + MCTS，进而覆盖 Atari 与连续控制（Hubert et al., 2021，见 sampled-muzero）、离线强化学习（Schrittwieser et al., 2021，见 muzero-unplugged）。但这一整条谱系共享同一套 MCTS 动作选择机制——PUCB（Rosin, 2011）变体 + 根节点 Dirichlet 噪声——这套机制是启发式拼凑出来的，并没有理论上的策略提升（policy improvement）保证。

论文用一个反例说明这一点（Example 1）：设 Q 值 q=(0, 0, 1)，策略网络给出 π=(0.5, 0.3, 0.2)，策略网络自身的价值 Σπ(a)q(a)=0.2。若用 n=2 次模拟，“只搜最可能的 top-n 个动作再取最优”这一常见启发式只会看到动作 {0, 1}，最终选到的动作期望价值 E[q]=0，比什么都不做（直接用策略网络）还差——即根本不是策略提升。AlphaZero 的 Dirichlet 噪声根节点探索也是同理：给潜在最优的策略网络加噪声，并不能保证之后的搜索结果比原策略更好。论文进一步指出，根节点的搜索性能只取决于最终推荐的动作（不依赖搜索过程中间访问了哪些节点，因为根节点没有”祖先”依赖它的评估结果——Bubeck et al., 2011），因此根节点关心的是简单遗憾（simple regret）而非 PUCB 优化的累积遗憾，这也是换用 Sequential Halving（Karnin et al., 2013）的理论依据。

本文围绕 AlphaZero/MuZero 搜索里的五个启发式机制逐一提出有策略提升保证的替代方案：(1) 根节点探索——Dirichlet 噪声 → Gumbel-Top-k 无放回采样；(2) 根节点动作选择——PUCB → 用同一组 Gumbel 值做 Sequential Halving（Algorithm 2）；(3) 环境中实际执行的动作——退火后的访问计数分布采样 → 直接取 Sequential Halving 剩下的最优动作；(4) 策略网络训练目标——访问计数交叉熵 → 基于”补全 Q 值”构造的策略提升目标（KL 蒸馏）；(5) 非根节点动作选择——PUCT → 基于补全 Q 值的确定性动作选择（对应 mctx 库里默认实现的 “Full Gumbel MuZero”）。这套改动同样适用于任何”策略网络 + 代价高昂的 Q 网络”的智能体，对 MuZero 本身即插即用。同期基准工具箱 lightzero 已把 Gumbel MuZero 列为其统一 MCTS/MuZero 算法族里的标准变体之一。

模型架构

本文严格来说是一个搜索/规划算法（root 与非 root 节点的动作选择机制），复用 AlphaZero/MuZero 的网络骨架（表示/动力学/预测函数三件套），架构改动仅限于为提速做的模块替换：

  • 网络骨干（棋类）：把 AlphaZero 原版的稠密残差块换成瓶颈残差块（bottleneck block, He et al., 2016），并每隔 8 层插入一个”广播残差块”（broadcast block，类似 squeeze-and-excitation/全局池化，作者发现比池化版更稳定）。9x9 围棋用 6 层网络，大规模 19x19 围棋用 32 层网络；均为 256 个隐藏 plane、128 个瓶颈 plane，每 8 层一个广播块。架构搜索在固定 300ms/步 的计算预算下比较（如 256-plane/32-block 瓶颈+广播网络可达到 104 次推理/秒，在该预算内相当于 3200 次模拟），同等推理速度下瓶颈块+广播/池化块相对稠密块显著提升 value MSE 与 policy 准确率；但该架构搜索本身是用基线 AlphaZero/MuZero 做的，Gumbel AlphaZero 尚未用于架构搜索（作者自陈）。
  • 网络骨干（Atari）：网络规模与超参直接沿用 MuZero Reanalyse 的配置（Schrittwieser et al., 2021，见 muzero-unplugged），未做改动。
  • 核心算法模块（本文贡献）：
    • Gumbel-Top-k 采样：对 logits 加一组独立 Gumbel(0) 噪声 g，取 argtop(g+logits, n) 得到 n 个无放回采样动作，等价于对策略网络分布做无放回采样（Gumbel-Max trick 的推广）。
    • Sequential Halving with Gumbel（Algorithm 2）：先用 Gumbel-Top-k 从 k 个动作里采样 m 个候选（Go 用 m=min(n,16)，Atari 用 m=min(n,18)，即动作总数），再用 Sequential Halving 把 n 次模拟预算分给 log2(m) 个阶段，每阶段结束淘汰一半候选，比较分数 g(a)+logits(a)+σ(q̂(a))；σ 取线性形式 σ(q̂(a))=(c_visit + max_b N(b))·c_scale·q̂(a)（Eq. 8），该缩放形式借鉴自 MPO 的策略更新（Abdolmaleki et al., 2018）。
    • 补全 Q 值（completed Q-values）：未访问动作的 Q 值用近似的 v_π 补全（Eq. 10），构造改进后的策略 π’=softmax(logits+σ(completedQ))（Eq. 11），用 KL(π’, π) 蒸馏回策略网络（Eq. 12），替代 MuZero 原来”访问计数分布”的交叉熵目标。
    • v_mix 混合价值估计器（Appendix D，Eq. 33）：v_π 的近似值——用价值网络输出 v̂_π 与”已访问动作按先验加权平均的 Q 值”做凸组合插值，权重取决于已访问动作的访问次数总和与先验概率质量；作者自陈这只是一个”不算精细”的估计器，明确留出空间给后续工作探索其他可能性。
    • 非根节点确定性动作选择（“Full Gumbel MuZero”）：推导出与 π’ 均方误差最小的确定性选择规则 argmax_a[π’(a) − N(a)/(1+ΣN(b))]（Eq. 14），取代 PUCT；mctx 库里 gumbel_muzero_policy 默认实现的就是这一”Full”版本。
  • 默认超参（来自官方 mctx 实现）：max_num_considered_actions=16（根节点候选上限）、gumbel_scale=1.0（Gumbel 噪声缩放，论文 Appendix F Figure 8b 显示围棋这类完美信息博弈评测时用 0 噪声即可，且不比有噪声差）、qtransform_completed_by_mix_value 默认 value_scale=0.1、maxvisit_init=50.0（即论文的 c_scale、c_visit，对应 Atari 设置；围棋/国际象棋论文实际用的是 c_scale=1.0，见训练方法节）。
  • 与标准 MuZero 的一个直接简化：Gumbel MuZero 不需要设置任何 Dirichlet 噪声超参数（因为根本不用 Dirichlet 噪声做探索）。

数据

本文是纯自对弈/在线交互式强化学习方法论文，没有预先采集的数据集，“数据”即自对弈生成的训练帧数：

  • 9x9 围棋：训练模拟次数 n∈{2,4,16,32,200}（自对弈时用），评测统一用 800 次模拟；每组实验 2 个随机种子，训练曲线跑到 8 亿帧（Figure 2/3/6/7/8）。参照对手用 Pachi（Baudiš & Gailly, 2011，10k 模拟/步），Elo 锚定 Pachi=1000。
  • 19x19 围棋（大规模）：训练模拟数 n=400，训练曲线到 8 亿帧（Figure 4a），同样锚定 Pachi=1000 Elo。
  • 国际象棋（大规模）：训练模拟数 n=400，从零自对弈训练到 16 亿帧（Figure 4b）；用 AlphaZero（非 MuZero）做对比，因为”国际象棋上 AlphaZero 比 MuZero 学得快”（论文原话）。
  • Atari：Arcade Learning Environment（Bellemare et al., 2013），启用 sticky actions（Machado et al., 2018）；ms pacman 训练模拟数 n∈{2,4,16,18,50}，跑到 2 亿帧，10 个随机种子；beam rider 用于 c_scale 消融，同样 2 亿帧、10 个种子。指标为最近 200 个训练 episode 的平均回报。评测阶段不额外增加模拟次数（Atari 上训练和评测都用同一 n）。
  • 无跨领域/跨模态数据混合，无人工示范数据，无 sim-to-real 概念（本文全部是在各自模拟器/游戏引擎里在线自对弈）。

训练方法

目标函数框架与标准 MuZero 一致（n-step value/reward 回归 + 策略蒸馏），核心改动集中在根节点的动作选择与探索机制、策略网络训练目标、非根节点动作选择三处：

  • 根节点探索 + 选择：不再对策略网络加 Dirichlet 噪声，改为采样 k 个 Gumbel 变量，取 m 个无放回候选（围棋 m=min(n,16)，Atari m=min(n,18)），再用 Sequential Halving 把 n 次模拟预算分配到 log2(m) 个阶段，比较 g(a)+logits(a)+σ(q̂(a))；每阶段至少给新一轮候选分 1 次新访问（max(1, ⌊n/(⌈log2(m)⌉·m)⌋)），预算用完即停止搜索，从”当前访问次数最多的一批动作”里选 g+logits+σ(q̂) 最高者作为 A_{n+1}。Atari 上 n≤18 时候选数等于模拟数，Sequential Halving 根本不会触发，Gumbel MuZero 直接从已访问的 n 个动作里选——论文特别指出这证明了”planning with Gumbel”（Gumbel-Top-k 采样 + 同一 Gumbel 值做最终 argmax）本身才是小模拟数下策略提升的关键，而非 Sequential Halving。
  • σ 缩放常数：σ(q̂(a))=(c_visit+max_b N(b))·c_scale·q̂(a)；围棋和国际象棋用 c_visit=50, c_scale=1.0（对完美信息博弈不敏感，c_visit≥50 结果相近，见 Figure 8a 消融）；Atari 用 c_visit=50, c_scale=0.1（Atari 奖励尺度差异大，Q 值先按标准 MuZero 方式归一化到 [0,1] 再乘 c_visit·c_scale；c_scale 过大会让搜索过度依赖 q̂(a) 而忽视先验 logits，在 partially-observable 的 beam_rider 上明显更差，见 Figure 5b 消融）。
  • 策略网络训练目标：先构造补全 Q 值 completedQ(a)=q(a)（已访问）或 v_π̂（未访问，用 v_mix 近似），再得 π’=softmax(logits+σ(completedQ))，用 KL(π’, π) 作为策略损失（Eq. 12），替代 MuZero 原来对访问计数分布做交叉熵的目标；论文证明这一目标同样满足策略提升（Appendix C）。消融显示：完整的 completed-Q 损失优于”只回归 argmax 动作”的简单损失（Lsimple=-logπ(A_{n+1})，Figure 3a）；且这一 completed-Q 机制是 TRPO MuZero/MPO MuZero（用 TRPO/MPO 正则项做搜索内正则化的两个消融基线）能work的必要条件——不用 completed Q 值时这两者都无法产生策略提升。
  • 非根节点动作选择：“Gumbel MuZero” 沿用标准 MuZero 的确定性 PUCT 式选择；“Full Gumbel MuZero” 换成基于补全 Q 值推导出的、与 π’ 均方误差最小的确定性规则 argmax_a[π’(a)-N(a)/(1+ΣN(b))]（Eq. 14，推导见 Appendix E）。作者明确只在非根节点推荐确定性选择（根节点的 Gumbel 噪声才是不同 episode 间产生探索多样性的来源）。9x9 围棋上 Full 版本相对普通 Gumbel MuZero 只有”小幅”提升（Figure 7）。
  • 探索：训练时自对弈前 30 步用探索性动作（按访问计数比例采样，与 AlphaGo Zero 相同做法），Gumbel MuZero 同样从中受益（Figure 8b）。
  • 关键消融汇总：策略损失 simple vs completed-Q（Fig 3a，completed-Q 更优）；采样数 m 敏感性（Fig 3b，m=4 太小、学习变慢）；v_mix vs 直接用价值网络原始输出补全（Fig 6，v_mix 更好）；非根节点选择 Gumbel MuZero vs Full Gumbel MuZero vs 从 π’ 随机采样（Fig 7，确定性>随机采样，Full 略优于普通）；c_visit 敏感性（Fig 8a，围棋上 c_visit≥50 结果相近）；评测时 Gumbel 噪声 0 vs 随机（Fig 8b，完美信息博弈下零噪声不差）；用小 n 训练的网络在不同评测模拟数下的表现（Fig 9，n=2 训练出的网络即使评测时只给 1 次模拟也能给出合理策略）。

Infra（训练 / 推理工程）

  • 硬件：全部实验使用 Google Cloud TPU v3；论文未披露具体 TPU 芯片总数或 TPU-核时（GPU/TPU-hours）总量。
  • 自对弈/训练算力配比：9x9 围棋上”用 3 倍于训练所用的 TPU 数量做自对弈”以避免训练被数据供给瓶颈限制；模拟数越小，自对弈所需 TPU 数越能成比例减少。
  • 模拟数→训练吞吐加速比（Table 1，9x9 围棋，固定 TPU 数下的训练步速度，以 MuZero/Gumbel MuZero n=200 为基准 1.0×）：
配置相对 n=200 的加速比
MuZero n=2001.0×
Full Gumbel MuZero n=2001.0×
Gumbel MuZero n=2001.0×
Gumbel MuZero n=325.9×
Gumbel MuZero n=1611.3×
Gumbel MuZero n=816.2×
Gumbel MuZero n=424.3×
  • 架构搜索的推理速度基准：固定 300ms/步计算预算下，256-plane/32-block 瓶颈+广播网络可达到 104 次推理/秒，对应约 3200 次 MCTS 模拟；这是论文里唯一给出的具体推理吞吐数字，用于网络架构选择而非 Gumbel MuZero 本身的延迟评测。
  • 推理延迟/控制频率(Hz)/边缘硬件：论文未披露（棋类/Atari 云端训练评测场景，不涉及边缘部署）。

评测 benchmark

  • 9x9 围棋（Figure 2，纵轴 Elo 0–3000，横轴自对弈帧数 0–8 亿，2 个随机种子）：对比 Gumbel MuZero、Replacement（Gumbel MuZero 但改为有放回采样 m 个动作，对应 sampled-muzero 的采样方式）、TRPO MuZero、MPO MuZero、标准 MuZero 五者，在训练模拟数 n∈{2,4,16,32,200} 下的学习曲线。标准 MuZero 在 n≤16 时学不动；Gumbel MuZero 在 n=2 时依然能可靠学习，是论文最核心的实证结果。
  • 19x19 围棋（大规模）（Figure 4a，n=400，锚定 Pachi=1000 Elo）：Gumbel MuZero 达到或超过标准 MuZero 的表现（MuZero 本身在 19x19 围棋上已是很强的基线）。
  • 国际象棋（大规模）（Figure 4b，n=400，训练到 16 亿帧）：对比 Full Gumbel AlphaZero / Gumbel AlphaZero / AlphaZero；论文摘要与结论均明确宣称在围棋、国际象棋、Atari 三个领域都匹配了当时的最优水平。
  • Atari ms pacman（Figure 5a，n∈{2,4,16,18,50}，2 亿帧，10 种子）：标准 MuZero 在 n≤4 时学不好；Gumbel MuZero 在 n=2 时依然学得好；MuZero 本身用 n=50 模拟已是 Atari 当时的最优结果（Schrittwieser et al., 2021）。
  • Atari beam rider（Figure 5b，c_scale∈{0.01,0.1,1,10,100} 对比 MuZero n=50 基线，2 亿帧，10 种子）：c_scale 过大时 Gumbel MuZero 表现明显变差（过度依赖 Q 值、忽视先验）。
  • 论文正文以学习曲线图（Elo / mean return vs 帧数）而非数值表格呈现绝大多数结果；除 Table 1（训练加速比）外没有给出其他数值表格意义上的最终得分，故本页不臆造具体端点数值，如需精确读数需查阅论文原图。

创新点与影响

  • 核心贡献：把 AlphaZero/MuZero 搜索里五个启发式机制（根节点探索、根节点动作选择、实际执行动作选择、策略网络训练目标、非根节点动作选择）逐一替换为有严格策略提升证明的机制，核心工具是 Gumbel-Top-k 无放回采样（用于根节点探索并避免”双重计数偏差”）与 Sequential Halving（用于根节点上的简单遗憾优化）。
  • 消除超参：不再需要 Dirichlet 噪声的 fraction/alpha 超参数。
  • 开源影响：配套发布 JAX 库 mctx（https://github.com/deepmind/mctx，现 https://github.com/google-deepmind/mctx），成为 MCTS+深度 RL 社区的标准参考实现，被 lightzero、Pgx、muax、mctx-az、a0-jax 等下游项目直接复用；据 Semantic Scholar（2026-07-16 抓取）统计，引用数 90，高影响力引用数 26。
  • 作者自陈的局限：(1) 非根节点的确定性动作选择（“Full Gumbel MuZero”）在 9x9 围棋上只带来小幅提升（Figure 7），改进幅度有限；(2) v_mix 混合价值估计器被作者自己评价为设计粗糙、不够精细，明确留出空间给后续工作探索其他可能性；(3) 论文的网络架构搜索（瓶颈块/广播块）是用基线 AlphaZero/MuZero 做的，Gumbel AlphaZero 尚未用于架构搜索；(4) Atari 部分作者明确指出未来可尝试按优势估计器的标准差归一化、以及像 Muesli（Hessel et al., 2021）一样裁剪归一化后的优势值，均属未做的方向。

原始链接

一手源存档（sources/）

  • gumbel-muzero—openreview.md —— OpenReview 论坛页摘要/元数据快照（经 Wayback Machine 20260302083649 快照获取；实时 OpenReview 页面被 Cloudflare Turnstile 人机验证拦截，等待 30 秒仍未通过）
  • gumbel-muzero—github-readme.md —— mctx 官方 README + 核心实现（policies.py/qtransforms.py/seq_halving.py）关键片段
  • 论文全文 PDF 经同一 Wayback Machine 快照获取并在本地读取用于撰写本页六个维度（不入 git，遵循仓库策略；见上方原始链接获取权威 URL）

关系图谱

反向链接