ArthJAX: The Differentiable Economy Built to Stress-Test a Central Bank
A JAX-powered agent-based model turns households, banks, markets, and policy rules into one compiled simulation, then adds a world model so the economy can be learned as well as run.
- ArthJAX frames macroeconomic simulation as compiled computation, so policy can be explored like model parameters instead of hand-tuned scenarios.
- Its core trick is functional JAX plumbing, where state moves through a deterministic pipeline instead of a slow Python loop.
- The model keeps economic intuition intact by pairing vectorized agents with familiar macro rules like Taylor-style policy and market feedback.
- A neural world model extends the simulator from running scenarios to learning transition dynamics.
Most agent-based macro models live in the slow lane. ArthJAX does something stranger: it treats an economy like a JAX program, then asks what happens when households, banks, markets, shocks, and policy all become one compiled state machine. That shift matters because it changes the question from "Can we simulate this?" to "Can we search this?"
A Central Bank Sandbox You Can Compile
The project is a GPU-accelerated agent-based model for macroeconomic stress testing. Instead of a loose pile of loops and mutable objects, it uses JAX primitives to make the whole simulation fast, reproducible, and amenable to optimization. The practical payoff is simple: you can run scenarios quickly enough to explore policy without treating every experiment like a batch job.
The repository positions itself as a differentiable sandbox. That phrase is doing real work here. If the full trajectory is expressible as a computational graph, then policy settings stop being just knobs and start looking like parameters.
Why JAX Changes the Shape of the Model
JAX is not just an implementation detail. It changes how the simulation is built. With jit, the transition code gets compiled. With vmap, behavioral variation becomes batchable. With lax.scan, time itself becomes a compiled loop instead of a Python one.
# Conceptual shape of the engine
state = init_state(config)
for t in range(T):
state = step(state, config)
# In ArthJAX, the temporal loop is pushed into JAX
trajectory = lax.scan(step, state0, time_steps)That structure matters because it collapses the usual distance between simulation and optimization. The model is not only executing a scenario. It is building a reusable graph over the scenario, which is the prerequisite for gradient-based policy search and other machine-learning-style workflows.
The Economy Runs as a State Machine
The engine threads one state dictionary through a fixed sequence of updates. Households act first, then companies, then banks, markets, shocks, contagion, and macro policy. The sequence is explicit, which makes the system easier to reason about and easier to test.
This functional design is the quiet superpower of the repo. The config is frozen, the state is passed forward, and each function returns a new state. That makes the model reproducible in the way researchers actually need: same inputs, same trajectory, no hidden mutation lurking in the background.
Agents Are Not Little Objects, They Are Vectors
The agent layer is designed for batch computation, not object-oriented drama. Households are grouped into behavioral types like Value, Momentum, Panic, and Saver, and the code uses vectorized tensor operations to apply different rules to different subsets in one pass. Banks track leverage and bad loans, which creates the familiar feedback loop where distress tightens credit and credit tightening deepens distress.
| Layer | How it behaves | Why it matters |
|---|---|---|
| Households | Vectorized behavioral types | Different rules can coexist without giving up GPU throughput |
| Banks | Tracks leverage and bad loans | Credit conditions feed back into the real economy |
| Companies | Respond to demand and financing conditions | Production and default dynamics become part of the same loop |
| Markets | Price discovery plus volatility structure | Asset moves look stylized instead of purely random |
The trick is that the behavior can be branchy without becoming slow. JAX lets the model keep the heterogeneity that agent-based modeling needs while still batching the work in a way GPUs like.
Macro Policy as a Thermostat
The top-down layer gives the economy a policy spine. A Taylor-style rule and a Phillips-curve-style relationship anchor the model in familiar macro logic, while the market layer adds mean reversion and volatility clustering. The result is not realism in the full empirical sense, but something more useful for experimentation: economically legible feedback.
That is the key design choice. ArthJAX does not try to replace macro theory. It uses macro theory as a control surface inside a richer simulation, which makes the policy levers easier to interpret when the system starts to wobble.
Contagion Is the Cleverest Part
This is where the model stops feeling like a generic ABM and starts feeling like a computation graph with economic meaning. Contagion is handled with a sector dependency matrix and a dot product, which means panic spreads through structure rather than bespoke scenario code. That is a small implementation choice with a big editorial payoff: the mechanism is simple, but the effect can be rich.
The visual intuition is strong because the math is strong. A dense linear operation can move stress across many links at once, and that makes the model feel closer to graph computation than to a hand-coded crisis script.
The World Model Makes the Simulator Useful for Learning
ArthJAX is not only a simulator. It also includes a neural world model that learns next-state transitions, so the simulated economy can become training data for a surrogate predictor. That matters because it bridges the gap between simulation and learning: once you can predict the trajectory, you can study stability, rollouts, and longer-horizon behavior in a different way.
The multistep loss is the important detail. One-step accuracy is easy to fake. Multi-step stability is harder, and it is what keeps a learned dynamics model from drifting too far when you roll it forward through time.
Where It Sits in the Macro Software Landscape
| Approach | Speed | Differentiability | Behavioral richness | Policy search readiness | Interpretability |
|---|---|---|---|---|---|
| ArthJAX | High on GPU | High in principle | High | High | Medium |
| Traditional ABM | Often low to medium | Low | High | Low | High |
| DSGE model | High | Medium to high | Low to medium | Medium | High |
| Plain Python loops | Low | Low | Medium | Low | Medium |
ArthJAX is not trying to win every comparison. It is trying to occupy a rare intersection: behavior-rich, fast enough to explore, and compatible with optimization workflows. That is a narrow but valuable niche, especially for people who want a central bank sandbox that feels less like spreadsheet archaeology and more like compiled computation.