LightZero：以 MCTS 为帆，航向决策 AI 的星辰大海

引言

图1：AlphaGO对李世石五番棋

2016 年 AlphaGo [1] 与李世石的世纪对决可谓是开启了近十年来人工智能技术探索和发展的浪潮。在围棋这一人类最经典优美的智力游戏上，传统的搜索算法往往迷失在浩如烟海的搜索空间中，朴素的神经网络设计也难以捕捉复杂多变的决策序列数据分布，但是，巧妙地将二者相结合，化用蒙特卡洛树搜索算法（Monte Carlo Tree Search, MCTS）和神经网络算法各自的优势，就孕育出了媲美甚至超过人类围棋职业棋手水准的 AlphaGo 智能体。而这背后所对应的 MCTS 系列技术流派，也从围棋一路拓展到象棋 [2]、视频游戏 [3]、视频编解码 [4] 和高性能计算 [5] 等等不同决策智能领域，进而成为决策智能领域数据利用率最高，训练稳定性最强的系列算法之一。

图2:决策智能相关业态发展‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍

但另一方面，由于 MCTS 系列技术涉及了多种经典算法的拓展与组合，其运行模块的设计就需要包含多方面的内容，进而导致这样的窘境——即虽然算法的能力在不断进步，但是算法的应用难度却逐年升高。具体来说，其中包含强化学习分布式系统的构建，树搜索算法的高效实现和并行化，经验回放中数据的统一管理与高效增删改查操作的执行，深度神经网络的高效推理与数据并行训练等技术栈。因此，开源社区对于相关算法工具的关注度也日渐水涨船高，相继出现了 KataGo [6]、AlphaZero-General [7]、MuZero-General [8] 等一系列工具项目，不断降低 MCTS 技术的使用门槛和应用开销。不过，已有的开源项目更多偏向于棋类游戏或是入门级示例环境，仍然限制了 MCTS 类算法的应用范围，无法拓展到更多有趣或有意义的用途，研究界最新的技术成果和前沿视角也未能完全融入现有工具，导致相关技术仍然是少数研究者手中的“屠龙之术”。

为了充分发挥 MCTS 系列算法技术的潜力，提升相关技术在各个决策智能领域的易用性与实用价值，上海人工智能实验室开源决策智能平台（OpenDILab）团队开源了 LightZero 项目，致力于打造一统 MCTS 的集大成之术。

欢迎体验LightZero（点击文末阅读原文可直达）：

https://github.com/opendilab/LightZero 

起源：从 AlphaGo 到 MuZero 全家桶‍‍‍‍‍‍

图3：MCTS 算法族发展历史‍‍

🌟MCTS：强大且高效的算法神器‍‍‍‍‍

蒙特卡洛树搜索（MCTS）是一种高效的在线规划（Planning）方法，对大规模的、机制复杂的决策问题拥有极高的效率与稳定性，在棋类、视频游戏等诸多领域展现了卓越性能。这种方法借助模拟当前决策状态的期望结果，评估决策搜索树中的各个分支，选出最优方案。与传统搜索方法相比，MCTS 具有以下三大优势：

1.高效：MCTS 是一种高度选择性的、最强优先的搜索方法，能迅速关注搜索空间中最有潜力的区域。它并不会考虑决策空间所有可能的情形，而是通过合理的采样与分析选择某些局部，从而克服维数诅咒。

2.易用：MCTS 只需一个黑盒的决策问题模拟器，就可以高效地构建搜索树，从而解决各种各样过大、过复杂、以至难以显式建模概率分布的问题。

3.强大：MCTS 利用随机模拟来评估决策行为的长期潜力，从而实现长时间周期的规划。通常情况下，无需任何启发式搜索方法或先验领域知识。配合一些高效的探索策略，MCTS 能够稳定地收敛到最优策略。

前文提到的 AlphaGo 背后的核心技术就包括蒙特卡洛树搜索、深度神经网络（Deep Neural Networks）以及强化学习（Reinforcement Learning）。通过这些技术，AlphaGo 可以评估围棋棋盘的局势，并自主地选择最佳的下一步行动。在 AlphaGo 出现之前，人们普遍认为围棋领域的人工智能还需要很长时间才能达到顶级水平。然而，AlphaGo 的成功打破了这一观念示了人工智能在处理复杂问题上的潜力。

