IRL for Restless Multi-Armed Bandits with Applications in Maternal and Child Health
Gauri Jain, Pradeep Varakantham, Haifeng Xu, Aparna Taneja, Prashant Doshi, Milind Tambe
TL;DR
The paper tackles learning rewards for restless RMABs under resource constraints in public health by introducing WHIRL, a gradient-based IRL framework that learns per-arm rewards from aggregate expert directives using a differentiable soft-top-$k$ policy and Whittle indices. It couples a trajectory-conversion mechanism to translate high-level goals into scalable expert-like demonstrations, enabling learning across thousands of beneficiaries. Empirical results on synthetic data and a real maternal health telehealth program in India show WHIRL outperforms baselines in both accuracy and computation, and can align interventions with expert risk-based objectives. The work supports practical deployment by providing what-if analyses, a public code release, and a discussion of deployment and ethical considerations for large-scale health resource allocation.
Abstract
Public health practitioners often have the goal of monitoring patients and maximizing patients' time spent in "favorable" or healthy states while being constrained to using limited resources. Restless multi-armed bandits (RMAB) are an effective model to solve this problem as they are helpful to allocate limited resources among many agents under resource constraints, where patients behave differently depending on whether they are intervened on or not. However, RMABs assume the reward function is known. This is unrealistic in many public health settings because patients face unique challenges and it is impossible for a human to know who is most deserving of any intervention at such a large scale. To address this shortcoming, this paper is the first to present the use of inverse reinforcement learning (IRL) to learn desired rewards for RMABs, and we demonstrate improved outcomes in a maternal and child health telehealth program. First we allow public health experts to specify their goals at an aggregate or population level and propose an algorithm to design expert trajectories at scale based on those goals. Second, our algorithm WHIRL uses gradient updates to optimize the objective, allowing for efficient and accurate learning of RMAB rewards. Third, we compare with existing baselines and outperform those in terms of run-time and accuracy. Finally, we evaluate and show the usefulness of WHIRL on thousands on beneficiaries from a real-world maternal and child health setting in India. We publicly release our code here: https://github.com/Gjain234/WHIRL.