用JAX打造的麻将强化学习模拟器，支持GPU并行训练。
Mahjax: A GPU-Accelerated Mahjong Simulator for Reinforcement Learning in JAX
- 基于JAX实现全向量化麻将环境，支持大规模并行计算。
- 在8张A100 GPU上分别达到每秒200万和100万步的吞吐量。
- 适合研究从零开始的强化学习算法，尤其适合对战智能体训练。
立直麻将是一种多玩家、不完全信息博弈，具有随机性和高维状态空间，其特性与强化学习中的复杂现实决策问题高度相似。以往研究多依赖人类对局日志进行监督预训练，而能从零学习（tabula rasa）的算法更具通用潜力，如AlphaZero系列。为推动此类研究，我们提出**Mahjax**，一个基于JAX实现的完全向量化立直麻将环境，可在图形处理器（GPU）上实现大规模推理并行化。同时提供高质量可视化工具，便于调试与交互训练好的智能体。实验表明，在8张NVIDIA A100 GPU下，无红规则和红规则场景分别达到每秒200万和100万步的处理速度。此外，通过实验证明该环境可有效训练智能体，显著提升其对基线策略的排名表现。
原文摘要 · Abstract (English)
Riichi Mahjong is a multi-player, imperfect-information game characterized by stochasticity and high-dimensional state spaces. These attributes present a unique combination of challenges that mirror complex real-world decision-making problems in reinforcement learning. While prior research has heavily relied on supervised learning from human play logs to pre-train the policy, algorithms capable of learning \textit{tabula rasa} (from scratch) offer greater potential for general applicability, as evidenced by the AlphaZero lineage. To facilitate such research, we introduce \textbf{Mahjax}, a fully vectorized Riichi Mahjong environment implemented in JAX to enable large-scale rollout parallelization on Graphics Processing Units (GPUs). We also provide a high-quality visualization tool to streamline debugging and interaction with trained agents. Experimental results demonstrate that Mahjax achieves throughputs of up to \textbf{2 million} and \textbf{1 million steps per second} on eight NVIDIA A100 GPUs under the no-red and red rules, respectively. Furthermore, we validate the environment's utility for reinforcement learning by showing that agents can be trained effectively to improve their rank against baseline policies.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。