Lecture 10: Sequence-space DEQNs

# Run-mode switch (smoke = CPU-bounded for CI, teaching = laptop figures, production = full reproduction).
RUN_MODE = "smoke"  # one of: "smoke", "teaching", "production"
SEED = 0

import matplotlib
matplotlib.use('Agg')  # non-blocking backend
import tensorflow as tf

import numpy as np

from scipy.interpolate import BSpline

import matplotlib.pyplot as plt

import time



plt.rcParams['font.size'] = 13



# ---------- tutorial / production toggle ----------
# Derived from the kit-mandated RUN_MODE in cell 1: smoke/teaching -> quick, production -> full.
TUTORIAL_MODE = "full" if RUN_MODE == "production" else "quick"



print(f"TensorFlow {tf.__version__}")

print(f"Mode: {TUTORIAL_MODE}")

print(f"GPUs: {tf.config.list_physical_devices('GPU')}")

# ---- Model parameters ----
ALPHA = 0.36   # capital share
DELTA = 0.025  # depreciation
BETA = 0.93    # discount factor
SIGMA = 1.0    # CRRA (log utility)

# Idiosyncratic shocks
EPS_VALS = np.array([0.5, 1.5])
PI_EPS = np.array([[0.9, 0.1],  # P(eps'=L | eps=L) = 0.9
                   [0.1, 0.9]])

# Aggregate shocks
Z_VALS = np.array([0.93, 1.07])
PI_Z = np.array([[0.7, 0.3],  # persistence = 0.7
                 [0.3, 0.7]])

N_EPS = len(EPS_VALS)
N_Z = len(Z_VALS)

# Deterministic steady state
R_ss = 1.0 / BETA
K_ss = ((R_ss - 1.0 + DELTA) / ALPHA) ** (1.0 / (ALPHA - 1.0))
print(f"Steady-state: K_ss = {K_ss:.4f}, R_ss = {R_ss:.4f}")

# TF constants for use in the graph
EPS_VALS_tf = tf.constant(EPS_VALS, dtype=tf.float32)
PI_EPS_tf = tf.constant(PI_EPS, dtype=tf.float32)
PI_Z_tf = tf.constant(PI_Z, dtype=tf.float32)
Z_VALS_tf = tf.constant(Z_VALS, dtype=tf.float32)

def compute_prices(K, L, Z):
    """Cobb-Douglas factor prices (numpy, for simulation)."""
    KL = max(K, 0.01) / max(L, 0.01)
    R = 1.0 - DELTA + ALPHA * Z * KL ** (ALPHA - 1.0)
    w = (1.0 - ALPHA) * Z * KL ** ALPHA
    return R, w

print(f"At K_ss with Z=1: R={compute_prices(K_ss, 1.0, 1.0)[0]:.4f}, w={compute_prices(K_ss, 1.0, 1.0)[1]:.4f}")

GRID_CONFIG = {

    "quick": {"N": 60, "K_MAX": 60.0},

    "full": {"N": 100, "K_MAX": 100.0},

}[TUTORIAL_MODE]



N = GRID_CONFIG["N"]

K_MAX = GRID_CONFIG["K_MAX"]

K_GRID_SHIFT = 0.5

BASIS_SHIFT = 0.5



x_grid = np.linspace(np.log(K_GRID_SHIFT), np.log(K_MAX + K_GRID_SHIFT), N + 1)

k_grid_np = np.exp(x_grid) - K_GRID_SHIFT

dk_np = np.diff(k_grid_np)

x0 = float(x_grid[0])

delta_x = float((x_grid[-1] - x_grid[0]) / N)



# TF constants

k_grid_tf = tf.constant(k_grid_np, dtype=tf.float32)

dk_tf = tf.constant(dk_np, dtype=tf.float32)



# Visualize

fig, ax = plt.subplots(figsize=(10, 1.5))

ax.scatter(k_grid_np, np.zeros(N + 1), s=6, alpha=0.7)

ax.set_xlabel("Capital")

ax.set_title(f"Log-spaced grid: {N+1} points in [0, {K_MAX}]")

ax.set_yticks([])

ax.axvline(K_ss, color='red', linestyle='--', alpha=0.5,

           label=f"deterministic $K_{{ss}}={K_ss:.1f}$")

ax.legend(fontsize=9)

plt.tight_layout()

plt.close("all")



print(f"Grid: {N+1} points, k_min={k_grid_np[0]:.4f}, k_max={k_grid_np[-1]:.2f}")

print(f"Smallest interval: {dk_np[0]:.4f}, largest: {dk_np[-1]:.2f}")

N_INTERIOR = 26 if TUTORIAL_MODE == "quick" else 56

ORDER = 4  # cubic I-splines



