Life-cycle OLG with borrowing constraints — the two-stage-loss model¶
The fourth examples/ stop, and the one that motivated a new capability in the
framework. This is Day 2, Exercise 4 of the Geneva 2026 Deep Learning for
Economics & Finance course: a 6-generation overlapping-generations economy
where households live $H = 6$ deterministic periods (one period $\approx$ 10
years, ages 20–80), save in capital subject to a borrowing constraint
$k^h_t \ge 0$, and supply age-dependent labor (less in the last two periods —
retirement).
Each working cohort $h \in \{0, \dots, H-2\}$ chooses a saving rate; the last cohort consumes everything. The optimality condition is an intertemporal Euler that holds with equality only when the borrowing constraint is slack:
$$\frac{1}{c^h_t} \;\ge\; \beta\,\mathbb{E}_t\!\left[\frac{1-\delta+r_{t+1}}{c^{h+1}_{t+1}}\right], \qquad k^{h+1}_{t+1}\ge 0,\qquad \text{complementary slackness.}$$
We encode the complementarity with a Fischer-Burmeister residual, exactly as the course notebook does.
Why this model needed a new loss¶
The FB nonlinearity wraps an expectation. And $\mathbb{E}[f^{FB}(\cdot)] \ne f^{FB}(\mathbb{E}[\cdot])$: averaging the residual over shocks and then squaring (the standard DEQN path) puts the nonlinearity in the wrong place and leaves a bias floor. So this model is trained with the framework's two-stage loss:
inside_fnreturns the shock-dependent continuation terms $(1-\delta+r')/c'^{\,j}$, which the loss averages to $\mathbb{E}[\cdot]$;combine_fnapplies the Fischer-Burmeister after the expectation.
The standard $(\mathbb{E}[r])^2$ path is the special case combine = identity.
This is the architectural prize from the port: occasionally-binding constraints
under uncertainty, solved MC-correctly.
Outline
- 1 — Inspect the model
- 2 — Train (two-stage loss, no closed-form steady state)
- 3 — Loss curve
- 4 — Reproduce the Exercise-4 diagnostic panels
- 5 — Ergodic accuracy + the borrowing-constraint corner
- 6 — Summary
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
import numpy as np
from deqn_jax.config import TrainConfig
from deqn_jax.models.olg_lifecycle import MODEL
from deqn_jax.models.olg_lifecycle.equations import _cohort_block, combine_fn, inside_fn
from deqn_jax.models.olg_lifecycle.variables import CONSTANTS, H, L_CYCLE
from deqn_jax.plots import plot_loss_curve
from deqn_jax.training.trainer import train_from_config
WARNING:2026-06-29 13:49:13,873:jax._src.xla_bridge:876: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
1. Inspect the model¶
State $\mathbf{x}_t = (Z_t,\, k^0_t, \dots, k^5_t)$ — TFP plus the capital held by each age group (7 states). Policy is the saving rate out of cash-at-hand for the five cohorts that still save (sigmoid-bounded to $(0,1)$, so consumption and saved capital are both positive by construction). Five Euler conditions, one per saving cohort.
There is no closed-form steady state; the cross-sectional capital distribution and TFP are solved jointly over the ergodic distribution, seeded from a random $\exp(\mathcal{U}(0,1))$ init — just like the course notebook.
print(f"states : {MODEL.n_states} {MODEL.state_names}")
print(f"policies : {MODEL.n_policies} {MODEL.policy_names}")
print(f"equations: {len(MODEL.equation_names)} {MODEL.equation_names}")
print(f"two-stage hooks present: inside_fn={MODEL.inside_fn is not None}, "
f"combine_fn={MODEL.combine_fn is not None}")
print(f"steady_state_fn: {MODEL.steady_state_fn} (None -> trained from random init)")
print()
print(f"alpha={CONSTANTS['alpha']}, beta={CONSTANTS['beta']:.4f} (=0.99^10), "
f"delta={CONSTANTS['delta']}, rho_z={CONSTANTS['rho_z']:.4f}, sigma_z={CONSTANTS['sigma_z']}")
print(f"age-dependent labor l_cycle = {L_CYCLE} (retirement in the last two periods)")
states : 7 ('Z', 'k0', 'k1', 'k2', 'k3', 'k4', 'k5')
policies : 5 ('s0', 's1', 's2', 's3', 's4')
equations: 5 ('euler_0', 'euler_1', 'euler_2', 'euler_3', 'euler_4')
two-stage hooks present: inside_fn=True, combine_fn=True
steady_state_fn: None (None -> trained from random init)
alpha=0.36, beta=0.9044 (=0.99^10), delta=0.8, rho_z=0.3487, sigma_z=0.1
age-dependent labor l_cycle = (1.0, 1.8, 2.3, 2.5, 1.6, 1.25) (retirement in the last two periods)
2. Train¶
Recipe from configs/olg_lifecycle.yaml: a [70, 70] ReLU MLP with sigmoid
output (Simon's 10 * n_input width), Adam at 3e-4, MC expectation with
antithetic shocks. The two-stage path is selected automatically because the
model declares inside_fn + combine_fn. A few thousand episodes converge in
seconds on CPU.
cfg = TrainConfig.from_yaml("../configs/olg_lifecycle.yaml")
params, history = train_from_config(cfg)
============================================================ DEQN-JAX Training ============================================================ Model: olg_lifecycle Optimizer: adam (lr=3e-04) Precision: float32 Network: 7 → 70 → 70 → 5 Parameters: 5,885 Batch size: 128 Expectations: 8 MC samples Warm start: no ============================================================ Schedule: 1 rollout (24×128=3072 states) → 1 epoch(s) × 1 minibatch(es) of 128 = 1 grad updates/cycle (3000 total over 3000 cycles)
[ 200/3000] loss=2.25e-03 | grad=1.51e-02 | 58 ep/s
euler_0 9.00e-03 euler_2 1.11e-04 euler_4 9.28e-04
euler_1 1.18e-03 euler_3 3.85e-05
[ 400/3000] loss=4.07e-04 | grad=4.95e-03 | 86 ep/s
euler_0 1.60e-03 euler_2 6.90e-05 euler_4 4.97e-06
euler_1 3.32e-04 euler_3 3.15e-05
[ 600/3000] loss=1.03e-04 | grad=2.02e-03 | 102 ep/s
euler_0 4.35e-04 euler_2 8.28e-06 euler_4 9.69e-06
euler_1 5.36e-05 euler_3 9.28e-06
[ 800/3000] loss=3.37e-05 | grad=6.28e-04 | 113 ep/s
euler_0 1.46e-04 euler_2 6.63e-06 euler_4 5.17e-06
euler_1 8.42e-06 euler_3 2.75e-06
[1000/3000] loss=1.48e-05 | grad=8.06e-04 | 121 ep/s
euler_0 6.53e-05 euler_2 4.57e-06 euler_4 1.74e-06
euler_1 1.31e-06 euler_3 9.45e-07
[1200/3000] loss=1.01e-05 | grad=8.36e-04 | 126 ep/s
euler_0 4.15e-05 euler_2 3.76e-06 euler_4 1.48e-06
euler_1 2.52e-06 euler_3 1.20e-06
[1400/3000] loss=6.40e-06 | grad=1.18e-03 | 131 ep/s
euler_0 2.41e-05 euler_2 3.52e-06 euler_4 1.29e-06
euler_1 1.75e-06 euler_3 1.32e-06
[1600/3000] loss=4.59e-06 | grad=2.23e-03 | 134 ep/s
euler_0 1.56e-05 euler_2 2.41e-06 euler_4 1.64e-06
euler_1 2.44e-06 euler_3 8.73e-07
[1800/3000] loss=4.21e-06 | grad=8.88e-04 | 137 ep/s
euler_0 1.33e-05 euler_2 2.52e-06 euler_4 1.28e-06
euler_1 2.76e-06 euler_3 1.22e-06
[2000/3000] loss=2.91e-06 | grad=6.03e-04 | 139 ep/s
euler_0 9.13e-06 euler_2 2.72e-06 euler_4 8.66e-07
euler_1 1.01e-06 euler_3 8.44e-07
[2200/3000] loss=2.59e-06 | grad=1.01e-03 | 141 ep/s
euler_0 7.81e-06 euler_2 2.15e-06 euler_4 7.46e-07
euler_1 1.63e-06 euler_3 6.07e-07
[2400/3000] loss=2.31e-06 | grad=1.70e-03 | 143 ep/s
euler_0 6.01e-06 euler_2 2.49e-06 euler_4 9.75e-07
euler_1 1.26e-06 euler_3 8.18e-07
[2600/3000] loss=2.26e-06 | grad=1.58e-03 | 145 ep/s
euler_0 5.07e-06 euler_2 2.31e-06 euler_4 6.76e-07
euler_1 2.04e-06 euler_3 1.20e-06
[2800/3000] loss=2.00e-06 | grad=2.17e-03 | 146 ep/s
euler_0 4.13e-06 euler_2 3.20e-06 euler_4 8.87e-07
euler_1 1.10e-06 euler_3 6.61e-07
[3000/3000] loss=1.77e-06 | grad=8.67e-04 | 147 ep/s
euler_0 3.98e-06 euler_2 2.29e-06 euler_4 8.92e-07
euler_1 1.07e-06 euler_3 6.14e-07
============================================================
Training complete in 20.4s (147 ep/s)
Final loss: 1.77e-06
euler_0 3.98e-06 euler_2 2.29e-06 euler_4 8.92e-07
euler_1 1.07e-06 euler_3 6.14e-07
============================================================
3. Loss curve¶
The loss is the mean squared Fischer-Burmeister residual across the five Euler conditions, with the expectation taken inside the FB.
fig, ax = plt.subplots(figsize=(7.5, 4))
plot_loss_curve(history, ax=ax, log_y=True)
ax.set_title("olg_lifecycle — two-stage FB-Euler loss (log scale)")
plt.show()
print(f"loss: {float(history['loss'][0]):.3e} -> {float(min(history['loss'][-50:])):.3e}")
loss: 1.226e-01 -> 1.367e-06
4. Reproduce the Exercise-4 diagnostic panels¶
The course notebook reads a trained policy off by plotting, by age group, the
cross-sectional capital, consumption, and cash-at-hand profiles; the return
distribution; the signed relative Euler error errREE; the Euler LHS vs
RHS check; and the saving policy. We reproduce that exact panel set on the
deqn-jax solution, using Simon's own quantities.
First, simulate the ergodic distribution under the trained policy, then compute
the per-cohort quantities with the model's own _cohort_block. The relative
Euler error is the MC-correct two-stage residual — average the continuation
terms over shocks, then apply Fischer-Burmeister (reusing inside_fn +
combine_fn):
$$\text{err}_{REE}^h = f^{FB}\!\Big(\underbrace{\tfrac{1}{c^h\,\beta\,\mathbb{E}[(1-\delta+r')/c'^{\,h+1}]} - 1}_{a:\ \text{Euler gap}},\ \underbrace{\tfrac{\text{sav}^h}{c^h}}_{b:\ \text{normalized slack}}\Big),$$
with the Euler LHS $= 1/c^h$ and RHS $= \beta\,\mathbb{E}[(1-\delta+r')/c'^{\,h+1}]$.
This is the exact errREE / LHS / RHS triple from the Day 2 Exercise 4 cost
function.
Reading the LHS/RHS panel. A visible LHS $>$ RHS gap at the youngest cohort ($h=0$) is expected, not a solver error: the borrowing constraint binds there ($k'\approx 0$), so the Euler inequality holds with slack rather than equality — faithful to Simon's reference. LHS tracks RHS for the older, unconstrained cohorts.
def simulate(params, key, n_traj=1024, T=80, burn=30):
"""Roll the trained policy forward; collect post-burn-in ergodic states."""
s = MODEL.init_state_fn(key, n_traj, CONSTANTS)
visited = []
for t in range(T):
key, sk = jr.split(key)
eps = jr.normal(sk, (n_traj, 1))
if t >= burn:
visited.append(s)
s = MODEL.step_fn(s, params(s), eps, CONSTANTS)
return jnp.concatenate(visited, axis=0)
def ergodic_diagnostics(params, X, key, n_shocks=64):
"""MC-correct Exercise-4 diagnostics: average the H continuation terms over
shocks, then form the SIGNED relative Euler error errREE = fb(a, b) plus the
Euler LHS / RHS, exactly as the course cost function does.
a = 1/(c^h * beta * E[(1-delta+r')/c'^{h+1}]) - 1 (Euler gap)
b = sav^h / c^h (normalized slack)
LHS = 1/c^h RHS = beta * E[(1-delta+r')/c'^{h+1}]
"""
insides = []
for _ in range(n_shocks):
key, sk = jr.split(key)
eps = jr.normal(sk, (X.shape[0], 1))
ns = MODEL.step_fn(X, params(X), eps, CONSTANTS)
ins = inside_fn(X, params(X), ns, params(ns), CONSTANTS)
insides.append(jnp.stack([ins[f"inside_{j}"] for j in range(H)], axis=1))
E = jnp.mean(jnp.stack(insides), axis=0) # [n, H]
Edict = {f"inside_{j}": E[:, j] for j in range(H)}
res = combine_fn(X, params(X), Edict, CONSTANTS)
errREE = jnp.stack([res[f"euler_{h}"] for h in range(H - 1)], axis=1) # signed
blk = _cohort_block(X[:, :1], X[:, 1 : 1 + H], params(X), CONSTANTS)
LHS = 1.0 / blk["c"][:, : H - 1] # 1/c^h
RHS = CONSTANTS["beta"] * E[:, 1:H] # beta E[(1-d+r')/c'^{h+1}]
return np.asarray(errREE), np.asarray(LHS), np.asarray(RHS)
X = simulate(params, jr.PRNGKey(1))
Z, k = X[:, :1], X[:, 1 : 1 + H]
blk = _cohort_block(Z, k, params(X), CONSTANTS)
c, sav, cah, r = (np.asarray(blk[n]) for n in ("c", "sav", "cah", "r"))
k_np = np.asarray(k)
errREE, LHS, RHS = ergodic_diagnostics(params, X, jr.PRNGKey(2))
print(f"ergodic sample: {X.shape[0]} states")
ergodic sample: 51200 states
def band(ax, data, title, ylabel, xl="age group"):
"""mean/min/max life-cycle profile on its natural (level) scale."""
a = np.arange(data.shape[1])
ax.plot(a, data.mean(0), "o-", label="mean")
ax.plot(a, data.min(0), "--", alpha=0.4, label="min")
ax.plot(a, data.max(0), "--", alpha=0.4, label="max")
ax.set_title(title); ax.set_xlabel(xl); ax.set_ylabel(ylabel); ax.legend(fontsize=8)
fig, ax = plt.subplots(2, 4, figsize=(18, 8))
# life-cycle VALUE profiles -- levels, shown on their natural scale
band(ax[0, 0], k_np, "capital by age group", "capital held")
band(ax[0, 1], c, "consumption by age group", "consumption")
band(ax[0, 2], cah, "cash-at-hand by age group", "cash-at-hand")
# return distribution
ax[0, 3].hist(r[:, 0], bins=40, color="steelblue")
ax[0, 3].set_title("return distribution")
ax[0, 3].set_xlabel("return r_t"); ax[0, 3].set_ylabel("count")
# SIGNED relative Euler error: an error centered on zero -> keep it near zero
# with a zero reference line (do NOT autoscale a log|.| into noise).
a = np.arange(errREE.shape[1])
ax[1, 0].axhline(0.0, color="grey", lw=0.8, zorder=0)
ax[1, 0].plot(a, errREE.mean(0), "o-", label="mean")
ax[1, 0].plot(a, errREE.min(0), "--", alpha=0.4, label="min")
ax[1, 0].plot(a, errREE.max(0), "--", alpha=0.4, label="max")
ax[1, 0].set_title("relative Euler error (signed)")
ax[1, 0].set_xlabel("age group / cohort h")
ax[1, 0].set_ylabel("relative Euler error"); ax[1, 0].legend(fontsize=8)
# Euler equation LHS vs RHS (levels, natural scale): should overlap where slack
ax[1, 1].plot(a, LHS.mean(0), color="k", label="LHS: 1/c")
ax[1, 1].plot(a, LHS.min(0), color="k", ls="--", alpha=0.3)
ax[1, 1].plot(a, LHS.max(0), color="k", ls="--", alpha=0.3)
ax[1, 1].plot(a, RHS.mean(0), color="r", label="RHS: beta E[(1-d+r')/c']")
ax[1, 1].plot(a, RHS.min(0), color="r", ls="--", alpha=0.3)
ax[1, 1].plot(a, RHS.max(0), color="r", ls="--", alpha=0.3)
ax[1, 1].set_title("Euler equation: LHS vs RHS")
ax[1, 1].set_xlabel("age group / cohort h"); ax[1, 1].legend(fontsize=8)
# saving policy: next-period capital carried by cohort h vs capital held now
for h in range(H - 1):
ax[1, 2].scatter(k_np[:, h], sav[:, h], s=2, alpha=0.3, label=f"h={h}")
ax[1, 2].set_title("saving policy by cohort")
ax[1, 2].set_xlabel("capital held (current)")
ax[1, 2].set_ylabel("capital saved (next period)"); ax[1, 2].legend(fontsize=8)
ax[1, 3].axis("off")
fig.tight_layout()
plt.show()
The panels match the Exercise-4 reference: a hump-shaped capital profile ($k^0=0$ for the newborn, rising through working life, peaking near retirement, then dissaved), a smooth rising consumption profile, a saving policy that clusters by cohort, a signed relative Euler error centered on zero (drawn on its own small scale, with the zero line — that small magnitude is the accuracy, not noise), and a Euler LHS that tracks the RHS wherever the borrowing constraint is slack. This is the standard life-cycle picture the course notebook produces — recovered here by the JAX port under the two-stage loss.
5. Ergodic accuracy certificate + the borrowing-constraint corner¶
We summarize accuracy with quantiles of the dimensionless residual $|\text{err}_{REE}|$ (the Fischer-Burmeister residual is already unit-free) — the median, 90th, and 99th percentiles over the ergodic sample. Quantiles are honest about the tail: a low median with a fat p99 would flag a few hard states rather than hiding behind a mean. The youngest cohort is hardest (it sits closest to the borrowing constraint, where the FB kink lives); older cohorts are tighter.
abs_err = np.abs(errREE) # dimensionless FB residual -> a clean accuracy certificate
print("Ergodic accuracy certificate -- |relative Euler error| (dimensionless FB residual)")
print(f" median : {np.median(abs_err):.2e}")
print(f" p90 : {np.quantile(abs_err, 0.90):.2e}")
print(f" p99 : {np.quantile(abs_err, 0.99):.2e}")
print()
print("per-cohort median |errREE| (h=0..4):", np.round(np.median(abs_err, 0), 4))
print()
print("mean capital by age :", np.round(k_np.mean(0), 3))
print("mean consumption :", np.round(c.mean(0), 3))
near_bind = (sav[:, : H - 1] < 1e-2).mean(0)
print("frac near borrowing constraint (k' < 1e-2), cohort 0..4:", np.round(near_bind, 3))
Ergodic accuracy certificate -- |relative Euler error| (dimensionless FB residual) median : 1.08e-03 p90 : 2.43e-03 p99 : 4.11e-03 per-cohort median |errREE| (h=0..4): [0.0017 0.001 0.0012 0.0006 0.0008] mean capital by age : [0. 0.001 0.137 0.377 0.632 0.508] mean consumption : [0.331 0.461 0.575 0.717 0.894 1.115] frac near borrowing constraint (k' < 1e-2), cohort 0..4: [1. 0. 0. 0. 0.]
On comparing to the reference. This is a faithful reproduction of the Geneva Day 2 Ex 4 method and diagnostics: same equilibrium conditions, same FB complementarity, same decade calibration, same diagnostic panels. On a like-for- like ergodic $|\text{err}_{REE}|$ the JAX port currently trails the original TensorFlow reference by roughly a part in ten on the log scale (a known gap tracked in the project notes, not a ship blocker) — the shapes and the economics agree; the last fraction of a decimal of accuracy is the open item.
The model exposes nothing exotic to get here: standard MLP, sigmoid-bounded
saving rates for $0 < s < 1$ (hence $c>0$ and $k'\ge 0$ by construction), and the
two-stage inside_fn/combine_fn hooks that put the Fischer-Burmeister
after the expectation.
Summary¶
- Model: 6-generation life-cycle OLG with borrowing constraints (Geneva Day 2 Ex 4). 7 states, 5 saving-rate policies, 5 Fischer-Burmeister Euler conditions, no closed-form steady state.
- Capability it unlocked: the two-stage expectation-inside-residual loss.
Because the FB wraps an expectation, $\mathbb{E}[f^{FB}] \ne f^{FB}(\mathbb{E})$;
inside_fnaverages the continuation terms andcombine_fnapplies the FB afterward. The standard $(\mathbb{E}[r])^2$ path is recovered ascombine = identity. - Result: textbook life-cycle behaviour (capital hump, consumption smoothing) with ergodic per-cohort Euler errors in the $10^{-2}$–$10^{-3}$ range, reproducing the course exercise's diagnostic panels.
- No special tooling: plain MLP + sigmoid bounds + the two model hooks; the
trainer, loss, and
deqn_jax.plotssuite are all stock.
The companion model with a closed-form oracle is examples/olg_analytic_6.ipynb
(Krueger-Kubler 2004); the constrained-labor sibling is
bm_labor_constrained (Day 2 Ex 3).