EconoJax: A Fast & Scalable Economic Simulation in Jax
This provides a practical tool for researchers in economics and AI to run large-scale experiments efficiently, though it is incremental as it builds on existing AI economist frameworks.
The paper tackles the problem of slow training times in multi-agent economic simulations by introducing EconoJax, a fast and scalable simulation written in JAX, which reduces training from days to 15 minutes for 100 agents.
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.