EconoJax: A Fast & Scalable Economic Simulation in Jax

TL;DR

EconoJax is a fast, scalable economic simulation framework built in JAX that enables rapid training of reinforcement learning agents in multi-agent economic environments, significantly reducing training time.

Contribution

The paper introduces EconoJax, a JAX-based economic simulation platform that accelerates training and experimentation in multi-agent economic models.

Findings

Training with EconoJax takes minutes instead of days.

Real-world economic behaviors emerge within 15 minutes of training.

Different multi-agent methods do not significantly affect behavioral diversity.

Abstract

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.