Two-country IRBC with irreversible investment — multipliers as policies¶
Day 3 of the Geneva 2026 Deep Learning for Economics & Finance course: a two-country International Real Business Cycle model solved as a planner problem. Three method lessons beyond the Brock-Mirman family:
- KKT multipliers are network outputs. Investment in each country is
irreversible, $i_j \ge 0$. In
bm_labor_constrainedthe complementarity wedge was an analytic expression; here the multipliers $\mu_j$ (and the aggregate shadow price $\lambda$) are learned policies, paired with their constraints through Fischer-Burmeister residuals $f^{FB}(\mu_j, i_j) = \mu_j + i_j - \sqrt{\mu_j^2 + i_j^2}$. - Consumption is pinned, not learned. Complete markets + Pareto weights $\tau_j$ give marginal-utility equalization $c_j = (\lambda/\tau_j)^{-1/\gamma_j}$ — the network never outputs consumption; the FOC structure does.
- Quadrature expectations over 3 shocks. Two country-specific TFP innovations plus one aggregate, integrated with a $3^3 = 27$-node Gauss-Hermite tensor grid instead of Monte Carlo.
The economics showcase is risk sharing under heterogeneous risk aversion: country 0 has $\gamma_0 = 0.25$, country 1 has $\gamma_1 = 1.0$. Efficient sharing loads aggregate risk onto the less risk-averse country, so $c_0$ is 4× as volatile (in logs) as $c_1$. Note what kind of check this is: because consumption is pinned from $\lambda$, the 4:1 ratio holds exactly by construction — it validates that the FOC structure is wired correctly, not that training converged. Training quality is certified separately in §9.
Outline
- 1 — Inspect the model
- 2 — Train (quadrature expectations)
- 3 — Training loss
- 4 — Risk sharing: who absorbs the volatility?
- 5 — Policy functions
- 6 — Irreversibility: the complementarity scatter
- 7 — Euler-equation accuracy (relative Euler error)
- 8 — Ergodic distribution
- 9 — Ergodic accuracy certificate (printed table)
- 10 — Summary
Sections 5, 7, and 8 mirror the result figures of the Geneva 2026 Day 3
IRBC exercise (policy scatter, LHS-vs-RHS Euler check, signed relative
Euler error, and the ergodic state cloud), reproduced with deqn_jax
quantities.
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
from deqn_jax.models import load_model
model = load_model("irbc")
print("states: ", model.state_names)
print("policies:", model.policy_names, " (consumption pinned by lambda)")
print("equations:", model.equation_names)
print("shocks: ", model.shock_names)
{k: v for k, v in model.constants.items()}
WARNING:2026-06-29 13:48:26,227: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.
states: ('k_0', 'k_1', 'z_0', 'z_1')
policies: ('k_0_next', 'k_1_next', 'lam', 'mu_0', 'mu_1') (consumption pinned by lambda)
equations: ('euler_0', 'euler_1', 'arc', 'fb_0', 'fb_1')
shocks: ('eps_0', 'eps_1', 'eps_agg')
{'beta': 0.99,
'delta': 0.01,
'zeta': 0.36,
'kappa': 0.5,
'rho_z': 0.95,
'sigma_eps': 0.01,
'A_tfp': 0.055836,
'gamma_0': 0.25,
'gamma_1': 1.0,
'tau_0': 0.5,
'tau_1': 0.5,
'fb_eps': 1e-08}
2 — Train¶
Recipe from configs/irbc.yaml: a BK-anchored residual network
(linear_plus_mlp: the policy is the stable linearized rule plus a learned
correction) trained with the composite loss, whose anchor/Jacobian terms pin
the policy's steady-state tangent to the Blanchard-Kahn-stable linearization.
Expectations by Gauss-Hermite quadrature (27 nodes over 3 shocks) — exact
for smooth integrands, no Monte Carlo noise or bias.
Why the anchoring matters here: a plain MLP trained on the same residuals reaches a small training loss while landing on a closed-loop unstable policy (spectral radius 1.23 at SS — simulations drift out of the training domain and market clearing fails on the states the economy actually visits). With the BK anchor the same training certifies at ρ = 0.98. Residual accuracy at visited states does not imply stable dynamics; the anchor supplies the equilibrium-selection discipline the loss alone lacks.
from deqn_jax.config import load_config
from deqn_jax.training.trainer import train_from_config
config = load_config("../configs/irbc.yaml")
config = config.model_copy(update={"verbose": False})
params, history = train_from_config(config)
print(f"final loss: {history['loss'][-1]:.3e}")
final loss: 8.647e-07
3 — Loss curve¶
The aggregate-resource-constraint residual (arc) starts much larger than the
Eulers — consumption pinned from an untrained $\lambda$ over/undershoots
aggregate output badly at first — and the mean-over-equations aggregation keeps
it from drowning the others.
fig, ax = plt.subplots(figsize=(7, 4))
ax.semilogy(history["loss"], lw=0.8)
ax.set_xlabel("episode"); ax.set_ylabel("training loss")
ax.set_title("IRBC training (quadrature expectations)")
fig.tight_layout()
4 — Risk sharing: who absorbs the volatility?¶
Simulate the trained economy and compare consumption volatility across countries. Efficient risk sharing with $\gamma_0 = 0.25 < \gamma_1 = 1.0$ predicts $\mathrm{std}(\log c_0) \approx 4 \times \mathrm{std}(\log c_1)$: marginal-utility equalization $\gamma_0\,d\log c_0 = \gamma_1\,d\log c_1$ along every shock path.
def simulate(net, T=4000, burn=500, n_paths=16, seed=0):
ss_state, _ = model.steady_state_fn(model.constants)
s = jnp.tile(jnp.asarray(ss_state)[None, :], (n_paths, 1))
shocks = jax.random.normal(jax.random.PRNGKey(seed), (T, n_paths, model.n_shocks))
def step(s, e):
nxt = model.step_fn(s, net(s), e, model.constants)
return nxt, nxt
_, traj = jax.lax.scan(step, s, shocks)
return traj[burn:]
traj = simulate(params)
states_flat = traj.reshape(-1, model.n_states)
defs = model.definitions_fn(states_flat, params(states_flat), model.constants)
lc0, lc1 = jnp.log(defs["c_0"]), jnp.log(defs["c_1"])
ratio = float(jnp.std(lc0) / jnp.std(lc1))
print(f"std(log c_0)/std(log c_1) = {ratio:.2f} (efficient sharing: gamma_1/gamma_0 = 4.0)")
fig, axes = plt.subplots(1, 2, figsize=(11, 4))
t = np.arange(600)
axes[0].plot(t, np.asarray(jnp.log(defs["c_0"][:600])), lw=0.7, label="log c_0 (gamma=0.25)")
axes[0].plot(t, np.asarray(jnp.log(defs["c_1"][:600])), lw=0.7, label="log c_1 (gamma=1.0)")
axes[0].set_title("consumption paths: the brave country absorbs the risk")
axes[0].legend()
axes[1].hist(np.asarray(lc0 - lc0.mean()), bins=80, alpha=0.6, label="log c_0 (demeaned)")
axes[1].hist(np.asarray(lc1 - lc1.mean()), bins=80, alpha=0.6, label="log c_1 (demeaned)")
axes[1].set_title(f"volatility ratio = {ratio:.2f} (theory: 4.0)")
axes[1].legend()
fig.tight_layout()
std(log c_0)/std(log c_1) = 4.00 (efficient sharing: gamma_1/gamma_0 = 4.0)
5 — Policy functions¶
The result figures of the Geneva 2026 Day 3 IRBC exercise, reproduced on our ergodic set. Four scatters over the visited states:
- capital policy $k_{t+1}$ vs capital $k_t$, with the $k_{t+1}=k_t$ diagonal — points sit just above it (the economy accumulates a little), consistent with a stable rule that pulls toward the steady state;
- aggregate shadow price $\lambda$ vs capital — the single price that pins both countries' consumption under complete markets;
- consumption policy $c_j$ vs capital — never a network output; it is read off $\lambda$ by $c_j=(\lambda/\tau_j)^{-1/\gamma_j}$. Plotted on a log $y$-axis: the heterogeneous risk aversion makes $c_0\sim10^{-3}$ and $c_1\sim10^{-1}$ differ by ~100×, so a shared linear axis would flatten country 0 to a line at zero — the log axis keeps both legible;
- investment policy $i_j$ vs capital, with the $i=0$ floor drawn — every point stays at or above it (irreversibility respected).
Apart from the (deliberately) log consumption axis these are values, so the axes show the value — no zoom into a flat line.
N = 2 # countries
rng = np.random.default_rng(1)
m = states_flat.shape[0]
sub = rng.choice(m, min(4000, m), replace=False)
S = np.asarray(states_flat[sub])
P = np.asarray(params(states_flat[sub]))
D = model.definitions_fn(states_flat[sub], params(states_flat[sub]), model.constants)
fig, axes = plt.subplots(2, 2, figsize=(12, 9))
# capital policy: k' vs k, with the 45-degree (k_{t+1}=k_t) line
for j in range(N):
axes[0, 0].scatter(S[:, j], P[:, j], s=3, alpha=0.3, label=f"country {j}")
lo = float(S[:, :N].min()); hi = float(S[:, :N].max())
axes[0, 0].plot([lo, hi], [lo, hi], "k--", lw=1, alpha=0.6, label="$k_{t+1}=k_t$")
axes[0, 0].set_xlabel("capital $k_t$"); axes[0, 0].set_ylabel("next-period capital $k_{t+1}$")
axes[0, 0].set_title("capital policy"); axes[0, 0].legend()
# aggregate shadow price lambda vs capital (country 0)
axes[0, 1].scatter(S[:, 0], P[:, N], s=3, alpha=0.3)
axes[0, 1].set_xlabel("capital $k_0$"); axes[0, 1].set_ylabel(r"shadow price $\lambda$")
axes[0, 1].set_title("aggregate shadow price")
# consumption policy (pinned by lambda) vs capital
for j in range(N):
axes[1, 0].scatter(S[:, j], np.asarray(D[f"c_{j}"]), s=3, alpha=0.3, label=f"country {j}")
axes[1, 0].set_yscale("log") # c_0 ~ 1e-3, c_1 ~ 1e-1: shared linear axis would flatten country 0
axes[1, 0].set_xlabel("capital $k_t$"); axes[1, 0].set_ylabel("consumption $c$ (log scale)")
axes[1, 0].set_title(r"consumption policy (pinned by $\lambda$)"); axes[1, 0].legend()
# investment policy vs capital, with the irreversibility floor i=0
for j in range(N):
axes[1, 1].scatter(S[:, j], np.asarray(D[f"i_{j}"]), s=3, alpha=0.3, label=f"country {j}")
axes[1, 1].axhline(0, color="k", lw=0.8, alpha=0.6, label="$i=0$ floor")
axes[1, 1].set_xlabel("capital $k_t$"); axes[1, 1].set_ylabel("investment $i$")
axes[1, 1].set_title(r"investment policy ($i \geq 0$)"); axes[1, 1].legend()
fig.suptitle("Policy functions on the ergodic set")
fig.tight_layout()
6 — Irreversibility: the complementarity scatter¶
The defining picture of a learned KKT system: on the ergodic set each $(i_j, \mu_j)$ pair should lie on the complementarity L — either investment is interior ($i_j > 0,\ \mu_j \approx 0$) or the constraint binds ($i_j \approx 0,\ \mu_j > 0$). At this calibration ($\sigma = 0.01$, quadratic adjustment costs) irreversibility never binds on the ergodic set: the per-panel binding fractions printed below are $0.0\%$, so the cloud is the horizontal leg only — there is no vertical-leg mass.
The fit to that leg is good but not exact. A thin band of residual
complementarity slack sits at interior investment ($i_0 \in [0.012, 0.020]$)
with $\mu_0$ up to $\sim 0.011$ — points genuinely off the L. This is the
learned Fischer–Burmeister residual, whose magnitude is certified in §9
(fb_0, fb_1 quantiles); it is small but nonzero, not the perfect
complementarity an exact KKT solution would show. What the figure does
validate is that the corner is respected ($i_j \ge 0$ everywhere, no
negative-investment mass), which the unanchored recipe failed. Crisis
calibrations with larger shocks move mass onto the vertical leg with no
change to the loss.
The two country panels share a $y$-axis so the slack is read on one scale: country 1's multiplier is genuinely an order of magnitude smaller, not merely drawn on a tighter axis.
fig, axes = plt.subplots(1, 2, figsize=(11, 4), sharey=True)
for j, ax in enumerate(axes):
i_j = np.asarray(defs[f"i_{j}"])
mu_j = np.asarray(defs[f"mu_{j}"])
ax.scatter(i_j, mu_j, s=2, alpha=0.2)
ax.axhline(0, color="k", lw=0.5); ax.axvline(0, color="k", lw=0.5)
bind = float((i_j < 1e-4).mean())
ax.set_xlabel(f"investment $i$ (country {j})")
ax.set_ylabel(rf"irreversibility multiplier $\mu$ (country {j})")
ax.set_title(f"country {j}: binds on {bind:.1%} of ergodic mass")
# Honest: mass sits on the horizontal leg; the off-leg spread is FB slack, not exact zero.
fig.suptitle("complementarity scatter: i>=0 respected; off-leg band is residual FB slack")
fig.tight_layout()
7 — Euler-equation accuracy (relative Euler error)¶
The diagnostic the Geneva course reports for every DEQN solution: the signed relative Euler error
$$\mathrm{errREE}_j \;=\; \frac{\mathrm{RHS}_j}{\mathrm{LHS}_j} - 1, \qquad \mathrm{LHS}_j = \lambda\,(1 + \Gamma'_{k'}),\quad \mathrm{RHS}_j = \mu_j + \beta\,\mathbb{E}\!\left[\lambda'\,M'_j - (1-\delta)\,\mu_j'\right],$$
with the expectation taken by the same 27-node Gauss-Hermite quadrature used in training (no Monte Carlo noise). This is signed and plotted on a linear axis: it lives near zero, so the panels are drawn near zero with a horizontal zero reference — that small scale is the result, not a zoom into noise. (Plotting $\log_{10}|\mathrm{residual}|$ instead would spike to $-\infty$ at every sign change and read as jagged; the certificate in §9 uses logs deliberately, on magnitudes.)
The companion LHS vs RHS scatter is Simon's headline check: every point on the $45^\circ$ line means the capital FOC holds state-by-state.
from deqn_jax.training.loss import gauss_hermite_nd, compute_residuals
# Expectation of the (signed) Euler residual by 27-node quadrature.
nodes, weights = gauss_hermite_nd(3, model.n_shocks)
rng_e = np.random.default_rng(0)
idx = rng_e.choice(states_flat.shape[0], min(4096, states_flat.shape[0]), replace=False)
eul_sample = states_flat[idx]
acc_eul = None
for nidx in range(nodes.shape[0]):
shock = jnp.broadcast_to(jnp.asarray(nodes[nidx])[None, :],
(eul_sample.shape[0], model.n_shocks))
r = compute_residuals(model, params, eul_sample, shock)
contrib = {k: weights[nidx] * r[k] for k in ("euler_0", "euler_1")}
acc_eul = (contrib if acc_eul is None
else {k: acc_eul[k] + contrib[k] for k in contrib})
# Our framework's euler residual is the *signed* (RHS - LHS); Simon's
# euler_lhs = lambda * (1 + dGamma/dk'). errREE = (RHS - LHS)/LHS = residual/LHS.
pol = np.asarray(params(eul_sample))
k_state = np.asarray(eul_sample[:, :N])
kp = pol[:, :N]
lam = pol[:, N]
kappa = float(model.constants["kappa"])
fig_e, axes_e = plt.subplots(2, 2, figsize=(12, 9))
for j in range(N):
d_adj_today = kappa * (kp[:, j] / k_state[:, j] - 1.0)
euler_lhs = lam * (1.0 + d_adj_today)
resid = np.asarray(acc_eul[f"euler_{j}"]).reshape(-1)
euler_rhs = euler_lhs + resid
errREE = resid / np.maximum(euler_lhs, 1e-12)
med = float(np.median(np.abs(errREE)))
# signed relative Euler error vs capital -- linear, near zero, honest axis
axes_e[0, j].scatter(k_state[:, j], errREE, s=3, alpha=0.3)
axes_e[0, j].axhline(0, color="k", lw=0.8, alpha=0.6)
axes_e[0, j].set_xlabel(f"capital $k_{j}$")
axes_e[0, j].set_ylabel(r"relative Euler error $\mathrm{errREE}$")
axes_e[0, j].set_title(f"country {j}: signed errREE (median |errREE| = {med:.1e})")
# LHS vs RHS Euler check -- 45-degree line
axes_e[1, j].scatter(euler_lhs, euler_rhs, s=3, alpha=0.3)
lim = [float(min(euler_lhs.min(), euler_rhs.min())),
float(max(euler_lhs.max(), euler_rhs.max()))]
axes_e[1, j].plot(lim, lim, "k--", lw=1, alpha=0.6, label="$45^\\circ$")
axes_e[1, j].set_xlabel("Euler LHS"); axes_e[1, j].set_ylabel("Euler RHS")
axes_e[1, j].set_title(f"country {j}: LHS vs RHS"); axes_e[1, j].legend()
fig_e.suptitle("Capital-Euler accuracy on the ergodic set (quadrature expectation)")
fig_e.tight_layout()
8 — Ergodic distribution¶
Where the trained economy actually lives. The left panel is Simon's capital cloud $k_0$ vs $k_1$ — the two countries' capital stocks co-move through the shared aggregate shock but spread along the idiosyncratic ones; the right panel shows the driving log-TFP cloud $z_0$ vs $z_1$. The policy figures above, and the accuracy certificate below, are evaluated on exactly this set.
fig_g, axes_g = plt.subplots(1, 2, figsize=(11, 4.5))
k0 = np.asarray(states_flat[:, 0]); k1 = np.asarray(states_flat[:, 1])
z0 = np.asarray(states_flat[:, 2]); z1 = np.asarray(states_flat[:, 3])
axes_g[0].scatter(k0, k1, s=2, alpha=0.1)
axes_g[0].set_xlabel("capital, country 0"); axes_g[0].set_ylabel("capital, country 1")
axes_g[0].set_title("ergodic distribution: capital")
axes_g[1].scatter(z0, z1, s=2, alpha=0.1)
axes_g[1].set_xlabel("log TFP, country 0"); axes_g[1].set_ylabel("log TFP, country 1")
axes_g[1].set_title("ergodic distribution: TFP shocks")
fig_g.tight_layout()
9 — Ergodic accuracy certificate¶
A printed table (not a figure): per-equation residual quantiles on the
ergodic set, expectation by the same 27-node quadrature used in training.
Eulers are scaled by $\lambda$ (marginal utility units), the resource
constraint by aggregate output; FB residuals are reported raw (they are
already in multiplier units). The fb_0 / fb_1 rows are the magnitude
of the complementarity slack seen off the L in §6.
from deqn_jax.training.loss import compute_loss, gauss_hermite_nd, compute_residuals
nodes, weights = gauss_hermite_nd(3, model.n_shocks)
sample = states_flat[
np.random.default_rng(0).choice(states_flat.shape[0], 2048, replace=False)
]
acc = None
for nidx in range(nodes.shape[0]):
shock = jnp.broadcast_to(jnp.asarray(nodes[nidx])[None, :],
(sample.shape[0], model.n_shocks))
r = compute_residuals(model, params, sample, shock)
acc = ({k: weights[nidx] * v for k, v in r.items()} if acc is None
else {k: acc[k] + weights[nidx] * v for k, v in r.items()})
d = model.definitions_fn(sample, params(sample), model.constants)
scale = {
"euler_0": d["lam"], "euler_1": d["lam"],
"arc": d["y_0"] + d["y_1"],
"fb_0": 1.0, "fb_1": 1.0,
}
print(f"{'equation':<10} {'median':>8} {'p90':>8} {'p99':>8} (log10 |scaled residual|)")
for eq, r in acc.items():
e = jnp.abs(r) / scale[eq]
qs = [float(jnp.log10(jnp.maximum(jnp.quantile(e, p), 1e-300)))
for p in (0.5, 0.9, 0.99)]
print(f"{eq:<10} {qs[0]:>8.2f} {qs[1]:>8.2f} {qs[2]:>8.2f}")
equation median p90 p99 (log10 |scaled residual|) euler_0 -4.32 -3.76 -3.52 euler_1 -4.36 -3.80 -3.47 arc -2.89 -2.49 -2.19 fb_0 -6.30 -6.21 -6.16 fb_1 -6.30 -6.21 -6.16
10 — Summary¶
- Multipliers as policies: $\lambda,\ \mu_0,\ \mu_1$ are network outputs, tied to their constraints by Fischer-Burmeister residuals — the pattern for any occasionally-binding constraint whose wedge has no closed form.
- Structure does the work: consumption never enters the network; complete markets pin it from $\lambda$, and the simulated volatility ratio reproduces the $\gamma_1/\gamma_0$ risk-sharing prediction.
- Quadrature expectations: with 3 shock dimensions, a 27-node tensor grid replaces Monte Carlo — exact expectations, no estimator bias.
The FB trilogy in this gallery: bm_labor_constrained (analytic wedge,
cap slack by construction) → irbc (learned multipliers, occasionally
binding) → olg_lifecycle (FB wrapping an expectation: the two-stage loss).