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

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:

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].

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:

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]

1. There exists θ∗ ∈ Θ which is optimal for every instance (ˆµ, λ).
2. Θ = B2(D, Rd ), the ℓ2 ball of radius D in Rd .
3. |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:

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:

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

• 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:

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:

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.

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.

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.

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.

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.