Decision-Focused Learning in Restless Multi-Armed Bandits with Application to Maternal and Child Care Domain

Restless multi-armed bandits (RMABs) (weber1990index; tekin2012online) are composed of a set of heterogeneous arms and a planner who can pull multiple arms under budget constraint at each time step to collect rewards. Different from the classic stochastic multi-armed bandits (gittins2011multi; bubeck2012regret), the state of each arm in an RMAB can change even when the arm is not pulled, where each arm follows a Markovian process to transition between different states with transition probabilities dependent on arms and the pulling decision. Rewards are associated with different arm states, where the planner’s goal is to plan a sequential pulling policy to maximize the total reward received from all arms. RMABs are commonly used to model sequential scheduling problems where limited resources must be strategically assigned to different tasks sequentially to maximize performance. Examples include machine maintenance (glazebrook2006some), cognitive radio sensing problem (bagheri2015restless), and healthcare (mate2022field).

In this paper, we study offline RMAB problems with unknown transition dynamics but with given arm features. The goal is to learn a mapping from arm features to transition dynamics, which can be used to infer the dynamics of unseen RMAB problems to plan accordingly. Prior works (mate2022field; sun2018cell) often learn the transition dynamics from the historical pulling data by maximizing the predictive accuracy. However, RMAB performance is evaluated by its solution quality derived from the predicted transition dynamics, which leads to a mismatch in the training objective and the evaluation objective. Previously, decision-focused learning (wilder2019melding) has been proposed to directly optimize the solution quality rather than predictive accuracy, by integrating the one-shot optimization problem (donti2017task; perrault2020end) or sequential problems (wang2021learning; futoma2020popcorn) as a differentiable layer in the training pipeline. Unfortunately, while decision-focused learning can successfully optimize the evaluation objective, it is computationally extremely expensive due to the presence of the optimization problems in the training process. Specifically, for RMAB problems, the computation cost of decision-focused learning arises from the complexity of the sequential problems formulated as Markov decision processes (MDPs), which limits the applicability to RMAB problems due to the PSPACE hardness of finding the optimal solution (papadimitriou1994complexity).

Our main contribution is a novel and scalable approach for decision-focused learning in RMAB problems using Whittle index policy, a commonly used approximate solution in RMABs. Our three key contributions are (i) we establish the differentiability of Whittle index policy to support decision-focused learning to directly optimize the RMAB solution quality; (ii) we show that our approach of differentiating through Whittle index policy improves the scalability of decision-focused learning in RMAB; (iii) we apply our algorithm to an anonymized maternal and child health RMAB dataset previously collected by  armman to evaluate the performance of our algorithm in simulation.

We establish the differentiability of Whittle index by showing that Whittle index can be expressed as a solution to a full-rank linear system reduced from Bellman equations with transition dynamics as entries, which allows us to compute the derivative of Whittle index with respect to transition dynamics. On the other hand, to execute Whittle index policy, the standard selection process of choosing arms with top-k Whittle indices to pull is non-differentiable. We relax this non-differentiable process by using a differentiable soft top-k selection to establish differentiability. Our differentiable Whittle index policy enables decision-focused learning in RMAB problems to backpropagate from final policy performance to the predictive model. We significantly improve the scalability of decision-focused learning, where the computation cost of our algorithm $O(NM^{\omega+1})$ scales linearly in the number of arms $N$ and polynomially in the number of states $M$ with $\omega\approx 2.373$, while previous work scales exponentially $O(M^{\omega N})$. This significant reduction in computation cost is crucial for extending decision-focused learning to RMAB problems with large number of arms.

In our experiments, we apply decision-focused learning to RMAB problems to optimize importance sampling-based evaluation on synthetic datasets as well as an anonymized RMAB dataset about a maternal and child health program previously collected by (armman) – these datasets are the basis of comparing different methods in simulation. We compare decision-focused learning with the two-stage method that trains to minimize the predictive loss. The two-stage method achieves the best predictive loss but significantly degraded solution quality. In contrast, decision-focused learning reaches a slightly worse predictive loss but with a much better importance sampling-based solution quality evaluation and the improvement generalizes to the simulation-based evaluation that is built from the data. Lastly, the scalability improvement is the crux of applying decision-focused learning to real-world RMAB problems: our algorithm can run decision-focused learning on the maternal and child health dataset with hundreds of arms, whereas state of the art is a 100-fold slower even with 20 arms and grows exponentially worse.

