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

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.

한 줄 요약

**[강화학습 환경 / GPU 가속]** Mahjax는 JAX로 vectorize한 Riichi 마작 시뮬레이터, A100 8장에서 no-red 200만 step/s·red 100만 step/s 처리량 달성, tabula rasa RL 연구 가능하게 함.

핵심 기여도

Riichi 마작이 다인용·불완전 정보·확률성·고차원 상태공간의 게임으로 복잡한 실세계 RL 의사결정 문제를 미러링하는 unique combination 보유.
선행 연구는 주로 인간 기보로부터의 supervised learning에 의존해 정책을 pre-train — tabula rasa 학습 가능 알고리즘이 AlphaZero 계보처럼 일반 적용성에 더 큰 가능성 보유.
Mahjax 도입 — JAX로 fully vectorize된 Riichi 마작 환경, GPU 대규모 rollout 병렬화 지원.
디버깅·학습 에이전트 상호작용을 streamline하는 고품질 시각화 도구 제공.

핵심 아이디어

복잡한 불완전정보 게임 RL 연구를 인간 기보 의존에서 해방하려면 GPU 대규모 병렬화 가능한 fully vectorized 환경이 필수이며, JAX 구현으로 8 GPU에서 초당 백만~수백만 step의 throughput을 달성해 tabula rasa 학습을 실현 가능하게 한다.

기술적 접근법

**방법론**: Mahjax — JAX 기반 fully vectorized Riichi 마작 시뮬레이터.
**핵심 기법**: (1) JAX의 jit·vmap·pmap으로 환경을 fully vectorize, (2) GPU 대규모 rollout 병렬화 지원, (3) No-red·red 두 규칙 변종 구현, (4) 학습 에이전트 디버깅·시각화 도구 통합, (5) Tabula rasa RL 학습 검증 — baseline 정책 대비 ranking 개선 시연.

주요 결과

NVIDIA A100 8장에서 no-red 규칙 약 200만 step/s, red 규칙 약 100만 step/s throughput.
Tabula rasa RL 에이전트가 baseline 정책 대비 ranking 개선 학습 검증.
인간 기보 의존 없는 학습 시연.
고품질 시각화 도구로 학습 에이전트 상호작용 분석 가능.

의의 및 한계

**의의**: 불완전정보·다인용 게임 RL 연구의 환경 표준 제공, GPU 대규모 병렬 throughput으로 tabula rasa 가능, 시각화 도구로 분석 용이성 확보, AlphaZero 계보의 마작 확장 토대. **한계**: Riichi 마작 단일 게임에 특화로 다른 카드·보드게임 일반화 별도 작업 필요, 학습 정책의 전문가 수준 도달은 미입증, 8 GPU 가정의 cost.

실용적 활용

불완전 정보 게임 RL 연구.
대규모 self-play 학습 벤치마크.
다인용 의사결정 알고리즘 평가.