A Comparative Tutorial of Bayesian Sequential Design and Reinforcement Learning

Constrained backward induction is an algorithm consisting of three simple steps. The first step is forward simulation. Our implementation here uses the assumption that the sequential nature is limited to sequential stopping, so trajectories can be generated assuming no stopping independently from decisions. Throughout we use $D_{t}=0$ to denote continuation. Other actions, $D_{t}\neq 0$, indicate stopping the study and choice of a terminal decision. The second step is constrained backward induction, which implements (6) using decisions restricted to depend on the history $H_{t}$ indirectly only through $S_{t}$. The third step simply keeps track of the best action and iterates until convergence. We first briefly explain these steps and then provide an illustration of their application in Example 1 and additional implementation considerations.

Step 1.

Forward simulation: Simulate many trajectories, say $M$, until some maximum number of steps $T_{\text{max}}$ (e.g. cohorts in a trial). To do this, each $m=1,\ldots,M$ corresponds to a different prior draw $\theta^{(m)}\stackrel{{\scriptstyle\mathclap{\tiny\mbox{iid}}}}{{\sim}}p(\theta^{(m)})$ and samples $Y^{(m)}_{t}\stackrel{{\scriptstyle\mathclap{\tiny\mbox{iid}}}}{{\sim}}p(Y^{(m)}_{t}\mid\theta^{(m)})$, $t=1,\ldots,T_{\text{max}}$. For each $m$ and $t$, we evaluate and record the summary statistic $S^{(m)}_{t}$ discretized over a grid.

Step 2.

Backward induction: For each possible decision $d$ and each grid value $S=j$, the algorithm approximates $\widehat{U}(S,d)\approx U(S,d)$ using the forward simulation and Bellman equation as follows. Denote with $A_{j}=\{(m,t_{m})\mid S^{(m)}_{t_{m}}=j\}$ the set of forward simulations that fall within grid cell $j$. Then,

$$\widehat{U}(S=j,d)=\begin{cases}\frac{1}{|A_{j}|}\sum_{(m,t_{m})\in A_{j}}\widehat{U}(S^{(m)}_{t_{m}+1},D^{\star}(S^{(m)}_{t_{m}+1}))&d=0\\ \frac{1}{|A_{j}|}\sum_{(m,t_{m})\in A_{j}}u(S_{t_{m}}^{(m)},D_{t_{m}}=d,\theta^{(m)})&d\neq 0.\end{cases}$$

(7)

The evaluation under $d=0$ requires the optimal actions $D^{\star}_{t+1}(S_{t_{m}+1})$. We use an initial guess (see below), which is then iteratively updated (see next).

Step 3.

Iteration: Update the table $D^{\star}(S)=\operatorname*{arg\,max}_{d}\,\widehat{U}(S,d)$ after step 2.

Repeat steps 2 and 3 until updating in step 3 requires no (or below a minimum number of) changes of the tabulated $D^{\star}(S)$. \cbend

Figure 3(a) shows the estimated utility function $\widehat{U}(S,d)$ in Example 1 using $M=1000$, $T_{\text{max}}=50$ and 100 grid values for the running average $p_{t}=\sum_{k=1}^{t}Y_{t}/t$ in $S=(t,p_{t})$. Optimal actions $D^{\star}_{t}(S)$ are shown in Figure 3(b). The numerical uncertainty due to the Monte Carlo evaluation of the expectations is visible. If desired, one could reduce it by appropriate smoothing (MacDonald et al., 2015). One can verify, however, that the estimates are a close approximation to the analytic solution which is available in this case (Müller et al., 2007).

We explain Step 2 by example. Consider Figure 3(b) and assume, for example, that we need the posterior expected utility of $S=(t,p_{t})=(20,0.25)$. In this stylized representation, only three simulations, $A=\{m_{1},m_{2},m_{3}\}$ pass through this grid cell. In this case, $t_{m}=t=20$ for the three trajectories since $t$ is part of the summary statistic. We evaluate $\widehat{U}(S,d=1)$ and $\widehat{U}(S,d=2)$ as averages $\frac{1}{3}\sum_{i\in A}u(S_{20}^{(m)},d,\theta^{(m)})$. For $\widehat{U}(S,d=0)$, we first determine the grid cells in the next period for each of the three trajectories. $S^{(m)}_{t_{m}+1}=(21,p^{(m)}_{21})$. We then look up the optimal decisions $D^{\star}(S_{21},p^{(m)}_{21})$ (using in this case $t_{m}+1=21$) and average $\frac{1}{3}\sum\widehat{U}(21,p^{(m)}_{21},D^{\star}(21,p^{(m)}_{21}))$, as in (7).

