AI Research 技术调研

LightZero: A Unified Benchmark for Monte Carlo Tree Search in General Sequential Decision Scenarios

18分钟阅读

一句话定位

LightZero 是首个统一的开源 MCTS/MuZero 算法基准与工具箱——把 9 个算法变体(AlphaZero、MuZero、MuZero w/SSL、EfficientZero、Sampled EfficientZero、Stochastic MuZero、Gumbel MuZero 等)和 20 余种环境(棋类、Atari、MuJoCo、MiniGrid、GoBigger、2048)统一到同一套解耦的四模块训练流水线里,NeurIPS 2023 Datasets and Benchmarks Track Spotlight。

背景与定位

muzero 之后,MCTS+学习模型这条路线分裂成多个各自独立发展的分支:sampled-muzero 用采样动作+重要性加权把 MuZero 扩展到高维/连续动作空间,stochastic-muzero 引入 afterstate 建模环境的内在随机性,efficientzero 用自监督一致性损失+value_prefix 大幅提升样本效率,muzero-unplugged 提出 Reanalyze 机制统一在线/离线训练,Gumbel MuZero 用 Gumbel-Top-k trick 在低模拟预算下保证策略改进。这些工作各自验证了单一改进方向,但缺乏统一实现:Sampled MuZero 官方从未开源;AlphaZero-General、MuZero-General 各自只支持单一算法且没有分布式实现;EfficientZero 虽支持多卡但代码仅限自身算法;KataGo 专精 AlphaZero+围棋且训练算力门槛很高。LightZero 是第一个把这些分支代码统一进同一套接口、可以任意组合各家技巧的基准库:论文先系统总结了通用 MCTS 求解器要面对的六个环境挑战维度(多模态观测空间、复杂动作空间、先验知识依赖、内在随机性、模拟成本、难探索),再把过去紧耦合的”算法设计”与”系统设计”解耦为数据采集器(Data Collector)、数据整理器(Data Arranger)、智能体学习器(Agent Learner)、智能体评估器(Agent Evaluator)四个可插拔子模块。这是一篇基准/工具箱论文而非新算法论文,贡献主要在”统一实现 + 系统解耦 + 两个衍生洞察”,而非提出新的 world model 结构。

模型架构

LightZero 本身不引入新的骨干网络,而是在统一接口下复现并整合各分支论文的架构,核心仍是标准 muzero 的三网络结构:

  • Representation Network $h_\theta$:把历史观测 $o_1,\dots,o_t$ 编码为隐状态 $s^0$。图像输入(Atari)用卷积骨干(Figure 23),向量输入(连续控制等)把卷积层替换为对应的全连接层,其余结构不变。
  • Dynamics Network $g_\theta$:给定隐状态和动作,输出转移后隐状态与预测奖励 $r^k, s^k = g_\theta(s^{k-1}, a^k)$(Figure 24 左)。
  • Prediction Network $f_\theta$:给定隐状态,输出策略与价值 $p^k, v^k = f_\theta(s^k)$(Figure 24 右)。
  • MCTS 核心公式沿用标准 PUCT:$a^* = \arg\max_a Q(s,a) + P(s,a)\cdot[c_1 + \log(\frac{\sum_b N(s,b)+c_2+1}{c_2})]\cdot\frac{\sqrt{\sum_b N(s,b)}}{1+N(s,a)}$,默认 $c_1=1.25$, $c_2=19652$(Dirichlet 噪声 alpha=0.3、权重 0.25,加在根节点)。
  • Value/Reward 表征:categorical 分布(离散化的支撑集),默认支撑集缩放 scale=300(Table 7,Atari 设置)。
  • 树搜索实现是本文的一个工程贡献:提供纯 Python(ptree)和 C++/Cython 扩展(ctree)两套实现,后者用于消除 Python getattr/setattr 等基础操作在高频调用下的开销。
  • 各分支算法在这套骨干上叠加各自的改动:EfficientZero 加自监督一致性损失(投影头+预测头,stop-gradient,类似 SimSiam);Sampled (Efficient)Zero 把先验替换成采样版 $\hat\pi_\beta$;Stochastic MuZero 加 afterstate dynamics/prediction 网络与两层 MLP 的 chance encoder;Gumbel MuZero 不改网络结构,只改动作选择与策略改进的计算方式。

数据

