NeurIPS 2023 Spotlight | 面向通用决策场景的MCTS基准框架：LightZero

2023-10-17 04:10

引言

人工智能研究中的强化学习（RL）算法已在棋类博弈、即时战略类游戏等领域表现出色，但仍然存在数据利用效率低和算法通用程度低等问题。目前，已有诸如 DQN，PPO，R2D2，SAC，IMPALA 等各种专门针对特定决策环境和学习需求的算法，但创建一个通用且具有高数据效率的决策算法方案仍面临挑战。蒙特卡洛树搜索（Monte Carlo Tree Search，MCTS）是一种强大的搜索方法，AlphaGo 通过将其与神经网络结合，诞生出首个击败人类顶尖围棋选手的 RL 智能体。

但由于 MCTS 对决策动作空间和环境规则等条件的依赖，其应用范围最初仅限于一部分问题。近些年来，AlphaGo 的后继者，如 MuZero 和 Sampled MuZero，试图通过利用世界模型和基于采样的机制来减少对环境特性和先验知识的依赖，将 MCTS 应用于更复杂的环境中。这些基于 MCTS 的技术已被广泛应用于各种决策智能领域，如象棋、视频游戏、视频编解码和高性能计算等。

尽管 MCTS 系列算法技术的能力持续提升，但是由于这类算法涉及强化学习、树搜索、深度神经网络等诸多方面，从而大大增加了算法的应用难度。开源社区中目前也出现了一些工具项目，比如 KataGo、AlphaZero-General 和 MuZero-General 等，试图努力降低 MCTS 技术的门槛，拓宽应用范围。然而，这些开源工具的设计和实现相当复杂，仍然需要较高的学习成本，此外，这些工具并未完全吸收和融合最新的研究成果，因此难以支持广泛的应用，只能主要专注于一小部分的算法和环境。

为了充分发挥 MCTS 系列算法技术的潜力，提升相关技术在各个决策智能领域的通用性和易用性，进一步挖掘这类算法的实用价值。本文将介绍由上海人工智能实验室开源决策智能平台（OpenDILab）团队发表在 2023 NeurIPS Datasets and Benchmarks Track 上的工作 LightZero，期待能够吸引更多的开源社区开发者参与讨论和实践，共同探索 MCTS 系列算法的新篇章。

论文标题：

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

论文链接：

https://arxiv.org/abs/2310.08348

代码链接：

https://github.com/opendilab/LightZero

概述

本文将介绍名为 LightZero 的通用算法 Benchmark，它集成了 MCTS/MuZero 算法的各个衍生分支，涵盖了 9 种算法和对超过 20 种决策环境的实践评估。LightZero 深入分析了 MCTS 方法的发展历史和 RL 环境的多样性，解读了开发通用决策算法面临的六大挑战性维度。为了解决由算法与系统架构的高度耦合性引发的各类问题，LightZero 设计了一个模块化的训练流程，以便灵活地定制和优化各种 MCTS 系列算法组件。具体地，LightZero 将这类方法的训练方案划分为四个子模块：数据收集器（Data collector）、数据整理器（Data arranger）、智能体学习器（Agent learner）和智能体评估器（Agent evaluator），致力于帮助开发者专注于环境和算法的研究。

另外，基于在框架设计和基准测试中发现的经验，LightZero 进一步探讨了将 model-based RL 的相关最新方法与 MCTS 方法结合的优势。例如引入基于 self-consistency 的自监督训练方法可以在部分环境上能显著加速算法收敛性，但不当的使用也会显著影响另一部分环境的训练稳定性。最终，详尽的基准测试和实验表明，LightZero 中整合并构建的最佳算法变体具有极高的数据利用效率和广泛的应用范围。

该论文的主要贡献如下：

推出了 LightZero，这是首个系统性评估 MCTS/MuZero 系列算法的通用 Benchmark 框架。
详述了设计通用决策算法时需要面临的关键挑战，为解决这些问题，LightZero 对 MCTS 算法和系统设计进行了解耦，设计了模块化的训练流程，其良好的可扩展性有助于更便捷地设计和集成强大的通用决策智能体。
展示了 LightZero 中构建的 MCTS 算法变体可作为一种通用序列决策问题解决方案的强大能力和潜在的未来价值。

