JaxMARL-BC: Multi-Agent Reinforcement Learning for Macroeconomic Modelling, in Seconds

Benchmarks are crucial in the development of machine learning algorithms, with available environments significantly influencing reinforcement learning (RL) research. Traditionally, RL environments run on the CPU, which limits their scalability with typical academic compute. However, recent advancements in JAX have enabled the wider use of hardware acceleration, enabling massively parallel RL training pipelines and environments. While this has been successfully applied to single-agent RL, it has not yet been widely adopted for multi-agent scenarios. In this paper, 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 and popular baseline algorithms. Our experiments show that, in terms of wall clock time, our JAX-based training pipeline is around 14 times faster than existing approaches, and up to 12500x when multiple training runs are vectorized. This enables efficient and thorough evaluations, potentially alleviating the evaluation crisis in the field. We also introduce and benchmark SMAX, a JAX-based approximate reimplementation of the popular StarCraft Multi-Agent Challenge, which removes the need to run the StarCraft II game engine. This not only enables GPU acceleration, but also provides a more flexible MARL environment, unlocking the potential for self-play, meta-learning, and other future applications in MARL. The code is available at https://github.com/flairox/jaxmarl.

Taxation is a fundamental instrument of economic policy and has long been studied through economic models. However, most existing taxation models focus on interactions within a single economic group, typically with one government and multiple households. Such a setting overlooks interactions among independent economic groups, where governments may compete for economic resources through fiscal policies. To capture these interactions, we formulate taxation as a hierarchical multi-group tax game. Within each group, a government sets tax policies while households respond, forming a hierarchical government--household interaction. Across groups, governments compete through fiscal policies, coupling the economic dynamics of distinct groups. This coupled structure poses challenges for standard multi-agent reinforcement learning (MARL) methods. To this end, we propose Multi-Group PPO (MGPPO), a bilevel MARL algorithm designed for hierarchical multi-group interactions. MGPPO incorporates Hierarchical Sampling to coordinate learning across agent levels and Curriculum Learning to improve training stability. We further develop a multi-group taxation simulation environment grounded in classical economic models, supporting the evaluation of fiscal policies under inter-group competition. Simulations show that in the two-group settings, MGPPO reduces the income and wealth Gini coefficients by 5.0% and 4.9%, respectively, and increases mean GDP by 4.1% compared with IPPO.

Power markets are a natural testbed for multi-agent reinforcement learning (MARL), where multiple self-interested participants repeatedly submit bids. A market-clearing mechanism then determines dispatch and prices subject to power grid constraints and market settlement rules. However, existing MARL environments typically focus on a single market setting, implement simplified clearing mechanisms, or rely on CPU-based optimization solvers that slow large-scale training and limit the systematic study of bidding strategies and market behavior. We introduce PowerMarketJax, a benchmark suite for MARL across five power markets: day-ahead wholesale, real-time balancing, ancillary services, peer-to-peer double auctions, and local flexibility. Each environment implements its own clearing, pricing, and settlement rules while providing a common framework for learning and evaluation. We find that learned bidding behavior depends strongly on the market design: independent learners can miss better strategies when gains require many agents to change together, when more profitable strategies lie beyond a region of lower profit, or when profits disappear as more agents adopt the same strategy. PowerMarketJax implements both market simulation and policy training in JAX, allowing the entire pipeline to run on the GPU with 1,024 X 1,200 parallelisms across both environments and market participants, achieving up to 33X speedup over CPU-based baselines. Our open-source benchmark is available at: https://github.com/powermarketjax/PowerMarketJax.