Learn to Intervene: An Adaptive Learning Policy for Restless Bandits in Application to Preventive Healthcare

1 Introduction

Preventive timely intervention (e.g., to check adherence or to provide medicines) can be used to significantly alleviate many public health issues such as diabetes [Newman et al., 2018], hypertension [Brownstein et al., 2007], tuberculosis [Chang et al., 2013; Rahedi Ong'ang'o et al., 2014], HIV [Christopher et al., 2011; Kenya et al., 2013], depression [L¨owe et al., 2004; Mundorf et al., 2018], etc. In this paper, we focus on the area of maternal health (well-being of women during pregnancy, childbirth and the post-natal period) where preventive interventions can impact millions of women. A key challenge in such preventive intervention problems is the limited amount of resources for intervention (e.g., availability of health care workers). Furthermore, the human behavior (w.r.t. taking medicines or adhering to a protocol) changes over time and to interventions, thereby requiring a strategic assignment of the limited resources to the beneficiaries most in need.

∗The research was conducted when A. Biswas worked at Google.

We are specifically motivated towards improving maternal health among low-income communities in developing countries, where maternal deaths remain unacceptably high due to not having timely preventive care information [Thaddeus and Maine, 1994]. We work with a non-profit organization, ARMMAN 2015, that provides a free call-based program for around 2.3 million pregnant women. This is similar to other programs, such as Mom-connect (https://tinyurl.com/ mom-connect-jnj). Each enrolled woman receives 140 automated calls to equip them with critical life-saving healthcare information across 80 weeks (pregnancy and up to one year after childbirth). Unfortunately, the engagement behavior (overall time spent listening to automated calls) changes and, for most women, the overall engagement decreases. This can have serious implications on their health. We ask the question: how do we systematically choose whom to provide interventions (personal visit by a healthcare worker) in order to maximize the overall engagement of the beneficiaries?

Preventive intervention problems of interest in this paper are challenging, owing to multiple key reasons: i) number of interventions are budgeted and much smaller than the total number of beneficiaries; ii) beneficiary's engagement is uncertain and may change after a few months; iii) postintervention improvement of beneficiary's engagement is uncertain; iv) intervention decisions at a time step have an impact on the state of beneficiaries and decisions to be taken at the next step. A relevant model for this setting is restless multi-armed bandits (RMABs). RMABs with prior knowledge of uncertainty model have been studied for health interventions [Lee et al., 2019; Mate et al., 2020; Mate et al., 2021; Killian et al., 2021; Bhattacharya, 2018], sensor monitoring tasks [Iannello et al., 2012; Glazebrook et al., 2006], anti-poaching patrols [Qian et al., 2016], and uplift modeling in eCommerce platforms [Gubela et al., 2019]. Due to the unpredictability of human beneficiaries, it is unrealistic to know the uncertainty model a priori. Our key contribution is in designing an intervention scheme under limited budget, without relying on a priori knowledge of the uncertainty model.

Contributions. First, we represent the preventive intervention problem as a Restless Multi-Arm Bandit (RMAB), where the uncertainty model associated with beneficiary behaviors with and without intervention are not known a priori, thereby precluding the direct application of Whittle Index based methods [Whittle, 1988]. Second, we develop a model-free learning method based on Q-learning, referred to as Whittle Index based Q-Learning (WIQL) that executes actions based on the difference between Q-values of active (intervention) and passive actions. We show that WIQL converges to the optimal solution asymptotically. Finally, to show that WIQL is a general approach and applicable to multiple domains, we evaluate WIQL on various examples. We then simulate the intervention problem in the context of maternal healthcare. We demonstrate that the model assumptions we make capture the problem of concern and the intervention scheme employed by WIQL significantly improves the engagement among beneficiaries compared to the existing benchmarks.

2 RelatedWork

The restless multi-armed bandit (RMAB) problem was introduced by Whittle 1988. The main result involves formulating a relaxation of the problem and solving it optimally using a heuristic called Whittle Index policy. This policy is optimal when the underlying Markov Decision Processes satisfy indexability, which is computationally intensive to verify. Moreover, Papadimitriou and Tsitsiklis 1994 established that solving RMAB is PSPACE-hard.