方法

这一部分，本文首先会概述讲解 LightZero，然后深入讨论设计通用 MCTS 算法将会面临的环境挑战，并提出模块化的训练流程设计方案和高度可扩展的 MCTS 工具链。

3.1 Benchmark简介

▲ 图1：LightZero 概览。左侧部分描绘了蒙特卡罗树搜索（MCTS）的发展过程，而右侧则展示了各类强化学习（RL）环境。LightZero 融合并拓展了 MCTS/MuZero 子领域的最新研究成果，并将这些成果有效地应用于多种不同的环境中。

如图 1 所示，LightZero 是首个汇集了 MCTS/MuZero 领域几乎全部最新算法进展的 Benchmark 测试平台。更准确地说，LightZero 汇集了源自 AlphaZero/MuZero 的 9 项关键算法和超过 20 种决策环境，为其设计了标准化的训练和部署接口，以适应多种决策环境的需要。与原有其他各类定制化算法框架相比，LightZero 提供了一种统一的视角和用户接口。这种统一性使得对比和分析各类技术组合成为可能，为 MCTS 相关算法和应用的研究提供一个强大且通用的基准平台。

3.2 如何评估一个通用MCTS算法：6种环境挑战

▲ 图2：MCTS+RL 方法和 model-free RL 方法（例如 PPO）在六个环境挑战及数据效率维度上的定性对比结果。LightZero 中将通用决策求解器的关键能力进行了分类，包括：多模态观察空间、复杂动作空间、环境固有随机性、对先验知识的依赖程度、仿真成本、探索难度和数据效率。图中的每条曲线表示一个算法在这六个类别上的评分。评分为1表示该算法在此维度上表现欠佳，只适用于有限场景，而更高的评分则意味着更广泛的应用范围和更优的性能。特别地，model-free RL 方法不依赖仿真，对先验知识的依赖度较小，因此在这些领域得分较高。请注意，LightZero 在此上下文中指的是在 LightZero 框架内，将各种技术和超参数设置进行最优组合的特殊算法变体。关于定性评分规则的详细信息，请参见原论文附录 D。

LightZero 对各类 RL 环境进行了统一的分析，总结出设计通用高效 MCTS 算法的六个关键挑战（图2）。具体来说，这六种环境挑战分别是：

1）多模态观察空间（Multi-modal observation spaces），需要智能体能够提取并有效融合不同形式的信息如低维向量、视觉图像和复杂结构化观察空间等。

2）复杂动作空间（Complex action space），需要智能体能在生成多样化决策信号上表现出色，包括离散动作选择、连续控制和混合结构的动作空间等。

3）依赖先验知识（Reliance on prior knowledge），这是 AlphaZero 等方法的一个主要缺点。这些方法本质上需要能够访问完美的模拟器和环境的具体规则。相反，MuZero 及派生方法通过学习环境模型来替代模拟器和相关先验，解决了这一限制。

4）环境固有随机性（Inherent stochasticity）在基于树搜索的规划方法中提出了根本性挑战。环境动力学的不确定性和部分可观测的状态空间都可能导致规划轨迹的错位，产生大量无用或冲突的搜索结果。

5）模拟成本（Simulation cost）是 MCTS-style 方法的时间消耗的主要来源。同时，如果算法在模拟过程中未能访问所有必要的动作，算法的性能会大幅下降。

6）探索困难（Hard exploration）是一个经常被忽视的关键挑战。虽然搜索树可以通过减少探索范围来提高效率，但 MCTS-style 的方法在有大量非终止轨迹（即没有有效的奖励信号的轨迹）的环境中容易遇到困难，如迷宫等。

3.3 如何简化一个通用MCTS算法：模块化训练流程

▲ 图3：LightZero 训练流程中的四个核心子模块。上下文交换器（Context Exchanger）负责在各个子模块之间传输配置、模型和轨迹。