图4：决策搜索树

近年来，研究者通过不断的算法创新，将 AlphaGo 进行了一系列扩展，以解决各种实际问题。例如：

• AlphaZero  [2]：不依赖于除游戏规则之外的任何先验知识，通过完全无人参与的自我对弈（self-play）进行训练，无需借助人类围棋棋谱，从随机策略出发，可在国际象棋、将棋和围棋等游戏中达到超越人类的水平。

• MuZero  [3]：即使不知道游戏规则（比如环境状态转移函数），也可通过学习抽象的环境模型，在 Atari 视频游戏和传统棋类上达到最先进的水平。

• EfficientZero  [10]  和 Gumbel MuZero [11] ：结合深度学习、强化学习、最优控制等多方面技术，通过各种技巧提高样本效率和算法性能，进一步扩展 MCTS 技术在各种场景的效率和稳定性。

💡MCTS 的困境（1）：在多样化动作空间中的挑战‍‍‍‍‍‍‍

图5：基于视频游戏的竞赛和学习环境中采取的行动空间类型汇总

在现实世界中，许多关键问题的动作空间是高维和连续的，甚至是复杂的结构化空间，这使得决策智能体完全枚举所有可能的动作变得不切实际。关于各种环境中具体的动作空间类型，可以参考上图，其中总结了各种基于视频游戏的竞赛和学习环境中采取的行动空间类型 [17]。

尽管 MCTS 结合深度强化学习具备强大的建模能力，但长期以来，其应用仍主要局限于低维离散动作空间，例如各种棋类游戏。一旦涉及高维或连续动作空间（比如机器人控制），计算效率将受到极大影响。

近年来，Sampled MuZero [12] 提出了一种基于采样的方法，通过在动作空间子集上计算改进策略，当采样动作数量趋近整个动作空间时，采样的改进策略按概率收敛于整个动作空间上的改进策略。这使得模型能够在具有任意复杂动作空间的环境中学习。然而，由于算法细节和系统设计的复杂性，整个开源社区难以找到完整且高质量的实现，导致 MCTS 系列技术最新的研究进展依然难以应用于实际现实任务。

💡MCTS 的困境（2）：在工具化和应用化中的迷茫

另一方面，尽管 MCTS + RL 方法在算法性能和样本利用效率方面都拥有极高水平，但其整体训练系统构建十分复杂，核心算法设计模块细节颇多，对于应用社区来讲，直接将相关算法推广到其他决策智能领域会遇到诸多难题，包括但不限于：仿真器的获取难度与运行效率，算法超参数及其调整策略，训练系统的数据吞吐量等等。因此，整个社区仍然在期待更高效的改进算法和更完善的开源工具。

值得欣慰的是，近年来的研究正朝着这一方向发展。MuZero Unplugged [13]  推出了重分析（Reanalyze）技术，一种简单高效的在线和离线下都能取得良好性能的 MuZero 增强技术。ROSMO  [14]  针对离线强化学习（Offline RL）场景下 MuZero 的潜在问题进行了研究，并建议采用正则化的一步展望方法解决相关问题。

在工具方面，许多算法在社区中缺乏高效且全面的开源实现。例如，Sampled MuZero [12]  尚无相应的社区实现。AlphaZero-General [7] 和 MuZero-General  [8]  分别只支持单一算法，且不具备分布式实现。虽然 EfficientZero [18] 支持分布式实现，但它仅限于单一的 EfficientZero 算法。KataGo [6] 主要关注于 AlphaZero 算法和围棋 AI 研究，但其训练和推理过程对计算资源要求较高，可能给普通用户带来一定的硬件门槛。

目前，尚无统一框架涵盖 AlphaZero/MuZero 系列算法。现有的单一实现通常较为复杂且难以理解，导致学术论文中的改进难以应用于实际任务。近年来的算法创新主要集中在 MCTS+RL 核心算法逻辑的某一方面，很少有工作全面分析算法实现细节中的关键要素。这些因素都使得 MCTS 在工具化和应用化方面面临着挑战。

开拓：LightZero--全新的 MCTS 系列算法工具库‍‍‍‍‍‍‍‍‍‍

