---
title: "Towards a Pretrained Model for Restless Bandits via Multi-arm Generalization"
authors:
  - "Yunfan Zhao"
  - "Nikhil Behari"
  - "Edward Hughes"
  - "Edwin Zhang"
  - "Dheeraj Nagaraj"
  - "Karl Tuyls"
  - "Aparna Taneja"
  - "Milind Tambe"
doi: "10.24963/ijcai.2024/36"
url: "https://ijcai.org/proceedings/2024/36"
pdf: "https://ijcai.org/proceedings/2024/0036.pdf"
venue: "Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence"
year: 2024
pages: "321-329"
track: "Main Track"
keywords:
  - "Agent-based and Multi-agent Systems: MAS: Multi-agent learning"
  - "Agent-based and Multi-agent Systems: MAS: Multi-agent planning"
  - "Agent-based and Multi-agent Systems: MAS: Resource allocation"
  - "Machine Learning: ML: Few-shot learning"
---

# Towards a Pretrained Model for Restless Bandits via Multi-arm Generalization

Yunfan Zhao, Nikhil Behari, Edward Hughes, Edwin Zhang, Dheeraj Nagaraj, Karl Tuyls, Aparna Taneja, Milind Tambe

*Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence*, pages 321-329. DOI: [10.24963/ijcai.2024/36](https://doi.org/10.24963/ijcai.2024/36)

## Abstract

Restless multi-arm bandits (RMABs) is a class of resource allocation problems with broad application in areas such as healthcare, online advertising, and anti-poaching. We explore several important question such as how to handle arms opting-in and opting-out over time without frequent retraining from scratch, how to deal with continuous state settings with nonlinear reward functions, which appear naturally in practical contexts. We address these questions by developing a pre-trained model (PreFeRMAB) based on a novel combination of three key ideas: (i) to enable fast generalization, we use train agents to learn from each other's experience; (ii) to accommodate streaming RMABs, we derive a new update rule for a crucial $\lambda$-network; (iii) to handle more complex continuous state settings, we design the algorithm to automatically define an abstract state based on raw observation and reward data. PreFeRMAB allows general zero-shot ability on previously unseen RMABs, and can be fine-tuned on specific instances in a more sample-efficient way than retraining from scratch.  We theoretically prove the benefits of multi-arm generalization and empirically demonstrate the advantages of our approach on several challenging, real-world inspired problems.

## Full text

## 1 Introduction

Restless multi-arm bandits (RMABs), a class of resource allocation problems involving multiple agents with a global resource constraint, have found applications in various scenarios, including resource allocation in multi-channel communication, machine maintenance, and healthcare [Hodge and Glazebrook, 2015; Mate et al., 2022]. RMABs have recently been studied from a multi-agent reinforcement learning perspective.

The usual RMAB setting considers a fixed number of arms, each associated with a known, fixed MDP with finite state and action spaces; the RMAB chooses K of N arms every round to optimize some long term objective. Even in this setting, the problem has been shown to be PSPACE hard [Papadimitriou and Tsitsiklis, 1999]. Several approximation algorithms have been proposed in this setting [Whittle, 1988; Hawkins, 2003], particularly when MDP transition probabilities are fully specified, which are successful in practice. State-of-the-art approaches for binary action RMABs commonly provide policies based on the Whittle index [Whittle, 1988], an approach that has also been generalized to multiaction RMABs[Hawkins, 2003; Killian et al., 2021b]. There are also linear programming-based approaches to both binary and multi-action RMABs [Zhang and Frazier, 2021]. Reinforcement learning (RL) based techniques have also been proposed as state-of-the-art solutions for general multi-action RMABs [Xiong and Li, 2023].

In this work, we focus on RL-based methods that provide general solutions to binary and multi-action RMABs, without requiring ground truth transition dynamics, or special properties such as indexability as required by other approaches [Wang et al., 2023]. Unfortunately, several limitations exist in current RMAB solutions, especially for state of the art RL-based solutions, making them challenging or inefficient to deploy in real-world resource allocation problems.

The first limitation arises when dealing with arms that constantly opt-in (also known as streaming RMABs), which happens in public health programs where new patients (arms in RMABs) arrive asynchronously. Existing solutions either require ground truth transition probabilities [Mate et al., 2021], which are often unknown in practice, or else require an entirely new model to be trained repeatedly, which can be extremely computationally costly and sample inefficient.

A second limitation occurs for new programs, or existing programs experiencing a slight change in the user base. In these situations, existing approaches do not provide a pretrained RMAB model that can be immediately deployed. In deep learning, pretrained models are the foundation for contemporary, large-scale image and text networks that generalize well across a variety of tasks [Bommasani et al., 2021]. For real-world problems modeled with RMABs, establishing a similar pretrained model is essential to reduce the burden of training new RMAB policies from scratch.

The third limitation occurs in handling continuous state multi-action RMABs that have important applications [Sinha and Mahajan, 2022; Dusonchet and Hongler, 2003]. Naturally continuous domain state-spaces, such as patient adherence, are often binned into manually crafted discrete state spaces to improve model tractability and scalability [Mate et al., 2022], resulting in the loss of crucial information about raw observations.

In this work we present PreFeRMAB, a Pretrained Flexible model for RMABs. Using multi-arm generalization, PreFeRMAB enables zero-shot deployment for unseen arms as well as rapid fine-tuning for specific RMAB settings.

## Our main contributions are:

- To the best of our knowledge, we are the first to develop a pretrained RMAB model with zero shot ability on entire sets of unseen arms.
- Whereas a general multiagent RL system could suffer from sample complexity exponential in the number of agents N [Gheshlaghi Azar et al., 2013], we prove PreFeRMAB benefits from larger N, via multi-arm generalization and better estimation of the population distribution of arm features.
- Our pretrained model can be fine-tuned on specific instances in a more sample-efficient way than training from scratch, requiring less than 12.5% of samples needed for training a previous multi-action RMAB model in a healthcare setting [Verma et al., 2023].
- We derive an update rule for a crucial λ-network, allowing changing numbers of arms without retraining. While streaming bandits received considerable attention [Liau et al., 2018], we are the first to handle streaming multiaction RMABs with unknown transition dynamics.
- Our model accommodates both discrete and continuous states. To address the continuous state setting, where real-world problems often require nonlinear rewards [Riquelme et al., 2018], we providing a StateShaping module to automatically define an abstract state.

## 2 RelatedWork

RMABs with binary and multiple actions. Solving an RMAB problem, even with known transition dynamics, is known to be PSPACE hard [Papadimitriou and Tsitsiklis, 1999]. For binary action RMABs, [Whittle, 1988] provides an approximate solution, using a Lagrangian relaxation to decouple arms and choose actions by computing so-called Whittle indices of each arm. It has been shown that the Whittle index policy is asymptotically optimal under the indexability condition [Weber and Weiss, 1990; Akbarzadeh and Mahajan, 2019]. The Whittle index was extended to a special class of multi-action RMABs with monotonic structure [Hodge and Glazebrook, 2015]. A method for more general multi-action RMABs based on Lagrangian relaxation was proposed by [Killian et al., 2021b]. Weakly coupled Markov Decisions Processes (WCMDP), which generalizes multi-action RMABs to have multiple constraints, was studied by [Hawkins, 2003], who proposed a Langrangian decomposition approach. WCMDP was subsequently studied by [Adelman and Mersereau, 2008], who proposed improvements in solution quality at the expense of higher computational costs. While above methods developed for WCMDP require knowledge of ground truth transition dynamics, our algorithm handles unknown transition dynamics, which is more common in practice[Wang et al., 2023]. Additionally, the above works in multi-action settings do not provide algorithms for continuous state RMABs.

Multi-agent RL and RL for RMABs. RMABs are a specific instance of the powerful multi-agent RL framework used to model systems with multiple interacting agents in both competitive and co-operative settings [Shapley, 1953; Littman, 1994], for which significant strides have been made empirically [Jaques et al., 2019; Yu et al., 2022] and theoretically [Jin et al., 2021; Xie et al., 2020]. [Nakhleh et al., 2021] proposed a deep RL method to estimate the Whittle index. [Fu et al., 2019] provided an algorithm to learn a Q-function based on Whittle indices, states, and actions. [Avrachenkov and Borkar, 2022] and [Biswas et al., 2021] developed Whittle index-based Q-learning methods with convergence guarantees. While the aforementioned works focus on binary action RMABs, [Killian et al., 2021a] generalized this to multi-action RMABs using tabular Q-learning. A subsequent work [Killian et al., 2022], which focussed on robustness against adversarial distributions, took a deep RL approach that was more scalable. However, existing works on multi-action RMABs do not consider streaming RMABs and require training from scratch when a new arm opts-in. Additionally, works built on tabular Q-learning [Fu et al., 2019; Avrachenkov and Borkar, 2022; Biswas et al., 2021; Killian et al., 2021a] may not generalize to continuous state RMABs without significant modifications. Our pretrained model addresses these limitations, and enables zero-shot ability on a wide range of unseen RMABs.

Streaming algorithms. The streaming model, pioneered by [Alon et al., 1996], considers a scenario where data arrives online and the amount of memory is limited. The model is adapted to multi-arm bandits (MAB), assuming that arms arrive in a stream and the number of arms that can be stored is limited [Liau et al., 2018; Chaudhuri and Kalyanakrishnan, 2020]. The streaming model has recently been adapted to binary action RMABs with known transition probabilities[Mate et al., 2021], but not studied in the more general and practical settings of multi-action RMABs with unknown transition dynamics. We aim to close this gap.

Zero-shot generalization and fine-tuning. Foundation models that have a strong ability to generalize to new tasks in zero shot and efficiently adapt to new tasks via fine-tuning have received great research attention [Bommasani et al., 2021]. Such models are typically trained on vast data, such as internet-scale text [Devlin et al., 2018] or images [Ramesh et al., 2021]. RL has seen success in the direction of foundation models for decision making, using simulated [Team et al., 2023] and real-world [Yu et al., 2020] environments. A pretrained model for RMABs is needed [Zhao et al., 2024]. To our knowledge, we are the first to realize zero-shot generalization and efficient fine-tuning in the setting of RMABs.

## 3 Problem Statement

We study multi-action RMABs with system capacity N, where existing arms have the option to opt-out (that is, the state-action-rewards corresponding to them are disregarded by the model post opt-out), and new, unseen arms can request to opt-in (that is, these arms are considered only post the optin time). Such requests will be accepted if and only if the system capacity permits. A vector ξt ∈ {0, 1}N represents the opt-in decisions:

<!-- formula-not-decoded -->

Notice that existing arms must opt-in in each round t to remain in the system. For each arm i ∈ [N], the state space Si can be either discrete or continuous, and the action space Ai is a finite set of discrete actions. Each action a ∈ Ai has an associated cost Ci(a), with Ci(0) denoting a no-cost passive action. The reward at a state is given by a function Ri : Si → R. We let β ∈ [0, 1) denote a discount factor. Each arm has a unique feature vector zi ∈ Rm that provides useful information about the arm. Notice our model directly utilizes feature information in its policy network, without requiring intermediate steps to extract transition dynamics information from features.

When the state space is discrete, each arm i ∈ [N] follows a Markov Decision Process (Si, Ai, Ci, Ti, Ri, β, zi), where Ti : Si ×Ai ×Si → [0, 1] is a transition matrix representing the probability of transitioning from the current state to the next state given an action. In contrast, when the state space is continuous, each arm i ∈ [N] follows a Markov Decision Process (Si, Ai, Ci, Γi, Ri, β, zi), where Γi is a set of parameters encoding the transition dynamics. For example, in the case that the next state moves according to a Gaussian distribution, Γi may denote the mean and variance of the Gaussian.

For simplicity, we assume that Si,Ai, Ci, and Ri are the same for all arms i ∈ [N] and omit the subscript i. Note that our algorithms can also be used in the general case where rewards and action costs are different across arms. For ease of notation, we let s ∈ RN denote the state over all arms, and we let A ∈ {0, 1}N×|A| denote one-hot-encoding of the actions taken over all arms. The agent learns a policy π that maps states s and features z to actions A, while satisfying a constraint that the sum cost of actions taken is no greater than a given budget B in every timestep t ∈ [H], where H is the length of the horizon.

Our goal is to learn an RMAB policy that maximizes the following Bellman equation The key difficulty in learning such a policy is how to utilize features z and address optin decisions ξ. These are important research questions not addressed in previous works [Killian et al., 2022].

<!-- formula-not-decoded -->

<!-- formula-not-decoded -->

<!-- formula-not-decoded -->

where cj ∈ C is the cost of j th action, and Aij = 1 if action j is chosen on arm i and Aij = 0 otherwise. Further, we assume that the rewards R are uniformly bounded by Rmax.

## 4 GeneralizedModel for RMABs

We first provide an overview of key ideas and then discuss each of the ideas in more detail. (See Figure 3 in Appendix for an overview of the training procedure.)

## 4.1 Key Algorithmic Ideas

Several key algorithmic novelties are necessary for our model to address limitations of existing works:

A pretrained model via multi-arm generalization: We train agents to learn from each others' experience. Whereas a general multiagent RL system could suffer from sample complexity exponential in the number of arms N [Gheshlaghi Azar et al., 2013], we prove PreFeRMAB benefits from a larger N, via generalization across arms.

A novel λ-network updating rule for opt-in: The opt-in and opt-out of arms induce a more complex form of the Lagrangian and add randomness to actions taken by agents. We provide a new λ-network update rule and train PreFeRMAB with opt-in and opt-out of arms, to enable zero-shot performance across various opt-in rates and accommodate streaming RMABs.

Handling continuous states with StateShaping subroutine: In the continuous state setting, real-world problems often require nonlinear rewards [Riquelme et al., 2018], and naively using raw observations to train models may result in poor performance (see Figure 5). To tackle this challenge, we design the algorithm to automatically define an abstract state based on raw observation and reward data.

## 4.2 A PretrainedModel viaMulti-arm Generalization

To enable multi-arm generalization, we introduce featurebased Q-values, together with a Lagrangian relaxation with features zi and opt-in decisions ξi:

<!-- formula-not-decoded -->

<!-- formula-not-decoded -->

<!-- formula-not-decoded -->

where Q is the Q-function, ai is the action of arm i, s′i is the state transitioned to from si under action ai, and π(λ) is the optimal policy under a given λ. Notice that this relaxation decouples the Q-functions of the arms, and therefore Qi can be solved independently for a given λ.

Now we discuss how we use feature-based Q-values and how agents could learn from each other. During pretraining, having received arms' opt-in and out decisions (line 5), Algorithm 1 samples an action-charge λ based on updated opt-in decisions ξ and features zi (line 6). Next, from opt-in arms we collect trajectories (lines 7-14), which are later used to train a single pair of actor/critic networks for all arms, allowing the policy for one arm to benefit from other arms' data. After that, we update the policy network θ and the critic network ϕ (Line 16), using feature-based Q-values to compute advantage estimates for the actor in PPO update. Critically, feature-based Q-values updated with one arm's data, improves the policy for other arms. In real-world problems with missing feature entries or less informative features, it is more important for agents to learn from each other (see Table 1 in Sec 5.2). Intuitively, if a model only learns from homogeneous arms, then we should expect this model to perform poorly when used out-of-the-box on arms with completely different behaviors.

## Algorithm 1 PreFeRMAB (Training)

- 1: Input: n epochs, n steps, λ-update frequency K ∈ N+, and system capacity N 2: Initialize actor θ, critic ϕ, λ-network Λ, buffer = [], state s ∈ RN , and features zi ∈ Rm 3: Initialize StateShaping, and set ¯ s ←StateShaping(s) 4: for epoch = 1, 2, . . . , n epochs do 5: Receive opt-in/out requests and update ξ and zi. 6: Compute λ = Λ(¯ s, {zi}Ni=1 , ξ) 7: for timestep t = 1, . . . , n steps do 8: for Arm i = 1, . . . , N do 9: if Arm i is opt-in (i.e. ξi = 1) then 10: Sample an action ai ∼ θ(¯ si, λ, zi) 11: s′i , ri = Simulate(si, ai) 12: ¯ s′i = StateShaping(s′i ) 13: Add tuple (si, ¯ si, ai, ri, ¯ s′i , zi) to buffer 14: si ← s′i , ¯ si ← ¯ s′i 15: Add tuple (λ, ξ) to buffer 16: Update the (θ, ϕ) pair using buffer. 17: if epoch // K = 0 then 18: Update Λ via Prop 2 using trajectories in buffer 19: Update ˆ r(·) in StateShaping using (s, r)-tuples in buffer

We will now give a theoretical guarantee of multi-arm generalization by considering the following simplified setting and assumptions where we do not consider opt-in and optout. That is, in each epoch we draw a new set of N arms. We let the distribution of arm features (i.e, zi) to be denoted by the probability measure µ∗ . Each 'sample' for our policy network training consists of N features corresponding to N arms (z1, . . . , zN), drawn i.i.d. from the distribution µ∗ . Call the empirical distribution of (zi) to be ˆµ. During training, we receive nepochs i.i.d. draws of N arm features each, denoted by ˆµ1, . . . , ˆµnepochs

Let Θ denote the space of neural network weights of the policy network (for clarity, we shorten the (θ, ϕ) in Algorithm 1 to θ). The neural network inputs are Lagrangian multiplier λ, state of an arm s, its feature z and the output is a ∈ A. Let V (s, θ, λ, ˆµ) denote the discounted reward, averaged over N arms with features ˆµ obtained with the neural network with parameter θ, starting from the state s (cumulative state of all arms). The proposition below shows the generalization properties of the output of Algorithm 1. The proof and a detailed discussion of the assumptions and consequences are given in Section D.

Proposition 1. Suppose the following assumptions hold:

1. Algorithm 1 learns neural network weights ˆ θ ∈ Θ, whose policy is optimal for each (ˆµi, λ) for 1 ≤ i ≤

nepochs and λ ∈ [0, λmax]

2. There exists θ∗ ∈ Θ which is optimal for every instance (ˆµ, λ).
3. Θ = B2(D, Rd ), the ℓ2 ball of radius D in Rd .
4. |V (s, θ1, λ, ˆµ) - V (s, θ2, λ, ˆµ)| ≤ L∥θ1 - θ2∥ and |V (s, θ, λ1, ˆµ) - V (s, θ, λ2, ˆµ)| ≤ L|λ1 - λ2| for all θ1, θ2, θ ∈ Θ and λ1, λ2, λ ∈ [0, λmax].

Then, the generalization error over unseen arms (ˆµ) satisfies:

<!-- formula-not-decoded -->

Here, ˜O hides polylogarithmic factors in nepochs,N and constants depending on d, D, L, β, B N , cj , Rmax and λmax

The assumption of existence of θ∗ is reasonable: This means that there exists a neural network which gives the optimal policy for a family of single-arm MDPs indexed by (z, λ). Proposition 1 shows that when nepochs and N are large, the Lagrangian relaxed value function of the learned network is close to that of the optimal network.

An important insight is that the generalization ability of the PreFeRMAB network becomes better as the number of arms per instance becomes larger. This is counter intuitive since a system with a larger number of agents are generally very complex. Jointly, the arms form an MDP with |S|N states and |A|N actions. General multi-agent RL problems with N arms thus can suffer from an exponential dependence on N in their sample complexity for learning (see sample complexity lower bounds in [Gheshlaghi Azar et al., 2013]). However, due to the structure of RMABs and the Lagrangian relaxation, we achieve a better generalization with a larger N. Our proof in the appendix shows that this is due to the fact that a larger number of arms helps estimate the population distribution µ∗ of the arm features better. We show in Table 1 that indeed having more number of arms helps the PreFeRMAB network generalize better over unseen instances.

## 4.3 ANovel λ-network Updating Rule

In real-world health programs, we may observe new patients constantly opt-in [Mate et al., 2021]. The opt-in / opt-out decisions render the updating rule in [Killian et al., 2022] unusable and add additional randomness to actions taken by the agent. To overcome this challenge and to stabilize training, we develop a new λ-network updating rule.

Proposition 2. [λ-network updating rule] The equation for gradient descent for the objective (Eq 2) with respect to λ, with step size α is:

<!-- formula-not-decoded -->

where ci,t is the cost of the action taken by the optimal policy on arm i in round t.

Critically, this update rule allows PreFeRMAB to handle streaming RMABs, accommodating a changing number of arms without retraining and achieving strong zero-shot performance across various opt-in rates (see Table 3 and 14). Having established an updating rule, we provide a convergence guarantee. The proofs are relegated to Appendix E.

Proposition 3 (Convergence of λ-network). Suppose the arm policies converge to the optimal Q-function for a given Λt, then the update rule (in Prop 2) for the λ-network converges to the optimal as the number of training epochs and the number of actions collected in each epoch go to infinity.

## 4.4 Handling Continuous States with StateShaping

Real-world problems may require continuous states with nonlinear rewards [Riquelme et al., 2018]. Existing RMAB algorithms either use a human-crafted discretization or fail to address challenging nonlinear rewards [Killian et al., 2022]. Discretization may result in loss of information and fail to generalize to different population sizes. For example, the popular SIS epidemic model [Yaesoubi and Cohen, 2011] is expected to scale to a continuum limit as the population size increases to infinity, and a continuous state-space model can better handle scaling by using proportions instead of absolute numbers. Under nonlinear rewards, naively using raw observations in training may result in poor performance (see Figure 5). We provide a StateShaping module to improve model stability and performance.

## Algorithm 2 StateShaping Subroutine

- 1: Input: estimator ∈ {Isotonic Regression,KNN}, states s ∈ RN , data D of (s, r) tuples
- 2: Output ¯ s = s if no normalization is desired
- 3: Compute

<!-- formula-not-decoded -->

- 4: Compute ˆ r(si) using the choice of Estimator.
- 5: Output ¯ s, where ¯ si = ˆ r(si)-rmin rmax-rmin (smax - smin), ∀i

In Algorithm 2, users can choose whether to obtain abstract state [Abel et al., 2018] from raw observations (lines 2). We compute ranges of reward and raw observations, and obtain an reward estimate (lines 3-4). After that, we automatically refine the raw observation such that reward is a linear function of the abstract state (line 4), improving model stability for challenging reward functions. Here a key assumption is that reward is an increasing function of the raw observation, which is common in RMABs [Killian et al., 2022]. Notice as we collect more observations, the accuracy of the reward estimator ˆ r(·) will improve (it is updated in line 24 of Algorithm 1).

StateShaping instantiates the idea of state abstraction, which is shown to improve generalizability and robustness [Li et al., 2006], in the RMAB context for continuous states. Applying Theorem 3 in [Li et al., 2006] to the Lagrangian relaxation (Eq. 2), we have that an optimal policy learnt using the abstract state space is guaranteed to be also optimal in the ground MDP (defined by raw observations).

## 4.5 Inference Using PretrainedModel

An important difference between training and inference is that during inference time, we strictly enforce the budget constraint on the trained model, by greedily selecting highest probability actions until the budget is reached. The rest of the inference components are similar to the training component.

## Algorithm 3 PreFeRMAB (Inference)

- 1: Input: : States s, costs C, budget B, features zi ∈ Rm, opt-in decisions ξ, agent actor θ, λ-network, StateShaping routine with trained estimator ˆ r(·) 2: Compute λ = Λ(¯ s, {zi}Ni=1 , ξ) 3: for Arm i = 1, . . . , N do 4: if Arm i is opt-in (i.e. ξi = 1) then 5: ¯ si = StateShaping(si) 6: Compute pi ∼ θ(¯ si, λ, zi) ▷ Greedily select highest
- 7: a = GreedyProba(p,C,B) probability actions until budget B is reached

## 5 Experimental Evaluation

We provide experimental evaluations of our model in three separate domains, including a synthetic setting, an epidemic modeling setting, as well as a maternal healthcare intervention setting. We first describe these three experimental domains. Then, we provide results for PreFeRMAB in a zeroshot evaluation setting, demonstrating the performance of our model on new, unseen test arms drawn from distributions distinct from those in training. Here, we demonstrate the flexibility of PreFeRMAB, including strong performance across domains, state representations (discrete vs. continuous), and over various challenging reward functions. Finally, we demonstrate the strength of using PreFeRMAB as a pretrained model, enabling faster convergence for fine-tuning on a specific set of evaluation arms. Code is available at https://github.com/yzhao3685/PreFeRMAB

In Appendix B, we provide ablation studies over (1) a wider range of opt-in rates (2) different feature mappings (3) DDLPO topline with and without features (4) more problem settings. For hyperparameter details, we refer to Appendix A.

## 5.1 Experimental Settings

Features: In all experiments, we generate features by projecting parameters that describe the ground truth transition dynamics into features using randomly generated projection matrices. The dimension of feature equals the number of parameters required to describe the transition dynamics. In Appendix B, we provide results on different feature mappings.

Synthetic: Following [Killian et al., 2022], we consider a synthetic dataset with binary states and binary actions. The transition probabilities for each arm i are represented by matrices T (i) s=0 and T (i) s=1 for arm i at states 0 and 1 respectively:

<!-- formula-not-decoded -->

Each pjk corresponds to the probability of transitioning from state j to state 0 when action k is taken. These values are sampled uniformly from the intervals:

<!-- formula-not-decoded -->

SIS Epidemic Model: Inspired by the vast literature on agent-based epidemic modeling, we adapt the SIS model given in [Yaesoubi and Cohen, 2011], following a similar experiment setup as described in [Killian et al., 2022]. Arms p represent a subpopulation in distinct geographic regions; states s are the number of uninfected people within each arm's total population Np; the number of possible states is S. Transmission within each arm is guided by parameters: κ, the average number of contacts within the arm's subpopulation in each round, and rinfect, the probability of becoming infected after contact with an infected person.

In this setting, there is a budget constraint over interventions. There are three available intervention actions a0, a1, a2 that affect the transmission parameters: a0 represents no action; a1 represents messaging about physical distancing; a2 represents the distribution of face masks. We discuss additional details in Appendix A.

ARMMAN: Similar to the set up in [Biswas et al., 2021; Killian et al., 2022], we model the real world maternal health problem as a discrete state RMAB. We aim to encourage engagement with automated health information messaging. There are three possible states, presenting self-motivated, persuadable, and lost cause. The actions are binary. There are 6 uncertain parameters per arm, sampled from uncertainty intervals of 0.5 centered around the transition parameters that align with summary statistics given in [Biswas et al., 2021].

Continuous State Modeling: Continuous state restless bandits have important applications [Lef´evre, 1981; Sinha and Mahajan, 2022; Dusonchet and Hongler, 2003]. By not explicitly having a switch in the model (switching between discrete and continuous state space), we enable greater model flexibility. To demonstrate this, we consider both a Continuous Synthetic and a Continuous SIS modeling setting. We provide details of these settings in Appendix A.3.

We present additional details, including hyperparameters and StateShaping illustration in Appendix A.

## 5.2 PreFeRMAB Zero-Shot Learning

We first consider three challenging datasets in the discrete state space. After that, we present results on datasets with continuous state spaces with more complex reward functions and transition dynamics.

Pretraining. For each pretraining iteration, we sample from a binomial with mean 0.8 to determine which arms will be opted-in given system capacity N. For new arms, we sample new transition dynamics to allow the model to see a wider range of arm features.

Evaluation. We compare PreFeRMAB to Random Action and No Action baselines. In every table in this subsection, we present the reward per arm averaged over 50 trials, on new, unseen arms arm sampled from the testing distribution.

Multi-arm Generalization: Table 1 on Synthetic illustrates that PreFeRMAB, learning from multi-arm generalization, achieves stronger performance when the number of unique arms (i.e. arms with unique features) seen during pretraining increases. Additionally, in practice arm features may be missing or not always reliable, such as in real-world ARMMAN data [Mate et al., 2022]. Our results demonstrate that when features are masked, arms could learn from similar arms' experience.

Table 1: Multi-arm generalization results on Synthetic (opt-in 100%). With the same total amount of data, PreFeRMAB achieves stronger performance when pretrained on more unique arms, especially when input arm features are masked.

| System capacity N = 21. Budget B = 7. - # Unique training arms   | System capacity N = 21. Budget B = 7. - 45   | System capacity N = 21. Budget B = 7. - 33   | 21        |
|------------------------------------------------------------------|----------------------------------------------|----------------------------------------------|-----------|
| No Action                                                        | 2.88±0.17                                    | 2.88±0.17                                    | 2.88±0.17 |
| Random Action                                                    | 3.25±0.22                                    | 3.25±0.22                                    | 3.25±0.22 |
| PreFeRMAB (2/4 Feats. Masked)                                    | 3.81±0.23                                    | 3.79±0.22                                    | 3.59±0.21 |
| PreFeRMAB (1/4 Feats. Masked)                                    | 3.92±0.24                                    | 3.70±0.21                                    | 3.58±0.20 |
| PreFeRMAB (0/4 Feats. Masked)                                    | 4.02±0.26                                    | 3.80±0.22                                    | 3.78±0.21 |

Table 2: Results on Synthetic (opt-in 100%). For each system capacity, we pretrain a model and present zero-shot results under different amounts of distributional shift.

| Wasserstein Distance - System capacity N = 48. Budget B = 16.   | 0.05 - System capacity N = 48. Budget B = 16.   | 0.10 - System capacity N = 48. Budget B = 16.   | 0.15 - System capacity N = 48. Budget B = 16.   | 0.20 - System capacity N = 48. Budget B = 16.   | 0.25 - System capacity N = 48. Budget B = 16.   |
|-----------------------------------------------------------------|-------------------------------------------------|-------------------------------------------------|-------------------------------------------------|-------------------------------------------------|-------------------------------------------------|
| No Action RandomAction PreFeRMAB                                | 3.07±0.10 3.49±0.09 4.50±0.09                   | 2.89±0.08 3.25±0.09 4.30±0.10                   | 2.68±0.07 2.99±0.16 3.81±0.17                   | 2.49±0.09 2.80±0.17 3.79±0.18                   | 2.35±0.07 2.57±0.17 3.46±0.12                   |
| System capacity N = 96. Budget B = 32.                          | System capacity N = 96. Budget B = 32.          | System capacity N = 96. Budget B = 32.          | System capacity N = 96. Budget B = 32.          | System capacity N = 96. Budget B = 32.          | System capacity N = 96. Budget B = 32.          |
| No Action RandomAction PreFeRMAB                                | 3.09±0.08 3.44±0.14 4.44±0.13                   | 2.88±0.04 3.23±0.09 4.26±0.13                   | 2.74±0.05 3.05±0.09 4.12±0.13                   | 2.62±0.06 2.90±0.10 3.97±0.16                   | 2.49±0.06 2.70±0.11 3.75±0.12                   |

Table 3: Results on SIS (N = 20, B = 16, S = 150). We pretrain a model and present zero-shot results on various distributions. During training, aeff 1 , aeff 1 are uniformly sampled from [1, 7].

| Number of arms System capacity                                | 80%                                                           | 85%                                                           | 90%                                                           | 95%                                                           | 100%                                                          |
|---------------------------------------------------------------|---------------------------------------------------------------|---------------------------------------------------------------|---------------------------------------------------------------|---------------------------------------------------------------|---------------------------------------------------------------|
| Parameters aeff 1 , aeff 1 are uniformly sampled from [2, 8]. | Parameters aeff 1 , aeff 1 are uniformly sampled from [2, 8]. | Parameters aeff 1 , aeff 1 are uniformly sampled from [2, 8]. | Parameters aeff 1 , aeff 1 are uniformly sampled from [2, 8]. | Parameters aeff 1 , aeff 1 are uniformly sampled from [2, 8]. | Parameters aeff 1 , aeff 1 are uniformly sampled from [2, 8]. |
| No Action RandomAction PreFeRMAB                              | 5.23±0.17 6.94±0.15 7.64±0.27                                 | 5.27±0.16 7.00±0.16 7.75±0.25                                 | 5.28±0.16 7.03±0.15 7.96±0.18                                 | 5.26±0.14 6.97±0.14 7.80±0.16                                 | 5.28±0.13 6.99±0.12 7.82±0.11                                 |
| Parameters aeff 1 , aeff 1 are uniformly sampled from [3, 9]. | Parameters aeff 1 , aeff 1 are uniformly sampled from [3, 9]. | Parameters aeff 1 , aeff 1 are uniformly sampled from [3, 9]. | Parameters aeff 1 , aeff 1 are uniformly sampled from [3, 9]. | Parameters aeff 1 , aeff 1 are uniformly sampled from [3, 9]. | Parameters aeff 1 , aeff 1 are uniformly sampled from [3, 9]. |
| No Action RandomAction PreFeRMAB                              | 5.29±0.16 7.21±0.15 7.77±0.29                                 | 5.30±0.17 7.28±0.18 7.87±0.28                                 | 5.29±0.15 7.26±0.15 7.90±0.22                                 | 5.26±0.14 7.22±0.13 7.95±0.16                                 | 5.28±0.13 7.22±0.12 7.95±0.11                                 |

Discrete State Settings with Different Distributional Shifts: Results on Synthetic (Table 2) shows PreFeRMAB consistently outperforms under varying amounts of distributional shift, measured in Wasserstein distance. Results on SIS (Table 3) shows PreFeRMAB performs well in settings with large state space S = 150 and multiple actions, under various testing distributions and opt-in rates. Results on ARMMAN (Table 4) shows PreFeRMAB could handle more challenging settings that mimics the scenario of a real-world non-profit organization using RMABs to allocate resources.

Continuous State Settings: Our results (Figure 1) show that StateShaping is crucial in handling continuous states, where the reward function can be more challenging. We provide additional evaluations in Table 5, showing PreFeRMAB outperforms in complex transition dynamics. More details are provided in Appendix A.

Table 4: Results on ARMMAN (N = 25, B = 7, S = 3). We pretrain a model and present zero-shot results on various testing distributions. During training, the proportion of self-motivated, persuadable, and lost cause arms are 20%, 20%, annd 60% respectively.

| Number of arms System capacity                      | 80%                                                 | 85%                                                 | 90%                                                 | 95%                                                 | 100%                                                |
|-----------------------------------------------------|-----------------------------------------------------|-----------------------------------------------------|-----------------------------------------------------|-----------------------------------------------------|-----------------------------------------------------|
| 40% motivated, 20% persuadable, and 40% lost cause. | 40% motivated, 20% persuadable, and 40% lost cause. | 40% motivated, 20% persuadable, and 40% lost cause. | 40% motivated, 20% persuadable, and 40% lost cause. | 40% motivated, 20% persuadable, and 40% lost cause. | 40% motivated, 20% persuadable, and 40% lost cause. |
| No Action RandomAction PreFeRMAB                    | 2.39±0.30 3.04±0.40 5.47±0.41                       | 2.32±0.28 3.06±0.38 5.00±0.37                       | 2.16±0.25 3.00±0.36 4.95±0.29                       | 2.26±0.28 3.14±0.36 5.34±0.27                       | 2.25±0.30 3.24±0.32 5.03±0.37                       |
| 40% motivated, 40% persuadable, and 20% lost cause. | 40% motivated, 40% persuadable, and 20% lost cause. | 40% motivated, 40% persuadable, and 20% lost cause. | 40% motivated, 40% persuadable, and 20% lost cause. | 40% motivated, 40% persuadable, and 20% lost cause. | 40% motivated, 40% persuadable, and 20% lost cause. |
| No Action RandomAction PreFeRMAB                    | 2.07±0.29 3.04±0.37 5.06±0.36                       | 2.19±0.28 3.02±0.33 4.81±0.35                       | 2.19±0.30 2.99±0.30 5.13±0.34                       | 2.05±0.24 3.12±0.31 5.01±0.26                       | 2.17±0.28 3.15±0.29 5.00±0.28                       |

Figure 1: Results for Continuous Synthetic domain (N=21,B=7.0) with challenging rewards r(s).

![Image](https://ijcai.org/proceedings/2024/fulltext/36/image_000000_918c39846e0623e734178ff7a37152eacfb44dcd25ced8ab8649d15a3200ed0a.png)

Table 5: Results on continuous states. For each problem instance, we pretrain a model.

| Number of arms System capacity          | 80%                                     | 85%                                     | 90%                                     | 95%                                     | 100%                                    |
|-----------------------------------------|-----------------------------------------|-----------------------------------------|-----------------------------------------|-----------------------------------------|-----------------------------------------|
| Continuous Synthetic (N=21, B=7.0, S=2) | Continuous Synthetic (N=21, B=7.0, S=2) | Continuous Synthetic (N=21, B=7.0, S=2) | Continuous Synthetic (N=21, B=7.0, S=2) | Continuous Synthetic (N=21, B=7.0, S=2) | Continuous Synthetic (N=21, B=7.0, S=2) |
| No Action RandomAction PreFeRMAB        | 0.70±0.25 3.44±0.48 3.94±0.31           | 0.71±0.25 3.43±0.45 3.76±0.33           | 0.70±0.23 3.45±0.41 4.01±0.31           | 0.70±0.23 3.37±0.41 4.02±0.29           | 0.66±0.21 3.20±0.38 3.67±0.27           |
| Continuous SIS Model (N=20, B=16)       | Continuous SIS Model (N=20, B=16)       | Continuous SIS Model (N=20, B=16)       | Continuous SIS Model (N=20, B=16)       | Continuous SIS Model (N=20, B=16)       | Continuous SIS Model (N=20, B=16)       |
| No Action RandomAction PreFeRMAB        | 5.64±0.25 7.11±0.22 7.91±0.17           | 5.68±0.19 7.31±0.22 8.08±0.11           | 5.57±0.21 7.23±0.22 7.95±0.14           | 5.48±0.18 7.24±0.21 7.98±0.13           | 5.62±0.17 7.18±0.17 7.82±0.12           |

Comparison with an Additional Baseline: DDLPO [Killian et al., 2022] could not handle distributional shifts or various opt-in rates, the more challenging settings that PreFeRMAB is designed for. Nevertheless, we provide comparisons with DDLPO in settings with no distributional shift (Table 6, see also Appendix B.4). Notably, PreFeRMAB zero-shot performance on unseen arms is near that of DDLPO, which is trained and tested on the same set of arms.

## 5.3 PreFeRMAB Fast Fine-Tuning

Having shown the zero-shot results of PreFeRMAB, we now demonstrate finetuning capabilities of the pretrained model.

Table 6: Comparison of PreFeRMAB zero-shot performance on unseen arms against that of DDLPO trained and tested on the same set of arms. For each problem instance, we pretrain a model.

| Number of arms System capacity                   | 80%                                   | 85%                                   | 90%                                   | 95%                                   | 100%                                    |
|--------------------------------------------------|---------------------------------------|---------------------------------------|---------------------------------------|---------------------------------------|-----------------------------------------|
| Synthetic with N = 96, B = 32, S = 2.            | Synthetic with N = 96, B = 32, S = 2. | Synthetic with N = 96, B = 32, S = 2. | Synthetic with N = 96, B = 32, S = 2. | Synthetic with N = 96, B = 32, S = 2. | Synthetic with N = 96, B = 32, S = 2.   |
| No Action RandomAction PreFeRMAB DDLPO (topline) | 3.22±0.12 3.62±0.13 4.63±0.12 n/a     | 3.24±0.12 3.66±0.13 4.71±0.12 n/a     | 3.19±0.11 3.58±0.13 4.53±0.13 n/a     | 3.18±0.11 3.60±0.12 4.47±0.12 n/a     | 3.18±0.11 3.60±0.12 4.61±0.10 4.58±0.13 |
| SIS with N = 20, B = 16, S = 150.                | SIS with N = 20, B = 16, S = 150.     | SIS with N = 20, B = 16, S = 150.     | SIS with N = 20, B = 16, S = 150.     | SIS with N = 20, B = 16, S = 150.     | SIS with N = 20, B = 16, S = 150.       |
| No Action RandomAction PreFeRMAB DDLPO (topline) | 5.33±0.16 7.03±0.17 8.35±0.12 n/a     | 5.30±0.15 7.13±0.16 8.38±0.11 n/a     | 5.31±0.14 7.02±0.14 8.26±0.11 n/a     | 5.29±0.13 7.11±0.13 8.10±0.11 n/a     | 5.28±0.13 7.06±0.13 8.00±0.10 8.09±0.11 |
| ARMMAN with N = 25, B = 7, S = 3.                | ARMMAN with N = 25, B = 7, S = 3.     | ARMMAN with N = 25, B = 7, S = 3.     | ARMMAN with N = 25, B = 7, S = 3.     | ARMMAN with N = 25, B = 7, S = 3.     | ARMMAN with N = 25, B = 7, S = 3.       |
| No Action RandomAction PreFeRMAB DDLPO (topline) | 2.12±0.26 2.86±0.32 5.06±0.34 n/a     | 2.30±0.29 3.27±0.40 5.26±0.33 n/a     | 2.29±0.27 3.01±0.30 4.68±0.33 n/a     | 2.19±0.23 3.09±0.35 4.75±0.35 n/a     | 2.26±0.25 2.96±0.31 4.61±0.27 4.68±0.09 |

In Figure 2, we compare the number of samples required to train DDLPO from scratch vs. the number of samples for fine-tuning PreFeRMAB starting from a pre-trained model (additional results in Appendix A.5). Results suggests the cost of pretraining can be amortized over different downstream instances. A non-profit organization using RMAB models may have new beneficiaries opting in every week, and training a new model from scratch every week can be 3-20 times more expensive than fine-tuning our pretrained model.

Figure 2: Comparison of samples per arm required by DDLPO and PreFeRMAB (fine-tuning using a pretrained model) to achieve maximum DDLPO reward across different environments. PreFeRMAB achieves the maximum topline reward with significantly fewer samples than DDLPO. Averages across training seeds are reported as interquartile means.

![Image](https://ijcai.org/proceedings/2024/fulltext/36/image_000001_bcfe95b46144a085c6864bb5f033dc16582402dc2b0943d6341ead393c574bb3.png)

## 6 Conclusion

Our pretrained model (PreFeRMAB) leverages multi-arm generalization, a novel update rule for a crucial λ-network, and a StateShaping module for challenging reward functions. PreFeRMAB demonstrates general zero-shot ability on unseen arms, and can be fine-tuned on specific instances in a more sample-efficient way than training from scratch.

## Ethical Statement

The presented methods do not carry direct negative societal implications. However, training reinforcement learning models should be done responsibly, especially given the safety concerns associated with agents engaging in extreme, unsafe, or uninformed exploration strategies. While the domains we considered such as ARMMAN do not have these concerns, the approach may be extended to extreme environments; in these cases, ensuring a robust approach to training reinforcement models is critical.

## Acknowledgments

The work is partially supported by Harvard HDSI funding.

## Contribution Statement

The first two authors are equal contribution.

## References

1. [Abel et al., 2018] David Abel, Dilip Arumugam, Lucas Lehnert, and Michael Littman. State abstractions for life- long reinforcement learning. In International Conference on Machine Learning, pages 10-19. PMLR, 2018.
2. [Adelman and Mersereau, 2008] Daniel Adelman and Adam J Mersereau. Relaxations of weakly coupled stochastic dynamic programs. Operations Research, 56(3):712-727, 2008.
3. [Akbarzadeh and Mahajan, 2019] N Akbarzadeh and A Ma- hajan. Restless bandits with controlled restarts: Indexabil- ity and computation of whittle index., 2019.
4. [Alon et al., 1996] Noga Alon, Yossi Matias, and Mario Szegedy. The space complexity of approximating the fre- quency moments. In Proceedings of the twenty-eighth an- nual ACM symposium on Theory of computing, pages 20- 29, 1996.
5. [Avrachenkov and Borkar, 2022] Konstantin E Avrachenkov and Vivek S Borkar. Whittle index based q-learning for restless bandits with average reward. Automatica, 139:110186, 2022.
6. [Biswas et al., 2021] Arpita Biswas, Gaurav Aggarwal, Pradeep Varakantham, and Milind Tambe. Learn to in- tervene: An adaptive learning policy for restless bandits in application to preventive healthcare. arXiv preprint arXiv:2105.07965, 2021.
7. [Bommasani et al., 2021] Rishi Bommasani, Drew A Hud- son, Ehsan Adeli, Russ Altman, Simran Arora, Syd- ney von Arx, Michael S Bernstein, Jeannette Bohg, An- toine Bosselut, Emma Brunskill, et al. On the oppor- tunities and risks of foundation models. arXiv preprint arXiv:2108.07258, 2021.
8. [Chaudhuri and Kalyanakrishnan, 2020] Arghya Roy Chaudhuri and Shivaram Kalyanakrishnan. Regret minimisation in multi-armed bandits using bounded arm memory. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 10085-10092, 2020.
9. [Devlin et al., 2018] Jacob Devlin, Ming-Wei Chang, Ken- ton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understand- ing. arXiv preprint arXiv:1810.04805, 2018. [Dusonchet and Hongler, 2003] Fabrice Dusonchet and M-O Hongler. Continuous-time restless bandit and dynamic scheduling for make-to-stock production. IEEE Transac- tions on Robotics and Automation, 19(6):977-990, 2003.
10. [Fu et al., 2019] Jing Fu, Yoni Nazarathy, Sarat Moka, and Peter G Taylor. Towards q-learning the whittle index for restless bandits. In 2019 Australian & New Zealand Con- trol Conference (ANZCC), pages 249-254. IEEE, 2019.
11. [Gheshlaghi Azar et al., 2013] Mohammad Ghesh- laghi Azar, Rémi Munos, and Hilbert J Kappen. Minimax pac bounds on the sample complexity of reinforcement learning with a generative model. Machine learning, 91:325-349, 2013.
12. [Hawkins, 2003] Jeffrey Thomas Hawkins. A Langrangian decomposition approach to weakly coupled dynamic opti- mization problems and its applications. PhD thesis, Mas- sachusetts Institute of Technology, 2003.
13. [Hodge and Glazebrook, 2015] David J Hodge and Kevin D Glazebrook. On the asymptotic optimality of greedy index heuristics for multi-action restless bandits. Advances in Applied Probability, 47(3):652-667, 2015.
14. [Jaques et al., 2019] Natasha Jaques, Angeliki Lazari- dou, Edward Hughes, Caglar Gulcehre, Pedro Ortega, DJ Strouse, Joel Z Leibo, and Nando De Freitas. Social influence as intrinsic motivation for multi-agent deep reinforcement learning. In International conference on machine learning, pages 3040-3049. PMLR, 2019.
15. [Jin et al., 2021] Chi Jin, Qinghua Liu, Yuanhao Wang, and Tiancheng Yu. V-learning-a simple, efficient, decen- tralized algorithm for multiagent rl. arXiv preprint arXiv:2110.14555, 2021.
16. [Killian et al., 2021a] Jackson A. Killian, Arpita Biswas, Sanket Shah, and Milind Tambe. Q-learning lagrange poli- cies for multi-action restless bandits. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, KDD '21, page 871-881, New York, NY, USA, 2021. Association for Computing Machinery.
17. [Killian et al., 2021b] Jackson A. Killian, Andrew Perrault, and Milind Tambe. Beyond "to act or not to act": Fast la- grangian approaches to general multi-action restless ban- dits. In AAMAS, pages 710-718, UK, 2021. AAMAS.
18. [Killian et al., 2022] Jackson A Killian, Lily Xu, Arpita Biswas, and Milind Tambe. Restless and uncertain: Ro- bust policies for restless bandits via deep multi-agent re- inforcement learning. In Uncertainty in Artificial Intelli- gence, pages 990-1000. PMLR, 2022.
19. [Lefévre, 1981] Claude Lefévre. Optimal control of a birth and death epidemic process. Operations Research, 29(5):971-982, 1981.
20. [Li et al., 2006] Lihong Li, Thomas J Walsh, and Michael L Littman. Towards a unified theory of state abstraction for mdps. In AI&M, 2006.
21. [Liau et al., 2018] David Liau, Zhao Song, Eric Price, and Ger Yang. Stochastic multi-armed bandits in constant space. In International Conference on Artificial Intelli- gence and Statistics, pages 386-394. PMLR, 2018.
22. [Littman, 1994] Michael L Littman. Markov games as a framework for multi-agent reinforcement learning. In Ma- chine learning proceedings 1994, pages 157-163. Else- vier, 1994.
23. [Mate et al., 2021] Aditya Mate, Arpita Biswas, Christoph Siebenbrunner, Susobhan Ghosh, and Milind Tambe. Efficient algorithms for finite horizon and streaming restless multi-armed bandit problems. arXiv preprint arXiv:2103.04730, 2021.
24. [Mate et al., 2022] Aditya Mate, Lovish Madaan, Aparna Taneja, Neha Madhiwalla, Shresth Verma, Gargi Singh, Aparna Hegde, Pradeep Varakantham, and Milind Tambe. Field study in deploying restless multi-armed bandits: As- sisting non-profits in improving maternal and child health. Proceedings of the AAAI Conference on Artificial Intelli- gence, 36:12017-12025, 06 2022.
25. [Nakhleh et al., 2021] Khaled Nakhleh, Santosh Ganji, Ping-Chun Hsieh, I Hou, Srinivas Shakkottai, et al. Neur- win: Neural whittle index network for restless bandits via deep rl. Advances in Neural Information Processing Systems, 34:828-839, 2021.
26. [Papadimitriou and Tsitsiklis, 1999] Christos H Papadim- itriou and John N Tsitsiklis. The complexity of optimal queuing network control. Mathematics of Operations Re- search, 24(2):293-305, 1999.
27. [Ramesh et al., 2021] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pages 8821-8831. PMLR, 2021.
28. [Riquelme et al., 2018] Carlos Riquelme, George Tucker, and Jasper Snoek. Deep bayesian bandits showdown: An empirical comparison of bayesian deep networks for thompson sampling. In International Conference on Learning Representations, 2018.
29. [Shalev-Shwartz and Ben-David, 2014] Shai Shalev- Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
30. [Shapley, 1953] Lloyd S Shapley. Stochastic games. Proceedings of the national academy of sciences, 39(10):1095-1100, 1953.
31. [Sinha and Mahajan, 2022] Amit Sinha and Aditya Maha- jan. Robustness of whittle index policy to model approxi- mation. Available at SSRN 4064507, 2022.
32. [Team et al., 2023] Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Bradley-Schmieg, Michael Chang, Natalie Clay, Adrian Collister, et al. Human-timescale adaptation in an open-ended task space. arXiv preprint arXiv:2301.07608, 2023.
33. [Verma et al., 2023] Shresth Verma, Aditya Mate, Kai Wang, Neha Madhiwalla, Aparna Hegde, Aparna Taneja, and Milind Tambe. Restless multi-armed bandits for maternal and child health: Results from decision-focused learning. In Proceedings of the 2023 International Conference on Autonomous Agents and Multiagent Systems, pages 1312- 1320, 2023.
34. [Vershynin, 2018] Roman Vershynin. High-dimensional probability: An introduction with applications in data sci- ence, volume 47. Cambridge university press, 2018.
35. [Wang et al., 2023] Kai Wang, Shresth Verma, Aditya Mate, Sanket Shah, Aparna Taneja, Neha Madhiwalla, Aparna Hegde, and Milind Tambe. Scalable decision-focused learning in restless multi-armed bandits with application to maternal and child health. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 12138-12146, 2023.
36. [Weber and Weiss, 1990] Richard R Weber and Gideon Weiss. On an index policy for restless bandits. Journal of applied probability, 27(3):637-648, 1990.
37. [Whittle, 1988] Peter Whittle. Restless bandits: Activity al- location in a changing world. Journal of applied probabil- ity, 25(A):287-298, 1988.
38. [Xie et al., 2020] Qiaomin Xie, Yudong Chen, Zhaoran Wang, and Zhuoran Yang. Learning zero-sum simultaneous-move markov games using function approx- imation and correlated equilibrium. In Conference on learning theory, pages 3674-3682. PMLR, 2020.
39. [Xiong and Li, 2023] Guojun Xiong and Jian Li. Finite- time analysis of whittle index based q-learning for restless multi-armed bandits with neural network function approx- imation. In Thirty-seventh Conference on Neural Informa- tion Processing Systems, 2023.
40. [Yaesoubi and Cohen, 2011] Reza Yaesoubi and Ted Co- hen. Generalized markov models of infectious dis- ease spread: A novel framework for developing dynamic health policies. European journal of operational research, 215(3):679-687, 2011.
41. [Yu et al., 2020] Wenhao Yu, Jie Tan, Yunfei Bai, Erwin Coumans, and Sehoon Ha. Learning fast adaptation with meta strategy optimization. IEEE Robotics and Automa- tion Letters, 5(2):2950-2957, 2020.
42. [Yu et al., 2022] Chao Yu, Akash Velu, Eugene Vinitsky, Ji- axuan Gao, Yu Wang, Alexandre Bayen, and Yi Wu. The surprising effectiveness of ppo in cooperative multi-agent games. Advances in Neural Information Processing Sys- tems, 35:24611-24624, 2022.
43. [Zhang and Frazier, 2021] Xiangyu Zhang and Peter I. Fra- zier. Restless bandits with many arms: Beating the central limit theorem, 2021.
44. [Zhao et al., 2024] Yunfan Zhao, Nikhil Behari, Edward Hughes, Edwin Zhang, Dheeraj Nagaraj, Karl Tuyls, Aparna Taneja, and Milind Tambe. Towards zero shot learning in restless multi-armed bandits. AAMAS (2-page extended abstract), 2024.