An instance of the restless multi-armed bandit (RMAB) problem is composed of a set of $N$ arms, each is modeled as an independent Markov decision process (MDP). The $i$-th arm in a RMAB problem is defined by a tuple $(\mathcal{S},\mathcal{A},R_{i},P_{i})$. $\mathcal{S}$ and $\mathcal{A}$ are the identical state and action spaces across all arms. $R_{i},P_{i}:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow\mathbb{R}$ are the reward and transition functions associated to arm $i$. We consider finite state space with $|\mathcal{S}|=M$ fully observable states and action set $\mathcal{A}=\{0,1\}$ corresponding to not pulling or pulling the arm, respectively. For each arm $i$, the reward is denoted by $R_{i}(s_{i},a_{i},s^{\prime}_{i})=R(s_{i})$, i.e., the reward $R(s_{i})$ only depends on the current state $s_{i}$, where $R:\mathcal{S}\rightarrow\mathbb{R}$ is a vector of size $M$. Given the state $s_{i}$ and action $a_{i}$, $P_{i}(s_{i},a_{i})=[P_{i}(s_{i},a_{i},s^{\prime}_{i})]_{s^{\prime}_{i}\in\mathcal{S}}$ defines the probability distribution of transitioning to all possible next states $s^{\prime}_{i}\in\mathcal{S}$.

In a RMAB problem, at each time step $t\in[T]$, the learner observes $\boldsymbol{s}_{t}=[s_{t,i}]_{i\in[N]}\in\mathcal{S}^{N}$, the states of all arms. The learner then chooses action ${\boldsymbol{a}}_{t}=[a_{t,i}]_{i\in[N]}\in\mathcal{A}^{N}$ denoting the pulling actions of all arms, which has to satisfy a budget constraint $\sum\nolimits_{i\in[N]}a_{t,i}\leq K$, i.e., the learner can pull at most $K$ arms at each time step. Once the action is chosen, arms receive action $\boldsymbol{a}_{t}$ and transitions under $P$ with rewards $\boldsymbol{r}_{t}=[r_{t,i}]_{i\in[N]}$ accordingly. We denote a full trajectory by $\tau=(\boldsymbol{s}_{1},\boldsymbol{a}_{1},\boldsymbol{r}_{1},\cdots,\boldsymbol{s}_{T},\boldsymbol{a}_{T},\boldsymbol{r}_{T})$. The total reward is defined by the summation of the discounted reward across $T$ time steps and $N$ arms, i.e., $\sum\nolimits_{t=1}^{T}\gamma^{t-1}\sum\nolimits_{i\in[N]}r_{t,i}$, where $0<\gamma\leq 1$ is the discount factor.

A policy is denoted by $\pi$, where $\pi(\boldsymbol{a}\mid\boldsymbol{s})$ is the probability of choosing action $\boldsymbol{a}$ given state $\boldsymbol{s}$. Additionally, we define $\pi(a_{i}=1\mid\boldsymbol{s})$ to be the marginal probability of pulling arm $i$ given state $\boldsymbol{s}$, where $\pi(\boldsymbol{s})=[\pi(a_{i}=1\mid\boldsymbol{s})]_{i\in[N]}$ is a vector of arm pulling probabilities. Specifically, we use $\pi^{*}$ to denote the optimal policy that optimizes the cumulative reward, while $\pi^{\text{solver}}$ to denote a near-optimal policy solver.

This paper studies the RMAB problem where we do not know the transition probabilities $P=\{P_{i}\}_{i\in[N]}$ in advance. Instead, we are given a set of features $\boldsymbol{x}=\{x_{i}\in\mathcal{X}\}_{i\in[N]}$, each corresponding to one arm. The goal is to learn a mapping $f_{w}:\mathcal{X}\rightarrow\mathcal{P}$, parameterized by weights $w$, to make predictions on the transition probabilities $P=f_{w}(\boldsymbol{x})\coloneqq\{f_{w}(x_{i})\}_{i\in[N]}$. The predicted transition probabilities are later used to solve the RMAB problem to derive a policy $\pi=\pi^{\text{solver}}(f_{w}(\boldsymbol{x}))$. The performance of the model $f$ is evaluated by the performance of the proposed policy $\pi$.

