EconoJax: A Fast & Scalable Economic Simulation in JAX

EconoJax is (loosely) a reimplementation of The AI Economist in JAX with a 1D observation space rather than the original 2D visual space. With GPU support, EconoJax's transition function is over 100x times faster and agents converge over 2000x times faster.

---

📦 Installation

For those using uv, it is possible to run a standard PPO implementation with default settings by directly running uv run main.py.

git clone git@github.com:ponseko/econojax.git
cd econojax
uv run main.py

Alternatively, install the project as an editable package in your favourite virtual environment software. E.g. using conda:

git clone git@github.com:ponseko/econojax.git
cd econojax
conda create -n econojax python=3.11
conda activate econojax
pip install -e .

python main.py

for CUDA support, additionally run pip install jax[cuda].

---

📑 Citing

If you use EconoJax in your research or projects, please cite:

@article{ponse2024econojax,
  title={EconoJax: A Fast \& Scalable Economic Simulation in Jax},
  author={Ponse, Koen and Plaat, Aske and van Stein, Niki and Moerland, Thomas M},
  journal={arXiv preprint arXiv:2410.22165},
  year={2024}
}