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

Problem Statement

Riichi Mahjong presents a unique challenge for Reinforcement Learning (RL) due to its nature as a multi-player, imperfect-information game characterized by high-dimensional state spaces and significant stochasticity. While prior research has achieved top human-level performance (e.g., Suphx), these successes predominantly rely on supervised learning (SL) from human play logs or offline RL for pre-training. The "tabula rasa" (from scratch) approach, exemplified by the AlphaZero lineage, offers greater potential for general applicability but requires vast amounts of trial-and-error experience.

Existing CPU-based simulators (e.g., Mjx, Libriichi) create computational bottlenecks for the massive parallelization required for large-scale self-play training. Although the RL community has shifted toward hardware-accelerated environments to achieve speedups exceeding 100× over CPU baselines, existing JAX-based frameworks like Pgx lack a comprehensive implementation of complex imperfect-information games like Riichi Mahjong.

Methodology

The authors introduce Mahjax, a fully vectorized Riichi Mahjong environment implemented natively in JAX. The system is designed to enable large-scale rollout parallelization on Graphics Processing Units (GPUs).

Key Design Choices

1. Functional Programming Paradigm: Adhering to JAX requirements, Mahjax strictly employs a functional architecture. The game state is stored in an immutable State dataclass containing all game information (hands, scores, winds, melds, masks) as JAX arrays. This contrasts with prior stateful, object-oriented simulators that hinder JAX implementation.
2. Vectorized Logic & Optimization: To handle complex conditional branching inherent in Mahjong logic without sacrificing GPU parallel performance, the authors employed:
   • Vectorized Logic: Replacing control flow divergence (e.g., if-else statements) with matrix operations where feasible.
   • Caching: Pre-computing and encoding statistics for all possible suit combinations (specifically for Yaku/hand value calculation) into bitmasks to avoid repetitive computation.
3. API Compatibility: Mahjax adopts the Application Programming Interface (API) design of Pgx, ensuring seamless compatibility with modern JAX-based RL pipelines.
4. Environment Configuration:
   • Rules: Supports both "Tenhou (Red)" rules (standard with red fives) and "No-Red" rules (simplified, removing complex rules like abortive draws for higher throughput).
   • Modes: Offers single (one round), east (up to 4 rounds), and half (up to 8 rounds) modes to vary difficulty and strategic depth.
   • Observations: Provides a structured dictionary observation suitable for Transformer-based agents, containing tokenized inputs (hand indices, action history) and scalar properties (shanten number, scores).

Key Contributions

The paper outlines four primary contributions:

1. Vectorized Environment: A high-performance Mahjong environment compatible with Pgx APIs and JAX pipelines.
2. Performance: Demonstrated scalability across multiple GPUs, achieving massive throughput.
3. Usability: Inclusion of an SVG-based visualization tool and web interface for debugging and interaction.
4. Validation: Successful training of RL agents to improve their rank against baseline policies, proving the environment's readiness for research.

Experimental Results

Computational Efficiency (Speed Benchmark)

The authors benchmarked Mahjax against Libriichi (a Rust-based CPU simulator) and Pgx (Shogi) on nodes equipped with eight NVIDIA A100 GPUs.

• Throughput: Mahjax achieved peak throughputs of 2 million steps per second (SPS) under "no-red" rules and 1 million SPS under "red" rules on eight GPUs.
• Scalability: While Libriichi plateaued at small batch sizes due to CPU limitations, Mahjax continued to scale effectively up to large batch sizes (2^14 environments).
• Comparison: Mahjax outperformed Libriichi by over 10× and surpassed Pgx Shogi in throughput.

Reinforcement Learning Validation

To validate the environment's stability for learning, the authors trained an agent using Proximal Policy Optimization (PPO) with Kullback-Leibler (KL) regularization.

• Setup: The agent was initialized via Behavioral Cloning (BC) on 500k samples from a heuristic agent, then fine-tuned using PPO with 1,024 parallel environments.
• Performance: The trained agent consistently achieved an average rank better (lower) than 2.5 (the theoretical expected rank for players of equal skill) when playing against three fixed BC opponents.
• Conclusion: This trajectory confirms that Mahjax provides a stable implementation capable of supporting deep RL training and policy improvement.

Significance and Claims

The paper positions Mahjax as a critical infrastructure tool for advancing pure RL research in complex, imperfect-information games. By enabling massive parallelization on GPUs, it addresses the data throughput bottleneck that has historically limited "from scratch" learning in Mahjong.

The authors modestly claim that their current work validates the environment's utility and reliability for RL training. They acknowledge that while the current RL evaluation is limited to single-round settings and relies on BC pre-training, the framework is designed to support future work toward learning from scratch and expanding to more complex game modes (e.g., 3-player variants). The code is made publicly available to facilitate further research in this domain.