mlsanityData to 14 Oct 2026

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

arXiv 2605.20577 ↗First posted 20 May 2026cs.AI cs.LGcited by 0

On the map: Deep reinforcement learning methods → RL training of language and agent models

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.

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.

Builds on

6 of 27 references.

Closest ancestors · 1

Same direction, newest first

  1. 2024NAVIX: Scaling MiniGrid Environments with JAX28 Jul

Other references · 5

  1. 2017Proximal Policy Optimization Algorithmscited by 15,215
  2. 2020Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problemscited by 1,920
  3. 2020Suphx: Mastering Mahjong with Deep Reinforcement Learningcited by 58
  4. 2022Building a 3-Player Mahjong AI using Deep Reinforcement Learningcited by 2
  5. 2025NashPG: A Policy Gradient Method with Iteratively Refined Regularization for Finding Nash Equilibriacited by 2

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2026BluffJAX: Adversarial Imperfect Information Games in JAX
  2. 2025Ludax: A GPU-Accelerated Domain Specific Language for Board Games
  3. 2025Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX
  4. 2026Factoriax: A GPU-Accelerated Factorio-Style Simulator for Reinforcement Learning
  5. 2025PixelBrax: Learning Continuous Control from Pixels End-to-End on the GPU
  6. 2026TABX: A High-Throughput Sandbox Battle Simulator for Multi-Agent Reinforcement Learning
  7. 2024GPUDrive: Data-driven, multi-agent driving simulation at 1 million FPS
  8. 2025Evolutionary Optimization of Deep Learning Agents for Sparrow Mahjong