Lecture 13, Notebook 08: Continuous-time Aiyagari — Steady-State PINN (marginal-value form) with an FD benchmark

Course: Deep Learning for Solving and Estimating Dynamic Models in Economics and Finance Script reference: Chapter 8, §§8.3–8.5 (HJB, KFE, stationary Aiyagari/Huggett equilibrium) and §8.7 (EMINNs). Notebook role: core Author: Simon Scheidegger and Pavel Ievlev

This notebook is a faithful, self-contained PyTorch port of the JAX ss/ soft-penalty solver behind the EMINN Krusell–Smith model (Gu, Laurière, Merkel, Payne, 2024, App. B.1), evaluated at its deterministic steady state. There is no finite-difference solve anywhere in training: the FD code is the validation reference only. The two neural nets run in float32 (the JAX original’s precision); the FD reference is float64 NumPy on CPU.

A Pedagogical Walkthrough: A Mesh-Free Steady-State PINN, Checked Against Finite Differences¶

This notebook is designed for teaching. It is intentionally explicit and heavily annotated.

We solve the stationary Aiyagari equilibrium two ways and compare them:

1. Finite Differences (FD), the classic Achdou et al. (2022a) upwind scheme with outer bisection on the interest rate. This is the trusted reference (the “FD gospel”), and nothing it computes ever enters PINN training.

2. A steady-state PINN, two small neural nets trained on the HJB and KFE residuals plus a handful of exact integral identities, with prices and market clearing built in by construction (no outer bisection).

By the end we compare the marginal value (W=\partial_a V), consumption (c(a,l)), saving (s(a,l)), the stationary density (g(a,l)), and the equilibrium aggregates (K, r, w), and we score the PINN with a single (L^\infty) saving gate against a dense FD policy.

In-class notebook (Chapter 8 / Lecture 13 deck Part II, the in-class mesh-free solver). Full training is about 9 min on CPU at RUN_MODE="production"; use "smoke" for a fast structural pass that intentionally does not clear the gate (the aggregate-saving ramp needs a converged policy, see the failure-modes section).

Learning Objectives¶

• Understand the stationary HJB + KFE + market-clearing fixed point, written in marginal-value form (W=\partial_a V).

• See how a PINN can make market clearing hold by construction: prices are the firm’s marginal products at the density net’s own aggregates, so there is no outer price loop.

• Learn the role of each loss term, the HJB residual, the KFE (in two forms), and the exact integral identities that close the gaps a pointwise residual leaves open.

• Understand the continuous Gauss–Seidel gradient gating (two detachments) that decouples the policy net from the distribution net.

• Diagnose solution quality with residuals, moments, and a saving (L^\infty) gate against a dense FD benchmark.

Economic Environment and Equilibrium¶

The EMINN Krusell–Smith model at its deterministic steady state: aggregate TFP is pinned at (z=\bar z) (no aggregate shocks), so the economy collapses to the classic Aiyagari/Huggett stationary equilibrium.

Households. Hold wealth (a\in[a_{\min},a_{\max}]=[10^{-6},20]), earn labor income (w,l) with (l\in{l_1,l_2}={0.3,1.7}) switching by a symmetric Poisson process ((\lambda_1=\lambda_2=0.4)), and choose consumption (c) with CRRA utility (u(c)=(c^{1-\gamma}-1)/(1-\gamma)), (\gamma=2.1), discount rate (\rho=0.05). The budget is (\dot a = s = w,l + r,a - c).

Borrowing constraint (soft form). Instead of a hard (a\ge 0), the flow utility carries the paper’s penalty (\psi(a)=-\tfrac12\kappa(a-a_{lb})^2) for (a\le a_{lb}) (zero above), with (\kappa=3,\ a_{lb}=1). Below (a_{lb}) the penalty’s gradient (\psi’(a)=-\kappa(a-a_{lb})>0) raises the marginal value of wealth, hence raises saving, so households endogenously avoid the region and the constraint never literally binds.

