Mctx: JAX中的蒙特卡洛树搜索

Mctx是一个使用JAX原生实现蒙特卡洛树搜索(MCTS)算法的库，包括AlphaZero、MuZero和Gumbel MuZero。为了提高计算速度，该实现完全支持即时编译(JIT)。Mctx中的搜索算法针对批量输入并行操作。这允许充分利用加速器，并使算法能够与由深度神经网络参数化的大型学习环境模型一起工作。

安装

你可以通过PyPI安装Mctx的最新发布版本：

pip install mctx

或者你可以从GitHub安装最新的开发版本：

pip install git+https://github.com/google-deepmind/mctx.git

动机

学习和搜索自人工智能研究早期以来就一直是重要的话题。用Rich Sutton的话说：

我们应该学到的一件事是[...]通用方法的强大力量，即那些随着可用计算能力的大幅增加而继续扩展的方法。似乎可以无限扩展的两种方法是搜索和学习。

最近，搜索算法已成功与由深度神经网络参数化的学习模型相结合，产生了迄今为止一些最强大和最通用的强化学习算法（例如MuZero）。然而，将搜索算法与深度神经网络结合使用需要高效的实现，通常用快速编译语言编写；这可能会牺牲可用性和可修改性，特别是对不熟悉C++的研究人员而言。反过来，这限制了这一关键主题的采用和进一步研究。

通过这个库，我们希望帮助各地的研究人员为如此令人兴奋的研究领域做出贡献。我们提供了核心搜索算法（如MCTS）的JAX原生实现，我们认为这在性能和可用性之间取得了很好的平衡，适合那些想用Python研究基于搜索的算法的研究人员。Mctx提供的搜索方法具有高度可配置性，允许研究人员在这个领域探索各种想法，并为下一代基于搜索的智能体做出贡献。

强化学习中的搜索

在强化学习中，智能体必须学会与环境互动以最大化标量奖励信号。在每一步中，智能体必须选择一个动作，并作为交换接收一个观察和一个奖励。我们可以称智能体用来选择动作的任何机制为智能体的策略。

传统上，策略直接由函数逼近器参数化（如在REINFORCE中），或者通过检查每个动作的学习估值集来推断策略（如在Q-学习中）。另外，搜索允许通过在每个状态下即时构建一个针对当前状态的局部策略或价值函数来选择动作，方法是使用学习到的环境模型进行搜索。

在任何非平凡环境中，对所有可能的未来行动路径进行穷尽搜索在计算上是不可行的，因此我们需要能够最好地利用有限计算预算的搜索算法。通常需要先验来指导应该扩展搜索树中的哪些节点（以减少我们构建的树的广度），并使用价值函数来估计树中不到达情节终点的不完整路径的价值（以减少搜索树的深度）。

快速入门

Mctx提供了一个低级通用search函数和高级具体策略：muzero_policy和gumbel_muzero_policy。

用户需要提供几个学习到的组件来指定MuZero使用的表示、动态和预测。在Mctx库的上下文中，根状态的表示由RootFnOutput指定。RootFnOutput包含策略网络的prior_logits、根状态的估计value，以及任何适合表示环境模型根状态的embedding。

动态环境模型需要由recurrent_fn指定。recurrent_fn(params, rng_key, action, embedding)调用接受一个action和一个状态embedding。该调用应返回一个元组(recurrent_fn_output, new_embedding)，其中包含一个RecurrentFnOutput和下一个状态的嵌入。RecurrentFnOutput包含转换的reward和discount，以及新状态的prior_logits和value。

在examples/visualization_demo.py中，你可以看到对策略的调用：

policy_output = mctx.gumbel_muzero_policy(params, rng_key, root, recurrent_fn,
                                          num_simulations=32)

policy_output.action包含搜索建议的动作。该动作可以传递给环境。为了改进策略，policy_output.action_weights包含可用于训练策略概率的目标。

我们建议使用gumbel_muzero_policy。Gumbel MuZero保证如果动作价值被正确评估，就能实现策略改进。策略改进在examples/policy_improvement_demo.py中得到了演示。

示例项目

以下项目演示了Mctx的使用：

• Pgx — 20多个向量化JAX环境的集合，包括双陆棋、国际象棋、将棋、围棋，以及一个AlphaZero示例。
• 使用Mctx的基础学习演示 — 随机迷宫上的AlphaZero。
• a0-jax — 四子棋、五子棋和围棋上的AlphaZero。
• muax — gym风格环境（CartPole，LunarLander）上的MuZero。
• 经典MCTS — 四子棋上的简单示例。
• mctx-az — 带有AlphaZero子树持久性的Mctx。

告诉我们你的项目。

引用Mctx

这个仓库是DeepMind JAX生态系统的一部分，要引用Mctx，请使用以下引用：

@software{deepmind2020jax,
  title = {The {D}eep{M}ind {JAX} {E}cosystem},
  author = {DeepMind and Babuschkin, Igor and Baumli, Kate and Bell, Alison and Bhupatiraju, Surya and Bruce, Jake and Buchlovsky, Peter and Budden, David and Cai, Trevor and Clark, Aidan and Danihelka, Ivo and Dedieu, Antoine and Fantacci, Claudio and Godwin, Jonathan and Jones, Chris and Hemsley, Ross and Hennigan, Tom and Hessel, Matteo and Hou, Shaobo and Kapturowski, Steven and Keck, Thomas and Kemaev, Iurii and King, Michael and Kunesch, Markus and Martens, Lena and Merzic, Hamza and Mikulik, Vladimir and Norman, Tamara and Papamakarios, George and Quan, John and Ring, Roman and Ruiz, Francisco and Sanchez, Alvaro and Sartran, Laurent and Schneider, Rosalia and Sezener, Eren and Spencer, Stephen and Srinivasan, Srivatsan and Stanojevi\'{c}, Milo\v{s} and Stokowiec, Wojciech and Wang, Luyu and Zhou, Guangyao and Viola, Fabio},
  url = {http://github.com/deepmind},
  year = {2020},
}