Lecture 03, Notebook 02: Stochastic Brock–Mirman DEQN - Deep Learning for Solving and Estimating Dynamic Models in Economics and Finance

Course: Deep Learning for Solving and Estimating Dynamic Models in Economics and Finance
Script reference: §2.4 (stochastic Brock–Mirman), §2.6 (Gauss–Hermite quadrature for the conditional expectation)
Notebook role: core
Author: Simon Scheidegger

RUN_MODE = "smoke"  # one of: "smoke", "teaching", "production"
SEED = 0

Simple Introduction to Deep Equilibrium Nets¶

Notebook 2: uncertainty and sampling states from the simulated path of the economy¶

Purpose of the notebook and economic model¶

The notebook is the second of three notebooks that should serve as a simple introduction to Deep Equilibirium Nets, a deep learning based method introduced in Azinovic et al. (2022).
To focus on the method, we are going to solve a simple optimal growth model with one representative agent, a simplified version of Brock and Mirman (1972).
The previous notebook introduced the method in a setting with a single state variable ( $\mathbf{X}_t = K_t$ ), without any uncertainty, and with sampling the states of the economy from an exogenously given interval $\mathcal{K}=[\underline{K}, \overline{K}]$ .
In this notebook we take the model a step further by introducing two new aspects.

We introduce aggregate uncertainty in form of an AR(1) process for the log of total factor productivity ( $\log(z_t)$ ). The evaluation of the equilibrium hence involves the evaluation of an expectation operator requiring numerical integration. Furthermore, the state space will become two dimensional ( $\mathbf{X}_t = [z_t, K_t]$ ).
The second new aspect we introduce is to sample the states, which we train the neural network on, from simulated paths of the economy. This is especially usefull when the state-space becomes high-dimensional. Since the state variables of economic models are often heavily correlated, training the neural network on a hyper-cubic domain is often exponentially wasteful (see Maliar et al. (2011) for more information on this topic).

Similarly to the previous notebook the planner aims to maximize her time-separable life time utility subject to her budget constraint:

\begin{align*} &\max_{\left\{c_{t}\right\}_{t=0}^{\infty}} \sum_{t=0}^{\infty} \beta^{t}E\left[ \ln\left(C_{t}\right) \right] \tag{1} \\ \text{s.t.} \quad & K_{t+1} + C_{t} = Y_t + (1 - \delta) K_t \tag{2} \end{align*}

where now

\begin{align*} Y_t &= z_t K_{t} ^ \alpha \tag{3} \\ \log(z_t) &= \varrho \log(z_{t - 1}) + \sigma \epsilon_t \tag{4} \\ \epsilon_t &\sim N(0, 1). \tag{5} \end{align*}

The difference to the previous notebook is that production now depends on a random variable, $z_t$ , which (in logs) follows an AR(1) process with persistence $\varrho$ and with a standard deviation of innovations given by $\sigma$ . The case, where $\varrho = 1$ and $\sigma = 0$ , corresponds to the certainty case solved in the previous notebook.

The above problem can again be formulated recursively, where the Bellman equation is given by

\begin{align*} V(z_t, K_t) &= \max_{K_{t+1}} \ln\left(C_{t}\right)+\beta E\left[V(z_{t+1}, K_{t+1})\right] \tag{6} \\ \text{subject to: }C_t &= Y_t + (1 - \delta) K_t - K_{t+1} \tag{7} \end{align*}

the state of the economy is now 2-dimensional and given by $\mathbf{X}_t = [z_t, K_t]$ as before the policy is 1-dimensional and dented with $f(\mathbf{X}_t) = K_{t+1}$ .

Again we are interested in approximating the policy $f(\cdot)$ with a neural network $\mathcal{N}(\cdot)$ , such that $\mathcal{N}(\mathbf{X}_t) = \mathcal{N}([z_t, K_t]) \approx f(\mathbf{X}_t) = K_{t+1}$ .
Taking the first order condition with respect to $K_{t+1}$ and applying the Envelope theorem (we follow the same steps as in Notebook 1), we obtain

\begin{align} 0 &= -\frac{1}{C_t} + \beta E\left[\frac{1}{C_{t+1}}\left(1 - \delta + r_{t+1}\right)\right] \end{align}

(8)

where

\begin{align*} Y_t &= z_t K_{t} ^ \alpha \tag{9} \\ r_{t} &= \alpha z_{t} K_{t}^{\alpha - 1} \tag{10} \\ C_t &= Y_t + (1 - \delta) K_t - K_{t+1} \tag{11} \\ \log(z_{t+1}) &= \varrho \log(z_{t}) + \sigma \epsilon_{t+1}. \tag{12} \end{align*}

The difference to the equation derived in the precvious notebook is the expectation operator on the right hand side, with the terms inside the expecation potentially varying with the realiztion of $z_{t+1}$ .

Our goal is again to find a policy function $f(\cdot)$ , such that this equation is fullfilled for $K_{t+1}= f(\mathbf{X}_t)$ and $K_{t+2}= f(\mathbf{X}_{t+1})$ .
In order to be able to interpret the remaining errors in the equilibrium condition, we again reformulate it such that deviations from 0 can be interpreted as relative consumption errors:

\begin{align*} 0 &= -\frac{1}{C_t} + \beta E\left[\frac{1}{C_{t+1}}\left(1 - \delta + r_{t+1}\right)\right] \tag{13} \\ \Leftrightarrow C_t &= \frac{1}{\beta E\left[\frac{1}{C_{t+1}}\left(1 - \delta + r_{t+1}\right)\right]} \tag{14} \\ \Leftrightarrow 0 &= \frac{1}{C_t \beta E\left[\frac{1}{C_{t+1}}\left(1 - \delta + r_{t+1}\right)\right]} - 1. \tag{15} \end{align*}

We will encode this equation as the loss function to train the neural network. I.e. we will train the neural network $\mathcal{N}(\cdot)$ such that the implied policy fullfills the equation above for given states $\mathbf{X}$ .
Three remarks are in order:

