Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
50,103 characters · 22 sections · 38 citation commands
Federated Offline Policy Learning
\twocolumn[ \icmltitle{Federated Offline Policy Learning}
\icmlsetsymbol{equal}{*}
\icmlaffiliation{sta}{Stanford University}
\icmlcorrespondingauthor{Aldo Gael Carranza}{[email removed]}
\icmlkeywords{Machine Learning, ICML, Offline Policy Learning, Contextual Bandits, Federated Learning}
\vskip 0.3in ]
\printAffiliationsAndNotice
Offline policy learning from observational bandit feedback data is an effective approach for learning personalized decision policies in applications where obtaining online, real-time data is impractical swaminathan2015batch, kitagawa2018should,athey2021policy. Typically, the observational data used in offline policy learning is assumed to originate from a single source distribution. However, in practice, we often have the opportunity to leverage multiple datasets collected from various experiments under different populations, environments, or logging policies kallus2021optimal. For instance, a healthcare policymaker may have access to data from multiple hospitals that conduct different types of clinical trials on distinct patient populations. Learning from multiple heterogeneous observational datasets, with their more diverse and extensive coverage of the decision space, may lead to better personalized decision policies, assuming sufficient similarity and generalization across data sources.
However, several practical constraints, such as privacy concerns, legal restrictions, proprietary interests, or competitive barriers, can hinder the consolidation of datasets across sources. Federated learning kairouz2021advances presents a potential solution to such obstacles by offering a framework for training machine learning models in a decentralized manner, thereby minimizing systemic risks associated with traditional, centralized machine learning. Although federated learning principles have been widely implemented on standard supervised learning tasks, they have not been extensively explored in more prescriptive machine learning tasks such as policy learning. Federated learning techniques applied to policy learning can enable platforms to learn targeted decision policies without centrally storing sensitive user information. It also has the potential to incentivize institutions to collaborate on developing policies that are more generalizable across diverse environments without having to share sensitive data, such as clinical patient data in hospitals.
In this work, we introduce the problem of learning personalized decision policies on observational bandit feedback data from multiple heterogeneous data sources. Our main technical contribution is the presentation of a novel regret analysis that establishes finite-sample upper bounds on distinguishing notions of global regret for all data sources on aggregate and of local regret for any given data source. We characterize these regret bounds by expressions of source heterogeneity and distribution shift. Moreover, we consider the additional complexities of this problem in the federated setting where central data aggregation across sources is constrained and data source heterogeneity is known to be a significant practical concern. We introduce an offline policy learning algorithm amenable to federation and we experimentally verify the effect of source heterogeneity on regret. Additionally, we present design choices to overcome local performance degradation due to distribution shift. Our analysis and supporting experimental results in this work help provide insights into the tradeoffs in the participation of heterogeneous data sources in offline policy learning.
\paragraph{Offline Policy Learning}
In recent years, there have been significant advancements in offline policy learning from observational bandit feedback data. swaminathan2015batch, kitagawa2018should established frameworks for structured decision policy learning using offline policy evaluation strategies. athey2021policy achieved optimal regret rates under unknown propensities through doubly robust estimators. zhou2023offline extended these results to the multi-action setting. kallus2018balanced found optimal weights for the target policy directly from the data. zhan2021policy ensure optimal regret guarantees under adaptively collected data with diminishing propensities. jin2022policy relaxed the uniform overlap assumption to partial overlap under the optimal policy. We also mention that many contextual bandit methods often utilize offline policy learning oracles when devising adaptive action-assignment rules bietti2021contextual,krishnamurthy2021adapting, simchi2022bypassing, carranza2022flexible. More relevant to our setting of policy learning under heterogeneous data sources, agarwal2017effective,he2019off,kallus2021optimal leveraged data from multiple historical logging policies, but they assume the same underlying populations and environments. Our problem setting is also very closely related to multi-task offline policy learning, where each task can be seen as a data source. hong2023multi study a hierarchical Bayesian approach to this problem where task similarity is captured using an unknown latent parameter. In our experimental section, we further discuss this method and compare results of this method with those of our approach.
\paragraph{Federated Learning}
kairouz2021advances, wang2019federated offer comprehensive surveys on federated learning and its challenges. mohri2019agnostic presented an agnostic supervised federated learning framework, introducing multiple-source learning concepts like weighted Rademacher complexity and skewness measures which we make use of in our work. wei2021federated established excess risk bounds for supervised federated learning under data heterogeneity. li2019convergence,li2020federated,karimireddy2020scaffold explore the impact of client heterogeneity on model convergence in federated learning. Contextual bandits in federated settings have been studied by agarwal2020federated,dubey2020differentially,huang2021federated,agarwal2023empirical. However, offline policy evaluation and learning in federated settings remain relatively underexplored. xiong2021federated investigated federated methods for estimating average treatment effects across heterogeneous data sources. zhou2022federated,shen2023distributed delved into federated offline policy optimization in full reinforcement learning settings but with significant limitations, including relying on strong linear functional form assumptions with highly suboptimal rates or a difficult saddle point optimization approach focusing more on policy convergence.
We introduce the problem of offline policy learning from observational bandit feedback data across multiple heterogeneous data sources. Throughout the paper, we refer to a heterogeneous data source as a client and the central planner that aggregates client models as the central server, following standard federated learning terminology.
Let $\mathcal{X}\subset\mathbb{R}^p$ be the context space, $\mathcal{A}=\{a_1,\dots,a_d\}$ be the finite action space with $d$ actions, and $\mathcal{Y}\subset\mathbb{R}$ be the reward space. A decision policy $\pi:\mathcal{X}\to\mathcal{A}$ is a mapping from the context space $\mathcal{X}$ to actions $\mathcal{A}$. We assume there is a central server and a finite set of clients $\mathcal{C}$, with each client $c\in\mathcal{C}$ possessing a local data-generating distribution $\mathcal{D}_c$ defined over $\mathcal{X}\times\mathcal{Y}^d$ which governs how client contexts $X^c$ and client potential reward outcomes $Y^c(a_1),\dots,Y^c(a_d)$ are generated. Moreover, the central server specifies a fixed distribution $\lambda$ over the set of clients $\mathcal{C}$ describing how clients will be sampled or aggregated\footnote{Clients are sampled in the cross-device FL setting and aggregated in the cross-silo FL settings.}, which we will simply refer to as the client sampling distribution.
The central server seeks to train a decision policy that performs well on the global data-generating mixture distribution defined by $\mathcal{D}_\lambda\coloneqq\sum_{c\in\mathcal{C}}\lambda_c\mathcal{D}_c$. At the same time, if there is a potential target client of interest, the central server may not want the personalized policy to perform poorly on the local distribution of this client, otherwise their locally trained policy may provide greater utility to the client, and thus their participation is disincentivized. In the following section, we introduce the exact policy performance measures that capture these two potentially opposing objectives.
We consider the immediate reward gained by a client by taking actions according to any given policy. Additionally, we extend this metric to a global version that captures the aggregate reward gained from the mixture of clients under the client sampling distribution.
The performance of a policy is typically characterized by a notion of regret against an optimal policy in a specified policy class $\Pi\subset\{\pi:\mathcal{X}\to\mathcal{A}\}$, which we assume to be fixed throughout the paper. We define local and global versions of regret based on their respective versions of policy values.
The objective of the central server is to determine a policy in the specified policy class $\Pi$ that minimizes global regret. On the other hand, the central server also aims to characterize the corresponding local regret of a target client under the obtained policy since this quantity captures the client's corresponding individual utility to a global policy.
We assume each client $c\in\mathcal{C}$ has a local observational data set $\{(X_{i}^{c},A_{i}^{c},Y_{i}^{c})\}_{i=1}^{n_c}\subset\mathcal{X}\times\mathcal{A}\times\mathcal{Y}$ consisting of $n_c\in\mathbb{N}$ triples of contexts, actions, and rewards collected using a local experimental stochastic policy $e_c:\mathcal{X}\to\Delta(\mathcal{A})$ in the following manner. For the $i$-th data point of client $c$,
We will let $\smash[b]{n\coloneqq\sumiMn_c}$ denote the total sample size across clients. Note that although the counterfactual reward outcomes $Y_{i}^{c}(a)$ for all $a\in\mathcal{A}\backslash\{A_{i}^{c}\}$ exist in the local data-generating process, they are not observed in the realized data. All clients only observe the outcomes associated to their assigned treatments. For this reason, such observational data is also referred to as bandit feedback data swaminathan2015batch.
Given these data-generating processes, it will be useful to introduce the data-generating distributions that also incorporate how actions are sampled. For each client $c\in\mathcal{C}$, the local historical policy $e_c$ induces a complete local data-generating distribution $\bar{\mathcal{D}}_c$ defined over $\mathcal{X}\times\mathcal{A}\times\mathcal{Y}^d$ that dictates how the entire local contexts, actions, and potential outcomes were sampled in the local data-generating process, i.e., $(X_{i}^{c},A_{i}^{c},Y_{i}^{c}(a_1),\dots,Y_{i}^{c}(a_d))\sim\bar{\mathcal{D}}_c$. Given this construction of the complete client distributions, we also introduce the complete global data-generating mixture distribution defined by $\bar{\mathcal{D}}_\lambda\coloneqq\sum_{c\in\mathcal{C}}\lambda_c\bar{\mathcal{D}}_c$.
We make the following standard assumptions on the data-generating process of any given client.
Note that the boundedness assumption is not essential and we only impose it for simplicity in our analysis. With additional effort, we can instead rely on light-tail distributional assumptions such as sub-Gaussian potential outcomes as in athey2021policy. Unconfoundedness ensures that action assignment is as good as random after accounting for measured covariates, and it is necessary to ensure valid policy value estimation using inverse propensity-weighted strategies. The uniform overlap condition ensures that all actions are taken sufficiently many times to guarantee accurate evaluation of any policy. This assumption may not be entirely necessary as recent work jin2022policy introduced an approach that does away with the uniform overlap assumption for all actions and only relies on overlap for the optimal policy. However, in our work, we made the above assumptions to simplify our analysis and maintain the focus of our contributions on the effects of data heterogeneity on policy learning. In any case, these stated assumptions are standard and they are satisfied in many experimental settings such as randomized controlled trials or A/B tests.
Next, we also impose the following local data scaling assumption on each client.
This assumption states that, asymptotically, the total sample size cannot increase without increasing across all data sources. We emphasize that this assumption is quite benign since $\nu_c$ could be any slowly increasing function (e.g., an iterated logarithm) and the asymptotic lower bound condition even allows step-wise increments. We only impose this assumption to ensure that the regret bounds in our analysis scale with respect to the total sample size with sensible constants. However, it does come at the cost of excluding scenarios in which a client always contributes $O(1)$ amount of data relative to the total data, no matter how much more total data is made available in aggregate, in which case one may expect it is better to exclude any such client.
The approach for the central server is to use the available observational data to construct an appropriate estimator of the global policy value and use this estimator to find an optimal global policy.
We define the following functions which we refer to as nuisance parameters as they will be required to be separately known or estimated in the policy value estimates.
In our estimation strategy, each client must estimate the conditional response and inverse conditional propensity functions when they are unknown. Following the literature on double machine learning chernozhukov2018double, we make the following high-level assumption on the estimators of these local nuisance parameters.
We emphasize this is a standard assumption in the double machine learning literature, and we can easily construct estimators that satisfy these rate conditions, given sufficient regularity on the nuisance parameters zhou2023offline. See Appendix (ref) for more details. They can be estimated with widely available out-of-the-box regression and classification implementations. Moreover, this condition is general and flexible enough to allow one to trade-off the accuracies of estimating the nuisance parameters. This is an important property in offline policy learning where distribution shift in the batch data can complicate consistent reward estimation.
Next, we define our policy value estimators. For any client $c\in\mathcal{C}$, we define the local augmented inverse propensity weighted (AIPW) score for each $a\in\mathcal{A}$ to be
where
are the centered outcomes and $(X^c,A^c,Y^c(a_1),\dots,Y^c(a_d))\sim\bar{\mathcal{D}}_c$. One can readily show that $Q_c(\pi)=\operatorname*{\mathbb{E}}_{\bar{\mathcal{D}}_c}[\Gamma^c(\pi(X^c))]$, and therefore $Q_\lambda(\pi)=\operatorname*{\mathbb{E}}_{\lambda}\operatorname*{\mathbb{E}}_{\bar{\mathcal{D}}_c}[\Gamma^c(\pi(X^c))]$ (see the proof in Lemma (ref)). Accordingly, our procedure is to estimate the local AIPW scores and appropriately aggregate them to form the global policy value estimator.
We assume we have constructed nuisance parameter estimates $\hat{\mu}_c$ and $\hat{w}_c$ that satisfy Assumption (ref). Then, for each data point $(X_{i}^{c},A_{i}^{c},Y_{i}^{c})$ in the local observational data set of client $c\in\mathcal{C}$, we define the approximate local AIPW score for each $a\in\mathcal{A}$ to be
where
are the approximate centered outcomes. Using these estimated scores, we introduce the following doubly robust global policy value estimate:
Our proposed estimator is a generalized aggregate version of the doubly robust policy value estimator introduced in the standard offline policy learning setting zhou2023offline. It is doubly robust in the sense that it is accurate as long as one of nuisance parameter estimates is accurate for each client. To ensure we can use the same data to estimate the nuisance parameters and to construct the policy value estimates, we utilize a cross-fitting strategy locally for each client. See Appendix (ref) for more details on this cross-fitting strategy.
The objective of the central server is to find a policy that maximizes the doubly robust global policy value estimate:
Note that in the centralized setting, this optimization can be done using standard policy optimization techniques athey2021policy on the centrally accumulated heterogeneous datasets with appropriate reweighting. However, as we discussed previously, the centralized collection of datasets can present difficulties in privacy sensitive settings.
For this reason, we seek to provide an optimization procedure that is also amenable to federated settings to overcome these challenges. In the federated setting, the central server does not have access to client raw data to estimate local policy values nor does it have access to the local policy values; only model updates can be shared through the network. In Section (ref), we discuss our optimization procedure for parametric policy classes that manages these constraints.
In this section, we establish regret bounds for the global policy solution $\hat{\pi}_\lambda$ to the optimization objective above. Refer to Appendices (ref), (ref), (ref), (ref) for a detailed discussion and proofs of the results in this section.
First, we introduce important quantities that appear in our regret bounds.
\paragraph{Policy Class Complexity} The following quantity provides a measure of policy class complexity based on a variation of the classical entropy integral introduced by dudley1967sizes, and it is useful in establishing a class-dependent regret bound. See Appendix (ref) for more details on its definition.
The entropy integral is constant for a fixed policy class, and rather weak assumptions on the class are sufficient to ensure it is finite such as sub-exponential growth on its Hamming covering number, which is satisfied by many policy classes including parametric and finite-depth tree policy classes zhou2023offline. In the binary action setting, the entropy integral of a policy class relates to its VC-dimension with $\kappa(\Pi)=\sqrt{\smash[b]{\text{VC}(\Pi)}}$, and for $D$-dimensional linear classes $\kappa(\Pi)=\mathcal{O}(\sqrt{\smash[b]{D}}\,)$.
\paragraph{Client Skewness} The following quantity measures the imbalance of the client sampling distribution $\lambda$ relative to the empirical distribution of samples across clients defined by $\bar{n}\coloneqq(n_c/n)_{c\in\mathcal{C}}$. This quantity naturally arises in the generalization bounds of weighted mixture distributions mansour2021theory.
The following result captures a root-$n$ finite-sample bound for the global regret that parallels the optimal regret bounds typically seen in the offline policy learning literature. In Appendix (ref), we provide a high-level sketch of the proof.
First, note that $V$ captures a notion of the worst-case AIPW score variance across clients. Next, we observe that root-$n$ rate is moderated by a skewness term which can also scale with the total sample size. For example, if $\lambda=\bar{n}$ then $\mathfrak{s}(\lambda\|\bar{n})/n=1/n$, and if $\lambda=(1,0,\dots,0)$ then $\mathfrak{s}(\lambda\|\bar{n})/n=1/n_1$. Thus, this skewness-moderated rate generalizes and smoothly interpolates between the rates one expects from the uniform weighted model and the single source model. Indeed, when clients are identical and $\lambda=\bar{n}$, we recover the best known rates from standard offline policy learning zhou2023offline.
From this observation, it may seem that the best design choice for the client sampling distribution is the empirical sample distribution, i.e., $\lambda=\bar{n}$. However, as we will observe in the next section, there are terms in the local regret bounds that introduce trade-offs on the choice of $\lambda$ when considering a specific target client.
In this next result, we capture the discrepancy in local and global regret due to client heterogeneity. This result is helpful in understanding the extent at which the global and local regret minimization objectives can be in conflict for a particular target client.
Note that the first term in this regret bound is inherently irreducible relative to the sample sizes and it is due to distribution shift between the complete local client distribution $\bar{\mathcal{D}}_c$ and the complete global mixture distribution $\bar{\mathcal{D}}_\lambda$. Thus, we can observe how the design choice on the client distribution $\lambda$ must balance a trade-off to achieve low skewness and low expected distribution shift across sources. Skewing towards the target client will reduce the distribution shift term, but it will further eat at the rates in the global regret bound. In our experiments, we take a heuristic approach to obtain a skewed $\lambda$, but we note that this may be chosen in a more principled manner, e.g., minimax skewness as in mohri2019agnostic. Observe that the constant $U$ in the distribution shift term is defined by the constants in the boundedness and overlap assumptions stated in Assumption (ref). However, in Appendix (ref), we provide an alternate upper bound that does not rely on bounded AIPW scores and instead is scaled by the worst-case AIPW variance, which can be smaller and also appears in our global regret bound.
Next, we demonstrate how the distribution shift term can be further tensorized into contributions due to distribution shift in the covariates, propensities, and potential outcomes.
This result directly reveals the contribution to local regret from each possible source of distribution shift. This implies that if we have prior knowledge that certain components of the data-generating distribution match, then we can claim tighter bounds on the local regret of clients.
In summary, the analysis we presented in this section help provide insights into understanding the tradeoffs and value of information in heterogeneous client participation in offline policy learning (see Appendix (ref) for further discussion).
We present FedOPL, a federated algorithm for finding the optimal global policy $\hat{\pi}_\lambda$ for the optimization problem stated in Section (ref). The standard approaches in federated learning are based on the federated averaging (FedAvg) algorithm konevcny2016federated,mcmahan2017communication, which works iteratively by selecting clients to participate in the training round, fine-tuning a local parametric model instance on each client using their own data, and then aggregating local model parameters on the central server via a weighted average. To enable standard federated learning strategies for policy optimization, we consider parametric policy classes $\Pi_\Theta=\{\pi_\theta:\mathcal{X}\to\mathcal{A}\mid\theta\in\Theta\}$, and we construct an iterative parametric policy optimization procedure for the local policy updates.
Specifically, we observe that the corresponding local policy optimization procedures
are equivalent to cost-sensitive multi-class classification (CSMC) problems beygelzimer2009error, dudik2011efficient, where the actions are the labels and the AIPW scores are the negative costs for the labels. Therefore, we can conduct iterative local policy model updates using online CSMC oracle methods beygelzimer2008machine, agarwal2017effective, often used for policy learning in contextual bandit algorithms agarwal2014taming, bietti2021contextual.
See Algorithms (ref) and (ref) for the server-side and client-side implementations of Fed-OPL based on online CSMC local model updates. Note that prior to policy learning, each client must estimate their local AIPW scores, as described in Section (ref). For data efficiency in constructing these scores, we use a cross-fitting strategy on local observational data to estimate local nuisance parameters. See Algorithm (ref) in Appendix (ref) for more details on the implementation of this AIPW estimation procedure.
\paragraph{Remark on Optimality Gap} Note that obtaining the optimal global policy is necessary for achieving the stated regret bounds above. FedOPL is guaranteed to converge to the optimal policy if the optimization problem is concave. However, this only holds for certain policy classes such as linear policy classes. Nevertheless, we can still establish useful regret bounds for nearly optimal policies under general policy classes. The regret bound is simply modified to include an additive term that captures the policy value optimality gap. In particular, if $\tilde\pi_\lambda$ is the approximate global policy, then we can easily extend the global regret bound to be
where the asymptotic term is the exact bound established above and the first term $\Delta_\lambda(\hat{\pi}_\lambda,\tilde\pi_\lambda)=Q_\lambda(\hat{\pi}_\lambda)-Q_\lambda(\tilde\pi_\lambda)$ is the policy value gap between the optimal global policy $\hat{\pi}_\lambda$ and the approximate global policy $\tilde\pi_\lambda$. Under a judicious choice of policy class and corresponding optimization procedure, the value gap can be rendered insignificant or of a similar order of magnitude as other terms in the regret bounds, especially under heterogeneous environments where there is significant irreducible distribution shift.
For our experiments, we compare empirical local and global regrets across different experimental settings involving homogeneous and heterogeneous clients.
We describe the experimental setup common to all of our experiments. For the environment, we consider the client set $\mathcal{C}=[C]$ with $C=3$, the action set $\mathcal{A}=\{a_1,\dots,a_d\}$ with $d=4$, and the context space $\mathcal{X}=\mathbb{R}^p$ with $p=d\times q$ where $q=10$. For every client $c\in\mathcal{C}$, we consider the following data-generating process:
where the noise parameters $\sigma_c, \rho_c$ and the true conditional reward functions $\mu_c$ are specified in each experiment below. For a given total sample size $n\in\mathbb{N}$, each client $c\in\mathcal{C}$ is allocated a local sample sample size determined by an increasing function $n_{c}=\nu_{c}(n)$, which we specify in each experiment below. By construction, this entire data-generating process satisfies our data assumptions stated in Section (ref).
For the parametric policy class, we consider the policy class induced by linear scores: $\pi_\theta(x)=\operatorname*{arg\,max}_{a\in\mathcal{A}}\phi(x,a)^\top \theta_a$ where $\theta\in\Theta=\mathbb{R}^{d\times q}$ and $\phi:\mathcal{X}\times\mathcal{A}\to\mathbb{R}^q$ is a feature extractor, which we simply set to $\phi(x,a)=x_a$. Then, for the local CSMC oracle, we employ the cost-sensitive one-against-all (CSOAA) implementation in Vowpal Wabbit vowpalwabbit. This method performs online multiple regressions of costs (i.e., negative rewards) on contexts for each action, and at inference time, it selects the action with the lowest predicted cost agarwal2017effective.
For our results, we consider a training sample size grid in the range up to $N$ samples, where $N=1$K for the homogeneous experiments and $N=10$K for the heteregeneous experiments (longer range due to slower convergence). For each sample size $n$ in this grid, we sample $n_{c}=\nu_{c}(n)$ training samples from each client distribution defined above, and we train a global policy $\hat{\pi}_\lambda$ to optimize the corresponding global policy value estimate. We will then compute and plot the empirical estimates of the global regret $R_\lambda(\hat{\pi}_\lambda)$ and local regret $R_c(\hat{\pi}_\lambda)$ for a fixed client $c$. Our plots will display the mean and the standard deviation bands of these empirical regrets over five different seed runs. We compute these metrics on 10K separate test samples from each client. The empirical regrets are computed against the optimal models trained on 100K total separate training samples.
For comparison against pure local training, we will also plot the local regret $R_c(\hat\pi_c)$ of the locally trained policy $\hat\pi_c$. The locally trained policy will be trained with the same number $n$ of total sample sizes across all clients. This will help us understand the benefits of federation compared to simply training a policy locally.
We will also compare against the HierOPO algorithm for multi-task offline policy learning introduced by hong2023multi. HierOPO is a hierarchical Bayesian approach where task similarity is captured using an unknown latent parameter. Although this approach can be applied to our setting, we note significant limitations. This approach assumes the reward model is known and, without access to an informative prior on the latent task parameter, the regret bounds only scale with the local sample size. Furthermore, this algorithm is less well-suited for federation since it relies on aggregation of summary statistics across tasks rather than model parameters, thereby increasing communication complexity for higher-dimensional problems and increasing privacy risks. In our experiments, we will let $\hat\pi_H$ denote the policy trained under the HierOPO algorithm.
For sanity check, we consider the homogeneous setting where all clients are identical. We should expect the local and global regrets to be the same. For every client $c\in\mathcal{C}$, we set $\mu_c(x;a)=\phi(x,a)^\top\theta_a$ where $\theta\sim\operatorname{Normal}(0,\omega^2 I_p)$ with $\omega^2=1$ and $\sigma_c^2=\rho_c^2=1$. Moreover, we simply divide the total sample size across all clients, i.e., $\nu_c(n)=\lfloor n/C\rfloor$. Since all clients are the same, we simply use the empirical mixture $\lambda=\bar{n}$ in our optimization procedure. Figure (ref) displays our results. As expected, each of these curves is nearly identical since the global and local regrets are identical in this scenario.
Next, we consider the more interesting setting where one client is different than all other clients. Specifically, we set
where $k=50$ and $\theta\sim\operatorname{Normal}(0,\omega^2I_p)$ with $\omega^2=5$. The above hyperparameters were simply chosen to better highlight the differences between regret curves, but we found our results to be robust to different hyperparameter choices. Under this data-generating process, beyond the distribution shift in the contexts and rewards due to different variances, there is also distribution shift due to different mean rewards. The idea behind this choice of reward function is that since the scaled sine function is nearly linear near zero, there is a wide range of contexts where the reward function for client $c=1$ is nearly identical to the linear reward function of all other clients. Therefore, there is distribution shift between client $c=1$ and all other clients, but there is some amount of similarity that can be exploited. Moreover, to illustrate the benefits of federation under sample size heterogeneity, we will also have client $c=1$ contribute significantly less data than the others with $n_1=\nu_1(n)=\lfloor\log n\rfloor$ and all other clients will evenly distribute the rest of the total sample size.
\paragraph{Empirical Client Sampling Distribution} The top figure in Figure (ref) plots the same type of regret curves as in the homogeneous experiment. In FedOPL, we use the empirical client sampling distribution $\lambda=\bar{n}$. We observe that $R_1(\hat{\pi}_\lambda)$ significantly suffers from distribution shift, as predicted. In fact, the locally trained policy $\hat\pi_1$ performs better than the globally trained policy $\hat{\pi}_\lambda$ at a sufficiently large sample size. Note that the HierOPO policy only outperforms the locally trained model at low sample sizes due to worse scaling.
\paragraph{Skewed Client Sampling Distribution} The bottom figure in Figure (ref) plots similar regret curves, but instead with the global policy trained with a skewed client sampling distribution $\lambda=\bar{n}+\bar{\varepsilon}$ where $\bar{\varepsilon}_c=-\alpha\cdot\bar{n}_c$ for $c\neq 1$ and $\bar{\varepsilon}_1=\alpha\cdot(1-\bar{n}_1)$ for $\alpha=0.2$. Here, we observe that $\hat{\pi}_\lambda$ still suffers some amount in terms of local regret, but not to such an extent that $\hat\pi_1$ beats it. Moreover, the local regret shift decreases with larger sample size. The idea is that this skewness upscales the distribution of client 1 to diminish the amount of distribution shift of $\bar{\mathcal{D}}_1$ from $\bar{\mathcal{D}}_\lambda$ especially at larger total sample sizes, at the cost of negatively affecting the other more homogeneous clients. Refer to the results in Appendix (ref) to see how the more similar clients are affected by this design choice modification to favor client 1. The takeaway in those experiments is that the other clients have less distribution shift from the average so their performance degradation is lesser under $\lambda=\bar{n}$, but their performance further slightly degrades under the skewed client sampling distribution $\lambda=\bar{n}+\bar{\varepsilon}$.
We studied the problem of offline policy learning from observational bandit feedback data across multiple heterogeneous data sources. Moreover, we considered the practical aspects of this problem in a federated setting to address privacy concerns. We presented a novel regret analysis and supporting experimental results that demonstrate tradeoffs of client heterogeneity on policy performance. In Appendix (ref), we provide additional discussion on the broader impact, limitations, and potential future directions to our work.