def build_ispline_basis(k_grid, n_interior=56, order=4, BASIS_SHIFT=0.1):

    """Precompute I-spline basis at the grid points (numpy/scipy)."""

    k_arr = np.asarray(k_grid)

    x = np.log(BASIS_SHIFT + k_arr)

    x_min, x_max = float(x[0]), float(x[-1])



    # Equidistant interior knots in x-space

    knots_int = np.linspace(x_min, x_max, n_interior + 2)[1:-1]

    knots = np.concatenate([np.repeat(x_min, order), knots_int, np.repeat(x_max, order)])

    n_basis = len(knots) - order



    # Compute I-splines (= integral of B-splines)

    basis = np.zeros((n_basis, len(k_arr)))

    for idx in range(n_basis):

        c = np.zeros(n_basis)

        c[idx] = 1.0

        bspl = BSpline(knots, c, order - 1, extrapolate=False)

        t_span = knots[idx + order] - knots[idx]

        if t_span < 1e-12:

            continue

        anti = bspl.antiderivative()

        basis[idx] = (order / t_span) * np.nan_to_num(anti(x) - anti(x_min), 0.0)



    return np.clip(basis, 0.0, 1.0)



ispline_basis_np = build_ispline_basis(k_grid_np, N_INTERIOR, ORDER, BASIS_SHIFT)

J = ispline_basis_np.shape[0]

print(f"Built {J} I-spline basis functions on {N+1} grid points")



# Store as TF constant for graph use

ispline_basis_tf = tf.constant(ispline_basis_np, dtype=tf.float32)



# Visualize

B = ispline_basis_np

print("Basis matrix shape:", B.shape)



fig, axes = plt.subplots(1, 2, figsize=(12, 4))

for j in range(J):

    axes[0].plot(k_grid_np, B[j], linewidth=0.8, alpha=0.5)

axes[0].set_xlabel("Capital $k$")

axes[0].set_ylabel("$I_j(k)$")

axes[0].set_title(f"I-spline basis in $k$-space ({J} functions)")

axes[0].grid(True, alpha=0.3)



x_vals = np.log(BASIS_SHIFT + k_grid_np)

for j in range(J):

    axes[1].plot(x_vals, B[j], linewidth=0.8, alpha=0.5)

axes[1].set_xlabel("$x = \\log(0.5 + k)$")

axes[1].set_ylabel("$I_j(x)$")

axes[1].set_title("I-spline basis in transformed space")

axes[1].grid(True, alpha=0.3)



plt.tight_layout()

plt.close("all")

H = 50            # history length (last H aggregate shocks)
HIDDEN_DIM = 64   # MLP hidden layer width
INPUT_DIM = H * (N_Z + 1)  # one-hot + values per lag

def init_actor(input_dim, hidden_dim, n_eps, J, seed=0):
    """Initialize a 2-hidden-layer MLP as a list of tf.Variables.

    Returns [W0, b0, W1, b1, W2, b2].
    """
    rng = np.random.default_rng(seed)
    output_dim = n_eps * (1 + J)
    W0 = tf.Variable(rng.normal(size=(input_dim, hidden_dim)).astype(np.float32)
                     * np.sqrt(2.0 / input_dim), name='W0')
    b0 = tf.Variable(np.zeros(hidden_dim, dtype=np.float32), name='b0')
    W1 = tf.Variable(rng.normal(size=(hidden_dim, hidden_dim)).astype(np.float32)
                     * np.sqrt(2.0 / hidden_dim), name='W1')
    b1 = tf.Variable(np.zeros(hidden_dim, dtype=np.float32), name='b1')
    W2 = tf.Variable(rng.normal(size=(hidden_dim, output_dim)).astype(np.float32)
                     * 0.01, name='W2')
    b2 = tf.Variable(np.zeros(output_dim, dtype=np.float32), name='b2')
    return [W0, b0, W1, b1, W2, b2]


@tf.function
def actor_c_grid(params_list, agg_input, R, w):
    """Forward pass: aggregate input -> consumption on the k-grid.

    Args:
        params_list: [W0, b0, W1, b1, W2, b2]
        agg_input:   (INPUT_DIM,) encoded Z-history
        R, w:        scalar factor prices (tf constants or python floats)

    Returns: c_grid of shape (N_EPS, N+1)
    """
    # MLP forward pass (tanh activations)
    h = tf.math.tanh(tf.matmul(agg_input[tf.newaxis, :], params_list[0]) + params_list[1])
    h = tf.math.tanh(tf.matmul(h, params_list[2]) + params_list[3])
    raw = tf.reshape(tf.matmul(h, params_list[4]) + params_list[5], [N_EPS, 1 + J])

    # Boundary MPC: sigmoid -> (0, 1)
    alpha = tf.nn.sigmoid(raw[:, 0])                             # (N_EPS,)

    # Phantom-zero softmax: guarantees weights >= 0 and sum < 1
    logits = raw[:, 1:]                                           # (N_EPS, J)
    augmented = tf.concat([tf.zeros([N_EPS, 1]), logits], axis=1)
    w_tilde = tf.nn.softmax(augmented, axis=1)[:, 1:]            # (N_EPS, J)

    # MPC on the grid: alpha * (1 - weighted I-splines)
    mpc = alpha[:, tf.newaxis] * (1.0 - tf.matmul(w_tilde, ispline_basis_tf))  # (N_EPS, N+1)

    # Cumulate consumption on the grid from the MPC schedule
    # Boundary: c(k_0) = MPC(k_0) * m(k_0), where m = w*eps + R*k
    c0 = mpc[:, 0] * (w * EPS_VALS_tf + R * k_grid_tf[0])       # (N_EPS,)
    # Increments: dc_n = MPC(k_n) * R * dk_n
    dc = mpc[:, 1:] * (R * dk_tf[tf.newaxis, :])                 # (N_EPS, N)
    c_grid = tf.concat([c0[:, tf.newaxis],
                        c0[:, tf.newaxis] + tf.cumsum(dc, axis=1)], axis=1)
    return c_grid