As before we can approximate the savings rate $s_t$ , such that $K_{t+1} = (1 - \delta) K_t + Y_t s_t =(1 - \delta) K_t + Y_t \mathcal{N}(K_t)$ . Since $Y_t$ is completely determined given the state $\mathbf{X}_t$ , which now also includes productivity $z_t$ , this formulation again encodes the policy $K_{t+1}$ , and we can again use a sigmoid activation function to ensure that $s_t\in(0, 1)$ and hence $K_{t+1}\geq 0$ and $C_t\geq 0$ .
In contrast to the previous notebook, our equilibrium condition now requires us to evaluate an expectation operator. There are many ways to do this, and which is optimal depends on the problem at hand. In the given problem, the shocks $\epsilon_{t+1}$ are normally distributed and hence we can conveniently use Gauss-Hermite quadrature, a form of Gaussian quadrature, which is designed to evaluate integrals of the form

\begin{align} \int_{-\infty}^{+\infty}h(x)e^{-x^2}dx \end{align}

(16)

by replacing the integral with a weighted sum over the function evaluated at $n$ points

\begin{align} \int_{-\infty}^{+\infty}h(x)e^{-x^2}dx \approx \sum_{i=1}^{n}w_i h(x_i). \end{align}

(17)

The $n$ integration nodes $\{x_i\}_{i = 1}^{n}$ and corresponding weights $\{w_i\}_{i = 1}^{n}$ are chosen from a (complicated) formula and the approximation is exact if $h$ is a polynomial of degree $2n - 1$ or lower. In our case, we do not know that the term inside the expectation is polynomial but it will be smooth enough for us to obtain a good approximation. In case we would solve models with kinks, for example, this assumption is violated and the accuracy of the numerical integration should be assessed. Alternative integration methods include, (Quasi) Monte Carlo, Monomial rules, and many more (see the great book Judd (1998)).
Alternatively one could discretize the AR(1) process into a Markov Chain and then use plain summation with the transition probabilities, as in Tauchen (1986) and Rouwenhorst (1995).
One final remark is that the weights and nodes returned by functions that implement Gauss-Hermite quadrature are designed to integrate over $e^{-x^2}$ and hence the weights add up to $\sqrt{\pi}$ . Since we want to integrate of the density $\frac{1}{\sqrt{\pi}}e^{-\left(\frac{x}{\sqrt{2}}\right)^2}$ , we need to divide the weights by $\sqrt{\pi}$ and multiply the nodes with $\sqrt{2}$ .

Again we want the equilibrium condition to hold for all possible states $\mathbf{X}_t$ , but in practice we need to focus on a finite space. As in the last notebook, we startout in the simplest possible way and sample the state from an exogenously chosen rectangle $\mathcal{X}=[\underline{z}, \overline{z}] \times [\underline{K}, \overline{K}]$ . Since capital and the tfp shock are correlated this will be wasteful. Therefore, in a second step, we will show how we can solve the model exlusively where it matters by sampling the states to train the neural network from simulated paths of the economy.

# Import necessary libraries
import numpy as np
import math
import tensorflow as tf
from tensorflow import keras

import matplotlib.pyplot as plt
from matplotlib import rc
plt.rcParams["font.size"] = 15

%matplotlib inline


# Reproducibility: fix seeds across numpy and TensorFlow.
SEED = 0
np.random.seed(SEED)
tf.random.set_seed(SEED)

print("Version of tensorflow is {}".format(tf.__version__))

2026-04-24 07:20:34.067506: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1777008034.091959   79031 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1777008034.099617   79031 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2026-04-24 07:20:34.120745: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.

Version of tensorflow is 2.18.0

Evaluating the expectation operator¶

In this section we define the integration nodes and weights we will use to evaluate the expectation operator.
We choose n_int integration points, denote the integration nodes by x_int_norm, and the weights by w_int and obtain them using np.polynomial.hermite.hermgauss.
The actual integration nodes we use later will be multiplied with the standard deviation of the innovations.

n_int = 5 # integration nodes
x_int_norm, w_int = np.polynomial.hermite.hermgauss(n_int) # obtain standard GH nodes and weights
w_int = w_int / np.sqrt(np.pi) # divide the weights by sqrt(pi)
x_int_norm = x_int_norm * np.sqrt(2) # multiply weights with 2**0.5

#convert to tensorflow
x_int_norm = tf.constant(x_int_norm, dtype = tf.float32)
w_int = tf.constant(w_int, dtype = tf.float32)

# for plotting the normal distribution
x_plot_norm = np.linspace(np.min(x_int_norm) - 0.7, np.max(x_int_norm) + 0.7, 200)
y_plot_norm = np.exp(- 0.5 * x_plot_norm ** 2) / np.sqrt(np.pi)

plt.bar(x_int_norm, w_int, width = 0.1, label = "GH nodes and weights")
plt.plot(x_plot_norm, y_plot_norm, color = "r", label = "std normal dist")
plt.xlabel("x")
plt.ylabel("y")
plt.legend()
plt.show()

# we should make sure the weights some to 1
print("sum(w_int) = ", np.sum(w_int))

2026-04-24 07:20:37.415623: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:152] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: UNKNOWN ERROR (303)

sum(w_int) =  1.0000001

Implementing the loss function¶

We define the economic parameters, which are constant throughout and will be globals in this notebook.

alpha = 0.36  # Capital share in the Cobb-Douglas production function
beta = 0.99  # Discount factor
delta = 0.1 # depreciation of capital
sigma_tfp = 0.04 # std. dev. for tfp process innvoations
rho_tfp = 0.9 # persistence of tfp process
x_int = x_int_norm * sigma_tfp # adjust the integration nodes

With $\delta=0.1$ the model does not have a closed-form policy, so we cannot benchmark the neural net against an analytical solution as in Notebook~1. We can, however, still pin down the deterministic steady state of capital, which is useful for choosing a sensible sampling range. At $z=1$ the Euler equation collapses to $1/\beta = 1 - \delta + \alpha k^{\alpha-1}$ , giving

k^{\star} = \left(\frac{1/\beta - 1 + \delta}{\alpha}\right)^{1/(\alpha-1)}.

(18)

The cell below computes this steady state.