Constrained backward induction requires iterative updates of $D^{\star}(S)$ and their values. The procedure starts with arbitrary initial values for $D^{\star}(S)$, recorded on a grid over $S$. For example, a possible initialization is $D^{\star}(S)=\max_{d\neq 0}\widehat{U}(S,d)$, maximizing over all actions that do not involve continuation. With such initial values, $\widehat{U}(S,d)$ can be evaluated over the entire grid. Then, for updating the optimal actions $D^{\star}(S)$ one should best start from grid values that are associated with the time horizon $T$, or at least high $t$. This is particularly easy when $t$ is an explicit part of $S_{t}$, as in Example 1 with $S_{t}=(t,p_{t})$. Another typical example arises in Example 2 with $S_{t}=(\mu_{t},\sigma_{t})$, the mean and standard deviation of some quantity of interest. For large $t$ we expect small $\sigma_{t}$, making it advisable to start updating in each iteration with the grid cells corresponding to smallest $\sigma_{t}$. We iterate until no more (or few) changes happen.

The algorithm can be understood as an implementation of (6).

Consider $f(S,d)$ as an arbitrary function over pairs $(S,d)$ and the function operator ${\mathcal{P}}f$ defined as $({\mathcal{P}}f)(S,d)=\max_{d^{\prime}}\mathbb{E}[f(S^{\prime},d^{\prime})|S]$ where $S^{\prime}$ is the summary statistic resulting from sampling one more data point $Y$ from the unknown $\theta$ and recompute $S^{\prime}$ from $S$. Then Bellman equation (6) can be written as $U={\mathcal{P}}U$. In other words, expected utility under the optimal decision $D^{\star}$ is a fixed point of the operator ${\mathcal{P}}$. Constrained backward induction attempts to find an approximate solution to the fixed-point equation. The same principle motivates the Q-learning algorithm in RL (Watkins and Dayan, 1992) (see Section 4). Backward induction is also closely related to the value iteration algorithm for Markov decision processes (Sutton and Barto, 2018), which relies on exact knowledge of the state transition function.

Inspection of Figure 3(b) suggests an attractive alternative algorithm. Notice the decision boundaries on $S=(t,p_{t})$ that trace a funnel with an upper boundary $\omega_{2}(t)$ separating $D^{\star}=2$ from $D^{\star}=0$, and a lower boundary $\omega_{1}(t)$ separating $D^{\star}=0$ versus $D^{\star}=1$.

Recognizing such boundaries suggests an alternative approach based on searching for optimal boundaries in a suitable family $\{\omega_{\phi,1},\omega_{\phi,2}\mid\phi\in\Phi\}$.

This approach turns the sequential decision problem of finding optimal $D^{\star}(S)$ into a non-sequential problem of finding an optimal $\phi^{\star}\in\Phi$. This method is used, for example, in Rossell et al. (2007).

In Example 1, we could use

using a single tuning parameter $\phi\in(0,1)$. Both functions are linear in $\sqrt{t-1}$, and mimic the funnel shape seen in Figure 3(b). The decision rules implied by these boundaries is

Here, the additional subscript _ϕ in $D_{\phi}(\cdot)$ indicates that the decision follows the rule implied by decision boundaries $\omega_{j}(t;\phi)$. Note that $\omega_{1}(1)=1$ and $\omega_{2}(1)=0$, ensuring continuation at $t=1$.