def encode_Z_history(z_hist):
    """Encode Z-history as one-hot + values (numpy, for buffer construction)."""
    onehot = np.eye(N_Z, dtype=np.float32)[z_hist].reshape(-1)
    values = Z_VALS[z_hist].reshape(-1).astype(np.float32)
    return np.concatenate([onehot, values])


def encode_Z_history_tf(z_hist_tf):
    """Encode Z-history as one-hot + values (TF, for use inside @tf.function)."""
    onehot = tf.reshape(tf.one_hot(z_hist_tf, N_Z), [-1])
    values = tf.reshape(tf.gather(Z_VALS_tf, z_hist_tf), [-1])
    return tf.concat([onehot, values], axis=0)


# Initialize and check
params_list = init_actor(INPUT_DIM, HIDDEN_DIM, N_EPS, J, seed=0)
n_params = sum(int(tf.size(p)) for p in params_list)
print(f"Actor parameters: {n_params}")
print(f"Input dim: {INPUT_DIM} (H={H} x {N_Z+1} features per lag)")
print(f"Output dim: {N_EPS} x (1 + {J}) = {N_EPS * (1 + J)}")

# Quick shape check with a dummy input
z_hist_dummy = np.zeros(H, dtype=np.int32)
agg_dummy = tf.constant(encode_Z_history(z_hist_dummy))
R_vis, w_vis = compute_prices(K_ss, 1.0, float(Z_VALS[0]))
c_test = actor_c_grid(params_list, agg_dummy, tf.constant(R_vis, tf.float32),
                      tf.constant(w_vis, tf.float32))
print(f"Output shape: {c_test.shape}  (expected ({N_EPS}, {N+1}))")

# Verify shape guarantees
c_np = c_test.numpy()
m_np = w_vis * EPS_VALS[:, None] + R_vis * k_grid_np[None, :]
print(f"\nShape check for initial random policy:")
print(f"  min dc across grid intervals      = {np.min(np.diff(c_np, axis=1)): .3e}  (should be >= 0)")
print(f"  max(c - cash on hand)             = {np.max(c_np - m_np): .3e}  (should be <= 0)")

# ---- Simulation config ----
SIM_CONFIG = {
    "quick": {"T": 100, "N_AGG": 8, "BUFFER_CAP": 128},
    "full":  {"T": 100, "N_AGG": 8, "BUFFER_CAP": 128},
}[TUTORIAL_MODE]

T = SIM_CONFIG["T"]
N_AGG = SIM_CONFIG["N_AGG"]
BUFFER_CAP = SIM_CONFIG["BUFFER_CAP"]


def distribution_aggregates(mu):
    """Compute aggregate K and L from distribution mu(eps, k).

    mu has shape (N_EPS, N+1) and sums to 1.
    K = sum k * mu(eps,k),  L = sum eps * mu(eps,k).
    """
    K = np.sum(mu * k_grid_np[None, :])
    L = np.sum(mu * EPS_VALS[:, None])
    return max(K, 0.01), max(L, 0.01)


def distribution_step(mu, c_grid_t, R, w):
    """Advance the distribution one period (Young 2010, numpy).

    For each (eps, k) grid point:
      1. Compute next-period capital: k' = w*eps + R*k - c(eps, k)
      2. Find the two neighboring grid points around k'
      3. Split the mass to those neighbors (lottery)
      4. Apply the idiosyncratic transition matrix PI_EPS
    """
    m_grid = w * EPS_VALS[:, None] + R * k_grid_np[None, :]
    k_next = np.clip(m_grid - c_grid_t, 0.0, k_grid_np[-1])

    # Find grid position of k' in log-space
    x_next = np.log(K_GRID_SHIFT + k_next)
    j_lo = np.clip(
        np.floor((x_next - x0) / delta_x).astype(np.int32), 0, N - 1)

    # Lottery weights
    k_lo = k_grid_np[j_lo]
    k_hi = k_grid_np[j_lo + 1]
    w_lo = (k_hi - k_next) / np.maximum(k_hi - k_lo, 1e-12)
    w_hi = 1.0 - w_lo

    # Scatter mass to neighboring grid points for each eps state
    mu_after = np.zeros_like(mu)
    for i in range(N_EPS):
        np.add.at(mu_after[i], j_lo[i], mu[i] * w_lo[i])
        np.add.at(mu_after[i], j_lo[i] + 1, mu[i] * w_hi[i])

    # Apply idiosyncratic transition
    mu_next = PI_EPS.T @ mu_after
    return mu_next