# ============================================================
# Training-budget switch -- one MODE controls all training cells
# ============================================================
# "smoke"      : ~30 s   on CPU, log10(loss) ~ -2  (sanity check)
# "teaching"  : ~3 min  on CPU, log10(loss) ~ -3.5 (default)
# "production" : ~20 min on CPU, log10(loss) ~ -5  (publication)
MODE = RUN_MODE  # aliased to the chapter-wide RUN_MODE switch in cell 1

if MODE == "smoke":
    NUM_EPISODES_UNIFORM, NUM_EPISODES_SIM = 1001, 501
elif MODE == "teaching":
    NUM_EPISODES_UNIFORM, NUM_EPISODES_SIM = 5001, 2001
elif MODE == "production":
    NUM_EPISODES_UNIFORM, NUM_EPISODES_SIM = 20001, 2001
else:
    raise ValueError(f"Unknown MODE: {MODE!r}")
print(f"MODE = {MODE}: uniform = {NUM_EPISODES_UNIFORM}, sim = {NUM_EPISODES_SIM}")

MODE = classroom: uniform = 5001, sim = 2001

def k_steady_state(alpha, beta, delta):
    """Deterministic steady state of capital under partial depreciation."""
    return ((1.0 / beta - 1.0 + delta) / alpha) ** (1.0 / (alpha - 1.0))

k_star = k_steady_state(alpha, beta, delta)
print("Deterministic steady state k* = {:.5f}".format(k_star))

Deterministic steady state k* = 6.36684

Deep neural network¶

In this section, we define the architecture of the deep neural net.
Our goal is for the neural network to approximate the savings rate $s_t$ , such that $K_{t+1} = Y_t s_t \approx Y_t \mathcal{N}\left(K_t \right)$ .
The neural network input is hence a now 2-dimensional state $\mathbf{X}_t = [z_t, K_t]$ and the output is the 1-dimesnional savings rate $s_t$ .

Hyper parameters¶

Hyper parameter defines the architecture of the deep neural net.
Note that the purpose of this notebook is in the demonstration of the deep neural net. Detail analyses of the choice of hyper parameters are omitted and we use a densely connected feed forward neural network with two hidden layers as in Azinovic et al. (2022). Since we are approximating the savings rate, we want our architecture to ensure that $s_t \in (0, 1)$ . We can do this by using a sigmoid activation function in the output layer. This is in the spirit of encoding prior knowledge based on economics directly into the neural network architecture (relatedly, Kahou et al. (2021) and Han et al. (2022) show how symmetry can be encoded into the neural-network architecture and Azinovic and Žemlička (2023) introduce market clearing neural network architectures.). We assume the following structure of layers and activation functions:
- Layer 1: the input layer, 2 neurons corresponding to the tfp value $z_t$ and the capital stock $K_t$
- Layer 2: the first hidden layer, 50 neurons and is activated by Relu
- Layer 3: the second hidden layer, 50 neurons and is activated by Relu
- Layer 4: the output layer, 1 neuron corresponding to the savings rate $s_t$ , such that $K_{t+1}=(1 - \delta) K_t + Y_t s_t$ . We use a sigmoid activation function, ensuring that $s_t \in (0, 1)$ .

# Layer setting
num_input = 2
num_hidden1 = 50
num_hidden2 = 50
num_output = 1
layers_dim = [num_input, num_hidden1, num_hidden2, num_output]
print("Dimensions of each layer are {}".format(layers_dim))

Dimensions of each layer are [2, 50, 50, 1]

We define the neural network¶

We use Keras to define the neural network

Hard vs. soft constraints — the central design choice in DEQNs¶

Two kinds of equilibrium conditions appear in any dynamic stochastic model:

Inequality / feasibility constraints — e.g.\ $C_t > 0$ , $K_{t+1} > 0$ , the resource constraint $C_t + K_{t+1} = Y_t$ . These must hold exactly.
Optimality conditions — e.g.\ the Euler equation. These hold in the equilibrium but not at every intermediate guess of the policy.

Azinovic, Gaegauf & Scheidegger (2022, §4.2.2; lecture script Fig. 2.3) make this distinction explicit and treat the two kinds very differently:

	Hard constraint (architecture)	Soft constraint (loss)
What	Built into the network output	Penalised in the cost function
How	Activation choice + algebraic identities	Squared residuals in $\mathcal{L}$
Cost	Always satisfied — even at random init	Only satisfied at convergence

Why this matters here. The cell below parameterises the savings share $s_t \in (0, 1)$ via a sigmoid output. Combined with the resource constraint $K_{t+1} = s_t Y_t$ and $C_t = (1 - s_t) Y_t$ , this guarantees $C_t > 0$ and $K_{t+1} > 0$ simultaneously, at every iteration of training. We never have to penalise infeasibility — the architecture rules it out. The Euler equation, by contrast, is enforced softly through the loss.

This split removes a whole class of bad local minima (network outputs that would imply $C_t < 0$ ) and is one reason DEQNs converge in regions where naive penalty methods do not.

# we use sigmod in the output layer so output is between 0 and 1
nn = keras.Sequential([
    keras.layers.Dense(num_hidden1, activation='relu', input_shape=(num_input,)),
    keras.layers.Dense(num_hidden2, activation='relu'),
    keras.layers.Dense(num_output, activation='sigmoid')
])

/usr/local/lib/python3.10/dist-packages/keras/src/layers/core/dense.py:87: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead.
  super().__init__(activity_regularizer=activity_regularizer, **kwargs)

print(nn.summary())

None

The batch dimension¶

Since neural networks are highly parallelizable, we typically do not want to evaluate the neural network for only a single state $\mathbf{X} \in \mathbb{R}^2$ , but for a matrix of $N_{data}$ different states $X \in \mathbb{R}^{N_{data}\times 2}$ . The output of the neural network is then a vector $\mathbf{s}\in \mathbb{R}^{N_{data}\times 1}$ .
The convention in deep learning is to have the different data samples on the 0-axis.

# test it
X_test = np.array([[1., 1.], [0.5, 2.]])

print("X_test = ", X_test)
print("nn prediction = ", nn(X_test))

X_test =  [[1.  1. ]
 [0.5 2. ]]