There are three main threads of relevant research. The first category focuses on specific classes of RMABs. Akbarzade et al. 2019 provide a class of bandits with "controlled restarts" and state-independent policies, which possess the indexability property and a Whittle index policy. Mate et al. 2020 consider two-state bandits to model a health intervention problem, where taking an intervention collapses the uncertainty about their current state. Bhattacharya 2018 models the problem of maximizing the coverage and spread of health information as an RMAB problem and proposes a hierarchical policy. Lee et al. 2019 study the problem of screening patients to maximize early-stage cancer detection under limited resource, by formulating it as a subclass of RMAB. Similarly, Glazebrook et al. 2006, Hsu 2018, Sombabu et al. 2020, Liu and Zhao 2010 give Whittle indexability results for different subclasses of (hidden) Markov bandits. This category of research assumes that the transition and observation models are known beforehand. Instead, our focus is on providing learning methods when the transition model is unknown a priori.

The second category contains different learning methods for RMABs. Fu et al. 2019 provide a Q-learning method where the Q value is defined based on the Whittle indices, states, and actions. However, they do not provide proof of convergence to optimal solution and experimentally, do not learn (near-)optimal policies. Along similar lines, Avrachenkov and Borkar 2020 provide a fundamental change to the Q-value definition for computing optimal whittle index policy. However, their convergence proof requires all homogeneous arms with same underlying MDPs. We provide a learning method that is not only shown to theoretically converge but also empirically outperforms the above mentioned methods on benchmark instances and also on a real (heterogeneous arms) problem setting.

The third relevant line of work is to predict adherence to a health program and effects of interventions, by formulating these as supervised learning problems. Killian et al. 2019 use the data from 99DOTS [Cross et al., 2019] and train a deep learning model to target patients at high risk of not adhering to the health program. On similar lines, Nisthala et al. 2020 use the engagement data of a health program and train deep learning model to predict patients who are at high risk of dropping out of the program. Son et al. 2010 use Support Vector Machine to predict adherence to medication among heart failure patients. There are many other papers [Howes et al., 2012; Lauffenburger et al., 2018] that train models on historical data for predicting adherence and provide interventions to patients who have low adherence probability. These works assume the training data to be available beforehand. In contrast, we consider the online nature of the problem where feedback is received after an intervention is provided, which in turn is used in making future decisions.

3 Preliminaries

RMAB models various stochastic scheduling problems, where an instance is a 3-tuple (N, M, {MDPi}i∈N) with N denoting the set of arms, M is the budget restriction denoting how many arms can be pulled at a given time, and MDPi is the Markov Decision Process for each arm i. An MDP consists of a set of states S, a set of actions A, transition probabilities P : S × A × S ↦→ [0, 1], and reward function R : S × A ↦→ R. The action set A of each MDP consists of two actions: an active action (1) and a passive action (0). At each time step t, an action Ai(t) is taken on an arm i, such that ∑i Ai(t) = M. Then, each arm i transitions to a new state and observes a reward, according to the underlying MDPi. Let Xi(t) ∈ S and RXi(t) i (Ai(t)) denote the current state and reward obtained at time t respectively. Now, policy per step π : X1(t) × . . . ×XN(t) ↦→ {Ai(t)}i∈N can be defined as a mapping from the current states of all beneficiaries to the actions to be taken on each arm. Thus, given a policy π, the action on an arm i is denoted as:

The goal is to find a policy π∗ that maximizes the expected reward until T time steps, subject to the budget.

To deal with the computational hardness of solving this problem, an index-based heuristic policy based on the Lagrangian relaxation of the RMAB problem was proposed by Whittle 1988-at each time step t, an index is computed for each arm depending on the current state of the arm, transition probabilities, and reward function of its MDP. Then, the top M arms with highest index values are selected.

Whittle Index-based policy. Whittle's relaxation is to replace the budget constraint of M on the number of arms to a time-averaged constraint, i.e.,

Further, using Lagrange's relaxation (with λ as the Lagrange's multipliers) and dropping other constants, the objective function can be rewritten as:

Whittle showed that this problem can be decoupled and solved for each arm by computing the index λi(Z) which acts like a subsidy that needs to be given to an arm i at state Z, so that taking the action 1 is as beneficial as taking the action 0. Assuming indexability, choosing arms with higher subsidy λi(Z) leads to the optimal solution for Equation 2.