def initial_distribution(K_target):
    """Create an initial distribution centered around K_target."""
    eps_weight = np.ones(N_EPS) / N_EPS
    sigma_k = K_target * 0.3
    density = np.exp(-0.5 * ((k_grid_np - K_target) / sigma_k) ** 2)
    density = density / np.sum(density)
    mu = eps_weight[:, None] * density[None, :]
    return mu / np.sum(mu)


def draw_Z_path(rng, T, z_init):
    """Draw T aggregate shocks (numpy). Returns (Z_path, z_terminal)."""
    Z_path = np.zeros(T, dtype=np.int32)
    z = z_init
    for t in range(T):
        Z_path[t] = z
        z = int(rng.uniform() >= PI_Z[z, 0])
    return Z_path, z


def make_agg_inputs(full_Z_seq, H, T):
    """Build encoded Z-history inputs for each simulation step (numpy)."""
    inputs = np.zeros((T, H * (N_Z + 1)), dtype=np.float32)
    for t in range(T):
        window = full_Z_seq[t:t + H]
        inputs[t] = encode_Z_history(window)
    return inputs


def simulate_young(params_list, mu_init, Z_path, agg_inputs):
    """Simulate T periods using Young's non-stochastic method (numpy outer loop).

    The actor_c_grid calls go through TF for the forward pass, but we
    immediately convert to numpy. No gradients are tracked here.
    """
    mu = mu_init.copy()
    T_sim = len(Z_path)
    for t in range(T_sim):
        Z_idx = int(Z_path[t])
        K_t, L_t = distribution_aggregates(mu)
        R_t, w_t = compute_prices(K_t, L_t, Z_VALS[Z_idx])
        agg_input_tf = tf.constant(agg_inputs[t])
        c_grid_t = actor_c_grid(params_list, agg_input_tf,
                                tf.constant(R_t, tf.float32),
                                tf.constant(w_t, tf.float32)).numpy()
        mu = distribution_step(mu, c_grid_t, R_t, w_t)
    return mu


def run_simulations(rng, params_list, buffer_z, buffer_mu, idx):
    """Generate shocks, simulate distributions forward, return terminal states."""
    results = []
    for i in idx:
        z_hist = buffer_z[i]
        mu_start = buffer_mu[i]
        z_init = int(z_hist[-1])

        Z_path, z_terminal = draw_Z_path(rng, T, z_init)

        # Build full Z-sequence for encoding (history + simulation path)
        # Z_path[0] = z_init (duplicates parent's terminal), skip with [1:]
        full_Z = np.concatenate([z_hist, Z_path[1:], [z_terminal]])
        agg_inputs = make_agg_inputs(full_Z, H, T)

        mu_final = simulate_young(params_list, mu_start, Z_path, agg_inputs)

        new_z_hist = full_Z[-H:].astype(np.int32)
        results.append((new_z_hist, mu_final, Z_path, z_terminal))
    return results


print(f"Simulation config: T={T}, N_AGG={N_AGG}, buffer={BUFFER_CAP}")
print("Using Young (2010) non-stochastic method -- no MC noise in aggregates.")

@tf.function
def interpolate_tf(vals_grid, eps_idx, k):
    """Piecewise-linear interpolation on the k-grid (differentiable).

    Args:
        vals_grid: (N_EPS, N+1) grid values (e.g. consumption)
        eps_idx:   (M,) integer indices into eps dimension
        k:         (M,) capital values to interpolate at
    Returns:
        (M,) interpolated values
    """
    x = tf.math.log(K_GRID_SHIFT + k)
    j = tf.clip_by_value(
        tf.cast(tf.floor((x - x0) / delta_x), tf.int32), 0, N - 1)
    k_lo = tf.gather(k_grid_tf, j)
    k_hi = tf.gather(k_grid_tf, j + 1)
    w_lo = (k_hi - k) / tf.maximum(k_hi - k_lo, 1e-12)
    # Gather from 2D grid using (eps_idx, j) pairs
    v_lo = tf.gather_nd(vals_grid, tf.stack([eps_idx, j], axis=-1))
    v_hi = tf.gather_nd(vals_grid, tf.stack([eps_idx, j + 1], axis=-1))
    return w_lo * v_lo + (1.0 - w_lo) * v_hi


