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

Daniele Montagnani; Andrea Coletta; Antonietta Mira; Aldo Glielmo

cs.MAOct 20, 2026

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

Authors: Daniele Montagnani, Andrea Coletta, Antonietta Mira, Aldo Glielmo

Organizations: Università di Pavia Italy · USI Lugano Switzerland · Banca d’Italia† Italy · Università dell’Insubria Italy

Abstract

Multi-agent reinforcement learning (MARL) offers a principled way of populating simulated economies with many heterogeneous decision makers. Its adoption, pioneered in finance, has nonetheless been slow in macroeconomics because of two main limitations: the lack of a shared frame of reference connecting MARL systems to dynamic stochastic general equilibrium (DSGE) models the profession already trusts, and the very long training times. The first limitation was recently addressed by the Multi-Agent Reinforcement Learning Business Cycle (MARL-BC), which places deep RL households in a real business cycle economy and reduces to classic DSGE models in specific limits. Here we address the second limitation. We present JaxMARL-BC, a JAX-accelerated simulator that trains MARL-BC economies in seconds (between 75x and 724x faster than the reference implementation), and scales to 200,000 independent learning households. We verify its correctness against the known DSGE limits, characterise its scaling through a simple cost law, and demonstrate its value by a proof-of-concept calibration of an extended model with heterogeneous returns on capital, in a two-hour experiment that would previously have taken over a month. We release JaxMARL-BC as open source.

Figures & tables

Panel of Figure 1 — Figure 1. Speedup illustration and resulting scaling with population size $n$ . Left : an illustration of the axes of parallelism. One rollout evaluation is a rectangular array of agent-transitions with $n$ agents along one axis (a plain array dimension) and $E$ environments along the other (with jax.vmap ). This is advanced sequentially over $T$ rollout periods, within a lax.scan , which is not parallelisable. Rollouts and training steps live inside one jit -compiled program. Centre and right : throughput and wall-clock time of training runs with $10^{5}$ periods each, as a function of the number of agents $n$ . Curves are shown for both CPU and GPU backends (dashed light blue and solid blue lines) and each plots the mean and standard deviation over three seeds. The right-most panel also reports a linear slope as a dashed grey line. The simulator is evaluated against comparable MARL-BC runs (orange and red lines) with times taken from the original paper.

Figure 2. The LLM-augmented dashboard for JaxMARL-BC.

Panel of Figure 3 — Figure 3. Correctness check on the Krusell–Smith economy. The Krusell–Smith economy of experiment (i) with $n=200$ (top) and $n=100{,}000$ (bottom) households. Left : one-step transitions of aggregate capital $(K_{t},K_{t+1})$ coloured by aggregate state (good/bad), with a law of motion fitted separately in each state, at four points in training. Centre : stationary distribution of individual capital, with its Gini coefficient, before and after training. Right : learned consumption fractions against wealth (cash-on-hand), by employment status, before and after training.

Figure 4. The capital distribution through training. The distribution of agents’ capital for $n=200$ agents, at evenly spaced training snapshots, with its mean overlaid. The economy quickly leaves its untrained state and settles onto the stationary distribution of Figure 3 , which then holds for the rest of training.

Panel of Figure 5 — Figure 5. Throughput and run-time results. Throughput (left) and run time (right) for a $9\times 9$ grid over $n$ (columns) and $E$ (rows), each ranging from $1$ to $500$ at fixed sequential depth, with colour on a log scale. Both panels are symmetric about the diagonal, as from Eq. ( 7 ), throughput saturates at $\approx 10^{7}$ agent-transitions/s near $W=E\times n\approx 10^{4}$ , past which run time grows linearly in $W$ .

$n$	GPU (s)	CPU (s)	(h)	( $\times$ )	(h)	( $\times$ )
	ours		PPO-CPU		SAC-CPU