Firms / prices. Cobb–Douglas (Y=e^z K^\alpha L^{1-\alpha}) ((\alpha=1/3), depreciation (\delta=0.1)) with competitive factor markets,

Equilibrium is a triple (policy, distribution, prices) such that the policy is optimal given prices (HJB), the distribution is stationary under that policy (KFE), and prices are the marginal products at the distribution’s own aggregates (K=\sum_j\int a,g_j,da), (L=\sum_j l_j\int g_j,da) (market clearing).

1) Calibration and Numerics¶

The calibration follows the paper’s Table 5 at the deterministic steady state ((z=\bar z=0)). The numerical hyperparameters follow the JAX ss/ recipe: two small nets (width 64, depth 4), a fixed 400-point trapezoid quadrature grid for normalization and aggregates, and fresh random collocation each step. Grid sizes and the training length are overridden from RUN_CFG so the smoke run stays fast.

2) The Equations the Solver Uses¶

2.1 HJB in marginal-value form¶

The solver never represents (V); it works with (W=\partial_a V) (as the whole repo does). Starting from the stationary HJB

the FOC gives (u’(c)=W), i.e. (c=W^{-1/\gamma}) (this is why (W>0) is structural). Differentiating the HJB in (a), the (\partial c/\partial a) terms cancel by the envelope theorem, leaving the HJB residual

All (z)-derivative terms and the distribution-coupling term of the full master equation vanish at the deterministic steady state, which is exactly what makes the SS problem self-contained. One extra structural fact used by the strong-form KFE is (\partial_a s = r - c’(a) = r + (1/\gamma)(c/W),\partial_a W).

2.2 Stationary KFE¶

For each labor state (j) (with (\tilde\jmath) the other state),

Two exact integral consequences (both used as losses): labor-mass balance (\lambda_1 M_1 = \lambda_2 M_2) with (M_j=\int g_j) (with (\lambda_1=\lambda_2) this forces the 50/50 split, hence (L=1)), and the aggregate-saving identity (,dK/dt = \int \sum_j s_j g_j,da = 0) (at stationarity the population saves nothing in aggregate).

The code below specializes the CRRA/switching equations of Chapter 8 to this calibration.

3) Finite-Difference Reference (validation only)¶

This is the trusted benchmark, the “FD gospel”: a standard Achdou et al. (2022a) upwind scheme on a uniform grid, with outer bisection on the interest rate (r). Nothing here ever enters PINN training; the FD objects are used only as ground truth for the diagnostics and the saving gate.

• _fd_inner: implicit-Euler upwind HJB at fixed prices, then the stationary KFE as the null vector of the generator transpose (one row replaced by the normalization).

• fd_solve: bisect (r) to clear the capital market (invert the firm FOC for (K_r), solve the partial-equilibrium problem, compare with (K_g=\int a,g)).

• fd_gate: the dense fixed-price FD saving policy at the FD-equilibrium prices, the reference for the (L^\infty) gate. Solving the HJB on the dense grid resolves the kink that interpolating a coarse policy would smear.

The FD reference deliberately uses n_a = n_quad, because the FD equilibrium (K) is grid-sensitive and the SS solver bakes its own equilibrium into (W); gating against a coarser FD would measure discretization mismatch, not solver error.

4) The Two Neural Nets¶

Both unknowns are plain MLPs in ((a_{\text{scaled}}, \text{onehot}(l))): 3 inputs, 4 tanh hidden layers of width 64, a linear head, Glorot/Xavier init, with (a) affinely scaled to ([-1,1]). Float32, seeds fixed at import. both(net, a) evaluates the two labor states in one batched forward (a.repeat(2) against stacked one-hots, reshaped to ((B,2))).

4.1 WNet, the policy side¶

