# Scalable Decision-Focused Learning in Restless Multi-Armed Bandits with Application to Maternal and Child Health

arXiv: https://arxiv.org/abs/2202.00916 · HTML: https://mlsanity.com/paper/2202.00916/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 44 references that are in the map.
- First posted on arXiv 2 Feb 2022.

## Builds on (most cited first, 16 of 16)

- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [Smart "Predict, then Optimize"](https://mlsanity.com/paper/1710.08005/) (2017)
- [Melding the Data-Decisions Pipeline: Decision-Focused Learning for Combinatorial Optimization](https://mlsanity.com/paper/1809.05504/) (2018)
- [Objective Mismatch in Model-based Reinforcement Learning](https://mlsanity.com/paper/2002.04523/) (2020)
- [Smart Predict-and-Optimize for Hard Combinatorial Optimization Problems](https://mlsanity.com/paper/1911.10092/) (2019)
- [A Refined Laser Method and Faster Matrix Multiplication](https://mlsanity.com/paper/2010.05846/) (2020)
- [Decision Trees for Decision-Making under the Predict-then-Optimize Framework](https://mlsanity.com/paper/2003.00360/) (2020)
- [Collapsing Bandits and Their Application to Public Health Interventions](https://mlsanity.com/paper/2007.04432/) (2020)
- [Faster Dynamic Matrix Inverse for Faster LPs](https://mlsanity.com/paper/2004.07470/) (2020)
- [Field Study in Deploying Restless Multi-Armed Bandits: Assisting Non-Profits in Improving Maternal and Child Health](https://mlsanity.com/paper/2109.08075/) (2021)
- [POPCORN: Partially Observed Prediction COnstrained ReiNforcement Learning](https://mlsanity.com/paper/2001.04032/) (2020)
- [Addressing the Loss-Metric Mismatch with Adaptive Loss Alignment](https://mlsanity.com/paper/1905.05895/) (2019)
- [Differentiable Top-k Operator with Optimal Transport](https://mlsanity.com/paper/2002.06504/) (2020)
- [An order optimal policy for exploiting idle spectrum in cognitive radio networks](https://mlsanity.com/paper/1709.00237/) (2017)

## Built on by (most cited first, 12 of 12)

- [A Survey of Contextual Optimization Methods for Decision Making under Uncertainty](https://mlsanity.com/paper/2306.10374/) (2023)
- [A Decision-Language Model (DLM) for Dynamic Restless Multi-Armed Bandit Tasks in Public Health](https://mlsanity.com/paper/2402.14807/) (2024)
- [Application-Driven Innovation in Machine Learning](https://mlsanity.com/paper/2403.17381/) (2024)
- [Towards a Pretrained Model for Restless Bandits via Multi-arm Generalization](https://mlsanity.com/paper/2310.14526/) (2023)
- [Towards Foundation-model-based Multiagent System to Accelerate AI for Social Impact](https://mlsanity.com/paper/2412.07880/) (2024)
- [Evaluating the Effectiveness of Index-Based Treatment Allocation](https://mlsanity.com/paper/2402.11771/) (2024)
- [Limited Resource Allocation in a Non-Markovian World: The Case of Maternal and Child Healthcare](https://mlsanity.com/paper/2305.12640/) (2023)
- [Online Decision-Focused Learning](https://mlsanity.com/paper/2505.13564/) (2025)
- [A Dual Perspective on Decision-Focused Learning: Scalable Training via Dual-Guided Surrogates](https://mlsanity.com/paper/2511.04909/) (2025)
- [Neural Index Policies for Restless Multi-Action Bandits with Heterogeneous Budgets](https://mlsanity.com/paper/2510.22069/) (2025)
- [SMaRT: Online Reusable Resource Assignment and an Application to Mediation in the Kenyan Judiciary](https://mlsanity.com/paper/2602.18431/) (2026)
- [Decisions and Deployment: The Five-Year SAHELI Project (2020-2025) on Restless Multi-Armed Bandits for Improving Maternal and Child Health](https://mlsanity.com/paper/2604.07384/) (2026)

## Abstract (arXiv)

This paper studies restless multi-armed bandit (RMAB) problems with unknown arm transition dynamics but with known correlated arm features. The goal is to learn a model to predict transition dynamics given features, where the Whittle index policy solves the RMAB problems using predicted transitions. However, prior works often learn the model by maximizing the predictive accuracy instead of final RMAB solution quality, causing a mismatch between training and evaluation objectives. To address this shortcoming, we propose a novel approach for decision-focused learning in RMAB that directly trains the predictive model to maximize the Whittle index solution quality. We present three key contributions: (i) we establish differentiability of the Whittle index policy to support decision-focused learning; (ii) we significantly improve the scalability of decision-focused learning approaches in sequential problems, specifically RMAB problems; (iii) we apply our algorithm to a previously collected dataset of maternal and child health to demonstrate its performance. Indeed, our algorithm is the first for decision-focused learning in RMAB that scales to real-world problem sizes.