为了进一步追踪上述问题，推动相关领域的发展，上海人工智能实验室开源决策智能平台（OpenDILab）团队推出了 LightZero：首个全面、高效且易于扩展的开源 MCTS 系列算法工具库。基于深度学习框架 PyTorch  [15] 和强化学习平台 DI-engine  [16] 实现。

🌟轻量级与可拓展性‍‍‍‍‍

LightZero 集成了具有连续与离散动作、向量与图像输入、是否完全信息博弈、是否包含环境随机性等不同属性的任务。在统一的框架下，LightZero 支持多种 MCTS 系列算法，可轻松解决各类属性的决策问题，具备优秀的扩展性，支持快速上手自定义环境，更易适应复杂决策问题中的各类需求。

目前，LightZero 支持的环境及算法如下表所示，其中"✔" 表示相应项目已完成并经过充分测试。"🔒" 表示相应项目正在等待列表中（进行中）。"---" 表示该算法不支持此环境。

Env./Alg.	AlphaZero	MuZero	EfficientZero	Sampled EfficientZero
Atari	---	✔	✔	✔
Tictactoe	✔	✔	🔒	🔒
Gomoku	✔	✔	🔒	🔒
Lunar Lander	---	✔	✔	✔
BipedalWalker	---	✔	✔	✔
Cartpole	---	✔	✔	✔
pendulum	---	✔	✔	✔

表1：LightZero 支持的环境和算法表（截止至2023年5月）

图5：LightZero 框架流程图‍‍‍‍‍‍‍

图5展示了 LightZero 的框架流程图，其中包含了三个核心模块：

• Policy：Policy  描述了网络更新方式以及与环境的交互方式，它包括三个过程：训练过程（learn）、采样过程（collect）和评估过程（evaluate）。

• Model：Model 负责定义网络结构，包含__init__ 函数用于初始化网络结构，以及forward 函数用于计算网络的前向传播。

• MCTS：MCTS 则定义了蒙特卡洛搜索树的结构以及与Policy的交互方式。MCTS 的实现有 python 和 cpp 两个版本，分别在 ptree 和 ctree 中实现。

LightZero 仓库的设计理念注重用户友好性，旨在让开发者能够轻松地将新环境集成到项目中，并便捷地使用各类算法。具体示例如下所示：

# 下面的 main_config, create_config 分别对应算法的主要参数设置和算法的各模块创建的设置，具体可以参见 cartpole_muzero_config


# 对于 LightZero zoo 里面已经集成的环境，通过下面的入口函数，便可以调用 LightZero 集成的算法
from lzero.entry import train_muzero
train_muzero([main_config, create_config], seed=0, max_env_step=int(1e5))


# 对于用户自定义的以标准 Gym 接口封装好的环境，通过下面的入口函数，便可以调用 LightZero 集成的算法
from lzero.entry import train_muzero_with_gym_env
train_muzero_with_gym_env([main_config, create_config], seed=0, max_env_step=int(1e5))

🌟高效性与易用性‍‍‍‍‍‍‍‍‍‍‍‍‍‍

为提升计算效率，LightZero 精心设计了针对 MCTS 系列算法的混合异构计算编程策略。通过采用基于 Python 的主运行流程，结合 Cython 的中间接口，调用 C++ 的 MCTS 实现，从而在保持主流程易读性的同时，大幅提高了算法执行效率。

此外，LightZero 为所有集成的算法提供了算法框架图（一个典型示例如下图所示）及详细文档，帮助用户理解算法内核，并在同一范式下比较算法异同。欢迎参见论文笔记。此外，LightZero 还为算法代码实现提供了函数调用图和网络结构图，便于用户快速定位关键代码。

论文笔记：

https://github.com/opendilab/LightZero/tree/main/assets/paper_notes

图6：LightZero 算法框架图示例‍‍‍‍‍‍‍‍‍‍‍‍‍‍

基于上述丰富的功能，LightZero 在具有连续与离散动作、向量与图像输入、棋类与非棋类等不同属性的任务上评估了各个算法的性能，并给出了详细的基准实验结果，具体参见基线结果。更进一步地，为深入探究算法核心要素，LightZero 还对关键超参数和核心设计进行了实验分析，提供了详细讨论，并最终总结为一个实用的 MCTS+RL 算法指南。

