mctx

Mctx: MCTS-in-JAX Mctx 是一个库,其中包含了 JAX 版本的蒙特卡洛树搜索 (MCTS) 算法,如 AlphaZero、MuZero 和 Gumbel MuZero。为了提高计算速度,该实现完全支持 JIT 编译。Mctx 中的搜索算法是针对批量输入定义的,并以并行方式运行。这使得我们能够充分利用加速器,并使算法能够与由深度神经网络参数化的大型学习环境模型一起工作。

自人工智能研究起始以来，学习和搜索一直是重要的主题。Rich Sutton 这样说道:

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

通过这个库,我们希望帮助世界各地的研究人员为这一令人兴奋的研究领域做出贡献。我们提供 JA