MCTS 系列方法虽然表现优秀，但其实现复杂性也相当高，这大大限制了其应用场景。与 DQN/PPO 等 model-free RL 算法不同，MCTS 方法在每次智能体-环境交互中都需要使用搜索树进行多步模拟。此外，为了提升训练数据质量，MuZero Unplugged 引入了数据重分析机制，通过使用新模型重新计算旧数据，以获取改进的训练目标。然而，这两种技术都需要多次调用模拟器或神经网络，增加了分布式训练和通信拓扑等方面的复杂性。因此，对算法进行集成以简化整个框架是必要的。

LightZero 中设计的训练流程（如图 3 所示）包括四个核心子模块。

首先，LightZero 同时支持 online RL 和 offline RL 训练。两者的主要区别在于是采用在线交互数据收集器还是直接利用离线数据集。

其次，LightZero 对相关训练流程进行了重构，将其分解为四个核心子模块，遵循高内聚、低耦合（high cohesion and low coupling）的设计原则。

数据收集器（Data collector）负责利用策略网络和搜索树进行有效的动作选择，同时包括各种探索策略，以及数据预处理和打包操作。
数据整理器（Data arranger）在 MCTS 中扮演着关键角色，它能有效地存储和准备用于训练的有价值数据。此模块涉及到数据重分析技术，用以校正 off-policy或离线数据。优先级采样的改良版确保了训练的小批量数据具有足够的多样性和较高的学习潜能。吞吐量限制器负责控制添加和采样数据的比例，以在固定的通信带宽内实现最优的数据利用效果。
智能体学习器（Agent learner）负责多个网络的训练，可以利用自监督表征学习，基于模型的rollout，基于值分布的强化学习和价值函数归一化等优化技术进行增强。这些技术有助于策略的提升，进而提高智能体的整体性能。
智能体评估器（Agent evaluator）提供各类针对强化学习的评价指标以监控训练过程并评估策略行为。同时，这个模块集成了一些模型推理时常用的技巧，如 beam search，以提升测试性能。

LightZero 基于上述设计提供一系列标准接口和工具，让研究者和工程师能更有效地开发算法或进行系统优化。其中，MCTS 的探索策略和世界模型的对齐问题至关重要，将在后续章节详细探讨。对于多向量环境和搜索树的并行方案也是值得深究的话题，一个简易的对比示意图如下所示，具体的数据流和开销分析请参考原论文附录 E。

▲ 图4：各种不同决策算法在数据收集时的并行化设计方案对比图。MCTS 系列算法需要更加深度优化的并行策略。

实验

4.1 Benchmark 结果

为了测试不同算法之间的性能差异和 LightZero中的算法变体作为通用决策求解器的能力，实验部分在各种强化学习环境中进行了广泛的对比。一共评估了下面这些算法变体，包括 AlphaZero，MuZero，EfficientZero，Sampled MuZero，Stochastic MuZero，Gumbel MuZero 和 LightZero 中综合改进的算法变体。下面的部分展示了其中一些 Benchmark 结果作为例子。有关环境和算法的详细设置，更详尽全面的 Benchmark 结果和相关分析，请参阅论文的附录 B，C，G。

▲ 图5：图示为在六大代表性 Atari 环境下（图像输入），LightZero 集成的各种算法的性能对比。横坐标代表环境步骤（Env Steps），纵坐标表示在 20 个 episode 中的平均收益（Return）。在此上下文中，MuZero w/ SSL 代表加入了自监督损失的原始 MuZero 算法。EfficientZero 则指通过自监督损失及 value_prefix 增强的 MuZero 算法。Sampled EfficientZero 则在 EfficientZero 的基础上引入了与采样相关的改进技术。

▲ 图6：图示为 AlphaZero 与 MuZero 在 Connect4 和 Gomoku 中的性能比较。AlphaZero 相比于 MuZero 表现出显著的样本效率优势，说明在环境模拟器可直接使用时，AlphaZero 将有明显优势。然而，即便在无模拟器的环境下，MuZero 仍可产生相当的结果，这充分体现了其广泛的适应性。

