Reinforcement learning in clinical medicine: a method to optimize dynamic treatment regime over time

Abstract

Precision medicine requires individualized treatment regime for subjects with different clinical characteristics. Machine learning methods have witnessed rapid progress in recent years, which can be employed to make individualized treatment regime in clinical practice. The idea of reinforcement learning method is to take action in response to the changing environment. In clinical medicine, this idea can be used to assign optimal regime to patients with distinct characteristics. In the field of statistics, reinforcement learning has been widely investigated, aiming to identify an optimal dynamic treatment regime (DTR). Q-learning is among the earliest methods to identify optimal DTR, which fits linear outcome models in a recursive manner. The advantage is its easy interpretation and can be performed in most statistical software. However, it suffers from the risk of misspecification of the linear model. More recently, some other methods not so heavily depend on model specification have been developed such as inverse probability weighted estimator and augmented inverse probability weighted estimator. This review introduces the basic ideas of these methods and shows how to perform the learning algorithm within R environment.

Introduction

Precision medicine is an important concept in the modern era, which dictates that the treatment regime should be individualized depending on genetic background, demographics and clinical characteristics. Many clinical conditions can be extremely heterogeneous that the treatment regime is highly variable. For example, sepsis is a heterogeneous syndrome which is characterized by different clinical outcomes and responses to fluid resuscitation (1). Treatment regimes should be individualized for different subphenotypes. Even for a certain subphenotype, the treatment regime can evolve over time. Thus, individualized dynamic treatment regime should be adopted to achieve the best potential outcome. Reinforcement learning (RL) is a powerful method in computer science that has been successfully used to teach an agent to learn to interact with the environment, aiming to achieve the best reward/return. This is quite similar to the situation in clinical practice, where a doctor needs to adopt an action (treatment regime) depending on a patient’s conditions (e.g., demographics, genetic background, clinical characteristics and previous response to a specific treatment regime). Thus, reinforcement learning can be used to solve a clinical decision problem, whereby the concept of precision medicine can be realized. In this review article, we will introduce (I) the concept of reinforcement learning, (II) how this concept can be adopted to clinical research, and (III) how to perform RL using R language.

Basic ideas behind reinforcement learning

The key elements of a RL system include a policy, a reward signal, a value function and a model of the environment (Figure 1). A policy is usually a function mapping from a state to an action. A policy defines how an agent takes action in response to the environment. A reward defines the goal of a RL problem. Every time (t) when an agent takes an action (a_t) on the environment, the environment sends a reward signal (r_t+1) to the agent. The goal of the agent is to maximize the reward by taking an appropriate action. While the reward signal indicates the immediate benefit after taking an action, the value function defines the long-term benefit. The value of a state is the total amount of reward an agent can expect to accumulate over time, starting from that state. For example, when a doctor resuscitates a patient with septic shock, the action space comprises vasoactive agents and fluid infusion, and the immediate reward can be a stabilized blood pressure, or a normalized serum lactate (2). Since the policy is a mapping from the state to an action, the doctor will make treatment decision depending on the environment state of the patient. The mortality outcome at hospital discharge can be considered as a final reward. The state space of septic shock can be any combinations of hemodynamics, electrolytes, serum lactate and demographics. However, this can be an extremely large space due to the curse of dimensionality, and dimension reduction is usually required (3).

Figure 1.

Schematic illustration of reinforcement learning system.

However, several key assumptions that are commonly used in RL in computer science may not be applicable in clinical trials. First, many RL algorithms such as dynamic programming requires that the complete knowledge of the environment is known (based on physical laws and domain knowledges), which however is usually unsatisfied in clinical data. In the septic shock example, the underlying mechanisms behind hospital mortality cannot be represented by a deterministic equation and must be estimated from samples. Second, the data collection can be cheap and easy to perform in some RL examples such as the blackjack, in which tens of thousands of episodes of the game can be simulated (i.e., the data generating mechanism is known). This is not true in clinical setting because the sample size is usually limited and collecting data is expensive (e.g., also for the confidential issues). Of note, the state space and action space can be larger than the sample size in most situations, and thus parametric models are usually required to model the data. Third, the Markov property is assumed to be true in computer science, which is defined that the probability of each possible value for S_t and R_t depends only on the immediately preceding state (S_t−1) and action (A_t−1), but not on earlier states and actions. However, there is no such assumption from the pathophysiology that earlier conditions have no causal association with final clinical outcome. Thus, the notion of history must be introduced in solving clinical problems. The history at a given stage includes all the present and past states, plus the present action. As a result, the history space grows exponentially with the number of stages (4).

