Mctx: MCTS-in-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）。然而，在将搜索算法与深度神经网络结合使用时，通常需要高效的实现，而这些实现往往用快速编译的语言编写；这可能会牺牲易用性和可 hack 性，尤其是对于不熟悉 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 = {DeepMind JAX 生态系统},
  author = {DeepMind 以及 Babuschkin, Igor、Baumli, Kate、Bell, Alison、Bhupatiraju, Surya、Bruce, Jake、Buchlovsky, Peter、Budden, David、Cai, Trevor、Clark, Aidan、Danihelka, Ivo、Dedieu, Antoine、Fantacci, Claudio、Godwin, Jonathan、Jones, Chris、Hemsley, Ross、Hennigan, Tom、Hessel, Matteo、Hou, Shaobo、Kapturowski, Steven、Keck, Thomas、Kemaev, Iurii、King, Michael、Kunesch, Markus、Martens, Lena、Merzic, Hamza、Mikulik, Vladimir、Norman, Tamara、Papamakarios, George、Quan, John、Ring, Roman、Ruiz, Francisco、Sanchez, Alvaro、Sartran, Laurent、Schneider, Rosalia、Sezener, Eren、Spencer, Stephen、Srinivasan, Srivatsan、Stanojevi\'{c}, Milo\v{s}、Stokowiec, Wojciech、Wang, Luyu、Zhou, Guangyao、Viola, Fabio},
  url = {http://github.com/deepmind},
  year = {2020},
}

Mctx 快速上手指南

Mctx 是一个基于 JAX 原生实现的蒙特卡洛树搜索（MCTS）算法库，支持 AlphaZero、MuZero 和 Gumbel MuZero 等算法。它完全支持 JIT 编译，能够并行处理批量输入，充分利用加速器性能，非常适合与深度神经网络结合进行强化学习研究。

环境准备

• 系统要求：Linux、macOS 或 Windows（推荐 Linux 以获得最佳 GPU/TPU 支持）。
• 前置依赖：
  • Python 3.8+
  • JAX：Mctx 强依赖于 JAX 生态。请确保已根据硬件环境（CPU/GPU/TPU）正确安装 JAX 及其对应的加速后端（如 jax[cuda12] 等）。

安装步骤

你可以选择从 PyPI 安装稳定版，或从 GitHub 安装最新开发版。

选项 1：安装最新稳定版（推荐）

pip install mctx

选项 2：安装最新开发版

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

提示：如果下载速度较慢，可配置 pip 使用国内镜像源（如清华源）： pip install mctx -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

Mctx 提供了底层的通用 search 函数以及高层的具体策略实现，如 muzero_policy 和 gumbel_muzero_policy。官方推荐使用 gumbel_muzero_policy，因为它在动作价值评估准确时能保证策略提升。

核心概念

在使用前，你需要定义两个关键组件：

1. 根状态表示 (RootFnOutput)：包含策略网络的先验 logits (prior_logits)、根状态的估计价值 (value) 以及用于环境模型的状态嵌入 (embedding)。
2. 递归动力学模型 (recurrent_fn)：一个函数 recurrent_fn(params, rng_key, action, embedding)，接收动作和当前状态嵌入，返回 (RecurrentFnOutput, new_embedding)。其中 RecurrentFnOutput 包含转移的奖励 (reward)、折扣因子 (discount) 以及新状态的先验 logits 和价值。

代码示例

以下展示了如何调用 Gumbel MuZero 策略进行搜索：

import mctx

# 假设 params, rng_key, root, recurrent_fn 已经定义好
# num_simulations: 模拟次数，控制搜索预算
policy_output = mctx.gumbel_muzero_policy(
    params, 
    rng_key, 
    root, 
    recurrent_fn,
    num_simulations=32
)

# 获取搜索建议的动作
action = policy_output.action

# 获取用于训练策略概率的目标权重
action_weights = policy_output.action_weights

执行后，policy_output.action 即为智能体应采取的动作，可直接传入环境交互；policy_output.action_weights 可用于后续的策略网络训练优化。

更多完整示例（包括可视化演示和策略提升验证）请参考仓库中的 examples/ 目录。