▲ 图7：上部：图示为 Sampled EfficientZero 在连续动作空间环境中，采用不同策略建模技术的性能比较。下部：图示为该算法在 MuJoCo 连续动作空间环境中，应用各种策略建模方法的性能比较。随着动作空间尺寸的增加，前者（连续空间离散化版本）表现逐渐下降，而后者表现则相对稳定。

▲ 图8：图示为在不同模拟成本下，Gumbel MuZero 和 MuZero 的性能比较。在模拟次数有限时，Gumbel MuZero 相对于 MuZero 表现出显著的性能优势，突显了其在设计低时延 MCTS 智能体上的潜力。针对 Gomoku（棋盘尺寸=6），评估了 sim={20, 10}；针对 LunarLander-v2，评估了 sim={20, 10, 5}；针对 Atari Games，评估了 sim={50, 16, 2}。

▲ 图9：图示为在不同随机性等级（num_chances=2 和 5）的 2048 环境中，Stochastic MuZero 和 MuZero 的性能比较。在环境的状态转移存在显著随机性的环境中，Stochastic MuZero 略超 MuZero。然而，随着随机性等级的升高，Stochastic MuZero 的性能也开始受限。

▲ 图10：左图：图示为在 MiniGrid-KeyCorridorS3R3-v0 环境中，不同探索策略的性能比较（收集阶段的回报）。利用好奇心机制探索状态空间的 IntrinsicExploration 策略表现出较高的样本效率。右图：图示为同样环境中，LightZero 实施的算法的性能比较。在特性为高维向量观察及稀疏奖励的环境中，自监督学习损失有助于模型对齐，但预测 value_prefix 可能带来挑战，对学习造成阻碍。

▲ 图11：图示为在多智能体对抗和协作环境 GoBigger 的 T2P2 和 T2P3 场景中，MuZero 和 EfficientZero（均在独立学习模式下进行训练）的性能比较。两种算法在与内置机器人对战时均能稳定收敛，其样本效率表现出显著优势，相对于其他非 MCTS 方法提升了约六倍。

▲ 图12：自监督一致性损失在不同环境中的影响，涉及到各种类型的观测。从左到右，性能比较涉及到标准图像输入、紧凑向量输入和棋盘图像输入，考虑了有和没有一致性损失的情况。实验显示，一致性损失对于标准图像输入来说是至关重要的。

4.2 关键结论和思考

通过对 LightZero 的统一设计和 Benchmark 结果的研究，得出了一些关于各种算法优缺点的关键结论，这些结论有助于全面理解这些算法的性能和潜在应用。

