# Neural Index Policies for Restless Multi-Action Bandits with Heterogeneous Budgets

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 17 of its 35 references that are in the map.
- First posted on arXiv 24 Oct 2025.

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

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [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)
- [Reinforcement Learning in Healthcare: A Survey](https://mlsanity.com/paper/1908.08796/) (2019)
- [Melding the Data-Decisions Pipeline: Decision-Focused Learning for Combinatorial Optimization](https://mlsanity.com/paper/1809.05504/) (2018)
- [Differentiable MPC for End-to-end Planning and Control](https://mlsanity.com/paper/1810.13400/) (2018)
- [Decision-Focused Learning: Foundations, State of the Art, Benchmark and Future Opportunities](https://mlsanity.com/paper/2307.13565/) (2023)
- [Differentiable Ranks and Sorting using Optimal Transport](https://mlsanity.com/paper/1905.11885/) (2019)
- [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)
- [Opportunistic Scheduling as Restless Bandits](https://mlsanity.com/paper/1706.09778/) (2017)
- [Scalable Decision-Focused Learning in Restless Multi-Armed Bandits with Application to Maternal and Child Health](https://mlsanity.com/paper/2202.00916/) (2022)
- [Contextual Restless Multi-Armed Bandits with Application to Demand Response Decision-Making](https://mlsanity.com/paper/2403.15640/) (2024)
- [Reinforcement learning with combinatorial actions for coupled restless bandits](https://mlsanity.com/paper/2503.01919/) (2025)
- [Asymptotically Optimal Policies for Weakly Coupled Markov Decision Processes](https://mlsanity.com/paper/2406.04751/) (2024)
- [Projection-based Lyapunov method for fully heterogeneous weakly-coupled MDPs](https://mlsanity.com/paper/2502.06072/) (2025)
- [Deep Index Policy for Multi-Resource Restless Matching Bandit and Its Application in Multi-Channel Scheduling](https://mlsanity.com/paper/2408.07205/) (2024)

## Abstract (arXiv)

Restless multi-armed bandits (RMABs) provide a scalable framework for sequential decision-making under uncertainty, but classical formulations assume binary actions and a single global budget. Real-world settings, such as healthcare, often involve multiple interventions with heterogeneous costs and constraints, where such assumptions break down. We introduce a Neural Index Policy (NIP) for multi-action RMABs with heterogeneous budget constraints. Our approach learns to assign budget-aware indices to arm--action pairs using a neural network, and converts them into feasible allocations via a differentiable knapsack layer formulated as an entropy-regularized optimal transport (OT) problem. The resulting model unifies index prediction and constrained optimization in a single end-to-end differentiable framework, enabling gradient-based training directly on decision quality. The network is optimized to align its induced occupancy measure with the theoretical upper bound from a linear programming relaxation, bridging asymptotic RMAB theory with practical learning. Empirically, NIP achieves near-optimal performance within 5% of the oracle occupancy-measure policy while strictly enforcing heterogeneous budgets and scaling to hundreds of arms. This work establishes a general, theoretically grounded, and scalable framework for learning index-based policies in complex resource-constrained environments.