def fb_loss_one_state_tf(params_list, z_history_np, mu_np):
    """Compute mean FB-squared on the full (eps, k) grid for one replay state.

    Prices and aggregates are numpy scalars (not traced by the tape).
    Only actor_c_grid calls contribute to the gradient.

    Args:
        params_list:   [W0, b0, W1, b1, W2, b2] tf.Variables
        z_history_np:  (H,) numpy int32 array
        mu_np:         (N_EPS, N+1) numpy float64 distribution
    Returns:
        scalar tf.Tensor -- mean FB-squared
    """
    z_t = int(z_history_np[-1])

    # Current prices from distribution (exact aggregates, numpy)
    K_t, L_t = distribution_aggregates(mu_np)
    R_t, w_t = compute_prices(K_t, L_t, Z_VALS[z_t])

    # Current policy on the full grid (TF -- traced)
    agg_input = tf.constant(encode_Z_history(z_history_np))
    c_grid_t = actor_c_grid(params_list, agg_input,
                            tf.constant(R_t, tf.float32),
                            tf.constant(w_t, tf.float32))

    # Next-period capital for each (eps, k) on the grid
    R_t_tf = tf.constant(R_t, tf.float32)
    w_t_tf = tf.constant(w_t, tf.float32)
    m_grid = w_t_tf * EPS_VALS_tf[:, tf.newaxis] + R_t_tf * k_grid_tf[tf.newaxis, :]
    k_next_grid = tf.clip_by_value(m_grid - c_grid_t, 0.0, k_grid_tf[-1])

    # Next-period aggregates from Young step (numpy, not traced)
    c_grid_t_np = c_grid_t.numpy()
    mu_next = distribution_step(mu_np, c_grid_t_np, R_t, w_t)
    K_next, L_next = distribution_aggregates(mu_next)

    # Continuation value: q = E[beta * R' * u'(c')]
    q = tf.zeros_like(c_grid_t)
    for z_next in range(N_Z):
        z_hist_next = np.concatenate([z_history_np[1:], [z_next]]).astype(np.int32)
        agg_next = tf.constant(encode_Z_history(z_hist_next))
        R_next, w_next = compute_prices(K_next, L_next, Z_VALS[z_next])

        c_grid_next = actor_c_grid(params_list, agg_next,
                                   tf.constant(R_next, tf.float32),
                                   tf.constant(w_next, tf.float32))

        # For each next-eps, interpolate c' at k_next_grid
        # c_next_all shape: (N_EPS, N_EPS, N+1) -- [next_eps, current_eps, k]
        c_next_parts = []
        for i_next in range(N_EPS):
            k_flat = tf.reshape(k_next_grid, [-1])           # (N_EPS*(N+1),)
            idx = tf.fill(k_flat.shape, i_next)              # (N_EPS*(N+1),)
            c_interp = interpolate_tf(c_grid_next, idx, k_flat)
            c_next_parts.append(tf.reshape(c_interp, [N_EPS, N + 1]))
        c_next_all = tf.stack(c_next_parts)  # (N_EPS, N_EPS, N+1)

        # u'(c') = 1/c  (log utility)
        mu_u = 1.0 / tf.maximum(c_next_all, 1e-10)
        # E_{eps'} [u'(c')] using PI_EPS transition
        exp_uc = tf.einsum('ij,jik->ik', PI_EPS_tf, mu_u)    # (N_EPS, N+1)
        q = q + PI_Z[z_t, z_next] * BETA * R_next * exp_uc

    # Euler-implied consumption and FB residual
    c_euler = 1.0 / tf.maximum(q, 1e-10)
    c_scale = tf.maximum(c_grid_t, 1e-10)
    g = (c_euler - c_grid_t) / c_scale          # relative Euler gap
    s = k_next_grid / c_scale                    # relative slack
    fb = tf.sqrt(g**2 + s**2 + 1e-12) - g - s   # Fischer-Burmeister residual

    return tf.reduce_mean(fb**2)


print("FB loss defined. Prices enter as numpy scalars (stop-gradient equivalent).")
print("Only actor_c_grid calls are traced by GradientTape.")

TRAIN_CONFIG = {

    "quick": {"LR": 5e-5, "FB_STEPS": 10, "FB_BATCH": 16, "N_EPOCHS": 50, "LOG_EVERY": 10},  # classroom: 50 (~7 min CPU); production: 5000

    "full":  {"LR": 5e-5, "FB_STEPS": 10, "FB_BATCH": 16, "N_EPOCHS": 10_000, "LOG_EVERY": 200},

}[TUTORIAL_MODE]



LR = TRAIN_CONFIG["LR"]

FB_STEPS = TRAIN_CONFIG["FB_STEPS"]

FB_BATCH = TRAIN_CONFIG["FB_BATCH"]