O1：在棋盘游戏环境中，AlphaZero 的样本效率显著优于 MuZero。这表明如果环境模拟器可用，直接使用 AlphaZero 是最佳推荐方案。然而，即使在没有模拟器的情况下，在经历足够的训练步数之后 MuZero 也可以取得满意的结果。
O2：自监督学习损失（SSL）在大多数涉及图像输入的 Atari 环境中显著提升了性能。如图 5 所示，配备 SSL的 MuZero 在 MsPacman 环境中与原始的 MuZero 表现相对等，但在其他五个环境中则超越了 MuZero。这一结果突显了 SSL 在提高世界模型对齐能力和加速图像输入环境学习过程中的重要作用。
O3：选择预测 value_prefix 而不是 reward 并不能始终保证性能的提升。例如，在图 5 中，EfficientZero 仅在 MsPacman 和 Breakout 环境中超过了配备 SSL 的 MuZero，而在其他环境中的表现与其相当。更具体地说，如图 12 所示，在稀疏奖励环境中，EfficientZero 的性能显著低于配备 SSL 的 MuZero。因此，在决定是否预测 value_prefix 时，应充分考虑环境的奖励函数特性。
O4：在大多数 Atari 环境和具有复杂结构化观察空间的环境中（如 GoBigger），配备 SSL 的 MuZero 和 EfficientZero 的表现相当。这一观察结果表明，具有复杂结构观察的环境可以从表征学习和对比学习技术中受益，从而提高样本效率和鲁棒性。
O5：在离散动作空间中，Sampled EfficientZero 的表现与动作空间维度密切相关。例如，Sampled EfficientZero 在 Breakout（动作空间维度为 4）中的表现与 EfficientZero 相当，但在 MsPacman（维度为 9）中的表现则有所下降。
O6：配备高斯策略表示的 Sampled EfficientZero 在连续动作空间中表现出更好的可扩展性。高斯版本在传统连续控制和 MuJoCo 环境中表现优异，而离散化版本则仅适用于较低维的连续动作空间。
O7：当模拟次数有限时，Gumbel MuZero 明显优于 MuZero，这揭示了其在设计低时间成本的蒙特卡洛树搜索（MCTS）智能体方面的潜力。
O8：在具有随机状态转换或部分可观察状态的环境中（如没有叠帧的 Atari），Stochastic MuZero 可以实现比 MuZero 更优秀的性能。
O9：有关 MCTS 系列算法提出的改进技术，如 EfficientZero 中的自监督损失，Sampled MuZero 中的采样技术，Gumbel MuZero 中利用 MCTS 搜索信息的计算改进，以及 Stochastic MuZero 中的环境随机性建模，这些方法都可以看作是相互正交的，之间的干扰非常小。LightZero 正在持续探索如何将这些技巧高效且插件化地集成在一起，以设计出通用的决策算法。
O10：如图 12 所示，对于标准图像输入，一致性损失显得至关重要。如果无视一致性损失，模型性能会显著下滑，这揭示了学习针对高维输入的动态模型的挑战性。然而，对于 LunarLander 这样的向量输入环境，一致性损失只提供了轻微的优势，这暗示在紧凑的向量观测上实施 dynamic model 学习相对较为简单。在特定的二维输入环境，如井字游戏（TicTacToe），一致性损失依然较大，凸显出在潜在状态输出间实现一致性的难度。此外，如果在一致性损失中使用不适当的超参数，可能会导致训练无法收敛。综合来看，LightZero 中的实验表明，一致性损失的有效性高度依赖于特定的观察属性。对于棋盘游戏，未来的研究方向应考虑探索适当的损失函数，以确保训练过程中的对齐性。

总结与未来工作

LightZero 是一个模块化整合了各种 MCTS 相关强化学习方法的统一算法基准框架。LightZero 系统性地剖析并应对了将 MCTS 作为通用且高效决策求解器部署过程中的挑战和机遇。通过模块化的系统设计、详尽全面的基准测试，揭示了 LightZero 作为一个可扩展、高效的决策问题工具对于研究社区的巨大潜力。

尽管 LightZero 展示了其变体算法在各类基准测试环境中的出色性能，但某些局限性仍需被关注。首先，尽管这种模块化的设计显著增强了 MCTS 系列算法的通用性和扩展性，但对于某些特定的决策环境仍需要部分定制化的调整。其次，由于 MCTS 算法的内在限制，在处理某些复杂（尤其是随机性强）的环境时可能遇到挑战。最后，MCTS 的高技术门槛可能会给初次接触相关技术的人带来困扰。未来的改进将会关注提高 LightZero 框架接口的易用性，丰富相关的文档和例子，并构建一个的活跃用户生态社区。

但对于未来 MCTS 系列算法技术的发展，仍充满无限可能，这里列出了一些未来的扩展方向：

扩大应用领域：期望更多的研究者和开发者将LightZero应用于更广泛的实践领域，包括但不限于自然语言处理，自动驾驶，以及复杂系统的控制和优化。
算法优化：欢迎社区贡献新的探索和优化策略，以进一步提升MCTS系列算法的运行效率以减小墙钟时间开销。
与其他前沿技术的整合：MCTS 系列算法和其他机器学习前沿方向的仍然结合十分重要。其中最值得关注的两个领域是：MCTS 与大型语言模型（LLM）的结合，如何学习更强大的世界模型来使用 MCTS 进行规划。
相信通过这些努力，有望进一步推动 MCTS 和相关强化学习技术的发展，为实际决策问题的解决提供更强大的工具。

LightZero 相关基准实验和训练框架已完全开源，可访问 GitHub 仓库查看更多信息：

https://github.com/opendilab/LightZero

更多阅读