nn prediction =  tf.Tensor(
[[0.46820828]
 [0.50078773]], shape=(2, 1), dtype=float32)

Implementing the cost function¶

As outlined above, we want to implement a cost-function which takes a vector of states $\mathbf{X}\in \mathbb{R}^{N_{data}\times 2}$ and a neural network $\mathcal{N}(\cdot)$ , and then returns a vector of errors in the equilibrium conditions implied by the policy, which is encoded by the neural network.
As explained above, the equilibrium condition is given by

\begin{align} 0 &= \frac{1}{C_t \beta E\left[\frac{1}{C_{t+1}}\left(1 - \delta + r_{t+1}\right)\right]} - 1. \end{align}

(19)

We use the tf.function decorator, which speeds up the evaluation of the cost function.

To compute the expecation operator we will iterate over the n_int states in the period $t+1$ . Hence we start out with programming a helper function that takes states $X_{t+1}$ and the neural network as input and returns the term inside the expectation $\frac{1}{C_{t+1}}\left(1 - \delta + r_{t+1}\right)$ .

@tf.function
def get_singleinside(X_tplus1, nn):
    n_data = X_tplus1.shape[0] # number of states is on the axis 0
    dim_state = X_tplus1.shape[1] # dimensionality of the state is on axis 1
    
    # read out the state 
    Z_tplus1 = X_tplus1[:, 0 : 1]
    K_tplus1 = X_tplus1[:, 1 : 2]
    
    # compute output
    Y_tplus1 = Z_tplus1 * K_tplus1 ** alpha
    
    # compute the return
    r_tplus1 = alpha * Z_tplus1 * K_tplus1 ** (alpha - 1.)
    
    # use the neural network to predict the savings rate
    s_tplus1 = nn(X_tplus1)
    
    # Compute consumpption
    C_tplus1 = Y_tplus1 - Y_tplus1 * s_tplus1
    
    # compute term inside the expectation
    ret = (1. / C_tplus1) * (1. - delta + r_tplus1)
    
    return ret

# let's try
X_tplus1 = tf.constant([[1., 1.], [2., 2.], [2., 1.3]])

print("ret = ", get_singleinside(X_tplus1, nn))

ret =  tf.Tensor(
[[2.3693488]
 [0.9419461]
 [1.2070402]], shape=(3, 1), dtype=float32)

Now we are ready to compute the cost function.

@tf.function
def compute_cost(X, nn):
    """ 
    Compute the mean squared error in the equilibrium conditions.
    """
    n_data = X.shape[0] # number of states is on the axis 0
    dim_state = X.shape[1] # dimensionality of the state is on axis 1
    
    # read out the state
    Z_t = X[:, 0 : 1]
    K_t = X[:, 1 : 2]
    
    # compute output today
    Y_t = Z_t * K_t ** alpha
    
    # compute return (not really needed)
    r_t = alpha * Z_t * K_t ** (alpha - 1.)
    
    # use the neural network to predict the savings rate
    s_t = nn(X)
    
    # get the implied capital in the next period
    K_tplus1 = (1. - delta) * K_t + Y_t * s_t

    # get consumption
    C_t = Y_t - Y_t * s_t
    
    # now we have to compute the expectation
    expectation = tf.zeros((n_data, dim_state))
    
    # we loop over the integration nodes
    for i in range(n_int):
        # integration weight
        weight_i = w_int[i]
        
        # innovation to the AR(1)
        innovation_i = x_int[i]
        
        # construct exogenous shock at t+1
        Z_tplus1 = tf.exp(rho_tfp * tf.math.log(Z_t) + innovation_i)
        
        # construct state at t+1
        X_tplus1 = tf.concat([Z_tplus1, K_tplus1], axis = 1)
        
        # compute term inside the expeectation
        inside_i = get_singleinside(X_tplus1, nn)
        
        # add term to the expectaion with the appropriate weight
        expectation = expectation + weight_i * inside_i

    # now we have all terms to construct the relative Euler error
    
    # Define the relative Euler error
    errREE = 1. - 1. / (C_t * beta * expectation)
    
    # compute the cost, i.e. the mean square error in the equilibrium conditions
    cost = tf.reduce_mean(errREE ** 2)
    
    # we return some more things for plotting
    LHS = 1. / C_t # LHS of Ee
    RHS = beta * expectation # RHS of Ee

    return cost, errREE, C_t, K_tplus1, r_t, LHS, RHS

# let's try
X = tf.constant([[0.5, 1.], [0.9, 2.], [1.1, 3.]])

print("cost = ", compute_cost(X, nn)[0])

cost =  tf.Tensor(0.001215832, shape=(), dtype=float32)

Gradients¶

Now we define a function that, for given data X, computes the gradient of the loss w.r.t. the neural network parameters.
This gradient will be used to update the neural network parameters into the direction which decreases the loss function.

def grad(X, nn):
    with tf.GradientTape() as tape:
        loss_value = compute_cost(X, nn)[0]
    return loss_value, tape.gradient(loss_value, nn.trainable_variables)

X = tf.constant([[1., 0.8], [2., 1.3], [0.5, 3.]])

loss, grads = grad(X, nn)

print("loss = ", loss)
print("grads = ", grads)

