Source-linked AI summary
JaxMARL-BC: Multi-Agent Reinforcement Learning for Macroeconomic Modelling, in Seconds
Daniele Montagnani, Andrea Coletta, Antonietta Mira, Aldo Glielmo
TL;DR
MARL has been slow to gain traction in macroeconomics because it lacks a trusted connection to DSGE models and can take hours to train. JaxMARL-BC uses JAX to accelerate MARL-BC and scale its learning households. The paper reports training in tens of seconds and brings populations of 200,000 households within reach of a single run.
Problem
Macroeconomic MARL adoption is limited by concerns about its connection to trusted DSGE models and by the time required to train RL systems.
Source: §1 · p. 1Method
JaxMARL-BC uses JAX to accelerate MARL-BC, parallelising independent economies while representing coupled households as an array and processing rollout periods sequentially.
Source: §3.2 · p. 3, §3.1 · p. 3Results
Training a MARL-BC economy takes tens of seconds, and populations of 2 × 10^5 independent learning households come within reach of a single run.
Source: §6 · p. 7Takeaways & Limitations
Large-scale runs and parameter calibration can become an integral part of model development.
Source: §6 · p. 7The reported economies are RBC and Krusell–Smith models, while wider model coverage remains a future extension.
Source: §6 · p. 7
Abstract
from arXiv · showhide
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.
1 Introduction
MARL offers a way to model heterogeneous economic decision makers, but its adoption in macroeconomics has lagged amid concerns about its connection to trusted DSGE models and its computational cost. JaxMARL-BC addresses the cost bottleneck by accelerating MARL-BC training and supporting larger populations.
1 Introduction
- Macroeconomic MARL adoption remains comparatively slow despite work on stylised economies, monetary policy, and tax design.The approach is not yet part of policymakers’ and practitioners’ toolkit.Source: §1 · p. 1
- A trusted DSGE frame of reference and fast solvers help make models practical, while MARL training can require multiple hours per run.MARL-BC connects deep RL households to a real business cycle economy and recovers established DSGE models in specific limits, but does not address the training-time bottleneck.Source: §1 · p. 1
- JaxMARL-BC fills the gap in JAX-accelerated tools for structural MARL-driven macroeconomic modelling.The introduction also describes an open-source release with a prototype natural-language dashboard for configuring simulations.Source: §1 · p. 2
- JaxMARL-BC uses JAX to train MARL-BC economies in seconds and scale to hundreds of thousands of independent learning agents.The contribution includes checking known DSGE limits, characterising scalability with an empirical cost law, and calibrating a model with heterogeneous capital returns in two hours rather than over a month.Source: §1 · pp. 1–2
2 Related work
JAX-based acceleration has been used in RL libraries and in finance, while economic applications include a stylised simulation. JaxMARL-BC targets structural heterogeneous macroeconomic modelling, a domain the cited packages do not cover.
2 Related work
- JAX RL libraries report speedups of one to four orders of magnitude over comparable CPU runs, and finance applications report substantial training-time reductions.JAX-LOB reports at least 7× speedup, while JaxMARL-HFT reports up to a 240× training-time reduction.Source: §2 · p. 2
- EconoJax trains populations of 100 agents in fifteen minutes, compared with several days for the original AI-Economist implementation.It ports the AI-Economist environment to JAX.Source: §2 · p. 2
- JaxMARL-BC targets structural heterogeneous macroeconomic modelling, which the cited JAX packages do not cover.ABIDES-Economist uses RL agents in agent-based economies but is not implemented in JAX and does not target established macroeconomic benchmarks.Source: §2 · p. 2
- JaxMARL-BC aims to support independent learning agents rather than compete with fast conventional methods on model classes they already solve well.The related work also distinguishes deep-learning methods without RL from the independent-agent modelling paradigm targeted here.Source: §2 · p. 2
3 The JaxMARL-BC simulator
JaxMARL-BC implements MARL-BC with JAX to accelerate training, while retaining its economic structure and providing a dashboard for launching simulations. Its design batches independent economies and parallelises household computations while preserving sequential rollout periods.
3.1 The underlying model
- Households use PPO as independent learners, with a shared policy network evaluated on each household’s private observations.The paper describes this parameter-sharing arrangement as a stylised form of social learning; household interactions make each learner’s problem non-stationary.Source: §3.1 · p. 3
- MARL-BC recovers representative-agent RBC with one agent and mean-field Krusell–Smith with many ex-ante identical agents.These limits provide known DSGE behaviours against which the implementation can be tested.Source: §3.1 · p. 3
3.2 The JAX speedup
- JAX represents coupled households as an array, batches independent economies with jax.vmap, and advances rollout periods sequentially with lax.scan.The policy network uses the combined batch width W = E × n, and the optimiser and environment updates run inside one jit-compiled program.Source: §3.2 · p. 3, Figure 1 · p. 3
- The run-time model separates fixed launch latency from width-dependent computation, yielding a near-constant-cost regime before execution time grows linearly with batch width.The transition occurs around W ≈ τP; latency and throughput are effective values averaged across rollout and gradient stages.Source: §3.2 · pp. 3–4
3.3 The LLM-augmented dashboard
- The dashboard converts natural-language requests into editable YAML configurations, launches runs, streams logs, and renders artefacts.The authors present it as a convenience for exploration and do not evaluate it in this work.Source: §3.3 · p. 4
4 Experimental setup
The experiments test JaxMARL-BC against the mean-field Krusell–Smith benchmark, measure its runtime across population and environment sizes, and calibrate an extension with heterogeneous capital productivity.
4.1 General setup
- Training runs use a P100 GPU, shared-parameter 64 × 64 tanh MLPs, and PPO with Adam; experiment (ii) also uses an Intel Xeon CPU.Updates use clipping 0.2, GAE decay 0.95, learning rate 3 × 10−4, four epochs, and 64 minibatches.Source: §4.1 · p. 4
- The simulator is tested on a mean-field Krusell–Smith economy and an extension with heterogeneous capital productivity.The benchmark uses α=0.36, β=0.95, and δ=0.025; the extension sweeps productivity dispersion σ from 0 to 0.5.Source: §4.1 · p. 4, §4.2 · p. 5
4.2 The three experiments
- Correctness is evaluated after 640 gradient updates, first with 200 households and then with 100,000 households.Each run uses 12 environments and 200-period episodes; evaluation averages the final 50 periods of a fresh simulation with a frozen policy.Source: §4.2 · p. 4
- Scaling tests compare CPU and GPU training times across six populations, then extend runs to 10,000 CPU households and 200,000 GPU households.A further 81-run sweep varies environment count E and population n across nine values each, using 10^5 rollout periods and averaging three runs.Source: §4.2 · p. 4
- The correctness figure compares results for 200 and 100,000 households across aggregate-capital transitions, capital distributions, and consumption policies.It shows aggregate transitions by good/bad state at four training points, plus capital distributions and consumption fractions before and after training.Source: Figure 3 · p. 5
- The calibration experiment sweeps ex-ante capital-productivity heterogeneity across 110 runs, with ten seeds at each of 11 dispersion levels.Productivity is sampled from a log-normal distribution parameterized to keep the population mean approximately 1; each run uses 200 households and 12 environments.Source: §4.2 · p. 5
5 Results
JaxMARL-BC reproduces key Krusell–Smith properties, substantially reduces training time, and supports a calibration of heterogeneous capital returns that would have taken over a month with the reference implementation.
5.1 Correctness
- At n=200, training reproduces the linear aggregate-capital law, bimodal wealth distribution, and declining consumption share used to validate the framework.The capital Gini is 0.22 in this run, while mean capital settles at approximately 13.Source: §5.1 · pp. 5–6
- At n=100,000, the aggregate-capital law reaches R²=1.000 in both states; at n=200, the wealth distribution settles within the first fifth of training.The n=200 population captures the reported phenomena and is trainable in minutes on a CPU.Source: §5.1 · p. 6, Figure 4 · p. 6
5.2 Scaling
- 75× at n=10 to 724× at n=500: JaxMARL-BC’s PPO training speedup over the reference implementation spans this range.The GPU run time rises from 22 s to 36 s across a fifty-fold increase in population size.Source: Table 1 · p. 6, §5.2 · p. 6
- Throughput saturates at ≈107 agent-transitions/s near W≈10,000, after which run time grows linearly with W.For low W, run time rises from ≈18 s at W≈1 to ≈48 s at W≈1,000; JaxMARL-BC also trains populations up to 200,000 households in tens of seconds.Source: §5.2 · pp. 6–7, Figure 5 · p. 7
5.3 Calibration of ex-ante heterogeneity
- Capital Gini rises from 0.20±0.05 at σ=0 to 0.72±0.05 at σ≈0.4, bringing wealth inequality closer to levels reported for advanced economies.The calibration used 110 runs and took about two hours, compared with a lower-bound estimate of more than a month for the reference implementation.Source: §5.3 · p. 7
6 Conclusions
JaxMARL-BC trains MARL-BC economies in tens of seconds and brings runs with 200,000 independent learning households within reach. The conclusions also show how heterogeneous capital productivity shapes wealth and returns, and identify broader model and algorithmic directions for future work.
6 Conclusions
- 200,000 independent learning households come within reach of a single run, while training a MARL-BC economy takes tens of seconds rather than hours.The authors argue that this speed makes large-scale runs and parameter calibration more feasible as part of model development.Source: §6 · p. 7
- Future work could extend the framework to Huggett, New Keynesian, and HANK economies, and to settings with richer heterogeneity or less stringent rational-expectations and general-equilibrium assumptions.The authors note that classical solution methods remain applicable to the rational-expectations counterparts of the economies studied here.Source: §6 · pp. 7–8
- At σ = 0.5, the bottom 80% of households hold under 20% of aggregate capital, and capital returns are right-skewed with a thin upper tail beyond 15%.Figure 6 reports capital Gini rising from about 0.2 to more than 0.7; Figure 7 shows that increasing σ shifts return mass lower and widens dispersion.Source: Figure 6 · p. 8, Figure 7 · p. 8
- PPO was adequate for the experiments, while richer environments may benefit from off-policy learners such as SAC; the largely differentiable simulator could also support gradient-based calibration.The authors suggest this as an alternative to grid search.Source: §6 · p. 8
- JaxMARL-BC is open source and includes an LLM-augmented dashboard; the authors expect its trusted model classes and rapid iteration to support wider MARL use in macroeconomic modelling.They frame this as a path from proof-of-concept demonstrations toward standard practice.Source: §6 · p. 8