Mahjax：一个用于JAX中强化学习的高性能GPU加速麻将模拟器
Mahjax是一个在JAX中实现的完全向量化立直麻将环境，可利用GPU进行大规模并行化，吞吐量达到在8块NVIDIA A100 GPU上每秒200万步（无红宝牌规则）和100万步（有红宝牌规则）。该环境支持从零开始（tabula rasa）的强化学习训练，并附有高质量可视化工具，实验验证了训练智能体可以有效提升排名。
文章情报
要点
- Mahjax是基于JAX的完全向量化立直麻将模拟器，支持GPU并行化。
- 在8块NVIDIA A100 GPU上，每秒可处理多达200万步（无红宝牌规则）。
- 无需人类数据，支持从零开始的强化学习训练。
- 包含可视化工具，便于调试与智能体交互。
为什么重要
这条新闻值得关注，因为Mahjax是基于JAX的完全向量化立直麻将模拟器，支持GPU并行化。
技术影响
可能影响 Agent 架构、工具调用、工作流自动化和产品集成。
近日，由Soichiro Nishimori等六位研究人员共同开发的Mahjax模拟器正式发布。该模拟器基于JAX框架，针对立直麻将这一复杂的不完全信息博弈进行了深度优化。立直麻将因其多人参与、信息不完全、随机性高以及状态空间极大等特点，成为强化学习领域极具挑战性的研究对象，其复杂性堪比现实世界中的许多决策问题。传统方法通常依赖人类对局日志进行监督学习来预训练策略，但Mahjax支持从头开始的“白板学习”（tabula rasa），类似于AlphaZero系列算法，具有更强的通用性和泛化能力。
Mahjax的核心创新在于其完全向量化的设计，这使得大规模并行仿真成为可能。通过在JAX中实现整个环境，研究人员能够充分利用GPU进行加速。在八块NVIDIA A100 GPU的测试平台上，Mahjax在无红宝牌规则下达到了每秒200万步的惊人吞吐量，即使在引入红宝牌规则后，吞吐量也保持在每秒100万步。这一性能为大规模强化学习实验提供了坚实基础，使得在合理时间内进行大量对局训练成为现实。
除了强大的计算性能，团队还开发了一款高质量的可视化工具，便于研究人员对训练过程进行调试，并与训练后的智能体进行直观交互。这大大降低了研究门槛，使得非专业用户也能轻松上手。为了验证环境的实用性，研究人员在该环境中训练了强化学习智能体，并将其与基线策略进行对比。实验结果表明，经过训练的智能体能够显著提升其在比赛中的排名，证明了Mahjax作为强化学习研究平台的有效性和价值。
Mahjax的出现，不仅为麻将AI研究提供了新的工具，更为在更广泛的复杂博弈中应用从零开始的强化学习开辟了道路。相关论文已提交至arXiv，编号为2605.20577，感兴趣的读者可以获取更多技术细节。