返回- 核心技术
- 以原创技术体系为根基，SenseCore商汤AI大装置为核心基座，布局多领域、多方向前沿研究，
快速打通AI在各个垂直场景中的应用，向行业赋能。
NeurIPS 2023 Spotlight | 面向通用决策场景的MCTS基准框架：LightZero
论文链接：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 evalsuator），致力于帮助开发者专注于环境和算法的研究。
另外，基于在框架设计和基准测试中发现的经验，LightZero进一步探讨了将model-based RL的相关最新方法与MCTS方法结合的优势。例如引入基于self-consistency的自监督训练方法可以在部分环境上能显著加速算法收敛性，但不当的使用也会显著影响另一部分环境的训练稳定性。最终，详尽的基准测试和实验表明，LightZero中整合并构建的最佳算法变体具有极高的数据利用效率和广泛的应用范围。
该论文的主要贡献如下：
1. 推出了LightZero，这是首个系统性评估MCTS/MuZero系列算法的通用Benchmark框架。
2. 详述了设计通用决策算法时需要面临的关键挑战，为解决这些问题，LightZero对MCTS算法和系统设计进行了解耦，设计了模块化的训练流程，其良好的可扩展性有助于更便捷地设计和集成强大的通用决策智能体。
3. 展示了LightZero中构建的MCTS算法变体可作为一种通用序列决策问题解决方案的强大能力和潜在的未来价值。
这一部分，本文首先会概述讲解LightZero，然后深入讨论设计通用MCTS算法将会面临的环境挑战，并提出模块化的训练流程设计方案和高度可扩展的MCTS工具链。
Benchmark 简介
图1：LightZero 概览。左侧部分描绘了蒙特卡罗树搜索（MCTS）的发展过程，而右侧则展示了各类强化学习（RL）环境。LightZero 融合并拓展了 MCTS/MuZero 子领域的最新研究成果，并将这些成果有效地应用于多种不同的环境中。
如图1所示，LightZero是首个汇集了MCTS/MuZero领域几乎全部最新算法进展的Benchmark测试平台。更准确地说，LightZero汇集了源自AlphaZero/MuZero的9项关键算法和超过20种决策环境，为其设计了标准化的训练和部署接口，以适应多种决策环境的需要。
与原有其他各类定制化算法框架相比，LightZero提供了一种统一的视角和用户接口。这种统一性使得对比和分析各类技术组合成为可能，为MCTS相关算法和应用的研究提供一个强大且通用的基准平台。
如何评估一个通用 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的方法在有大量非终止轨迹（即没有有效的奖励信号的轨迹）的环境中容易遇到困难，如迷宫等。
如何简化一个通用 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 evalsuator）提供各类针对强化学习的评价指标以监控训练过程并评估策略行为。同时，这个模块集成了一些模型推理时常用的技巧，如beam search，以提升测试性能。
LightZero基于上述设计提供一系列标准接口和工具，让研究者和工程师能更有效地开发算法或进行系统优化。其中，MCTS的探索策略和世界模型的对齐问题至关重要，将在后续章节详细探讨。对于多向量环境和搜索树的并行方案也是值得深究的话题，一个简易的对比示意图如下所示，具体的数据流和开销分析请参考原论文附录 E。
图4：各种不同决策算法在数据收集时的并行化设计方案对比图。MCTS 系列算法需要更加深度优化的并行策略。
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：自监督一致性损失在不同环境中的影响，涉及到各种类型的观测。从左到右，性能比较涉及到标准图像输入、紧凑向量输入和棋盘图像输入，考虑了有和没有一致性损失的情况。实验显示，一致性损失对于标准图像输入来说是至关重要的。
关键结论和思考
通过对LightZero的统一设计和Benchmark结果的研究，得出了一些关于各种算法优缺点的关键结论，这些结论有助于全面理解这些算法的性能和潜在应用。
01：在棋盘游戏环境中，AlphaZero的样本效率显著优于MuZero。这表明如果环境模拟器可用，直接使用AlphaZero是最佳推荐方案。然而，即使在没有模拟器的情况下，在经历足够的训练步数之后MuZero也可以取得满意的结果。
02：自监督学习损失（SSL）在大多数涉及图像输入的Atari环境中显著提升了性能。如图5所示，配备SSL的MuZero在 MsPacman环境中与原始的MuZero表现相对等，但在其他五个环境中则超越了MuZero。这一结果突显了SSL在提高世界模型对齐能力和加速图像输入环境学习过程中的重要作用。
03：选择预测value_prefix而不是reward并不能始终保证性能的提升。例如，在图5中，EfficientZero仅在MsPacman和Breakout环境中超过了配备SSL的MuZero，而在其他环境中的表现与其相当。更具体地说，如图12所示，在稀疏奖励环境中，EfficientZero的性能显著低于配备SSL的MuZero。因此，在决定是否预测 value_prefix时，应充分考虑环境的奖励函数特性。
04：在大多数Atari环境和具有复杂结构化观察空间的环境中（如GoBigger），配备SSL的MuZero和 EfficientZero的表现相当。这一观察结果表明，具有复杂结构观察的环境可以从表征学习和对比学习技术中受益，从而提高样本效率和鲁棒性。
05：在离散动作空间中，Sampled EfficientZero的表现与动作空间维度密切相关。例如，Sampled EfficientZero在Breakout（动作空间维度为4）中的表现与EfficientZero相当，但在MsPacman（维度为9）中的表现则有所下降。
06：配备高斯策略表示的Sampled EfficientZero 在连续动作空间中表现出更好的可扩展性。高斯版本在传统连续控制和MuJoCo环境中表现优异，而离散化版本则仅适用于较低维的连续动作空间。
07：当模拟次数有限时，Gumbel MuZero明显优于MuZero，这揭示了其在设计低时间成本的蒙特卡洛树搜索（MCTS）智能体方面的潜力。
08：在具有随机状态转换或部分可观察状态的环境中（如没有叠帧的Atari），Stochastic MuZero可以实现比MuZero更优秀的性能。
09：有关MCTS系列算法提出的改进技术，如 EfficientZero中的自监督损失，Sampled MuZero中的采样技术，Gumbel MuZero中利用MCTS搜索信息的计算改进，以及Stochastic MuZero中的环境随机性建模，这些方法都可以看作是相互正交的，之间的干扰非常小。LightZero正在持续探索如何将这些技巧高效且插件化地集成在一起，以设计出通用的决策算法。
10: 如图12所示，对于标准图像输入，一致性损失显得至关重要。如果无视一致性损失，模型性能会显著下滑，这揭示了学习针对高维输入的动态模型的挑战性。然而，对于LunarLander这样的向量输入环境，一致性损失只提供了轻微的优势，这暗示在紧凑的向量观测上实施dynamic model学习相对较为简单。在特定的二维输入环境，如井字游戏（TicTacToe），一致性损失依然较大，凸显出在潜在状态输出间实现一致性的难度。此外，如果在一致性损失中使用不适当的超参数，可能会导致训练无法收敛。综合来看，LightZero中的实验表明，一致性损失的有效性高度依赖于特定的观察属性。对于棋盘游戏，未来的研究方向应考虑探索适当的损失函数，以确保训练过程中的对齐性。
Hugging Face Model Zoo for LightZero
另外，为了进一步拓展决策智能算法和模型的应用范围，LightZero项目也已经在Hugging Face Model Zoo 中设立了专门的模型空间：Hugging Face Model Zoo for LightZero。这个专门为LightZero打造的模型库集合，旨在为研究者和开发者提供一个共享、交流研究成果，并进行快速接入和实验的模型平台，其中包含了最前沿MCTS系列算法模型，并针对特定应用场景进行优化。致力于让研究者能在一个标准化的环境下快速复现、测试最新的基于MCTS的决策智能算法，并进行深度的比较和评估。
Hugging Face Model Zoo for LightZero空间：
https://huggingface.co/OpenDILabCommunity
特色功能和资源
预训练模型：提供一系列预训练好的LightZero模型，覆盖了从基础到高级的各种算法变体，可以直接被快速下载和部署。
配置管理：每个模型都提供了详细的配置文件，包括了环境设置、网络架构和训练参数等，确保了实验的可重复性。
训练和评估工具：集成了便捷的训练和评估接口，研究者可以轻松地在自己的数据集上训练模型，或在标准测试集上评估模型的性能。
文档和教程：提供了丰富的文档资源和教程，帮助用户理解和使用这些模型，无论是新手还是有经验的研究者都能快速上手。
社区支持：通过与Hugging Face社区的紧密合作，用户可以获取到最新的相关资讯，参与社区讨论，或直接向模型库贡献自己的模型。
Model Card 与使用示例
图13：Hugging Face Model Zoo for LightZero 中一个典型示例CartPole-v0-MuZero 的Model Card界面。包含了模型描述、模型使用、模型信息、环境、评估渲染视频等5大部分。
总结与未来工作
LightZero是一个模块化整合了各种MCTS相关强化学习方法的统一算法基准框架。LightZero系统性地剖析并应对了将MCTS作为通用且高效决策求解器部署过程中的挑战和机遇。通过模块化的系统设计、详尽全面的基准测试，揭示了LightZero作为一个可扩展、高效的决策问题工具对于研究社区的巨大潜力。
尽管LightZero展示了其变体算法在各类基准测试环境中的出色性能，但某些局限性仍需被关注。首先，尽管这种模块化的设计显著增强了MCTS系列算法的通用性和扩展性，但对于某些特定的决策环境仍需要部分定制化的调整。其次，由于MCTS算法的内在限制，在处理某些复杂（尤其是随机性强）的环境时可能遇到挑战。最后，MCTS的高技术门槛可能会给初次接触相关技术的人带来困扰。未来的改进将会关注提高 LightZero框架接口的易用性，丰富相关的文档和例子，并构建一个的活跃用户生态社区。
尽管存在上述挑战，但对于未来MCTS系列算法技术的发展，仍充满无限可能，这里列出了一些未来的扩展方向：
扩大应用领域：期望更多的研究者和开发者将LightZero应用于更广泛的实践领域，包括但不限于自然语言处理，自动驾驶，以及复杂系统的控制和优化。
算法优化：欢迎社区贡献新的探索和优化策略，以进一步提升MCTS系列算法的运行效率以减小墙钟时间开销。
与其他前沿技术的整合：MCTS系列算法和其他机器学习前沿方向的仍然结合十分重要。其中最值得关注的两个领域是：MCTS与大型语言模型（LLM）的结合，如何学习更强大的世界模型来使用MCTS进行规划。