纯 online RL 自交互/自对弈,不依赖预训练数据集(接口上支持离线数据集输入,但主实验都是在线训练):

  • 棋类:TicTacToe、Connect4、Gomoku(主实验 board_size=6)、Chess、Go(棋类支持矩阵见 Table 5,Chess/Go 部分算法仍在”waiting list”标记为未完全测试)。
  • Atari:主实验选取 Pong、Qbert、Ms.Pacman、Breakout、UpNDown、Seaquest 六款代表性游戏(源自 Arcade Learning Environment 57 款游戏中的子集)。
  • MuJoCo:Hopper-v3、Walker2d-v3(主实验),附录另提及 HalfCheetah-v3、Humanoid-v3。
  • MiniGrid:KeyCorridorS3R3-v0、FourRooms-v0 等稀疏奖励网格世界。
  • GoBigger:多智能体合作-竞争环境,子环境按队伍数(t)和每队智能体数(p)命名(t2p2/t2p3/t3p2/t4p3 等),论文主实验用 t=p=2 和 t=2,p=3 两个场景,独立学习范式(每个智能体共享同一策略-价值网络,各自独立做搜索)。
  • 2048:随机性环境,论文比较 num_chances=2(随机块为{2,4})与 num_chances=5(随机块为{2,4,8,16,32})两种设定。
  • 具体规模数字(Table 7,MuZero w/SSL@Atari):replay buffer size(以 transition 计)= 1e6;每个 game segment 长度 = 400;TD steps = 5;unroll steps = 5;帧堆叠 = 4,帧跳 = 4。
  • 评测协议:每个算法独立跑 5 个随机种子,评测阶段取 20 个 episode 的平均 Return 作为指标,曲线阴影区为跨种子标准差。

训练方法

  • 优化器与学习率(Table 7,MuZero w/SSL@Atari 默认配置):Adam,学习率 3×10⁻³;discount factor 0.997(棋类/2048 环境用 1);policy/value/reward loss 权重 1/0.25/1;policy entropy loss 权重默认 0(不启用);SSL 一致性损失权重 2(而”MuZero”不加 SSL 时权重为 0);batch size 256;model update ratio 0.25;目标网络更新频率 每 100 步;weight decay 1e-4;max gradient norm 10;augmentation 开启;离散动作用 One-Hot 编码;归一化用 Batch Normalization;priority exponent/correction 系数 0.6/0.4;默认模拟次数(sim)50;reanalyze ratio 默认 0(基础配置未启用,SpeedyZero 等后续工作会调整该比例)。
  • 算法差异化超参(附录 G.4,Table 8-14):
    • Sampled EfficientZero(Atari):采样动作数 K 按环境不同(见 Fig.7 图注),policy loss 用交叉熵。
    • Sampled EfficientZero(MuJoCo 连续控制,Table 9):帧堆叠/帧跳=1,不裁剪奖励,K=20,policy entropy 权重 0.005,game segment 长度 200,sim=50,max grad norm=0.5。
    • Gumbel MuZero(Atari,Table 10):考虑动作数上限=动作空间大小,Gumbel scale=10,max visit init=50,value scale=0.1。
    • Stochastic MuZero(Atari,Table 11):chance space size = 16×可能的随机块数,帧堆叠/帧跳=1,discount=0.999,batch=512,game segment=200,TD steps=10,sim=100。
    • Stochastic MuZero(2048,Table 12):board_size=6,discount=1,policy/value loss 权重均为 1,sim=100,不使用 categorical 价值建模。
    • AlphaZero(Gomoku,Table 13):board_size=6,frames stacked=1,discount=1,SSL 权重=0(AlphaZero 无 dynamics 网络故不适用一致性损失),game segment/TD steps=18,不用 augmentation,sim=100,categorical 支撑集 scale=10;MuZero(Gomoku,Table 14)沿用同一 Gomoku 环境设置,但表格本身具体数值未在抓取的文本中完整呈现,不逐项列出。
  • 两个算法层面的 case study(第 5 节):
    1. 探索机制:在 MiniGrid-KeyCorridorS3R3/FourRooms 上对比 Naive(手工温度衰减)、NaiveDoubleSimulation(加倍模拟次数)、FixedTemperature、PolicyEntropyRegularization(权重 0.05/0.005)、EpsGreedy、IntrinsicExploration(基于 RND)七种策略。RND 用原始观测 $o_t$(而非隐状态,论文指出隐状态方案效果差,因为训练过程中隐状态分布持续漂移导致目标网络输出等同噪声)作为输入,内在奖励用 min-max 归一化,与外在奖励按 $r_t = r_t^e + \beta r_t^i$ 混合,权重 $\beta = 1/300$(保证单集内内在奖励总和小于最大外在奖励 1)。
    2. 环境模型对齐:用 MuZero 做基线,对比”MuZero w/SSL”(一致性损失权重 2)与”MuZero”(权重 0)在图像输入(Pong)、向量输入(LunarLander)、棋盘图像输入(TicTacToe)三类观测上的差异。一致性损失对标准图像输入最关键;向量输入增益较小;棋盘二值图像的余弦相似度始终偏低(训练 0.8M 步时仍只有 0.662),且用 SGD 优化器时加一致性损失反而拖慢 TicTacToe 早期学习。