Definition 1 (Indexability) Let Φ(λ) be the set of states for which it is optimal to take action 0 when taking an active action costs λ. An arm is indexable if Φ(λ) monotonically increases from ∅ to S when λ increases from -∞ to +∞. An RMAB problem is indexable if all the arms are indexable.

4 The Model

We formulate the problem of selecting M out of N statetransitioning arms at each time step, as an RMAB, with beneficiaries being the arms in preventive healthcare intervention scenarios. We represent the engagement pattern as an MDP (Figure 1) with three (abstract) "behavioral" state: (i) self motivated (S): in this state, beneficiary shows high engagement and there is no need to intervene in this state. (ii) persuadable (P): in this state, the beneficiary engages less frequently with a possibility of increasing engagement when intervened, which makes this the most interesting state, and (iii) lost Cause (L): The engagement is very low in this state and very likely to remain low irrespective of intervention.

These three states capture different levels of engagement as well as differences in terms of the benefit obtained by interventions, which is not possible to represent using other existing intervention models, such as the model described by Mate et al. 2020. Note that the more the states, slower is the convergence of any online algorithm. Thus, for shortterm intervention programs (a year or two), we recommend a three-states model. In Section 7.2, we provide a mechanism to obtain the states from real-data. Let Ai(t) ∈ {0, 1} be the action taken on beneficiary i at time t; 1 denotes an intervention and 0 otherwise. Also, ∑i Ai(t) = M for each time slot t. Depending on the action taken, each beneficiary changes its states according to the transition probabilities.

Though our model is inspired by a preventive intervention in healthcare, this formulation captures the intervention problem in many other domains, such as sensor monitoring, antipoaching patrols, and uplift modeling in eCommerce. Depending on the domain, a persuadable state is where the benefit of an intervention is maximum. When a beneficiary is in the persuadable state and intervention is provided, an arm is more likely to transition to the self-motivated state whereas, when no intervention is provided, it is more likely to transition to the lost-cause state. These transitions are denoted as pP S i,1 and pPL i,0 for action 1 and 0, respectively. Every transition generates a state-dependent reward. Since better engagements have higher rewards, we assume RS i > RP i > RL i . The transition probabilities pYZ i,a may be different for each beneficiary and thus, the effect of intervention would vary from one beneficiary to another. If the transition probabilities were known beforehand, one could compute the Whittle-index policy and select accordingly. However, here, transition probabilities are unknown, and hence, we consider the problem of learning the Whittle index while simultaneously selecting a set of best arms depending on the estimated Whittle Index.

5 Whittle Index based Q-learning (WIQL)

Q-Learning [Watkins and Dayan, 1992] is a well-studied reinforcement learning algorithm for estimating Q∗ (Z, a) for each state-action pair (Z, a) of an MDP.

where the optimal expected value of a state is given by

Q-Learning estimates Q∗ using point samples-at each time t, an agent (policy maker) takes an action a using estimated Q values at the current state Z, a reward R is observed, a new state Z is reached, and Q values are updated according to the following update rule:

Here, αt(·) ∈ [0, 1] is the learning parameter. When α = 0, the agent does not learn anything new and retains the value obtained at tth step, while α = 1 stores only the most recent information and overwrites all previously obtained rewards. Setting 0 < α < 1 strikes a balance between the new values and the old ones. With mild assumptions on α, the convergence of Q-Learning to the optimal Q∗ values has been established [Watkins and Dayan, 1992; Jaakkola et al., 1994; Borkar and Meyn, 2000]. We build on the results to provide a Q-Learning approach for learning Whittle-index policy.

Algorithm 1: Whittle Index based Q-learning(WIQL)