In this paper, instead of grappling with the optimal policy, we consider the Whittle index policy (whittle1988restless) – the dominant solution paradigm used to solve the RMAB problem. Whittle index policy is easier to compute and has been shown to perform well in practice. In this section we establish that it is also possible to backpropagate through the Whittle index policy. This differentiability of Whittle index policy allows us to run decision-focused learning to directly maximize the performance in the RMAB problem.

In this paper, we use two different variants of evaluation metric: importance sampling-based evaluation (sutton1998introduction) and simulation-based (model-based) evaluation.

We compare two-stage learning (TS) with our decision-focused learning (DF-Whittle) that optimizes importance sampling-based evaluation directly. We consider three different evaluation metrics including predictive loss, importance sampling evaluation, and simulation-based evaluation to evaluate all learning methods. We perform experiments on three synthetic datasets including $2$-state fully observable, $5$-state fully observable, and $2$-state partially observable RMAB problems. We also perform experiments on a real dataset on maternal and child health problem modelled as a $2$-state fully observable RMAB problem with real features and historical trajectories. For each dataset, we use $70\%$, $10\%$, $20\%$ of the RMAB problems as the training, validation, and testing sets, respectively. All experiments are averaged over $50$ independent runs.

This paper presents the first decision-focused learning in RMAB problems that is scalable for large real-world datasets. We establish the differentiability of Whittle index policy in RMAB by providing new method to differentiate through Whittle index and using soft-top-k to relax the arm selection process. Our algorithm significantly improves the performance and scalability of decision-focused learning, and is scalable to real-world RMAB problem sizes.

Aditya Mate was supported by the ARO and was accomplished under Grant Number W911NF-17-1-0370. Sanket Shah and Kai Wang were also supported by W911NF-17-1-0370 and ARO Grant Number W911NF-18-1-0208. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of ARO or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein. Kai Wang was additionally supported by Siebel Scholars.

We run both Decision Focused Learning and Two-Stage Learning for 50 epochs in 2-state and 5-state synthetic domain problems, 30 epochs in ARMMAN domain and 18 epochs in 2-state partially observable setting. The learning rate $r$ is kept at $0.01$ and $\gamma=0.59$ is used in all experiments. All the experiments are performed on an Intel Xeon CPU with 64 cores and 128 GB memory.

The large-scale quality improvement study conducted by armman contains 7668 beneficiries in the Round Robin Group. Over a duration of 7 weeks, $20\%$ of the beneficiaries receive at least one active action (LIVE service call). We randomly split the 7668 beneficiaries into 12 groups while preserving the proportion of beneficiaries who received at least one active action. There are 43 features available for every beneficiary which describe characteristics such as age, income, education level, call slot preference, language preference, phone ownership etc.

In this section, we provide information about consent related to data collection, analyzing data, data usage and sharing.

We have shown the computation cost of backpropagating through Whittle indices in Section 4.5. This section covers the remaining computation cost associated to other components, including the computation cost of Whittle indices in the forward pass, and the computation cost of constructing soft Whittle index policy using soft-top-k operator.

Unlike the synthetic datasets that we can produce multiple trajectories of an RMAB problem, in the real problem of service call scheduling problem operated by ARMMAN, there is only one trajectory available to us for every RMAB problem. Due to the specialty of the maternal and child health domain, it is unlikely to have the exactly same set of the pregnant mothers participating in the service call scheduling program at different times and under the same engagement behavior.

Given this restriction, we must evaluate the performance of a newly proposed policy using the only available trajectory. Unfortunately, the standard CWPDIS in Equation 12 does not work because the CWPDIS estimator is canceled out when there is only one trajectory:

which is fixed regardless what target policy $\pi$ is used and the associated importance sampling weights $\frac{\pi(a_{t,i}\mid s_{t})}{\pi_{\text{beh}}(a_{t,i}\mid s_{t})}$ and $\rho_{ti}=\prod\nolimits_{t^{\prime}=1}^{t}\frac{\pi(a_{t^{\prime},i}\mid s_{t^{\prime}})}{\pi_{\text{beh}}(a_{t^{\prime},i}\mid s_{t^{\prime}})}$. This implies that we cannot use CWPDIS to evaluate the target policy when there is only one trajectory.