Working example

A simulated dataset is created for the illustration purpose.

As noted above, a data frame is generated with two stages. On stage 1, there are three covariates denoted as x11, x12 and x13; and treatment was A1. On stage 2, the three covariates are x21, x22 and x23, and the treatment is A2. Finally, we observe an outcome which is a continuous variable. Higher values indicate a better outcome. For stage 1, the optimal decision rule dictates that A1 should take value 1 for subjects with x11 >-0.54 and x12 <0.54. However, clinicians choose A1 value by a logit function (the propensity function) in clinical practice. Thus, the A1 used in real world practice may not equal the value adopted by the optimal rule. The purpose of RL is to learn the optimal rule to maximize the outcome Y. In the same vein, the optimal decision rule for A2 is to assign A2=1 to subjects with x21 >0.3 and x23 <0.46. However, the exact form of the outcome function is not easy to specify but can be approximated.

Y=2+0.25·x11+0.25·x12-0.25·x13−0.5·[A1-I(x11 > -0.54)×I(x12 < 0.54)]²−[A2-I(x21 > 0.3)×I(x23 < 0.46)]²

There is an indicator function I(·) which returns 1 when the express is true, and 0 otherwise. The quadratic function suggests that the outcome Y can only be maximized when A1 and A2 is equal to the value determined by the optimal decision rule in both stage 1 and 2.

Q-learning algorithm

Q-learning is a temporal difference control algorithm that can be used to estimate optimal dynamic treatment regime from longitudinal clinical data. Suppose there are two stages, the Q function can be defined as (5):

$Q_{\text{2}}(h_{\text{2}},a_{\text{2}})\triangleq E(Y|H_2=h_2,A_2=a_2)$

$Q_1(h_1,a_1)\triangleq E(\underset{a_2\in \left\{-1,1\right\}}{\max }{Q}_2\left(h_2,a_2\right)|H_1=h_1,A_1=a_1)$

The Q function at t=2 measures is the value function of assigning $a_2$ to a patient with history $h_2$. In a similar sense, the Q function at t=1 measures the quality of assigning $a_1$ to a patient with history $h_1$, assuming an optimal regime at stage 2 is adopted. The first step of Q learning is the regress Y on H₁ and A₂ to obtain estimators of ${\widehat{\beta }}_2$:

${\widehat{Q}}_2(H_2,A_2,{\widehat{\beta }}_2)\triangleq H_{20}^T{\widehat{\beta }}_{20}\text{+}{A}_2{H}_{2\text{1}}^T{\widehat{\beta }}_{2\text{1}}$

Where H₂ includes history variables up to t=2; H₂₀ is the main effect model (treatment-free) consisting a subset of variables in H₂. Similarly, H₂₁ is also a subset of variables of H₂ for the contrast function. A contrast function defines how the expected outcome varies depending on treatment regime. The only means by which the treatment can influence outcome are via the contrast function (6).

The second step is to maximize the Q₂ function by varying A2. In the above example, if A2 takes value {-1, 1}, we can define $\tilde{Y}\triangleq \underset{a_2\in \left\{-1,1\right\}}{\max }{\widehat{Q}}_2\left(H_2,a_2,{\widehat{\beta }}_2\right)$, which gives $\tilde{Y}=H_{20}^T{\widehat{\beta }}_{20}\text{+}\left|H_{\text{21}}^T{\widehat{\beta }}_{\text{21}}\right|$. $\tilde{Y}$ is the potential outcome when optimal regime is adopted at stage 2. Finally, the estimated $\tilde{Y}$ is regressed on H₁ and A₂ to obtain ${\widehat{\beta }}_1$ in the model

${\widehat{Q}}_{\text{1}}(H_{\text{1}},A_{\text{1}},{\widehat{\beta }}_{\text{1}})\triangleq H_{\text{1}0}^T{\widehat{\beta }}_{\text{1}0}\text{+}{A}_{\text{1}}{H}_{\text{11}}^T{\widehat{\beta }}_{\text{11}}$