Input: N, M, α(c), and initial states Xi(0) ∈ S ∀ i ∈ [N]. 1 Initialize: Q0i (Z, a) ← 0 and λ0i (Z) ← 0 for each state Z ∈ Si, each action a ∈ {0, 1}, each arm i ∈ [N]. 2 for t in 1, . . . , T do // Select M arms using ϵ-decay policy 3 ϵ ← N N+t . 4 With probability ϵ, select M arms uniformly at random. Otherwise, select topM arms according to their values λti(Xi(t)). Store the selected arms in Ψ. // Take suitable actions on the arms 5 for i in 1 : N do 6 if i ∈ Ψ then 7 Take action Ai(t) = 1 on arm i. 8 else 9 Take action Ai(t) = 0 on arm i. 10 Observe reward r and next state Xi(t + 1). // Update Q, λ using Z =Xi(t), a = Ai(t) 11 chi,Z,a = t∑ h=0 I{Xi(h) = Z & Ai(h) = a}. 12 for i in 1 : N do 13 Qt+1 i (Z, a) ← (1 -α(chi,Z,a)) · Qti(Z, a) + α(chi,Z,a) · (r + max a′∈{0,1}Qti (Xi(t + 1), a′ )). 14 λt+1 i (Xi(t) = Z) ←Qt+1 i (Z, 1) -Qt+1 i (Z, 0)

We adopt Q-Learning for RMABs and store Q values separately for each arm. Q values of state-action pairs are typically used for selecting the best action for an arm at each state; however, for RMAB, the problem is to select M arms. Our action-selection method ensures that an arm whose estimated benefit is higher is more likely to get selected. Algorithm 1 describes the action-selection and the update rule. This algorithm is not specific to any particular state representation and can be used for any finite state RMAB instance. However, when the number of states are large (10 states), the convergence is typically slow and not suitable for short horizon.

We propose Whittle Index based Q-learning (WIQL), that uses an ϵ-decay policy to select M arms at each time step t. During early steps, arms are likely to be selected uniformly at random. As time proceeds, arms with higher values of their estimated λi(Xi(t)) gets more priority. The selected set of M arms (who receive interventions) is called Ψ. Each arm is restless, i.e., each arm transitions to a new state and observes a reward, with or without interventions. These observations are then used for updating the Q values in Step 13 (Eq. 3). While updating Qi(Z, a), we use a learning parameter α(ci,Z,a) that decreases with increase in cti,Z,a (number of times the arm i received action a at state Z); eg, α(cti,Z,a ) = 1/(cti,Z,a + 1) satisfies this criteria. These Q values are then used to estimate the Whittle index λ(Xi(t)).

6 Theoretical Results

In this section, we show that the WIQL does not alter the optimality guarantees of Q-learning. First, we show that taking intervention decisions based on the benefit, (i.e., difference in Q values of active and passive actions) is equivalent to optimizing joint Q over all arms subject to the budget on intervention actions.

Theorem 1 Taking action 1 on the top M arms according to (Q∗i (si, 1) - Q∗i (si, 0)) is equivalent to maximizing ∑i Q∗i (si, ai) over all possible action profiles (a1, . . . , aN) ∈ {0, 1}N such that ∑i ai = M.

Proof Sketch. For ease of explanation, we prove this for M=1. Let i∗ be the arm that maximizes the benefit of taking an intervention action (Q∗i∗(si∗, 1) -Q∗i∗(si∗, 0)) at its current state si∗ . Then, for any j ∈ {1, . . . , N} \ {i∗ }.

Eq. (5) shows that taking intervention action on i∗ and passive actions on all other arms would maximize ∑i Q∗ (si, ai) whenM=1. This argument holds true for anyM ≥ 1. □

Theorem 2 WIQL converges to the optimal with probability 1 when ∑h α(chi,Z,a ) = ∞ and ∑h α(chi,Z,a )2 < ∞.

Proof Sketch. This proof follows from (1) the convergence guarantee of Q-Learning algorithm [Watkins and Dayan, 1992], (2) ϵ-decay selection policy, and (3) theorem 1. It has been established in [Watkins and Dayan, 1992] that the update rule of Q-Learning converges to Q∗ whenever ∑t αt(Z, a) = ∞ and ∑t αt(Z, a)2 < ∞. These assumptions require that all state-action pairs be visited infinitely often, which is guaranteed by the ϵ-decay selection process, where each arm has a non-zero probability of being selected uniformly at random. Thus, Qti (Z, a) converges to Q∗i (Z, a) which implies that Qti (Z, 1) - Qti (Z, 0) converges to Q∗i (Z, 1)-Q∗i (Z, 0) (using series convergence operation). Finally, using Theorem 1, we claim that selecting arms based on highest values of Q∗i (Z, 1)-Q∗i (Z, 0) would lead to an optimal solution problem. This completes the proof that WIQL converges to the optimal. □