Accordingly, we use the following variant to evaluate the performance:

where the new importance sampling weights are defined by $\rho^{\prime}_{t,i}(\tau)=\frac{\pi(a_{t,i}\mid s_{t})}{\pi_{\text{beh}}(a_{t,i}\mid s_{t})}$, which is not multiplicative compared to the original ones.

The main motivation of this new evaluation metric is to segment the given trajectory into a set of length-1 trajectories. We can apply CWPDIS to the newly generated length-1 trajectories to compute a meaningful estimate because we have more than one trajectory now. The OPE formulation with segmentation is under the assumption that we can decompose the total reward into the contribution of multiple segments using the idea of trajectory segmentation (krishnan2017transition; ranchod2015nonparametric). This assumption holds when all segments start with the same state distribution. In our ARMMAN dataset, the data is composed of trajectories of the participants who have enrolled in the system a few weeks ago, which have (almost) reached a stationary distribution. Therefore, the state distribution under the behavior policy, which is a uniform random policy, does not change over time. Our assumption of identical distribution is satisfied and we can decompose the trajectories into smaller segments to perform evaluation. Empirically, we noticed that this temporal decomposition helps define a meaningful importance sampling-based evaluation with the consistency benefit brought by CWPDIS.

We provide the learning curves of fully observable 2-state RMAB, fully observable 5-state RMAB, partially observable 2-state RMAB, and the real ARMMAN fully observable 2-state RMAB problems in Figure 7,8, 9, 10, respectively. Across all domains, two-stage method consistently converges to a lower predictive loss faster than decision-focused learning in Figure 7(a),8(a), 9(a), 10(a). However, the learned model does not produce a policy with good performance in the importance sampling-based evaluation metric in Figure 7(b),8(b), 9(b), 10(b), and similarly in the simulation-based evaluation metric in Figure 7(c),8(c), 9(c), 10(c).

Now, this set of equations is typically hard to solve because of the $\max$ terms in Equation 16. Specifically, knowing whether $\arg\max_{a}=0$ or $\arg\max_{a}=1$ for some state $s$ is equivalent to knowing what the optimal policy is for this modified MDP; such equations are typically solved using Value Iteration or variations thereof. However, this problem is slightly more complicated than a standard MDP because one also has to determine the value of $m_{s}$. The way that this problem is traditionally solved in the literature is the following:

1. 1.

   One guesses a value for the subsidy $m_{s}$.

2. 2.

   Given this guess, one solves the Bellman Optimality Equations associated with the modified MDP.

3. 3.

   Then, one checks the resultant policy. If it is more valuable to act than to not act in state $s$, the value of the guess for the subsidy is increased. Else, it is decreased.

4. 4.

   Go to Step 2 and repeat until convergence.

Given the monotonicity and the ability to bound the values of the whittle index, Step 3 above is typically solved using binary search. However, even with Binary Search, this process is quite time-consuming.

In this paper, we provide a much faster solution method for our application of interest. We leverage the small size of our state space to search over the space of policies rather than over the correct value of $m_{s}$. Concretely, because $S=\{0,1\}$ and $A=\{0,1\}$, the whittle index equations for state $s=0$ above boil down to:

These are 3 equations in 3 unknowns ($V(0),V(1),m_{s_{0}}$). The only hiccup here is that Equation 17 has a $\max$ term and so this set of equations can not be solved as normal linear equations would be. However, we can ‘unroll’ Equation 17 into 2 different equations:

Each of these corresponds to evaluating the value function associated with the partial policies $s=1\rightarrow a=0$ and $s=1\rightarrow a=1$. Then to get the optimal policy, we can just evaluate both of the policies and choose the better of the two policies, i.e., the policy with the higher expected value $V(1)$. In practice, we pre-compute the Whittle index and value function using the binary search and value iteration approach studied by  qian2016restless. Therefore, to determine which equation is satisfied, we just use the pre-computed value functions to evaluate the expected future return of different actions, and use the one with higher value to form a set of linear equations.

This gives us a set of linear equations where Whittle index is a solution. We can therefore derive a closed-form expression of the Whittle index as a function of the transition probabilities, which is differentiable. This completes the differentiability of Whittle index. This technique is equivalent to saying that the policy does not change if we infinitesimally change the input probabilities.