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

TL;DR

A Neural Index Policy (NIP) is introduced for multi-action RMABs with heterogeneous budget constraints, which establishes a general, theoretically grounded, and scalable framework for learning index-based policies in complex resource-constrained environments.

Abstract

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.

Figures (5)

• Figure 1: Neural Index Policy (NIP) Architecture. The policy features a neural network that predicts arm indices $I_n$ based on features $\mathbf{x}_n$ and state $s_n$. These indices inform a differentiable Sinkhorn-relaxed Knapsack layer to select a feasible set of multi-actions under heterogeneous budgets $\mathbf{b}_a$, yielding the Transportation Plan $\Gamma_{na}$. The entire system is trained end-to-end using a decision-focused KL-divergence loss $\mathcal{L}_{\text{KL}}(\Gamma, \omega^*)$ against a target optimal occupancy measure $\omega^*$, allowing gradients to flow back through the optimization layer to the network parameters $\theta$.
• Figure 2: Visualization of the optimal transport plan under different values of the entropy regularization parameter $\epsilon$. Higher $\epsilon$ results in smoother transport plans, while lower $\epsilon$ yields more discrete solutions.
• Figure 3: Training and validation loss (KL divergence) and corresponding percentage reward gap over epochs. Lower values indicate better alignment with the oracle policy. Results are shown for $\epsilon = 0.1$ and $N = 500$.
• Figure 4: Percentage reward gap versus cohort size $N$. The proposed methods achieve consistent improvements over the random policy, approaching the 5% reference bound. Results are shown for $\epsilon = 0.1$.
• Figure 5: Effect of the Sinkhorn regularization parameter $\epsilon$ on the average percentage reward gap across different cohort sizes. Darker regions indicate lower reward gaps and closer alignment with the oracle policy.