Mahjax: A GPU-Accelerated Mahjong Simulator for Reinforcement Learning in JAX

2026-05-19

AI Summary

This work addresses the challenge of inefficient reinforcement learning (RL) training in Japanese Riichi Mahjong, a domain characterized by multi-agent interactions, imperfect information, and a high-dimensional state space. To overcome the lack of an RL-friendly environment, the authors present the first high-performance simulator designed specifically for pure reinforcement learning. Built with JAX, the simulator features fully vectorized operations enabling massive GPU-parallelized rollouts and includes interactive visualization tools to facilitate debugging. On a system with eight NVIDIA A100 GPUs, it achieves throughput of 2 million steps per second without red fives and 1 million steps per second with red fives. Using this infrastructure, the authors successfully train agents from scratch that significantly improve competitive rankings, thereby surpassing the limitations of conventional approaches that rely on pretraining from human gameplay logs.

Abstract

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.

Problem

• Riichi Mahjong
• reinforcement learning
• imperfect-information game
• high-dimensional state space
• stochasticity

Innovation

• GPU acceleration
• vectorized environment
• JAX
• reinforcement learning
• Riichi Mahjong