loss =  tf.Tensor(0.013667074, shape=(), dtype=float32)
grads =  [<tf.Tensor: shape=(2, 50), dtype=float32, numpy=
array([[ 0.0000000e+00,  3.0820281e-03,  0.0000000e+00,  5.7562208e-03,
         1.2723984e-02, -4.8042974e-04,  0.0000000e+00,  0.0000000e+00,
         1.5041227e-03,  0.0000000e+00,  0.0000000e+00, -7.8256950e-03,
         0.0000000e+00,  5.4290518e-05, -1.4148600e-03,  2.7571325e-03,
         1.1582919e-03, -1.9126746e-05,  0.0000000e+00,  0.0000000e+00,
         6.4880832e-04,  0.0000000e+00, -1.0917010e-03,  3.0080602e-03,
         5.2756863e-05, -4.0761549e-03,  9.5555035e-04,  3.4817047e-03,
         3.2763649e-03,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
         1.3548485e-03, -1.4205422e-03, -3.7496869e-04, -5.7002250e-03,
         1.1548721e-02,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
         9.2778535e-04,  5.9693493e-04,  3.5221130e-03, -4.5576273e-03,
         3.9401418e-04, -5.1415991e-05, -2.8362998e-04,  7.7686440e-03,
        -3.2507256e-03,  0.0000000e+00],
       [ 0.0000000e+00,  2.3115915e-03,  0.0000000e+00,  8.5461512e-04,
         8.6304434e-03,  4.2308364e-03,  0.0000000e+00,  0.0000000e+00,
         1.1352310e-04,  0.0000000e+00,  0.0000000e+00,  3.7430208e-03,
         0.0000000e+00, -2.5755540e-04, -2.6191911e-03,  4.1428395e-04,
        -1.1249429e-03, -1.6088597e-04,  0.0000000e+00,  0.0000000e+00,
        -5.5367802e-04,  0.0000000e+00, -1.2641185e-03, -6.0173990e-03,
        -2.6591588e-04,  3.6250763e-03,  1.7996349e-03, -2.4923831e-03,
        -4.4576675e-03,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        -3.1630003e-03, -9.6352748e-04, -9.6931035e-05, -4.6420656e-04,
         1.3182389e-02,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
         5.5439761e-03, -1.0579462e-03, -2.4395608e-03,  1.2268131e-03,
        -2.4303151e-03,  2.5188178e-04, -7.1195367e-04, -1.4863629e-03,
         2.0157322e-03,  0.0000000e+00]], dtype=float32)>, <tf.Tensor: shape=(50,), dtype=float32, numpy=
array([ 0.0000000e+00,  1.8945287e-03,  0.0000000e+00,  3.3356389e-03,
        7.5615030e-03,  4.2673951e-04,  0.0000000e+00,  0.0000000e+00,
        8.3703129e-04,  0.0000000e+00,  0.0000000e+00, -3.8044062e-03,
        0.0000000e+00, -1.9393768e-04, -1.2252710e-03,  1.7499779e-03,
        4.9770740e-04, -1.2155832e-04,  0.0000000e+00,  0.0000000e+00,
        2.4958298e-04,  0.0000000e+00, -7.5081683e-04,  8.1827259e-04,
       -1.9676052e-04, -1.8924754e-03,  7.1774086e-04,  1.6761324e-03,
        1.2559472e-03,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        3.3260928e-04, -8.4418792e-04, -2.1296809e-04, -3.3286596e-03,
        7.3421681e-03,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        1.8625467e-03,  1.9900850e-04,  1.5414150e-03, -2.2390606e-03,
       -2.8248876e-05,  1.8789712e-04, -2.4115024e-04,  4.2368062e-03,
       -1.6856943e-03,  0.0000000e+00], dtype=float32)>, <tf.Tensor: shape=(50, 50), dtype=float32, numpy=
array([[ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00, ...,
         0.0000000e+00,  0.0000000e+00,  0.0000000e+00],
       [-6.1806047e-04,  0.0000000e+00,  1.9905510e-04, ...,
         1.5742856e-03,  3.0515471e-04, -1.2499602e-04],
       [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00, ...,
         0.0000000e+00,  0.0000000e+00,  0.0000000e+00],
       ...,
       [-7.4027432e-04,  0.0000000e+00,  2.3841579e-04, ...,
         1.8425728e-03,  1.5021765e-03, -1.4629765e-04],
       [ 7.9818245e-05,  0.0000000e+00, -2.5706562e-05, ...,
        -1.4596945e-04,  1.8621939e-03,  1.1589786e-05],
       [ 0.0000000e+00,  0.0000000e+00,  0.0000000e+00, ...,
         0.0000000e+00,  0.0000000e+00,  0.0000000e+00]], dtype=float32)>, <tf.Tensor: shape=(50,), dtype=float32, numpy=
array([-1.0037005e-03,  0.0000000e+00,  3.2325607e-04, -2.6654673e-04,
        4.3113623e-03,  0.0000000e+00, -3.5289438e-03,  4.0892605e-04,
        0.0000000e+00, -1.2905553e-03,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  3.0032881e-03,  0.0000000e+00,
        0.0000000e+00,  2.6593888e-03,  0.0000000e+00, -1.3680644e-03,
        4.9801642e-04,  1.9348750e-03,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  2.4105171e-03,
        0.0000000e+00,  2.8948522e-05,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  3.5973138e-04, -2.5844993e-03,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  2.7459664e-03,  5.5970304e-04,
        3.5225567e-03,  0.0000000e+00, -1.5520444e-04,  0.0000000e+00,
        4.2449869e-04,  0.0000000e+00,  0.0000000e+00,  2.8726207e-03,
        5.2809883e-03, -2.2808195e-04], dtype=float32)>, <tf.Tensor: shape=(50, 1), dtype=float32, numpy=
array([[-8.12295638e-03],
       [ 0.00000000e+00],
       [ 5.14636282e-03],
       [ 3.43606854e-03],
       [-9.62987635e-03],
       [ 0.00000000e+00],
       [-4.21041809e-03],
       [-4.95478162e-05],
       [ 0.00000000e+00],
       [ 2.35846266e-03],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [-1.21107325e-02],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [-1.08039416e-02],
       [ 0.00000000e+00],
       [-7.22747389e-03],
       [ 7.37232715e-03],
       [-5.13390545e-03],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [-1.05799362e-02],
       [ 0.00000000e+00],
       [ 3.31182592e-03],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [ 3.64641682e-03],
       [ 3.63867916e-03],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [-3.51302186e-03],
       [ 1.20910201e-02],
       [ 5.83436899e-03],
       [ 0.00000000e+00],
       [ 8.13384540e-05],
       [ 0.00000000e+00],
       [ 3.94021161e-04],
       [ 0.00000000e+00],
       [ 0.00000000e+00],
       [-3.22613074e-03],
       [ 3.65258195e-04],
       [-9.12425108e-03]], dtype=float32)>, <tf.Tensor: shape=(1,), dtype=float32, numpy=array([-0.01388596], dtype=float32)>]

Optimizer¶