Softplus enforces (W>0) (marginal utility is positive; (c=W^{-1/\gamma}) requires it) and blocks the parametrization form of the (W\to 0) cheat. But (W\to 0) is still an attractor the optimizer can approach ((W) small (\Rightarrow c) huge (\Rightarrow s) very negative (\Rightarrow) the transport term washes out the HJB residual); what actually guards against it is the pretrain level (§7) and the shape penalty. Audit check D1: the trained net has (\min W \approx 0.19), nowhere near the corner, and (\partial_a W < 0) everywhere on a fine grid.

4.2 GNet, the distribution side¶

Three deliberate choices: log-space (positivity for free, and densities spanning 4+ orders of magnitude become (O(1)) quantities); a trainable tail slope (\beta_j) (the true density decays exponentially, but a bounded tanh-MLP cannot steepen its own tail fast enough where the KFE residual is proportional to the tiny local density, so the decay rate is made a single parameter per labor state, receiving gradient from every tail node, initialized at the pretrain prior’s rate (1/\bar a) via the softplus inverse); and normalization by construction via logZ = logsumexp over (nodes (\times) states) on the fixed trapezoid quadrature, so (\sum_j\int g = 1) identically (audit D2: total mass = 1.0000), removing the KFE’s scale invariance without a tradeable penalty.

4.3 Prices by construction¶

Every loss evaluation recomputes, from the g-net’s current weights, (f_q\to\log Z\to g_q\to K=\sum a,w_q,g,\ L=\sum l,w_q,g\to r,w). There is no stored price variable and no outer bisection: market clearing holds identically at every step because (r,w) are defined as marginal products at the g-net’s own aggregates. The classical nested algorithm (solve a full partial-equilibrium problem per candidate (r), then bisect, as fd_solve does) is replaced by one flat loop in which the distribution shifting is the price adjustment. This relies on the Aiyagari fixed point being stable ((K) too high (\Rightarrow r) low (\Rightarrow) saving down (\Rightarrow) mass flows down (\Rightarrow K) falls); an unstable GE feedback would break the joint scheme.

5) The Joint Loss, Term by Term¶

make_loss(p) builds a closure over the quadrature grid. Each call computes:

HJB residual (trains the W-net) at n_col=256 fresh random collocation points, both labor states. (\partial_a W) is taken by exact autograd: the batch is a.repeat(2) with requires_grad, one forward, one torch.autograd.grad(W.sum(), a, create_graph=True). The .sum() trick is valid because the Jacobian is diagonal (each output depends only on its own input), and create_graph=True makes the derivative itself differentiable so the loss can train through it. The policy inside is live (the net is asked to satisfy its own FOC-consistent equation); prices are detached. This term is fully mesh-free in both KFE modes.

Shape penalty (L_{\text{shape}}=\text{mean}[\max(\partial_a W,0)^2]): the monotonicity prior (V) concave. It costs nothing once satisfied (it reads 0.0 after warm-up) but blocks sign-flipped or oscillating (W) during the transient, and is one of the two guards against the (W\to 0) corner.

KFE, form “fv” (default): the paper’s conservative upwind finite-volume operator (kfe_drift) applied to the g-net masses on the quadrature grid. This is not an FD solve: no generator is assembled and no linear system is solved; kfe_drift is an algebraic map (shifts, clamps, multiplies) from the net’s values to the rate at which mass would flow, and the gradient flows through the masses into the g-net. Because it is conservative/upwind, total mass is conserved for any policy, the zero-flux BCs are built in, and the strong form’s structural cheats are impossible by construction. The μ/da scaling makes the residual a density-rate, so the loss magnitude is grid-resolution-independent, and the operator equals the FD generator (A^\top) to ~3e-10 (the FD solver enters only as the validation reference; scrambling it leaves training bit-identical).