The optimal regime at stage one can be obtained by maximizing the above equation.

Q-learning with R

In the example, the DynTxRegime package (v4.1) will be used to obtain optimal regime at each stage. The package iqLearn is also available in R that performs Q-learning algorithm (5).

The above code creates two components (e.g. the main effect and the contrast components) of the outcome model. Note that all covariates are specified to be predictive on the outcome, and only covariates collected at stage two determine the optimal decision rule. However, the functional form is mis-specified in the above equation.

The second stage can be fit by using the above code. We can create a confusion matrix to see whether the learning algorithm can assign optimal regime to more subjects in the training set, than that assigned by the physician (e.g., suppose the longitudinal data were collected from electronic healthcare records and the treatment was determined by physicians).

The result shows that the Q-learning method is not able to assign optimal regime to more subjects than that assigned by physicians. In the observational cohort (the first matrix), 217 patients received A2=0, and 71 received A2=1, which are optimal regimes. The Q-learning algorithm correctly assigned A2=0 to 357 patients, and A2=1 to 71 patients. Other patients fail to receive the optimal regime if treatment regime is determined by the Q-learning algorithm. This is due to the misspecification of the outcome linear function, which has been widely discussed in the literature (7,8).

Q-learning is implemented through a backward recursive fitting procedure based on a dynamic programming algorithm. The previous step fitted a linear regression model for stage two, and next we will regress the estimated outcome on covariates (X1) and treatment (A1) at stage one.

The specification of the regression model is similar to that at for stage two except that the response argument takes an object returned by previous fit. The coefficient for the first stage model can be examined by the following code:

Again, the confusion matrix can help to examine the correctness of Q-learning algorithm to assign optimal regime.

It appears that the Q-learning algorithm fails to assign optimal regime to more subjects than the physicians’ judgement. However, the linear model cannot correctly capture the function form as that for data generation.

Inverse probability weighted estimator (IPWE)

To circumvent the risk of model misspecification as demonstrated above, the inverse probability weighted estimator for the population mean outcome can be employed (9). The estimator for E{Y*(g_η)} is:

$\text{IPWE}\left(\text{η}\right)=n^{-1}{\displaystyle \sum _{i=1}^n\frac{C_{\eta ,i}·Y_i}{{\pi }_c\left(X_i;\eta ,\widehat{\gamma }\right)}}$

Where C_η=A·g(X,η)+(1−A)·[1−g(X,η)], and g(X,η)∈G_η is a specific treatment regime conditional on covariate X and coefficient η. The equation indicates that the observed outcome Y (under treatment A) is the counterfactual outcome Y*(g_η) (under treatment g(X,η)) when C_η=1. Otherwise, Y*(g_η) is not observed because a patient actually received treatment A is not equal to g(X,η). Suppose the treatment regime can only takes value 0 or 1. ${\pi }_c\left(X_i;\eta ,\widehat{\gamma }\right)$ is the probability of C_η=1 giving covariates X:

$\mathrm{Pr}\left(C_{\eta }=1|X\right)={\pi }_c\left(X_i;\eta ,\widehat{\gamma }\right)=\pi \left(X,\widehat{\gamma }\right)·g\left(X,\eta \right)+\left[1-\pi \left(X,\widehat{\gamma }\right)\right]·\left[1-g\left(X,\eta \right)\right]$

The estimator IPWE(η) is actually a weighted average of observed outcome under decision rule g(X,η). Observed outcome Y not following the decision rule was excluded (by the C_η,i factor in the numerator) from the calculation. For patients who actually receive A=0, by substituting A=0 into the IPWE(η) equation, the contrast function can be written as:

Similarly, the contrast function for A=1 can be written as:

${\widehat{C}}_{IPWE}\left(X_i\right)=-\frac{g\left(X,\eta \right)·Y_i}{\pi \left(X,\widehat{\gamma }\right)·g\left(X,\eta \right)}=-\frac{Y_i}{\pi \left(X,\widehat{\gamma }\right)}$

Collectively, the contrast function is written as:

${\widehat{C}}_{IPWE}\left(X_i\right)=-\frac{A_i·Y_i}{\pi \left(X,\widehat{\gamma }\right)}=-\frac{\left(1-A_i\right)·Y_i}{\left[1-\pi \left(X,\widehat{\gamma }\right)\right]}$

From the above discussion, we note that only the propensity model is required in the estimation, the mean outcome model is not required, which reduces the risk of model misspecification. We then define a function Z=I(C(X)>0), so that Z=1 indicate subjects who will benefit from treatment A=1 than A=0. The optimization problem can be regarded as minimization of a weighted misclassification error:

$g^{opt}=\underset{g\in G}{\arg \min }E\left\{\left|C\left(X\right)\right|·{\left[I\left\{C\left(X\right)>0\right\}-g\left(X\right)\right]}^2\right\}=\underset{g\in G}{\arg \min }{{\displaystyle \sum _{i-1}^n\widehat{W}\left[{\widehat{Z}}_i-g\left(X_i\right)\right]}}^2$

Where ${\widehat{W}}_i=\left|\widehat{C}\left(X_i\right)\right|$ is an estimate of contrast function for each subject. The problem can be considered as a classification problem with ${\widehat{Z}}_i$ as the binary response, X_i as the predictor and ${\widehat{W}}_i$ as the weight. Classification and regression trees (CART) can be used to solve this problem (10).

The following code builds a propensity score model in which the variables x21 and x22 determine the choice of treatment regime. Since treatment A is a binary variable, a generalized linear model with the response variable following binomial distribution.

Then, the rpart package (version 4.1-13) is employed to solve the classification problem. The CART model uses x21, x22 and x23 as the predictors.

The second-stage analysis using IPW can be done with the following code.

Recall that the optimal decision rule is to assign A=1 to subjects with x21 >0.3 and x23 <0.46. The CART assigns x21 ≥0.15 and x23 <0.49 to receive A=1; but there are still moderate misclassifications (Figure 2).

Figure 2.

Decision tree to assign treatment regime A=1 or A=0 to subjects with covariates X2 at stage two by the IPWE method.

The confusion matrix can be obtained with the following code.

The results showed that the IPW method can significantly improve the accuracy of assigning optimal regime to patients. There are only 50 patients who would benefit from A=0 but receive A=1 as determined by the IPW method.

Augmented Inverse Probability Weighted Estimators (AIPWE)

The IPWE method does not depend on parametric or semi-parametric at all, making it robust to model misspecification. However, the IPWE estimator of contrast values can be too noisy to successfully inform the classification regime. One method to circumvent the problem is to maximize on a doubly robust augmented inverse probability weighted estimator (AIPWE) for the population mean outcome over all possible regimes (9,11). The AIPWE for the mean outcome is given by:

$\text{AIPWE}\left(\text{η}\right)=n^{-1}{\displaystyle \sum _{i=1}^n\left\{\frac{C_{\eta ,i}·Y_i}{{\pi }_c\left(X_i;\eta ,\widehat{\gamma }\right)}-\frac{C_{\eta ,i}-{\pi }_c\left(X_i;\eta ,\widehat{\gamma }\right)}{{\pi }_c\left(X_i;\eta ,\widehat{\gamma }\right)}·m\left(X_i;\eta ,\widehat{\beta }\right)\right\}}$

The m(X_i;η,β)=μ(1,X,β)·g(X,η)+μ(0,X,β)·[1-g(X,η)] is a model for the potential outcome under treatment regime g_η that E(Y*(g_η)|X)=μ(1,X)·g(X,η)+μ(0,X)·[1−g(X,η)]. AIPWE possesses the doubly robust property that the estimate is consistent for E(Y*(g_η)) if either π(X,γ) or μ(A,X,β) is correctly specified. The contrast function for AIPWE estimator can be derived from the equation:

${\widehat{C}}_{AIPWE}\left(X_i\right)=\frac{A_i·Y_i}{\pi \left(X,\widehat{\gamma }\right)}-\frac{\left(1-A_i\right)·Y_i}{\left[1-\pi \left(X,\widehat{\gamma }\right)\right]}-\frac{A_i-\pi \left(X,\widehat{\gamma }\right)}{\pi \left(X,\widehat{\gamma }\right)}·\mu \left(1,X,\widehat{\beta }\right)-\frac{A_i-\pi \left(X,\widehat{\gamma }\right)}{\left[1-\pi \left(X,\widehat{\gamma }\right)\right]}·\mu \left(0,X,\widehat{\beta }\right)$

