AGENTIC-SYSTEMSMay 22, 20261 min read10 views

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

Researchers have introduced Mahjax, a GPU-accelerated Riichi Mahjong environment implemented in JAX for reinforcement learning. It enables large-scale parallelization, achieving up to 2 million steps per second and facilitating tabula rasa learning for AI agents.

Computer Science > Artificial Intelligence

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

View PDF HTML (experimental)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.

Bibliographic and Citation Tools

Code, Data and Media Associated with this Article

Demos

Recommenders and Search Tools

arXivLabs: experimental projects with community collaborators

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.

Source: arXiv cs.AI Recent

More in this category

agentic-systems

HG-RAG: Hierarchy-Guided Retrieval-Augmented Generation for Structured Knowledge Graphs

HG-RAG (Hierarchy-Guided Retrieval-Augmented Generation) is a novel framework designed to overcome the limitations of traditional flat RAG systems by traversing hierarchical knowledge graphs, significantly improving multi-hop reasoning and reducing hallucinations in LLMs.

agentic-systems

Enhancing Small Language Models Reasoning through Knowledge Graph Grounding

Researchers have proposed a neuro-symbolic framework that enhances the reasoning capabilities of Small Language Models (SLMs) like Gemma 3 and Llama 3.2 using Knowledge Graphs, achieving up to a 2x performance boost.

agentic-systems

Orchestrating Power Grid Studies with Multi-Agent AI and MCP Servers

This paper explores how Agentic AI and Model Context Protocol (MCP) can support power-grid studies by integrating LLMs with numerical simulation tools under human supervision.

agentic-systems

Human AI Construction of Bayesian Networks for Operational Decision Support -- A Virtual Survey Approach

Researchers propose a new methodology using Large Language Models and AI agents to construct Bayesian Belief Networks, bridging the gap between expert opinion and data-driven learning. The framework is demonstrated through a healthcare decision-making model, revealing key insights into customer behavior.

agentic-systems

Interpretable Language Model for Closed-Loop Type 1 Diabetes Control

Researchers have developed LLM-T1D, a novel approach combining reinforcement learning and large language models to automate insulin delivery for Type 1 Diabetes. The system not only outperforms traditional methods but also explains its decisions in plain language, addressing the 'black-box' challenge in medical AI.

agentic-systems

RegNetAgents: A Multi-Agent Framework for Cross-Network Regulatory Driver Identification in Cancer Genomics

Researchers have introduced RegNetAgents, an AI-oriented multi-agent framework designed to identify regulatory driver candidates across heterogeneous gene networks in cancer genomics. By integrating bulk tumor and single-cell data, the system offers a powerful tool for discovering therapeutic targets in breast and colorectal cancers.