For a given $\phi$, the forward simulations are used to evaluate expected utilities under the policy $D_{\phi}$. Let $T^{(m)}_{\phi}=\min\{t:\;D_{\phi}(S^{(m)}_{t}\neq 0\}$ denote the stopping time under $D_{\phi}$. Then the expected utility under policy $D_{\phi}$ is

where the expectation is with respect to data $Y_{T}$ and $\theta$. It is approximated as an average over all Monte Carlo simulations, stopping each simulation at $T^{(m)}_{\phi}$, as determined by the parametric decision boundaries,

Optimizing $\widehat{U}(\phi)$ w.r.t. $\phi$ we find the optimal decision boundaries $\phi^{\star}=\arg\max\widehat{U}(\phi)$. As long as the nature of the sequential decision is restricted to sequential stopping, the same set of Monte Carlo simulations can be used to evaluate all $\phi$, using different truncation to evaluate $\widehat{U}(\phi)$. In general, a separate set of forward simulations for each $\phi$, or other simplifying assumptions might be required.

Figure 3(e) shows the decision boundaries for the best parameter estimated at $\phi^{\star}=0.503$ in Example 1. The estimated values for $\widehat{U}(\phi)$ are in Figure 3(f). In Figure 3(e), the boundaries using constrained backward induction are overlaid in the image for comparison. The decision boundaries trace the optimal decisions under the backward induction well. The differences in expected utility close to the decision boundary are likely very small, leaving minor variations in the decision boundary negligible.

The same approach is applied to the (slightly more complex) Example 2.

Recall the form of the summary statistics $S=(s_{\delta},\bar{\delta})$, the posterior standard deviation and mean of the ED95 effect. We use the boundaries

parameterized by $\phi=(b_{1},b_{2},c)$. The implied decision rules are

The results are in Figure 4(a). Again, the sequential decision problem is reduced to the optimization problem of finding the optimal $\phi$ in (9). Since now $\phi\in\Re^{3}$ the evaluation of $\widehat{U}$ requires a 3-dimensional grid. To borrow strength from Monte Carlo evaluations of (9) across neighboring grid points for $\phi$ we proceed as follows. We evaluate $\widehat{U}(\phi)$ on a coarse $10\times 10\times 10$ grid, and then fit a quadratic response surface (as a function of $\phi$) to these Monte Carlo estimates. The optimal decision $\phi^{*}$ is the maximizer of the quadratic fit. We find $\phi^{*}=(b_{1}^{*},b_{2}^{*},c^{*})=(1.572,1.200,0.515)$ Instead of evaluating $\widehat{U}$ on a regular grid over $\phi$, one could alternatively select a random number of design points (in $\phi$).

The use of parametric boundaries is closely related to the notion of function approximation and the method of policy gradients in RL, which will be described next.

The Markov property for a decision process is defined by the conditions

That is, the next state $S_{t+1}$ and the reward $R_{t}$ depend on the history only indirectly through the current state and action. When the condition holds, the decision process is called an MDP. For MDPs, the optimal policy is only a function of the latest state $S_{t}$ and not of the entire history $H_{t}$ (Puterman, 2014). Many RL algorithms assume the Markov property. However, many sequential decision problems are more naturally characterized as partially observable MDP (POMDP) that satisfy the Markov property only conditional on some $\theta$ that is generated at the beginning of each episode. Such problems have

been studied in the RL literature under the name of Hidden-Parameter MDP (HiMDP) (Doshi-Velez and Konidaris, 2016).

There is a standard – Bayesian motivated – way to cast any POMDP as an MDP using so-called belief states (Cassandra et al., 1997). Belief states are obtained by including the posterior distribution of unobserved parameters as a part of the state. With a slight abuse of notation, we may write the belief states as $S_{t}=p(\theta\mid H_{t})$. While the belief state is, in general, a function, it can often be represented as a vector when the posterior admits a finite (sufficient) summary statistic. The reward distribution can also be written in terms of such belief states as

See Figure 5 for a graphical representation of an HiMDP and the resulting belief MDP.

We implement this approach for Example 1. The reward is chosen to match the definition of utility. It suffices to define it in terms of $\theta$ (and let the posterior take care of the rest, using (10)). We use

as in (3). Next we introduce the belief states. Recall the notation from Example 1. We have $p(\theta\mid S_{t})=\mbox{Bin}(p_{t},p_{t}(1-p_{t})/t)$. The summary statistic $S_{t}=(t,p_{t})$ is a two-dimensional representation of the belief state.

Considering Example 2, we note that the utility function (5) depends on the state $S_{t}$ only, and does not involve $\theta$. We define

While the reward is clearly Markovian, the transition probability is not necessarily Markovian. This is the case because in this example the posterior moments $S_{t}$ are not a sufficient statistic. In practice, however, a minor violation of the Markov assumption for the transition distribution does not seem to affect the ability to obtain good policies with standard RL techniques.

Q-learning (Watkins and Dayan, 1992; Clifton and Laber, 2020; Murphy, 2003) is an RL algorithm that is similar in spirit to the constrained backward induction described in Section 3.1. The starting point is Bellman optimality equation for MDPs (Bellman, 1966). Equation (6) for the optimal $D^{\star}$ and written for MDPs becomes

where the expectation is with respect to $R_{t}$ and $S_{t+1}$. The optimal policy is implicitly defined as the solution to (13).

Q-learning proceeds iteratively following the fixed-point iteration principle. Let $Q^{(k)}$ be some approximation of $Q^{\star}$. We assume that a set of simulated transitions $\{(s_{t},d_{t},r_{t},s_{t+1})\}_{t=1}^{n}$ is available. This collection is used like the forward simulations in the earlier discussion. In RL it is known as the “experience replay buffer”, and can be generated using any stochastic policy. And suppose, for the moment, that the state and action spaces are finite discrete, allowing to record $Q^{(k)}$ in a table. Q-learning is defined by updating $Q^{(k)}$ as

Note the moving average nature of the update. The procedure iterates until convergence from a stream of transitions.

Deep Q-networks (DQN) (Mnih et al., 2013) are an extension of Q-learning for continuous states. A neural network is used to represent $Q(\cdot)$. Let $\phi^{(k)}$ denote the parameters of the neural network at iteration $k$. Using simulated transitions from the buffer

DQN performs updates

In practice, exact minimization is replaced by a gradient step from mini-batches, together with numerous implementation tricks (Mnih et al., 2013).

We implemented DQN in Example 1 using the Python package Stable-Baselines3 (Raffin et al., 2021). The experience replay buffer is continuously updated. The algorithm uses a random policy to produce an initial buffer and then adds experience from an $\epsilon$-greedy policy, where the current best guess for the optimal policy is chosen with probability $(1-\epsilon)$, and otherwise fully random actions are chosen with probability $\epsilon$. Figure 3(c) shows $\hat{Q}$, the estimate of $Q^{\star}$, for each state and action $d\in\{0,1,2\}$. Figure 3(d) shows the corresponding optimal actions. Overall, the results are similar to the results with constrained backward induction, but much smoother. Also, notice that the solution under DQN is usually better in terms of expected utility as shown in Figure 3(f), even in (out-of-sample) evaluation episodes. This is likely due to to the flexibility and high-dimensional nature of the neural network approximation.

The better performance of RL comes at a price. First is sample efficiency (the number of simulations required by the algorithm to yield a good policy). The best $\hat{Q}$ is obtained after 2 million sampled transitions. Data efficiency is a known problem in DQN, and in RL in general (Yu, 2018). In many real applications investigators cannot afford such a high number of simulation steps. Another limitation is training instability. In particular, Figure 3(f) illustrates a phenomenon known as catastrophic forgetting, which happens when additional training decreases the ability of the agent to perform a previously learned task (Atkinson et al., 2021). This can happen because of the instability that arises from a typical strategy of evaluating performance periodically and keeping track of the best performing policy only. Several improvements over basic DQN have been proposed, with improved performance and efficiency (Hessel et al., 2018).

The approach is similar to the use of parametric boundaries discussed before. Policy gradient (PG) approaches start from a parameterization $\pi_{\bm{\phi}}$ of the policy. Again, consider a neural network (NN) with weights $\phi$. The goal of a PG method is to maximize the objective

This objective is the analogue to maximizing $U(\phi)$ in (9), except that here the stochastic policy $\pi_{\phi}$ is a probability distribution on decisions $S_{t}$. The main characteristic of PG methods is the use of gradient descent to solve (16).

The evaluation of gradients is based on the PG theorem (Sutton et al., 1999),

where the total return $G$ is a function $G(\tau)$ of the entire trajectory $\tau=(S_{1},D_{1},\ldots,S_{T},D_{T})$. Using gradients, PG methods can optimize over high-dimensional parameter spaces like in neural networks. In practice, estimates of the gradient are known to have huge variance, affecting the optimization. But several implementation tricks exist that improve the stability and reduce the variance. Proximal policy optimization (PPO) (Schulman et al., 2017) incorporates many of these tricks, and is widely used as a default for PG-based methods.

The PG theorem is essentially Leibniz rule for the gradient of $J(\phi)$.

With a slight abuse of notation, write $\pi_{\bm{\phi}}(\tau)$ for the distribution of $\tau$ induced by $\pi_{\bm{\phi}}$ for the sequential decisions. Then Leibniz rule for the gradient of the integral gives

where the expectations are with respect to the (stochastic) policy $\pi_{\phi}$ over $\tau$. The log probability in the last expression can be written as a sum of log probabilities, yielding (17).

PPO is implemented in Example 2 using Stable-Baselines3 (Raffin et al., 2021). The results are shown in Figure 4(b). Not surprisingly, the results are similar to those obtained earlier using parametric decision boundaries. Interestingly, the figure shows that neural networks do not necessarily extrapolate well to regions with low data. This behavior is noticeable on the lower left corner of the figure, where there could be data, but where it is never observed in practice because of the early stopping implied by the boundaries.