We now define an optimizer, essentially an improved version of SGD

learning_rate = 0.0003
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

Sampling data exogenously¶

we make a function to generate training data. Here we just sample exogenously from an interval.

def get_training_data(z_lb, z_ub, k_lb, k_ub, n_data):
    
    Z = tf.random.uniform(
        shape = [n_data, 1],
        minval=z_lb,
        maxval=z_ub,
        dtype=tf.dtypes.float32)
    
    
    K = tf.random.uniform(
        shape = [n_data, 1],
        minval=k_lb,
        maxval=k_ub,
        dtype=tf.dtypes.float32)
    
    X = tf.concat([Z, K], axis = 1)
    return X

Training¶

We iterataively generate training data and update the neural network

# Keep results for plotting
train_loss = []

num_episodes = NUM_EPISODES_UNIFORM
n_data_per_epi = 128
z_lb = 0.7
z_ub = 1.3
k_lb = 0.9
k_ub = 12.0


for ep in range(num_episodes):
    # generate training data
    X = get_training_data(z_lb, z_ub, k_lb, k_ub, n_data_per_epi)

    # compute loss and gradients
    loss, grads = grad(X, nn)
    
    # apply gradients
    optimizer.apply_gradients(zip(grads, nn.trainable_variables))
    
    # record loss
    train_loss.append(loss.numpy())
    
    # print progress
    if ep % int(0.05 * num_episodes) == 0:
        print("#=================================================================")
        print("episode = {}, loss [log10] = {}".format(ep, np.log10(loss.numpy())))
        
    if ep % int(0.2 * num_episodes) == 0 or ep == num_episodes - 1:
        cost, errREE, C_t, K_tplus1, r_t, LHS, RHS = compute_cost(X, nn)
        
        plt.title("loss function")
        plt.plot(np.log10(np.array(train_loss)))
        plt.xlabel("Training Episode")
        plt.ylabel("loss [log10]")
        plt.show()
        plt.close()

        plt.title("policy")
        plt.xlabel("K")
        plt.ylabel("Knext")
        plt.scatter(X[:, 1], X[:, 1], label = "diagonal")
        plt.scatter(X[:, 1], K_tplus1[:, 0], label = "Knext")
        plt.legend()
        plt.show()
        plt.close()

        plt.title("policy")
        plt.xlabel("Z")
        plt.ylabel("Knext")
        plt.scatter(X[:, 0], K_tplus1[:, 0], label = "Knext")
        plt.legend()
        plt.show()
        plt.close()

        plt.title("consumption policy")
        plt.xlabel("K")
        plt.ylabel("cons")
        plt.scatter(X[:, 1], C_t[:, 0], label="C_t")
        plt.legend()
        plt.show()
        plt.close()

        plt.title("consumption policy")
        plt.xlabel("Z")
        plt.ylabel("cons")
        plt.scatter(X[:, 0], C_t[:, 0], label="C_t")
        plt.legend()
        plt.show()
        plt.close()

        plt.title("Rel Ee")
        plt.xlabel("K")
        plt.ylabel("Rel Ee")
        plt.scatter(X[:, 1], errREE[:, 0], label="REE")
        plt.show()
        plt.close()
       
        plt.title("Rel Ee")
        plt.xlabel("Z")
        plt.ylabel("Rel Ee")
        plt.scatter(X[:, 0], errREE[:, 0], label="REE")
        plt.show()
        plt.close()
        
        plt.xlabel("K")
        plt.scatter(X[:, 1], LHS[:, 0], s = 100, label="LHS Ee")
        plt.scatter(X[:, 1], RHS[:, 0], s = 20, label="RHS Ee")
        plt.legend()
        plt.show()
        plt.close()
        
        plt.xlabel("Z")
        plt.scatter(X[:, 0], LHS[:, 0], s = 100, label="LHS Ee")
        plt.scatter(X[:, 0], RHS[:, 0], s = 20, label="RHS Ee")
        plt.legend()
        plt.show()
        plt.close()

#=================================================================
episode = 0, loss [log10] = -2.7412846088409424

#=================================================================
episode = 250, loss [log10] = -3.486482620239258

#=================================================================
episode = 500, loss [log10] = -3.5200955867767334

#=================================================================
episode = 750, loss [log10] = -3.811296224594116

#=================================================================
episode = 1000, loss [log10] = -4.103127479553223

#=================================================================
episode = 1250, loss [log10] = -4.120211601257324

#=================================================================
episode = 1500, loss [log10] = -4.407349109649658

#=================================================================
episode = 1750, loss [log10] = -4.387524127960205

#=================================================================
episode = 2000, loss [log10] = -4.614704132080078

#=================================================================
episode = 2250, loss [log10] = -4.779208183288574

#=================================================================
episode = 2500, loss [log10] = -4.790186882019043

#=================================================================
episode = 2750, loss [log10] = -5.060757160186768

#=================================================================
episode = 3000, loss [log10] = -5.168338298797607

#=================================================================
episode = 3250, loss [log10] = -4.982332229614258

#=================================================================
episode = 3500, loss [log10] = -5.116687297821045

#=================================================================
episode = 3750, loss [log10] = -5.556701183319092

#=================================================================
episode = 4000, loss [log10] = -5.446951866149902

#=================================================================
episode = 4250, loss [log10] = -5.481368541717529

#=================================================================
episode = 4500, loss [log10] = -5.223493576049805

#=================================================================
episode = 4750, loss [log10] = -5.6075053215026855

#=================================================================
episode = 5000, loss [log10] = -5.604979515075684

Simulating the model from the policy¶

Given a policy function, we can simulate the model forward. Say we start out with a state $\mathbf{X}_t=[z_t, K_t]$ .
We can use the neural network to obtain $K_{t+1} = Y_t \cdot \mathcal{N}(\mathbf{X}_t)$ .
Then, we can draw a random innovation $\epsilon_t\sim N(0, 1)$ using a pseudo-random number generator and obtain $\log Z_{t+1}=\varrho \log Z_t + \sigma \epsilon_t$ (so $Z_{t+1}=\exp(\varrho\log Z_t+\sigma\epsilon_t)$ ).
Now we have $\mathbf{X}_{t+1}=[Z_{t+1}, K_{t+1}]$ .
We can repeat the same procedure to obtain $\mathbf{X}_{t+2}$ and so on and so forth.
Next, we will implement a function that takes a batch of states, a batch of innovations, and the neural network and then simulates the states one period forward based on the policy encoded by the neural network.
Afterward, we will write a second function to simulate longer sequences.