近年来的算法创新主要集中在 MCTS+RL 核心算法逻辑的某一方面，很少有工作全面分析算法实现细节中的关键要素。这些因素都使得 MCTS 在工具化和应用化方面面临着挑战。

基线结果：

https://github.com/opendilab/LightZero/tree/main#benchmark

展望：Roadmap‍‍‍‍

优秀的开源项目一定是不断生长壮大的，LightZero 开发团队希望以初始版本的发布为契机，吸引广大开源社区爱好者不断探索、尝试并提升 LightZero，让 MCTS 系列技术再次掀起波澜。具体来说，LightZero 将首先作为一个全面且易用的基准工具，助力各类基于 MCTS/MuZero 系列算法的学术研究问题，探索相关技术对于不同数据质量，不同搜索条件，不同决策空间等方向下的适用性和通用性。

图7：LightZero 的开源愿景‍‍‍‍‍‍‍‍‍‍‍‍‍

另一方面，从应用的角度出发，LightZero 也将提供主流类型决策问题环境的 Hugging Face Model Zoo，支持开发者通过多种方式快速搭建自定义的决策智能应用，探索智能体在不同类型问题中的角色定位。也就是说，智能体不仅可以作为最强大的胜利者，还可以成为协作模式中最精彩的弄潮儿，又或是化身引人入胜的最佳拍档，万千种可能正待发掘与强化。

更多详细的上手指南、协作规范和未来开发计划，欢迎参阅 LightZero 开源贡献小贴士。演化即无限！

LightZero 开源贡献小贴士：

https://github.com/opendilab/LightZero/tree/main/assets/contributor_tutorial.pdf

参考文献

[1] Silver D, Huang A, Maddison C J, et al. Mastering the game of Go with deep neural networks and tree search[J]. nature, 2016, 529(7587): 484-489.

[2] Silver D, Hubert T, Schrittwieser J, et al. Mastering chess and shogi by self-play with a general reinforcement learning algorithm[J]. arXiv preprint arXiv:1712.01815, 2017.

[3] Schrittwieser J, Antonoglou I, Hubert T, et al. Mastering atari, go, chess and shogi by planning with a learned model[J]. Nature, 2020, 588(7839): 604-609.

[4] Mandhane A, Zhernov A, Rauh M, et al. Muzero with self-competition for rate control in vp9 video compression[J]. arXiv preprint arXiv:2202.06626, 2022.

[5] Fawzi A, Balog M, Huang A, et al. Discovering faster matrix multiplication algorithms with reinforcement learning[J]. Nature, 2022, 610(7930): 47-53.

[6] Wu D J. Accelerating self-play learning in go[J]. arXiv preprint arXiv:1902.10565, 2019.

[7] https://github.com/suragnair/alpha-zero-general

[8] https://github.com/werner-duvaud/muzero-general

[9] https://papers.nips.cc/paper_files/paper/2010/file/edfbe1afcf9246bb0d40eb4d8027d90f-Paper.pdf

[10] Ye W, Liu S, Kurutach T, et al. Mastering atari games with limited data[J]. Advances in Neural Information Processing Systems, 2021, 34: 25476-25488.

[11] Danihelka I, Guez A, Schrittwieser J, et al. Policy improvement by planning with Gumbel[C]//International Conference on Learning Representations. 2022.

[12] Hubert T, Schrittwieser J, Antonoglou I, et al. Learning and planning in complex action spaces[C]//International Conference on Machine Learning. PMLR, 2021: 4476-4486.

[13] Schrittwieser J, Hubert T, Mandhane A, et al. Online and offline reinforcement learning by planning with a learned model[J]. Advances in Neural Information Processing Systems, 2021, 34: 27580-27591.

[14] Liu Z, Li S, Lee W S, et al. Efficient Offline Policy Optimization with a Learned Model[J]. arXiv preprint arXiv:2210.05980, 2022.

[15] https://pytorch.org/

[16] https://github.com/opendilab/DI-engine

[17] Kanervisto A, Scheller C, Hautamäki V. Action space shaping in deep reinforcement learning[C]//2020 IEEE Conference on Games (CoG). IEEE, 2020: 479-486.

[18] https://github.com/YeWR/EfficientZero

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

·

·