# EconoJax: A Fast & Scalable Economic Simulation in Jax

arXiv: https://arxiv.org/abs/2410.22165 · HTML: https://mlsanity.com/paper/2410.22165/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Adversarial robustness and RL in finance, one of 15 directions in Deep reinforcement learning methods.
- Builds on 13 of its 32 references that are in the map.
- First posted on arXiv 29 Oct 2024.

## Placed in

- Deep reinforcement learning methods / [Adversarial robustness and RL in finance](https://mlsanity.com/deeprl/adversarial-robustness-and-rl-in-finance/)

## Builds on (most cited first, 13 of 13)

- [The Surprising Effectiveness of PPO in Cooperative, Multi-Agent Games](https://mlsanity.com/paper/2103.01955/) (2021)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Monotonic Value Function Factorisation for Deep Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2003.08839/) (2020)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Equinox: neural networks in JAX via callable PyTrees and filtered transformations](https://mlsanity.com/paper/2111.00254/) (2021)
- [Structured State Space Models for In-Context Reinforcement Learning](https://mlsanity.com/paper/2303.03982/) (2023)
- [Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning](https://mlsanity.com/paper/2402.16801/) (2024)
- [Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning](https://mlsanity.com/paper/2303.17503/) (2023)
- [Building a Foundation for Data-Driven, Interpretable, and Robust Policy Design using the AI Economist](https://mlsanity.com/paper/2108.02904/) (2021)
- [AI for Global Climate Cooperation: Modeling Global Climate Negotiations, Agreements, and Long-Term Cooperation in RICE-N](https://mlsanity.com/paper/2208.07004/) (2022)
- [NAVIX: Scaling MiniGrid Environments with JAX](https://mlsanity.com/paper/2407.19396/) (2024)
- [Mava: a research library for distributed multi-agent reinforcement learning in JAX](https://mlsanity.com/paper/2107.01460/) (2021)
- [Foragax: An Agent-Based Modelling Framework Based on JAX](https://mlsanity.com/paper/2409.06345/) (2024)

## Built on by (most cited first, 1 of 1)

- [EconGym: A Scalable AI Testbed with Diverse Economic Tasks](https://mlsanity.com/paper/2506.12110/) (2025)

## Abstract (arXiv)

Accurate economic simulations often require many experimental runs, particularly when combined with reinforcement learning. Unfortunately, training reinforcement learning agents in multi-agent economic environments can be slow. This paper introduces EconoJax, a fast simulated economy, based on the AI economist. EconoJax, and its training pipeline, are completely written in JAX. This allows EconoJax to scale to large population sizes and perform large experiments, while keeping training times within minutes. Through experiments with populations of 100 agents, we show how real-world economic behavior emerges through training within 15 minutes, in contrast to previous work that required several days. We additionally perform experiments in varying sized action spaces to test if some multi-agent methods produce more diverse behavior compared to others. Here, our findings indicate no notable differences in produced behavior with different methods as is sometimes suggested in earlier works. To aid further research, we open-source EconoJax on Github.