Mahjax: A GPU-Accelerated Mahjong Simulator for Reinforcement Learning in JAX
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866918513194565632 |
|---|---|
| author | Nishimori, Soichiro Okano, Shinri Habara, Keigo Koyamada, Sotetsu Yu, Eason Sugiyama, Masashi |
| author_facet | Nishimori, Soichiro Okano, Shinri Habara, Keigo Koyamada, Sotetsu Yu, Eason Sugiyama, Masashi |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2605_20577 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Mahjax: A GPU-Accelerated Mahjong Simulator for Reinforcement Learning in JAX Nishimori, Soichiro Okano, Shinri Habara, Keigo Koyamada, Sotetsu Yu, Eason Sugiyama, Masashi Artificial Intelligence Machine Learning 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. |
| title | Mahjax: A GPU-Accelerated Mahjong Simulator for Reinforcement Learning in JAX |
| topic | Artificial Intelligence Machine Learning |
| url | https://arxiv.org/abs/2605.20577 |