10	$21.97(3)$	$13.03(28)$	$0.46(14)$	$75$	$0.70(15)$	$115$
20	$23.38(10)$	$19.13(39)$	$0.70(15)$	$108$	$0.77(16)$	$119$
50	$26.49(11)$	$29.89(27)$	$1.49(24)$	$203$	$0.96(13)$	$130$
100	$27.98(12)$	$68.93(175)$	$2.72(26)$	$350$	$1.45(9)$	$187$
200	$29.63(8)$	$134.14(365)$	$4.00(28)$	$486$	$1.90(22)$	$231$

Table 1. Speedup with respect to the original implementation. The reference times are single-CPU times of the original implementation ( Gabriele et al., 2026 ) in the same mean-field Krusell–Smith limit and at the same six population sizes. Our times use

E=1

and the same budget of

10^{5}

periods, on a GPU and on a CPU backend; both are the mean over three seeds. Parenthesised digits give the standard deviation in units of the last digit. Speedups are relative to our GPU times.

Figure 6. Calibration of ex-ante heterogeneity in capital productivity. Individual capital productivities

\kappa^{i}

are drawn from a log-normal distribution with unit mean and dispersion

\sigma

, with

\sigma

swept from

0

(the homogeneous Krusell–Smith benchmark,

\kappa^{i}\equiv 1

) to

0.5

, with ten seeds per value. Every curve is a mean over seeds with a one-standard-deviation band. Left : the realised

\kappa

profiles by agent rank. Increasing

\sigma

tilts the flat benchmark profile into an increasingly skewed one. Centre : Lorenz curves of the resulting wealth distributions move away from the diagonal in the same fashion, the bottom

80\%

of households holding

70\%

of aggregate capital at

\sigma=0

and under

20\%

\sigma=0.5

. Right : the resulting capital Gini, rising from

\approx 0.2

to more than

0.7

\sigma=0.5

Figure 7. Heterogeneous return on capital. Density of percentage returns on capital across agents, in percentage points, for $\sigma=0.40$ and $\sigma=0.50$ . Heterogeneity produces right-skewed distributions in which most agents earn less than the average, with a thin upper tail extending beyond returns of 15%. Increasing $\sigma$ shifts mass towards lower returns and widens dispersion.

Multi-Agent Reinforcement Learning

Multi-agent reinforcement learning studies how interacting agents learn coordinated or competitive policies through shared experience and rewards.

See all 371 papers

JaxMARL: Multi-Agent RL Environments and Algorithms in JAX
Nov 16, 2023Alexander Rutherford, Benjamin Ellis, Matteo Gallici +18RL Benchmarks Reinforcement Learning
University of Oxford · Universitat Politècnica de Catalunya · University College London +2
Hierarchical Multiagent Reinforcement Learning for Multi-Group Tax Game
May 6, 2026Honglei Guo, Yexin Li, Chiyuan Wang +1Game Theory Hierarchical RL
College of Computer Science and Technology Zhejiang University State Key Laboratory of General Artificial Intelligence, BIGAI
PowerMarketJax: A JAX Benchmark Suite for Multi-Agent Reinforcement Learning in Power Markets
Sep 29, 2026Zhanhua Pan, Xin Qin, Xiao Liu +3RL Benchmarks Multi-Agent Reinforcement Learning
Nanyang Technological University, Singapore. · Cornell University, USA. · University of Bristol, UK.

Nov 16, 2023J/K move · Enter open · S save

JaxMARL: Multi-Agent RL Environments and Algorithms in JAX

Alexander Rutherford, Benjamin Ellis, Matteo Gallici, Jonathan Cook+17

University of Oxford · Universitat Politècnica de Catalunya · University College London+2

RL Benchmarks Reinforcement Learning Game-Playing Agents GPU-Accelerated RL Multi-Agent Reinforcement Learning

Open paper AskPDF arXiv

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.

Abstract

Multi-Agent Reinforcement Learning

Explore similar work

JaxMARL: Multi-Agent RL Environments and Algorithms in JAX

Hierarchical Multiagent Reinforcement Learning for Multi-Group Tax Game

PowerMarketJax: A JAX Benchmark Suite for Multi-Agent Reinforcement Learning in Power Markets