@tf.function
def simulate_single_step(X_t, eps_tplus1, nn):
    # function simulates the economy one step forward based on the neural network policy and the innovation
    
    n_data = X_t.shape[0] # number of states is on the axis 0
    dim_state = X_t.shape[1] # dimensionality of the state is on axis 1
    
    # read out the state
    Z_t = X_t[:, 0 : 1]
    K_t = X_t[:, 1 : 2]
    
    # compute output today
    Y_t = Z_t * K_t ** alpha
    
    # use the neural network to predict the savings rate
    s_t = nn(X_t)
    
    # get the implied capital in the next period
    K_tplus1 = (1. - delta) * K_t + Y_t * s_t
    
    # get tfp in the next period
    Z_tplus1 = tf.exp(rho_tfp * tf.math.log(Z_t) + sigma_tfp * eps_tplus1)    
    
    # construct the next step
    X_tplus1 = tf.concat([Z_tplus1, K_tplus1],axis = 1)
    
    return X_tplus1

def sim_periods(X_start, nn, num_periods):
    n_tracks = X_start.shape[0] # number of states is on the axis 0
    dim_state = X_start.shape[1] # dimensionality of the state is on axis 1
    
    # create an empty array to store the states
    X_simulation = np.empty((num_periods, n_tracks, dim_state)) # 0 axis: time period, 1 axis: the different trajectories, 2 axis: the different state variables
    
    # draw random innovation
    eps = tf.random.normal((num_periods, n_tracks), dtype = tf.float32)
    
    # set starting state
    X_simulation[0, :, :] = X_start
    
    X_old = X_start
    
    # simulate the periods
    for t in range(1, num_periods):
        eps_use = eps[t, :, tf.newaxis] # newaxis makes sure the shape is n_tracks x 1
        
        X_new = simulate_single_step(X_old, eps_use, nn)
        
        X_simulation[t, :, :] = X_new
        
        X_old = X_new
        
    return X_simulation

# let's pick a starting state
X_start = np.array([[0.5 * (z_ub + z_lb), 0.5 * (k_ub + k_lb)]], dtype = np.float32)

# and simulate it some periods forward using the neural network
num_periods = 200 
X_simulation = sim_periods(X_start, nn, num_periods)

plt.plot(X_simulation[:, 0, 0])
plt.xlabel("t")
plt.ylabel("Z_t")
plt.show()

plt.plot(X_simulation[:, 0, 1])
plt.xlabel("t")
plt.ylabel("K_t")
plt.show()

Like often in macro, the two variables seem to be moving together, so let’s try to get a sense of the ergodic distribution of the state space.

num_periods = 10000

# we simulate more periods
X_simulation = sim_periods(X_start, nn, num_periods)

# we make a scatter plot for the two state-variables
plt.scatter(X_simulation[:, 0, 0], X_simulation[:, 0, 1], alpha = 0.1)
plt.xlabel("Z_t")
plt.ylabel("K_t")
plt.show()

As we can see the model essentially lives on a cloud around the diagonal.
The model never reaches states with extremely low productivity and extremely high capital or states with extremely high productivity and extremely low capital.
This means all the effort we spent training the neural network in these areas (top left and bottom right) was essentially in vain.
With more than two variables, this issues becomes exponentially more extreme (see Maliar et al. (2011) for more information on this topic).
Moreover, we may not even know where the cloud will be, since it depends on the very policy we want to solve for!
Hence, it would be very good if we could simultaneously learn the policy on the cloud of interest while also moving the cloud as we update the policies.
Azinovic et al. (2022) address this issue by iterating between simulating new states from the neural network policy and then training the neural network on those simulated states. This is what we will do next.

Iterating between training an simulation¶

What changes compared to our previous approach? Only how we sample the states!

Instead of drawing exogenously from a given interval, we simulate the model using the policy encoded in the neural network. Since neural networks can be very efficiently executed on batches, and to increase independence of the data, we simulate more than one track in parallel, but each for a shorter period of times. For example, we can get 1000 simulated states by either simulating one track for 1000 periods of 100 tracks for 10 periods, the latter will be much faster.

def get_training_data_simulation(X_start, nn, n_periods):
    n_tracks = X_start.shape[0]
    n_dim = X_start.shape[1]
    
    # we simulate the model
    X_simulation = sim_periods(X_start, nn, n_periods)
    
    # we read out the last state (so that we can use it as starting poit for the next simulation)
    X_end = np.float32(X_simulation[-1, :, :])
    
    # we reshape the data into (n_tracks * n_periods) x 2 array 
    X_training = np.float32(np.reshape(X_simulation, (n_tracks * n_periods, n_dim)))
    
    return X_training, X_end

# let's try
# let's get a new neural network (re-seed for reproducibility, SEED from cell 2)
tf.random.set_seed(SEED)
nn = keras.Sequential([
    keras.layers.Dense(num_hidden1, activation='relu', input_shape=(num_input,)),
    keras.layers.Dense(num_hidden2, activation='relu'),
    keras.layers.Dense(num_output, activation='sigmoid')
])


n_tracks = 50
n_periods = 3
print("n_tracks = ", n_tracks)
print("n_periods = ", n_periods)

# start from a random state
X_start = get_training_data(z_lb, z_ub, k_lb, k_ub, n_tracks)
print("X_start.shape = ", X_start.shape)

print("With every simulation we get", n_tracks * n_periods, "new states")

X_training, X_end = get_training_data_simulation(X_start, nn, n_periods)
print("X_training.shape = ", X_training.shape)
print("X_end.shape = ", X_end.shape)

learning_rate = 0.0001
print("learning_rate = ", learning_rate)
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

print(compute_cost(X_start, nn)[0])