Infra(训练 / 推理工程)

  • 硬件配置:几乎所有基础实验和消融都部署在 Kubernetes 集群上,单实例分配 1 张 A100 40G GPU + 24 个 CPU 核心 + 100GB 内存。
  • 训练时长(同一份小资源实例下):Atari 训练 100K 步耗时 4 小时;Gomoku 自对弈 100K 步耗时 5 小时。
  • 多卡 scaling(EfficientZero,PongNoFrameskip-v4,PyTorch DDP,Table 4):训练到 1M env steps 所需时间——1 卡 844 分钟,2 卡 363 分钟,4 卡 152 分钟;论文原文表述为 4 卡相对 1 卡获得约 5 倍加速(“approximately five-fold”),同时性能水平相近。
  • 系统级优化(附录 E,Efficiency Analysis):
    • 环境延迟:用 JAX 的 vmap/jit + LRU cache 加速棋类环境的 step/reset(AlphaZero 类方法单次动作选择要调用环境 step 50+ 次,环境交互常占训练总时长一半以上)。
    • 树搜索:用 C++/Cython 扩展(ctree)重写 MCTS 核心操作,规避纯 Python 中 getattr/setattr 等原语在高频调用下的开销;指出朴素向量化环境方案对树搜索类方法增益有限,因为每个环境需要独立的搜索树,不适合批处理。
    • 模型推理:利用 PyTorch 2.0 特性、混合精度训练、大 batch 训练等社区工具降低推理成本。
    • 并行 MCTS:提出”分组批处理”采集方案(Figure 15e)——把 k 个环境+搜索树分成若干组交替执行 env step / 模型推理 / 树搜索三个阶段,使三者时间重叠,相比朴素批处理(Figure 15d)减少 GPU/CPU 之间的相互等待。
    • Reanalyze:借鉴 SpeedyZero 的做法,只对训练潜力高(状态-动作新颖度高、TD 误差大)且训练稳定的数据做高频 reanalyze,其余数据降低 reanalyze 频率以节省时间。
    • 通信:用 RDMA(Remote Direct Memory Access)实现推理 GPU 到训练 GPU 的 P2P 直传,绕过传统”GPU Tensor→CPU Tensor→压缩→Numpy→序列化→网络传输→反序列化→…”的多步流程(Figure 16);用 throughput limiter 把数据生产/消费比控制在固定 batch size 下的 0.8×-1.2× 区间内,避免过度采集或过拟合。
  • 推理 FPS / 控制 Hz / 边缘硬件指标:论文未披露——只给出训练时长和多卡加速比,没有以 FPS 或控制频率(Hz)为单位的推理延迟数字,记为未披露。

评测 benchmark