KFE, form “strong”: pure mesh-free, the KFE enforced by exact autodiff at the collocation points, with the policy pair ((s,\partial_a s)) detached. Here mesh-free means precisely that the PDE is enforced by exact autodiff at sampled points; the integral identities may still be evaluated by quadrature of net values (normalization already requires that), which involves no stencils. The pointwise residual alone is not well posed, so three exact identities close the gaps, each killing a failure mode diagnosed in the JAX runs:

1. Pointwise total-flux identity (\sum_j s_j g_j = 0). Summing the two KFE equations cancels the jump terms, so the strong residual only forces (\partial_a(\text{total flux})=0); total flux = any nonzero constant has zero residual, the family (g\sim 1/|s|) (JAX run 1 found exactly this: (K) drifted to ~11, (r<0), near-zero KFE loss). Forcing the flux to 0 pointwise excludes it ((F_1=-F_2=\text{const}\neq 0) would force (\lambda_1 g_1=\lambda_2 g_2) and (s_1 g_1=-s_2 g_2) pointwise, hence generically (g=0)). In fv form the boundary faces are hard-zero, so this is impossible.

2. Labor-mass balance (\lambda_1 M_1=\lambda_2 M_2). An MSE-imperfect pointwise residual does not pin the labor split (JAX run 3 settled at (L\approx 0.3), (r\approx -0.05), a pseudo-equilibrium). In fv form the balance telescopes exactly inside the operator.

Plus zero-flux endpoint penalties (L_{\text{bc}}) (the BCs the fv operator gets for free) and an extra top collocation band over the employed’s (s\approx 0) weak-transport zone, where the transport term (s,\partial_a g) barely transmits information so the region needs explicit sampling density.

Aggregate-saving identity (both forms, ramped) (L_{\text{agg}}=(\sum_{\text{nodes},j} w_q,s_j,g_j)^2): the first moment of the stationary KFE ((dK/dt=0)). The pointwise KFE MSE plateaus at a floor (~1.4e-7 in the JAX fv runs) below which a 10×-overweight tail, worth +0.2 on (K) through the (a\approx 15) lever arm, hides; the identity is exactly the (K)-relevant functional the pointwise residual barely sees, and its gradient (\propto s_i) moves mass from dissaving regions toward savers (surgical, unlike tail-reweighting the MSE). It is ramped (weight 0 until 20% of n_iter, linear to full at 50%) because switched on from step 0 it acts on the garbage pretrain policy and crushes all mass to the bottom ((K) locks at ~1.2–1.4). This is not folklore: both of this port’s n_iter=1000 smoke runs reproduced the crush exactly. By ~20% of training the policy is sane (HJB ~1e-3) and the identity is purely stabilizing.

with all weights 1 (the ramp factor is the only schedule) and, in fv mode, (L_{\text{flux}}=L_{\text{mass}}=L_{\text{bc}}=0) identically.

6) Gradient Gating: the Continuous Gauss–Seidel¶

There is one joint loss, one backward pass, one Adam step over both nets simultaneously, no alternation. The Gauss–Seidel character comes from exactly two detachments:

• (r,w) are detached right after computation, so (\partial L_{\text{hjb}}/\partial(\text{g-net})=0). The g-net cannot bend aggregates/prices to flatter the HJB residual; the W-net solves “the HJB at the current prices.”

• the policy values ((s_q) on the grid; (s,\partial_a s) at collocation in strong mode) are detached in every g-side term, so (\partial(L_{\text{kfe}}+L_{\text{flux}}+L_{\text{agg}}+L_{\text{bc}})/\partial(\text{W-net})=0). The W-net cannot bend the policy to flatter the KFE; the g-net solves “the KFE under the current policy.”

Each net descends only its own equation given the other’s current values, a continuous Gauss–Seidel iteration on the Aiyagari fixed point, which converges because that fixed point is economically stable. The HJB’s own policy stays live, that is the net’s optimality condition, not a coupling channel.

Audit Test E verified that both cross-gradients are exactly 0.0 at the trained point.

7) Training Pipeline¶