7 Experimental Evaluation

We compare WIQL against five benchmarks: (1) OPT: assumes full knowledge of the underlying transition probabilities and has access to the optimal Whittle Indices, (2) AB [Avrachenkov and Borkar, 2020], (3) Fu [Fu et al., 2019], (4) Greedy: greedily chooses the top M arms with the highest difference in their observed average rewards between actions 1 and 0 at their current states, and (5) Random: chooses M arms uniformly at random at each step.

We consider a numerical example and a maternal healthcare application to simulate RMAB instances using beneficiaries' behavioral pattern from the call-based program. For each experiment, we plot the total reward averaged over 30 trials, to reduce the effect of randomness in the actionselection policy.

7.1 Numerical Example: Circulant Dynamics

This example has been studied in the existing literature on learning Whittle Indices [Avrachenkov and Borkar, 2020; Fu et al., 2019]. Each arm has four states S = {0, 1, 2, 3} two actions A = {0, 1}. The rewards are R0 = -1, R1 = R2 = 0, and R3 = 1 for a ∈ {0, 1}. The transition probabilities for each action a ∈ {0, 1} are represented as a |S| × |S| matrix:

The optimal Whittle Indices for each state are as follows: λ∗ (0) = -0.5, λ∗ (1) = 0.5, λ∗ (2) = 1, and λ∗ (3) = -1. Whittle Index policy would prefer taking action 1 on arms who are currently at state 2, followed by those at state 1, then those at state 0, and lastly those at state 3.

Figure 2 demonstrates that WIQL gradually increases towards the OPT policy, for two sets of experiments-(1) N=5 and M=1 and (2) N=100 and M=20. We observe that AB converges towards an average reward of zero. This happens because it prioritizes arms who are currently at state 1 over other states. Since, the expected reward of taking action 1 at state 1 is P(1, 1, 1) · R1 + P(1, 1, 2) · R2 = 0, the total average reward also tends to zero. The result obtained by the algorithm Fu is the same as what is shown in Figure 2 of their paper [Fu et al., 2019] where the total average reward converges to a value of 0.08. As expected, Greedy and Random, one being too myopic and the other being too exploratory in nature, are unable to converge to the optimal value.

These observations show that WIQL outperforms the existing algorithms on the example that was considered in the earlier papers. Note that, while implementing the algorithms AB and Fu, we fixed the hyperparameters to the values specified for this example. However for the real-world application, that we consider next, it is not obvious how to obtain the best set of hyperparameters for their algorithms. Thus, we do not compare these algorithms for the maternal healthcare application. Next, we compare the performance of WIQL algorithm with Greedy, Random and a Myopic policy (defined in the subsequent paragraph).

7.2 Real-world Application: Maternal Healthcare

We now focus on the maternal healthcare intervention problem where only a small subset of beneficiaries can be selected for providing interventions every week. We use the data, obtained from the call-based preventive care program, that contains call-records of enrolled beneficiaries-how long they listened to the call, whether an intervention was given, and when. The data also contain the ID of a healthcare worker dedicated to provide personalized intervention to each enrolled beneficiary. The data was collected for an experimental study towards building up a robust intervention program (which is the focus of this work). During the experimental study, only one intervention (in-person visit by a health worker) was provided to a selected set of beneficiaries (1559 chosen from among 3031 beneficiaries) who were more likely to drop-out of the program. We call this the Myopic intervention and use it as a benchmark to compare our approach.

(a) Average reward: 1 out of 5 arms chosen.

• (b) Average reward: 20 out of 100 arms chosen.

Figure 2: Circulant Dynamics: Results averaged over 30 trials.

We use the three-state MDP model (Figure 1) to simulate the behavior of the beneficiaries. During a particular week, the beneficiaries listening to > 50% of the content of the automated calls are in state S (self-motivated), those listening to 5 - 50% are in state P (persuadable), and those listening to < 5% are in state L (lost cause). A reward of 2 is obtained when a beneficiary is in state S, a reward of 1 is obtained in state P, and a reward of 0 is obtained in state L. Thus, a high total reward accumulated per week implies that a large number of beneficiaries are at either state S or P .

