SRL Tutorial

This repo presents a pedagogical JAX implementation of policy-gradient methods for heterogeneous-agent macro models in general equilibrium, including with aggregate risk.

SRL goes with the paper Structural Reinforcement Learning for Heterogeneous Agent Macroeconomics by Yucheng Yang, Chiyuan Wang, Andreas Schaab, and Benjamin Moll. We try to make the repo and our codes as accessible as possible. Our hope is that you will be able to clone the repo, read the notebooks top to bottom, and easily follow along with both the economics and the method. Note this is the tutorial repo for SRL, not the final replication codes.

The tutorials

Start with the tutorials. They start with the simplest, canonical household problem and build up to our full implementation of a heterogeneous agent model with aggregate risk by adding one step at a time. Read them in order from notebook 0. TUTORIAL.md is the full syllabus.

#	Notebook	Open in Colab	Colab G4 runtime	Free Colab T4 runtime
0	Household problem, in NumPy		~1 min	~3 min
1	The same problem, in JAX		~6 s	~15 s
2	The same problem, by policy gradient		~30 s	~1 min
3	The household with moving prices		~45 s	~1.5 min
4	Huggett without aggregate risk		~4.5 min	~8 min
5	Huggett with aggregate risk		~35 s	~1.5 min

Each notebook runs top to bottom on a Colab GPU (click the badge) or locally on a CPU (but more slowly). On Colab the first cell clones the repo and installs the package, so there's no token and no setup notebook.

Install (local)

pip install -e .

Python ≥ 3.11. JAX is pinned (jax==0.10.0) so results reproduce exactly across machines.

Repo structure

• srl/ is the package: The policy-gradient solver (SPGSolver), a value-function-iteration baseline (VFISolver), and the discretization and plotting helpers (srl.utils).
• tutorials_new/: the notebooks, the shared calibration.py, and the SSJ reference data.
• TUTORIAL.md is the syllabus.

License & citation

MIT, see LICENSE. If you use SRL in your work, please cite our paper:

@misc{SRL,
      title={Structural Reinforcement Learning for Heterogeneous Agent Macroeconomics},
      author={Yucheng Yang and Chiyuan Wang and Andreas Schaab and Benjamin Moll},
      year={2025},
      eprint={2512.18892},
      archivePrefix={arXiv},
      primaryClass={econ.TH},
      url={https://arxiv.org/abs/2512.18892},
}