n_tracks =  50
n_periods =  3
X_start.shape =  (50, 2)
With every simulation we get 150 new states
X_training.shape =  (150, 2)
X_end.shape =  (50, 2)
learning_rate =  0.0001

tf.Tensor(0.00054157386, shape=(), dtype=float32)

The only thing that changes in the training loop is how we sample the data.

# Keep results for plotting
train_loss = []

num_episodes = NUM_EPISODES_SIM

# initialize X_start
X_start = X_end
print("n_tracks = ", X_start.shape[0])
print("n_periods = ", n_periods)

for ep in range(num_episodes):
    # generate training data, now by sampling
    X, X_end = get_training_data_simulation(X_start, nn, n_periods)

    # update X_start
    X_start = X_end
    
    # compute loss and gradients
    loss, grads = grad(X, nn)

    # apply gradients
    optimizer.apply_gradients(zip(grads, nn.trainable_variables))
    
    # record loss
    train_loss.append(loss.numpy())
    
    # print progress
    if ep % int(0.05 * num_episodes) == 0:
        print("#=================================================================")
        print("episode = {}, loss [log10] = {}".format(ep, np.log10(loss.numpy())))
        
    if ep % int(0.2 * num_episodes) == 0 or ep == num_episodes - 1:
        cost, errREE, C_t, K_tplus1, r_t, LHS, RHS = compute_cost(X, nn)
        
        plt.title("loss function")
        plt.plot(np.log10(np.array(train_loss)))
        plt.xlabel("Training Episode")
        plt.ylabel("loss [log10]")
        plt.show()
        plt.close()
        
        plt.title("simulated ergodic set")
        plt.scatter(X[:, 0], X[:, 1])
        plt.xlabel("Z")
        plt.ylabel("K")
        plt.show()
        plt.close()
        
        
        plt.title("policy")
        plt.xlabel("K")
        plt.ylabel("Knext")
        plt.scatter(X[:, 1], X[:, 1], label = "diagonal")
        plt.scatter(X[:, 1], K_tplus1[:, 0], label = "Knext")
        plt.legend()
        plt.show()
        plt.close()

        plt.title("policy")
        plt.xlabel("Z")
        plt.ylabel("Knext")
        plt.scatter(X[:, 0], K_tplus1[:, 0], label = "Knext")
        plt.legend()
        plt.show()
        plt.close()

        plt.title("consumption policy")
        plt.xlabel("K")
        plt.ylabel("cons")
        plt.scatter(X[:, 1], C_t[:, 0], label="C_t")
        plt.legend()
        plt.show()
        plt.close()

        plt.title("consumption policy")
        plt.xlabel("Z")
        plt.ylabel("cons")
        plt.scatter(X[:, 0], C_t[:, 0], label="C_t")
        plt.legend()
        plt.show()
        plt.close()

        plt.title("Rel Ee")
        plt.xlabel("K")
        plt.ylabel("Rel Ee")
        plt.scatter(X[:, 1], errREE[:, 0], label="REE")
        plt.show()
        plt.close()
       
        plt.title("Rel Ee")
        plt.xlabel("Z")
        plt.ylabel("Rel Ee")
        plt.scatter(X[:, 0], errREE[:, 0], label="REE")
        plt.show()
        plt.close()
        
        plt.xlabel("K")
        plt.scatter(X[:, 1], LHS[:, 0], s = 100, label="LHS Ee")
        plt.scatter(X[:, 1], RHS[:, 0], s = 20, label="RHS Ee")
        plt.legend()
        plt.show()
        plt.close()
        
        plt.xlabel("Z")
        plt.scatter(X[:, 0], LHS[:, 0], s = 100, label="LHS Ee")
        plt.scatter(X[:, 0], RHS[:, 0], s = 20, label="RHS Ee")
        plt.legend()
        plt.show()
        plt.close()

n_tracks =  50
n_periods =  3

#=================================================================
episode = 0, loss [log10] = -3.3395793437957764

#=================================================================
episode = 100, loss [log10] = -3.9759445190429688

#=================================================================
episode = 200, loss [log10] = -3.8828840255737305

#=================================================================
episode = 300, loss [log10] = -4.040897369384766

#=================================================================
episode = 400, loss [log10] = -4.162646770477295

#=================================================================
episode = 500, loss [log10] = -4.094448089599609

#=================================================================
episode = 600, loss [log10] = -4.180356502532959

#=================================================================
episode = 700, loss [log10] = -4.178821563720703

#=================================================================
episode = 800, loss [log10] = -4.199854373931885

#=================================================================
episode = 900, loss [log10] = -4.253724575042725

#=================================================================
episode = 1000, loss [log10] = -4.251932621002197

#=================================================================
episode = 1100, loss [log10] = -4.2519121170043945

#=================================================================
episode = 1200, loss [log10] = -4.539088726043701

#=================================================================
episode = 1300, loss [log10] = -4.39324426651001

#=================================================================
episode = 1400, loss [log10] = -4.528472423553467

#=================================================================
episode = 1500, loss [log10] = -4.536249160766602

#=================================================================
episode = 1600, loss [log10] = -4.417702674865723

#=================================================================
episode = 1700, loss [log10] = -4.449241638183594

#=================================================================
episode = 1800, loss [log10] = -4.467284202575684

#=================================================================
episode = 1900, loss [log10] = -4.707545280456543

#=================================================================
episode = 2000, loss [log10] = -4.752475261688232

Final but important remark on simulation based methods: Without a question, simulation based methods offer a huge advantage, exponentially so in high dimensions. However, they also introduce some fragility into the learning process, since the distribution of the training data changes when the policy changes. On the one hand, this is exactly what we want, but on the other hand, it can make training unstable when moving too quikly to unseen data. Hence some parameters may have to be optimized more carefully, especially the learning rate, so that the data distribution is not changing to quickly. Depending on the model, we may also need to take care of the fact that in the beginning of training the neural network policies are random and may hence predict infeasible states (like negative aggregate capital), which make useful learning impossible. Azinovic et al. (2022) address some of these points in the appendices and Azinovic and Žemlička (2023) introduce market-clearing neural network architectures, as well as a procedure to introduce multiple assets, which stabilized training.