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

arXiv: https://arxiv.org/abs/2605.20577 · HTML: https://mlsanity.com/paper/2605.20577/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in RL training of language and agent models, one of 15 directions in Deep reinforcement learning methods.
- Builds on 6 of its 27 references that are in the map.
- First posted on arXiv 20 May 2026.

## Placed in

- Deep reinforcement learning methods / [RL training of language and agent models](https://mlsanity.com/deeprl/rl-training-of-language-and-agent-models/)

## Builds on (most cited first, 6 of 6)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Suphx: Mastering Mahjong with Deep Reinforcement Learning](https://mlsanity.com/paper/2003.13590/) (2020)
- [NAVIX: Scaling MiniGrid Environments with JAX](https://mlsanity.com/paper/2407.19396/) (2024)
- [Building a 3-Player Mahjong AI using Deep Reinforcement Learning](https://mlsanity.com/paper/2202.12847/) (2022)
- [NashPG: A Policy Gradient Method with Iteratively Refined Regularization for Finding Nash Equilibria](https://mlsanity.com/paper/2510.18183/) (2025)

## Abstract (arXiv)

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 textittabula rasa (from scratch) offer greater potential for general applicability, as evidenced by the AlphaZero lineage. To facilitate such research, we introduce 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 2 million and 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.