Pre-training (model-implied targets, never FD): (f\to -a/\bar a) (a generic decreasing exponential with prior capital (K_0\approx 4.63)) and (W\to u’(w_0 l + r_0 a)) (the marginal value of the zero-saving policy at the prices implied by the pretrain density’s own ((K_0, L_0=1))). The level matters: the HJB near (a_{\min}) must cancel (\psi’(a_{\min})\approx \kappa a_{lb}=3) against ((\rho - r + \lambda)W\approx 0.3,W), which needs (W(a_{\min},l_1)\approx u’(w_0 l_1)\approx 10). Both nets are warm-started to equal quality so neither overfits the other’s noise. Pretraining is deliberately crude; residual training does all the work.

Collocation sampling (fresh every step): 50% uniform on ([a_{\min},a_{\max}]), 50% uniform on ([a_{\min},6]) (the focus band over the constraint and density peak), and in strong mode an extra 10% over ([a_{\max}-2,a_{\max}]). The quadrature grid is fixed; sampling never moves it.

Optimization: a single Adam over both nets, LR linearly decayed (3\times10^{-4}\to 10^{-6}), the agg-identity ramp over [0.2, 0.5] * n_iter. Every n_iter/20 steps the loss components and the validation gate are logged (the gate is a cheap no-grad forward and consumes no RNG, so monitoring cannot perturb training).

Pretraining is deliberately crude: the post-pretrain gate metric is ~44.5, and residual training does all the work (44.5 (\to) 3.9e-3). Full training is about 9 min on CPU in fv mode, ~13 min in strong (the extra autodiff through the g-net).

8) Diagnostics: FD vs PINN¶

The first dashboard overlays the PINN solution on the FD reference for the marginal value (W), consumption (c), saving (s) (evaluated at the matched FD-equilibrium prices, the gate convention), and the stationary density (g), and shows the joint loss and the running (L^\infty) saving gate. The PINN is never trained on these FD arrays; they are ground truth for comparison only.

9) The Saving (L^\infty) Gate¶

The gate is (\lVert s_{\text{NN}} - s_{\text{FD}}\rVert_\infty) on a dense grid at matched FD-equilibrium prices, with (s_{\text{NN}} = w_{\text{fd}},l + r_{\text{fd}},a - c_{\text{NN}}). Since (c_{\text{NN}} = W^{-1/\gamma}) does not depend on prices, passing requires both the right policy shape and the right equilibrium: if (r_{\text{nn}}\neq r_{\text{fd}}) the W-net solved a different problem and the comparison blows up. The pass threshold is (\epsilon=10^{-2}).

A full production run clears the gate (the JAX recipe reports fv ~3.9e-3, strong ~8.2e-3, with (K_{\text{nn}}\approx 5.07) against (K_{\text{fd}}\approx 5.10)). The smoke run will report FAIL by design: the aggregate-saving ramp engages while the policy is still unconverged, so it does not reach the threshold in 2000 steps.

Resolution matching matters: the FD equilibrium (K) is grid-sensitive (5.27 at 93 points, 5.10 at 400, 5.07 at 800), so the reference deliberately uses n_a = n_quad. Notably the trained (K_{\text{nn}}=5.0706) lands on FD(800) = 5.0730, closer to the continuum than its own 400-point reference: the mesh-free HJB side pulls the equilibrium toward the continuum solution.

Reference Results (CPU, seed 0 unless noted)¶

These are the numbers a full production run reproduces; the smoke run does not reach them.

The strong form’s gate oscillates 4e-3 ↔ 1.2e-2 over the last ~15k steps (briefly failing at intermediate checkpoints, final 8.2e-3), the price of losing the FV operator’s built-in conservation.

Failure Modes (all observed, all understood)¶

• K-crush by the agg identity (both forms): engage the identity on a garbage policy and it drains all mass to the bottom, (K\to) ~1.2–1.4, and the run never recovers. Cure: the ramp. Any smoke run short enough that the ramp hits an unconverged policy will fail the gate, by design.