N_EPOCHS = TRAIN_CONFIG["N_EPOCHS"]

LOG_EVERY = TRAIN_CONFIG["LOG_EVERY"]



optimizer = tf.keras.optimizers.Adam(learning_rate=LR, clipnorm=1.0)



# Seed the replay buffer

rng = np.random.default_rng(0)

buffer_z = []

buffer_mu = []



mu_init = initial_distribution(K_ss)

for _ in range(min(BUFFER_CAP, 16)):

    buffer_z.append(np.zeros(H, dtype=np.int32))

    buffer_mu.append(mu_init.copy())



print(f"Ready to train: {N_EPOCHS} epochs, lr={LR}, fb_steps={FB_STEPS}, fb_batch={FB_BATCH}")

print(f"Buffer seeded with {len(buffer_z)} distributions around K_ss={K_ss:.2f}")

history = []
t0 = time.time()

for epoch in range(N_EPOCHS):
    # 1. Sample N_AGG states from the buffer
    idx = rng.integers(0, len(buffer_z), size=N_AGG)

    # 2. Simulate forward to generate new terminal distributions
    sim_results = run_simulations(rng, params_list, buffer_z, buffer_mu, idx)

    # 3. Add terminal states to the buffer
    for new_z_hist, mu_final, _, _ in sim_results:
        buffer_z.append(new_z_hist)
        buffer_mu.append(mu_final)
    while len(buffer_z) > BUFFER_CAP:
        buffer_z.pop(0)
        buffer_mu.pop(0)

    # 4. Mini-batch FB gradient steps
    fb2 = 0.0
    for _ in range(FB_STEPS):
        mb = rng.integers(0, len(buffer_z), size=min(FB_BATCH, len(buffer_z)))
        # Accumulate loss over mini-batch then update
        with tf.GradientTape() as tape:
            batch_loss = tf.constant(0.0)
            for bi in mb:
                batch_loss = batch_loss + fb_loss_one_state_tf(
                    params_list, buffer_z[bi], buffer_mu[bi])
            batch_loss = batch_loss / float(len(mb))
        grads = tape.gradient(batch_loss, params_list)
        optimizer.apply_gradients(zip(grads, params_list))
        fb2 = float(batch_loss)

    if (epoch + 1) % LOG_EVERY == 0:
        fb2_val = fb2
        buf_K = float(np.mean([
            np.sum(m * k_grid_np[None, :]) for m in buffer_mu
        ]))
        wall = time.time() - t0
        history.append((epoch + 1, fb2_val, buf_K, wall))
        print(f"Epoch {epoch + 1:5d} | FB2 = {fb2_val:.2e} | K_bar = {buf_K:.2f} | {wall:.0f}s")

wall_total = time.time() - t0
print(f"\nDone! {N_EPOCHS} epochs in {wall_total:.0f}s ({wall_total/60:.1f} min)")
print(f"Final FB2 = {fb2:.2e}")

epochs_h = [h[0] for h in history]

fb2_h = [h[1] for h in history]



fig, ax = plt.subplots(figsize=(8, 4))

ax.semilogy(epochs_h, fb2_h, linewidth=2, color='#1f77b4')

ax.axhline(1e-6, color='red', linestyle='--', alpha=0.5, label='$10^{-6}$ target')

ax.set_xlabel("Epoch", fontsize=12)

ax.set_ylabel(r"FB$^2$ (log scale)", fontsize=12)

ax.set_title("Euler Equation Accuracy During Training", fontsize=14)

ax.legend(fontsize=11)

ax.grid(True, alpha=0.3)

plt.tight_layout()

plt.close("all")

# Pick a representative buffer state for plotting

buf_K = float(np.mean([

    np.sum(m * k_grid_np[None, :]) for m in buffer_mu

]))



fig, axes = plt.subplots(1, 2, figsize=(13, 5))

colors = ["#1f77b4", "#ff7f0e"]

labels_eps = [r"$\varepsilon_L$", r"$\varepsilon_H$"]

labels_Z = [r"$Z_L$", r"$Z_H$"]



for z_idx in range(N_Z):

    z_hist_plot = np.full(H, z_idx, dtype=np.int32)

    agg_input = tf.constant(encode_Z_history(z_hist_plot))

    R, w = compute_prices(max(buf_K, 0.01), 1.0, float(Z_VALS[z_idx]))

    c_grid_learned = actor_c_grid(params_list, agg_input,

                                   tf.constant(R, tf.float32),

                                   tf.constant(w, tf.float32)).numpy()



    ls = '-' if z_idx == 0 else '--'

    for e in range(N_EPS):

        # Consumption

        axes[0].plot(k_grid_np, c_grid_learned[e], color=colors[e], linestyle=ls,

                     linewidth=2, label=f"{labels_eps[e]}, {labels_Z[z_idx]}")



        # Cash-on-hand envelope

        m = float(w) * float(EPS_VALS[e]) + float(R) * k_grid_np

        if z_idx == 0:

            axes[0].plot(k_grid_np, m, color=colors[e], linestyle=':', alpha=0.4,

                         label=f"$m(k, {labels_eps[e]})$" if z_idx == 0 else None)



        # MPC

        mpc = np.zeros_like(k_grid_np)

        mpc[0] = c_grid_learned[e, 0] / max(m[0], 1e-10)

        mpc[1:] = np.diff(c_grid_learned[e]) / (float(R) * dk_np)

        axes[1].plot(k_grid_np, mpc, color=colors[e], linestyle=ls, linewidth=2)



