EconoJax: The Economic Simulator That Trades 2D Worlds for GPU Speed
A JAX-native rework of the AI Economist that strips away spatial clutter, vectorizes the whole world, and turns taxation research into something you can actually iterate on.
Unlike the AI economist, EconoJax does not feature a 2D grid world in which the population agents move.
- EconoJax argues that spatial detail is often overhead when the real research target is taxation, trade, and inequality.
- Its main trick is architectural, not cosmetic: it collapses the AI Economist’s 2D world into a 1D JAX state machine that can run much faster.
- The project makes fiscal policy part of the simulation loop, so taxes and equality are not outputs to inspect later but signals that shape the next step.
- Its niche is sharp and deliberate: it is narrower than a general economic testbed, but much better suited to fast policy iteration on GPUs.
Why EconoJax Deletes the Map
EconoJax makes a blunt claim: if your goal is to study economic behavior, the map may be the wrong place to spend your compute budget. The project strips the original AI Economist down from a 2D grid world to a 1D state representation, then uses JAX to push the whole system through GPU-friendly transforms.
That is the thesis. It is not a prettier simulator. It is a narrower instrument that tries to preserve the economic signal while discarding the spatial noise.
In practice, that means less time watching agents walk around and more time asking whether taxes, trade, specialization, and inequality are changing in useful ways. The simplification is the point.
What It Inherits from the AI Economist
EconoJax is a reimplementation of Salesforce Research’s AI Economist, not a separate genre of simulator. The inheritance is clear: agents gather resources, craft goods, trade through a market, and live under a government that sets tax policy.
| Dimension | AI Economist | EconoJax |
|---|---|---|
| World model | 2D grid with movement | 1D vectorized state |
| Primary bottleneck | Simulation overhead and slower training | Economic dynamics with JAX throughput |
| Research feel | Rich spatial toy world | Sharper policy experiment |
| Main question | How do agents behave in a world? | How fast can we study fiscal behavior without spatial clutter? |
| Execution style | Traditional RL stack | JAX-native, GPU-friendly loop |
The difference is not just speed. It is what speed enables. When the simulation becomes cheap enough to iterate on, policy research stops feeling like a long batch job and starts feeling like a live instrument.
# Core idea, simplified
state = EnvState(
inventories=..., # per-agent resources
skills=..., # gather/craft abilities
market_orders=..., # escrow and exchange state
tax_history=..., # fiscal trace
)
next_state, reward, done = env_step(state, action)
# JIT, VMAP, and GPU execution stay in play because the state is JAX-native.The World Is a JAX State Machine
The heart of the repository is a JAX-native state machine. `EnvState` stores inventories, skills, market orders, and tax history as arrays, which means the environment can be transformed with `jit`, `vmap`, and other JAX tools without breaking the functional model.
That matters because simulation throughput is the real product here. If every agent and every market action can stay inside the array world, the environment can run in parallel across many episodes without dragging data back and forth between CPU and GPU.
This is also why Equinox shows up. `eqx.Module` lets the project keep a complex nested state while remaining compatible with JAX’s immutable style. The code is built to be transformed, not merely executed.
Tax Policy, Inequality, and the Gini Loop
This is where the project becomes more than a fast environment. EconoJax uses inequality as part of the control problem. The government does not simply tax output and call it a day. The system measures distribution, computes Gini, and folds that signal back into the policy loop.
The repo’s fiscal machinery is specific enough to matter. It includes progressive brackets, escrowed trading, and a feedback structure that makes the economy feel like a live optimization problem rather than a static benchmark.
| Mechanism | What it does | Why it matters |
|---|---|---|
| Progressive tax brackets | Applies different rates at different income levels | Lets the environment study redistribution, not just flat extraction |
| Escrow | Holds resources and coin during exchange | Prevents trade from being treated as magic |
| Gini calculation | Measures inequality from agent wealth | Turns fairness into a state variable the policy can react to |
| Government update | Changes the next policy step | Makes fiscal choice endogenous to outcomes |
The important shift is causal, not cosmetic. Policy changes outcomes, outcomes reshape inequality, and inequality changes the next policy step. That loop is what researchers care about.
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.
Why the Math Helpers Matter
The utility module is easy to overlook, which is exactly why it is worth paying attention to. `util.py` contains the Gini helper and the Pareto-based skill distribution, and both are part of the scientific argument.
The Gini calculation is especially telling. The project uses an exact method for small populations and a faster approximation for larger ones. That is a practical tradeoff, but it is also a philosophy: the code adapts its precision to the scale of the experiment.
The Pareto skill distribution does something similar. It bakes in a familiar inequality shape, where a smaller number of agents end up highly specialized. That makes the economy feel less synthetic and gives trade and crafting real structure.
# Conceptually, util.py balances realism and speed
if population_size < 30:
gini = get_gini_exact(wealth)
else:
gini = get_gini_approx(wealth)
skills = get_pareto_skill_dists(population_size)
# The model is built to stay fast without pretending all populations look the same.Who Should Use This Instead of a Slower Simulator
EconoJax is not trying to be the universal answer. It is a specialization play for researchers who care about fiscal policy, multi-agent learning, and GPU-scale iteration more than spatial realism.
| Project | Best for | Tradeoff |
|---|---|---|
| EconoJax | Fast taxation and inequality research | Drops 2D spatial complexity |
| The AI Economist | High-fidelity origin model | Slower to iterate and heavier to run |
| TaxAI | Tax-focused MARL experiments | Less JAX-native scaling emphasis |
| ArthJAX | Macro shocks and contagion | Different problem space |
| Gymnax | JAX environment infrastructure | Not an economic model itself |
| EconGym | Modular economic testbeds | Broader, less specialized |
If your question is about emergent fiscal dynamics and you need to run many experiments quickly, EconoJax is the sharper tool. If your question depends on spatial movement or richer world geometry, the original AI Economist may still be the better fit.
That is the real editorial takeaway. EconoJax does not win by doing more. It wins by doing less, more cleanly, and in a way that lets the economics surface faster.