以下数值均来自论文正文/附录 B 的实验描述(多为学习曲线图,论文本身较少给出精确终值表格,如实按图注/正文描述转述):

  • 棋类(Fig. 8,Connect4/Gomoku board_size=6):AlphaZero 样本效率明显高于 MuZero(有模拟器优势);但 MuZero 即使没有模拟器依然给出令人满意的结果。
  • Atari 六款代表性游戏(Fig. 4/7):MuZero w/SSL 在其中 5 款上超过 MuZero,仅 MsPacman 上表现相近;EfficientZero(w/SSL+value_prefix)只在 MsPacman、Breakout 上超过 MuZero w/SSL,其余环境相近,而在 MiniGrid 稀疏奖励设置(Fig. 12)下 EfficientZero 明显更差——论文据此提出”是否预测 value_prefix 需视环境属性谨慎决定”。Sampled EfficientZero 在 Breakout(动作空间维度 4)上性能持平 EfficientZero,在 MsPacman(维度 9)上性能下降。
  • 连续控制(Fig. 9,Pendulum-v1/LunarLanderContinuous-v2/Hopper-v3/Walker2d-v3,动作维度依次 1/2/3/6,K=20、sim=50 固定):factored(离散化)策略表示随动作维度指数级增大(11/49/125/15625)而性能单调下降;Gaussian 策略表示保持相对稳定,但两者在 MuJoCo 上均不及 model-free 方法——论文推测因 MuJoCo 最优动作常在 ±1 极值附近,高斯采样难以命中极值。
  • Gumbel MuZero(Fig. 10,Gomoku sim={20,10}、LunarLander-v2 sim={20,10,5}、Atari sim={50,16,2}):模拟预算越受限,Gumbel MuZero 相对 MuZero 的优势越明显。
  • Stochastic MuZero(Fig. 11,2048 环境):num_chances=2 时略优于 MuZero;num_chances=5(随机性更强)时优势减弱。
  • MiniGrid 探索消融(Fig. 5,KeyCorridorS3R3-v0):仅 IntrinsicExploration(RND)取得近最优表现;单纯加倍模拟次数(NaiveDoubleSimulation)未见明显提升;EpsGreedy/FixedTemperature/PolicyEntropyRegularization 方差偏大且不能保证充分探索。
  • GoBigger 多智能体(Fig. 13,T2P2/T2P3,独立学习范式):MuZero/EfficientZero 均能对built-in bot 稳定收敛;论文指出 T2P2 场景下 MuZero 约 400K env steps 达到 return 150K,而根据 GoBigger 原论文数据,MAPPO 类方法需约 3M env steps 才能达到同等水平——相对非 MCTS 方法样本效率提升约 6 倍。
  • 六维雷达图定性打分(Fig. 2 / 附录 D):对 PPO(model-free)、AlphaZero、MuZero、“LightZero”(作者定义为整合了各子模块最优组合的特化版本,非单一算法)在多模态观测/复杂动作空间/先验依赖/内在随机性/模拟成本/难探索/数据效率七个维度做 1-5 分定性打分——这是作者主观评估的定性对比图,不是严格量化 benchmark,论文也做了明确说明。
  • 与 MCTX 库对比(附录 H,Table 5/6):LightZero(PyTorch)整合 9 种算法 × 20 余种环境,而 MCTX(JAX)及其衍生库中 Gumbel MuZero/Stochastic MuZero 的完整训练流水线仍处于开发中(“in development”状态),环境覆盖面也更窄。

创新点与影响

  • 首个统一的 MCTS/MuZero 开源基准:把此前分散、部分未开源(如 Sampled MuZero)、部分单算法单卡(AlphaZero-General、MuZero-General)、部分单算法多卡(EfficientZero)的实现统一进同一套代码库、同一套接口,支持技巧任意组合。
  • 系统解耦贡献:提出数据采集器 / 数据整理器(含优先级重算、data reanalyze、throughput limiter)/ 智能体学习器 / 智能体评估器四个高内聚低耦合的子模块划分,让 tree search、data reanalyze 这类原本与算法紧耦合的机制可以被独立替换、优化、组合,并配套 Python(ptree)+ C++/Cython(ctree)双实现降低树搜索开销。
  • 两个衍生算法洞察:(1)RND 内在奖励能以很低的实现代价嵌入 MuZero,有效缓解稀疏奖励环境下 MCTS 探索不足的问题;(2)EfficientZero 的自监督一致性损失收益高度依赖观测类型——对标准图像输入关键,对向量输入增益有限,对棋盘类图像甚至可能因优化设置不当而拖慢学习,论文认为棋类场景下的对齐损失设计仍是开放问题。
  • 后续影响:该代码库之后延伸出 UniZero(TMLR 2025)、ReZero(CoRL 2025 workshop)、ScaleZero(ICLR 2026,统一多任务 world model)等后续工作,并在原论文中就提出了”MCTS 作为 LLM 高层任务规划的底层执行器”以及吸收 RSSM/DreamerV2 式隐状态一致性损失等展望方向。
  • 作者自陈局限(附录 I):(1)当前所有 LightZero 算法都不直接支持混合(discrete+continuous)动作空间,需要借助 action representation 等技术适配;(2)受 MCTS 本身局限,Sampled MuZero、Stochastic MuZero 在高维动作空间或强随机性场景仍需进一步优化;(3)MCTS+RL 算法本身学习门槛高,对新手不友好,文档和社区建设仍是未来重点。

原始链接

一手源存档(sources/)

  • lightzero—github-readme — 官方 GitHub README 快照(算法×环境支持矩阵、安装方式、benchmark 图索引、后续论文 UniZero/ReZero/ScaleZero 的更新记录)
  • 论文全文见上方 arXiv 链接(arXiv 原文 PDF,不入 git)

关系图谱

反向链接