axes[0].set_xlabel("Capital $k$", fontsize=12)

axes[0].set_ylabel("Consumption $c$", fontsize=12)

axes[0].set_title("Learned consumption policy", fontsize=14)

axes[0].legend(fontsize=8, loc='lower right')

axes[0].grid(True, alpha=0.3)

axes[0].set_xlim(0, min(50, K_MAX))



axes[1].set_xlabel("Capital $k$", fontsize=12)

axes[1].set_ylabel("MPC", fontsize=12)

axes[1].set_title("Marginal propensity to consume", fontsize=14)

axes[1].set_ylim(-0.05, 1.05)

axes[1].grid(True, alpha=0.3)

axes[1].set_xlim(0, min(50, K_MAX))



plt.tight_layout()

plt.close("all")



# Shape check on the last plotted policy

m_check = float(w) * EPS_VALS[:, None] + float(R) * k_grid_np[None, :]

print(f"Shape check for trained policy:")

print(f"  min dc across grid intervals      = {np.min(np.diff(c_grid_learned, axis=1)): .3e}  (should be >= 0)")

print(f"  max(c - cash on hand)             = {np.max(c_grid_learned - m_check): .3e}  (should be <= 0)")

def euler_diagnostics_one_state(params_list, z_history_np, mu_np):
    """Per-grid-point Euler gap and FB residual for one buffer state."""
    z_t = int(z_history_np[-1])

    K_t, L_t = distribution_aggregates(mu_np)
    R_t, w_t = compute_prices(K_t, L_t, Z_VALS[z_t])

    agg_input = tf.constant(encode_Z_history(z_history_np))
    c_grid_t = actor_c_grid(params_list, agg_input,
                            tf.constant(R_t, tf.float32),
                            tf.constant(w_t, tf.float32)).numpy()

    m_grid = w_t * EPS_VALS[:, None] + R_t * k_grid_np[None, :]
    k_next_grid_np = np.clip(m_grid - c_grid_t, 0.0, k_grid_np[-1])

    # Next-period aggregates from Young step
    mu_next = distribution_step(mu_np, c_grid_t, R_t, w_t)
    K_next, L_next = distribution_aggregates(mu_next)

    q = np.zeros_like(c_grid_t)
    for z_next in range(N_Z):
        z_hist_next = np.concatenate([z_history_np[1:], [z_next]]).astype(np.int32)
        agg_next = tf.constant(encode_Z_history(z_hist_next))
        R_next, w_next = compute_prices(K_next, L_next, Z_VALS[z_next])

        c_grid_next = actor_c_grid(params_list, agg_next,
                                    tf.constant(R_next, tf.float32),
                                    tf.constant(w_next, tf.float32)).numpy()

        # Interpolate c' at k_next for each next-eps
        c_next_all = np.zeros((N_EPS, N_EPS, N + 1))
        for i_next in range(N_EPS):
            k_flat = k_next_grid_np.ravel()
            # Simple numpy interpolation for diagnostics
            x_kn = np.log(K_GRID_SHIFT + k_flat)
            j_lo = np.clip(np.floor((x_kn - x0) / delta_x).astype(np.int32), 0, N - 1)
            w_lo = (k_grid_np[j_lo + 1] - k_flat) / np.maximum(
                k_grid_np[j_lo + 1] - k_grid_np[j_lo], 1e-12)
            c_interp = (w_lo * c_grid_next[i_next, j_lo]
                        + (1.0 - w_lo) * c_grid_next[i_next, j_lo + 1])
            c_next_all[i_next] = c_interp.reshape(N_EPS, N + 1)

        mu_u = 1.0 / np.maximum(c_next_all, 1e-10)
        exp_uc = np.einsum('ij,jik->ik', PI_EPS, mu_u)
        q += PI_Z[z_t, z_next] * BETA * R_next * exp_uc

    c_euler = 1.0 / np.maximum(q, 1e-10)
    c_scale = np.maximum(c_grid_t, 1e-10)
    g_grid = (c_euler - c_grid_t) / c_scale
    s_grid = k_next_grid_np / c_scale
    fb_grid = np.sqrt(g_grid**2 + s_grid**2 + 1e-12) - g_grid - s_grid

    # MPC via finite-difference
    mpc = np.zeros_like(c_grid_t)
    mpc[:, 0] = c_grid_t[:, 0] / np.maximum(m_grid[:, 0], 1e-10)
    mpc[:, 1:] = np.diff(c_grid_t, axis=1) / (R_t * dk_np[None, :])

    return {
        "K": float(K_t), "c_grid": c_grid_t, "mpc_grid": mpc,
        "g_grid": g_grid, "fb_grid": fb_grid,
    }