• Const-flux cheat family (strong only): (g\sim 1/|s|), zero pointwise residual, wrong answer. Cure: the pointwise total-flux identity. Impossible in fv.

• Labor-split drift (strong only): pointwise MSE tolerates a wrong (M_1/M_2); the equilibrium drifts through (L). Cure: the mass-balance identity. Impossible in fv.

• Tail overweight / K bias (both): hides below the KFE MSE floor. Cure: the agg identity plus the trainable tail slope.

• (W\to 0) attractor: approached (never reached, softplus) if the W-net’s level near the constraint is initialized far too low. Cure: the pretrain level plus the shape penalty.

• Strong-form late-phase noise: the gate metric oscillates over the last ~15k steps, the price of losing the FV operator’s built-in conservation.

What Is and Isn’t Ported; Scope¶

Ported (behavior-faithful to the JAX ss/ solver, same calibration, recipe, ramp, targets, and identities): the soft-penalty baseline (fv finite-volume KFE) and the variant (strong mesh-free KFE).

Not ported from ss/: the --hard-constraint variant (atoms of mass at (a_{\min}), the boundary-singularity exponent, the weak/CDF-form KFE) and the EMA/Polyak stabilizers. The deterministic steady state collapses the EMINN master equation (the (z)-derivative and (L_g) distribution-coupling terms all vanish), which is exactly what makes this self-contained; aggregate shocks reintroduce them and call for the surrogate-density approach of §8.7 (EMINNs).

Economic Discussion (for Teaching)¶

What should students see?¶

1. The FD benchmark is the trusted reference. Standard upwind discretization with outer bisection on (r); the HJB and KFE residuals are tiny and the GE fixed point converges quickly.

2. The PINN recovers the same economic objects without an outer price loop. Market clearing holds by construction because prices are the firm’s marginal products at the density net’s own aggregates. The distribution shifting is the price adjustment.

3. Structure is what makes the mesh-free PINN work. Functional-form pretraining (level at the constraint, zero-saving start), the trainable tail slope, normalization by construction, and the ramped aggregate-saving identity together pin the parts of the solution a raw pointwise residual cannot see.

Economic interpretation¶

• The soft borrowing penalty makes households endogenously avoid low wealth, so mass concentrates above (a_{lb}) without a hard constraint ever binding.

• The employed state (l_2) saves more and populates the right tail more heavily; the tail carries a long lever arm into aggregate (K).

• The equilibrium interest rate reflects the intersection of precautionary-savings supply and firm capital demand, recovered here by the joint net rather than by bisection.

Suggested Classroom Use¶

1. Run sections in order once at RUN_MODE="smoke" to see the full pipeline quickly (the gate will FAIL by design), then at "teaching" or "production" to clear it.

2. Flip KFE_FORM from "fv" to "strong" and watch which loss terms switch on (flux, mass, bc) and how the late-phase gate gets noisier.

3. Ask students to change one parameter (the switching rate (\lambda), the soft-penalty (\kappa), or the agg-identity ramp window) and compare how convergence, policies, and the density tail move.

Notebook path: lectures/lecture_13_continuous_time_ha_numerics/code/lecture_13_08_Aiyagari_Continuous_Time_FD_and_PINN_PyTorch.ipynb

Takeaway¶

On the deterministic steady-state Aiyagari calibration, two small nets, the marginal value (W=\partial_a V) and a normalized log-density (g), reproduce the FD reference for the policy, the distribution, and the equilibrium aggregates, with market clearing built in by construction and no FD solve anywhere in training. The continuous Gauss–Seidel gating (two detachments) decouples the policy and distribution problems, and a handful of exact integral identities pin the (K)-relevant content the pointwise residuals miss. For aggregate-shock extensions, the same recipe carries over by replacing (g) with a finite-dimensional surrogate (Section 8.7, EMINNs; Gu et al., 2024).