The only difference between ${\widehat{C}}_{AIPWE}\left(X_i\right)$ and ${\widehat{C}}_{IPWE}\left(X_i\right)$ is that the former borrows information from specified parametric regression model of outcome μ(A,X,β), trying to strike a balance between two extremes that are non-parametric or completely depending on the parametric model (11). R code to perform AIPWE is very similar to that for IPWE, but we need to define the parametric outcome model. The outcome model is defined as

then the second stage for AIPWE analysis is:

Note that the only difference between AIPWE and IPWE is the use of parametric outcome model as specified by the moMain and moCont arguments. We know from the Q-learning section that the linear outcome model is mis-specified. The confusion matrix is examined in the following code:

The results show that the AIPWE significantly improve the accuracy of correctly assigning optimal regime to the patients. In this analysis, only 18 subjects are incorrectly assigned A=1 by the AIPW method, and 15 subjects are falsely assigned A=0. This number is better than the IPWE method.

The decision tree shows that x21 and x23 are used to split most nodes (Figure 3). Subjects with x21 <0.3 were correctly assigned A=0. The subjects with x23 ≥0.46 are also correctly assigned to receive A=0.

Figure 3.

Decision tree to assign treatment regime A=1 or A=0 to subjects with covariates X2 at stage two by the AIPWE method. AIPWE, Augmented Inverse Probability Weighted Estimators.

The first stage decision rule can be learned in a recursive manner as that for stage two. The only difference is that the predictors for both outcome model and propensity model are chosen from x11, x12 and x13. Furthermore, the response argument in the optimalClass() function receives the object returned in the stage two analysis.

Firstly, we define the propensity for treatment model and methods. Here the variable x11 and x12 are used to predict treatment assignment.

Secondly, we define the classification model in which the CART method is used. Other machine learning methods such as support vector machine are also allowed and can be passed to the buildModelObj() function by the solver.method argument.

Thirdly, we define the outcome model which is composed of the main effect model and contrast model.

Finally, the AIPWE model can be fit with the optimalClass() as before. Note that an optimalClass object is assigned to the response argument, which passes the counterfactual outcome under optimal treatment at stage two to the modeling strategy for the first stage. This is consistent with the idea behind reinforcement learning algorithm in computer science that the action-value function is updated by assuming all subsequent steps takes the optimal action to maximize the action-value function.

The data generating mechanism dictates that optimal regime A=1 should be given to subjects with x11 > −0.54 and x12 <0.54. The majority of subjects receiving A=1 following the AIPWE rule is those with x11 >−0.52, x12 <0.57 and x12 >−1 (Figure 4). Classification accuracy can be further examined with confusion matrix:

Figure 4.

Decision tree to assign treatment regime A=1 or A=0 to subjects with covariates X1 at stage one by the AIPWE method. AIPWE, Augmented Inverse Probability Weighted Estimators.

The above output shows that the treatment options made by physician misclassify 95 subjects to the A=1 and 140 to the A=0 group. The AIPWE rule is able to significantly improve the classification accuracy. Only 19 subjects are misclassified to A=1 group and 30 subjects are misclassified to the A=0 group.

Final remarks

Exploring dynamic treatment regime borrows ideas from the reinforcement learning algorithm from the computer science. The Q-learning algorithm is easy to interpret for domain experts, but it is limited by the risk of misspecification of the linear outcome model. IPWE overcomes this problem by non-parametric modeling that the mean outcome is estimated by weighting the observed outcome. AIPWE borrows information from both the propensity model and mean outcome model and enjoys the property of double robustness. However, these statistical method does not allow the action space to have multiple levels. In real clinical practice, the interest may focus on the combinations of treatment regime, giving rise to a multi-dimensional action space.

Ethical Statement: The author is accountable for all aspects of the work in ensuring that questions related to the accuracy or integrity of any part of the work are appropriately investigated and resolved.