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

• URL: http://arxiv.org/abs/2605.20577v1
• Date: Wed, 20 May 2026 00:33:28 GMT
• Title: Mahjax: A GPU-Accelerated Mahjong Simulator for Reinforcement Learning in JAX
• Abstract summary: Riichi Mahjong is a multi-player, imperfect-information game characterized by byity and high-dimensional state spaces.<n>We introduce textbfMahjax, a fully vectorized Riichi Mahjong environment implemented in JAX to enable large-scale rollout parallelization.<n>We show that Mahjax achieves throughputs of up to textbf2 million and textbf1 million steps per second.
• Score: 38.43966132249977
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Scaling Strategy, Not Compute: A Stand-Alone, Open-Source StarCraft II Benchmark for Accessible Reinforcement Learning Research [2.5959683852997806]
  We present the Two-Bridge Map Suite, a benchmark for StarCraft IIs mini-games.<n>By disabling economy mechanics, the environment isolates two core tactical skills: long-range navigation and micro-combat.<n>Preliminary experiments show that agents learn coherent maneuvering and engagement behaviors without imposing full-game computational costs.
  arXiv  Detail & Related papers  (2026-02-19T04:13:14Z)
• JaxMARL: Multi-Agent RL Environments and Algorithms in JAX [105.343918678781]
  We present JaxMARL, the first open-source, Python-based library that combines GPU-enabled efficiency with support for a large number of commonly used MARL environments. Our experiments show that, in terms of wall clock time, our JAX-based training pipeline is around 14 times faster than existing approaches. We also introduce and benchmark SMAX, a JAX-based approximate reimplementation of the popular StarCraft Multi-Agent Challenge.
  arXiv  Detail & Related papers  (2023-11-16T18:58:43Z)
• Accelerate Multi-Agent Reinforcement Learning in Zero-Sum Games with Subgame Curriculum Learning [65.36326734799587]
  We present a novel subgame curriculum learning framework for zero-sum games. It adopts an adaptive initial state distribution by resetting agents to some previously visited states. We derive a subgame selection metric that approximates the squared distance to NE values.
  arXiv  Detail & Related papers  (2023-10-07T13:09:37Z)
• Scaling Laws for Imitation Learning in Single-Agent Games [28.257046559127875]
  We investigate whether carefully scaling up model and data size can bring similar improvements in the imitation learning setting for single-agent games.<n>We first demonstrate our findings on a variety of Atari games, and thereafter focus on the extremely challenging game of NetHack.<n>We find that IL loss and mean return scale smoothly with the compute budget and are strongly correlated, resulting in power laws for training compute-optimal IL agents.
  arXiv  Detail & Related papers  (2023-07-18T16:43:03Z)
• SPRING: Studying the Paper and Reasoning to Play Games [102.5587155284795]
  We propose a novel approach, SPRING, to read the game's original academic paper and use the knowledge learned to reason and play the game through a large language model (LLM) In experiments, we study the quality of in-context "reasoning" induced by different forms of prompts under the setting of the Crafter open-world environment. Our experiments suggest that LLMs, when prompted with consistent chain-of-thought, have great potential in completing sophisticated high-level trajectories.
  arXiv  Detail & Related papers  (2023-05-24T18:14:35Z)
• JueWu-MC: Playing Minecraft with Sample-efficient Hierarchical Reinforcement Learning [13.57305458734617]
  We propose JueWu-MC, a sample-efficient hierarchical RL approach equipped with representation learning and imitation learning to deal with perception and exploration. Specifically, our approach includes two levels of hierarchy, where the high-level controller learns a policy to control over options and the low-level workers learn to solve each sub-task. To boost the learning of sub-tasks, we propose a combination of techniques including 1) action-aware representation learning which captures underlying relations between action and representation, 2) discriminator-based self-imitation learning for efficient exploration, and 3) ensemble behavior cloning with consistency filtering for
  arXiv  Detail & Related papers  (2021-12-07T09:24:49Z)
• Forgetful Experience Replay in Hierarchical Reinforcement Learning from Demonstrations [55.41644538483948]
  In this paper, we propose a combination of approaches that allow the agent to use low-quality demonstrations in complex vision-based environments. Our proposed goal-oriented structuring of replay buffer allows the agent to automatically highlight sub-goals for solving complex hierarchical tasks in demonstrations. The solution based on our algorithm beats all the solutions for the famous MineRL competition and allows the agent to mine a diamond in the Minecraft environment.
  arXiv  Detail & Related papers  (2020-06-17T15:38:40Z)
• Suphx: Mastering Mahjong with Deep Reinforcement Learning [114.68233321904623]
  We design an AI for Mahjong, named Suphx, based on deep reinforcement learning with some newly introduced techniques. Suphx has demonstrated stronger performance than most top human players in terms of stable rank. This is the first time that a computer program outperforms most top human players in Mahjong.
  arXiv  Detail & Related papers  (2020-03-30T16:18:16Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.