Some other observations on the call-records are as follows. The beneficiaries who were in state P on a particular week never transitioned to state S in the immediate next week, unless they received an intervention. On the other hand, a few beneficiaries who were in state L transitioned to state P even without intervention. Moreover, the fraction of times a transition from state L to P occurred is almost the same with and without the intervention at state L. Even though these transition probabilities at the level of all users are known based on this data, it is difficult to know the actual transitions for a given beneficiary a priori and hence, we simulate this behavior. We conduct two set of simulations, namely static (transition probabilities of each beneficiary remain same throughout their enrollment), and dynamic (transition probabilities change a few weeks after their enrollment).

Static State-Transition Model. We assume three categories of arms-(A) high chance of improvement: highly likely to improve their engagement on receiving an intervention, and deteriorate in absence of an intervention when they are at state P, i.e., pPS = 0.8 and pPL = 0.8, (B) medium chance of improvement: pPS = 0.4 and pPL = 0.6, and (C) low change improvement: pPS = 0.1 and pPL = 0.6. We assume 10 arms belong to category-A, 10 arms belong to category-B and 30 arms belong to category-C. This assumption helps us determine the efficacy of any learning algorithm; in particular, the most efficient algorithm would quickly learn to intervene the 10 arms of category-A whenever they are at state P. We compare WIQL with greedy, random, and Myopic algorithm.

Dynamic State-Transition Model. For the dynamic setting, we simulate the first 28 weeks as described in Section 7.2.1. Further, we assume that the arms in category (A) change to "medium" chance of improvement after they are enrolled for around 28 weeks, those in category (B) change to "low" chance of improvement. Also, there is a new set of 10 arms with "high" chance of improvement. Note that, in reality, the behavior change would happen at any arbitrary time; however, to check if WIQL adapts to dynamic change of transition probabilities, we set this simulation environment.

We run each simulation for 80 weeks (the length of the program). In Figure 3 we provide results obtained by considering various values of M ∈ {500, 1000, 1500}, where a value of M represents the total number of personalized visits made by the 100 health-care workers on a week. We observe that, for M = 500 (each health-care worker visits only 5 beneficiaries per week), WIQL performs (only) marginally better than the other benchmarks. In contrast, when M ∈ {1000, 1500}, the reward obtained by WIQL is higher than Greedy and significantly more than Myopic and Random. Comparing WIQL and Greedy based on their total reward per week, we observe that the per-week engagement outperforms Greedy by a significant margin. Observe that, the convergence of WIQL to a total reward of 5000 is quicker when M is higher. This is because more sample points are observed per week. Additionally, we observe that the myopic algorithm leads to overall low engagement among the beneficiaries, even under the static setting. These results show that WIQL is able to learn which arms should be intervened at which state without any prior knowledge about the transition probabilities, and also adapts better to the underlying dynamics of the transition probabilities.

8 Conclusion and Discussion

We focus on a limited-resource sequential decision problem and formulate it as an RMAB setting. We provide a mechanism to systematically learn and decide on whom to intervene, and hence, improve the overall benefit of intervention. Our method possesses the capacity of impacting and improving the well-being of millions, for example, in the maternal healthcare domain, as demonstrated in this paper.

WIQL is a general solution for learning RMABs, and we demonstrate this using examples from other domains, such as Circulant Dynamics. Additionally, WIQL can be used in a more general setting where new arms arrive over time. For the departing arms, we can assume that each arm ends in a new state "dropped" and never transitions to the earlier states. Moreover, our experiments show that WIQL adapts to the dynamic behavioral change. In practice, however, WIQL may not be directly applicable for domains where there are beneficiaries with extreme health risk. The non-zero probability of being selected randomly may come at a cost of a patient in critical need of intervention. One way to mitigate this issue is to assign priorities to beneficiaries depending on their comorbidities, possible complication during pregnancy and after childbirth, etc. If the high-risk class is small, we can target the intervention for all of them, and run WIQL on the remaining beneficiaries. This constraint may hamper the convergence guarantee of WIQL; however, it would benefit the enrolled women at large.

Going ahead, it would be interesting to solve the RMAB learning problem considering a large number of states and more than two actions, each with a different cost of operation.