# Evaluate on all buffer states
print(f"Running Euler diagnostics on {len(buffer_z)} buffer states...")
diag_results = []
for i in range(len(buffer_z)):
    d = euler_diagnostics_one_state(params_list, buffer_z[i], buffer_mu[i])
    diag_results.append(d)

all_K = np.array([d["K"] for d in diag_results])
all_c = np.stack([d["c_grid"] for d in diag_results])        # (n_states, N_EPS, N+1)
all_mpc = np.stack([d["mpc_grid"] for d in diag_results])    # (n_states, N_EPS, N+1)
abs_g = np.abs(np.stack([d["g_grid"] for d in diag_results]))
abs_fb = np.abs(np.stack([d["fb_grid"] for d in diag_results]))

print(f"K range: [{all_K.min():.2f}, {all_K.max():.2f}]")
print(f"Mean |g| = {abs_g.mean():.2e},  Mean FB2 = {(abs_fb**2).mean():.2e}")

# --- Diagnostic Plot 1: Consumption and MPC vs K ---



target_k_vals = [0.0, 1.0, 3.0, 7.0, 15.0, 30.0, float(k_grid_np[-1])]

target_k_vals = [tk for tk in target_k_vals if tk <= float(k_grid_np[-1])]

k_indices = [int(np.argmin(np.abs(k_grid_np - tk))) for tk in target_k_vals]



colors_eps = ["#1f77b4", "#ff7f0e"]

eps_labels = [r"$\varepsilon_L$", r"$\varepsilon_H$"]



fig, axes = plt.subplots(2, len(k_indices), figsize=(3.5 * len(k_indices), 8),

                         sharex=True)

for col, ki in enumerate(k_indices):

    for e in range(N_EPS):

        axes[0, col].scatter(all_K, all_c[:, e, ki], s=3, alpha=0.3,

                             color=colors_eps[e],

                             label=eps_labels[e] if col == 0 else None)

        axes[1, col].scatter(all_K, all_mpc[:, e, ki], s=3, alpha=0.3,

                             color=colors_eps[e])



    axes[0, col].set_title(f"$k = {k_grid_np[ki]:.1f}$", fontsize=9)

    axes[1, col].set_xlabel("$K$", fontsize=9)

    axes[0, col].grid(True, alpha=0.3)

    axes[1, col].grid(True, alpha=0.3)

    axes[0, col].tick_params(labelsize=7)

    axes[1, col].tick_params(labelsize=7)



axes[0, 0].set_ylabel("Consumption $c$")

axes[1, 0].set_ylabel("MPC")

axes[0, 0].legend(fontsize=7, markerscale=3)

fig.suptitle("Consumption and MPC vs aggregate capital $K$", fontsize=13, y=1.02)

plt.tight_layout()

plt.close("all")



print("Tight scatter clouds = approximate aggregation holds.")

print("Variation across K comes from the encoder mapping different Z-histories to different policy coefficients.")

# --- Diagnostic Plot 2: FB error bands ---



n_states = abs_g.shape[0]

fig, axes = plt.subplots(1, N_EPS, figsize=(6.2 * N_EPS, 8.2), sharex=True)



for i in range(N_EPS):

    fb_med = np.median(abs_fb[:, i, :], axis=0)

    fb_p10 = np.percentile(abs_fb[:, i, :], 10.0, axis=0)

    fb_p90 = np.percentile(abs_fb[:, i, :], 90.0, axis=0)



    # |FB| panel

    axes[i].fill_between(k_grid_np, np.maximum(fb_p10, 1e-12),

                            np.maximum(fb_p90, 1e-12),

                            color=colors_eps[i], alpha=0.18)

    axes[i].plot(k_grid_np, np.maximum(fb_med, 1e-12),

                    color=colors_eps[i], linewidth=2.0)

    axes[i].set_yscale("log")

    axes[i].set_ylabel(r"$|FB(g,s)|$")

    axes[i].set_xlabel("Capital $k$")

    axes[i].grid(True, alpha=0.3)



    axes[i].set_title(f"{eps_labels[i]}: FB error bands ({n_states} buffer states)",

                         fontsize=11)



plt.tight_layout()

plt.close("all")



print("The band shows the 10th-90th percentile range across buffer states.")

print("Wider bands indicate more variation in accuracy across aggregate states.")

print("High |g| near k=0 reflects the borrowing constraint; high |g| at high k")

print("is often on grid points with zero household mass.")