Paper: Federated Offline Policy Learning with Heterogeneous Observational Data (ICLR-2024)

Reviewer comment (Novelty):
It seems that the proposed algorithm is an application of FedAvg to CSMC. I wonder if there are some special challenges when extending offline policy learning algorithms to the federated setting with FedAvg, which have not been addressed in other federated learning or federated RL literature.

Author response:
Indeed, the algorithmic aspect of our work involves adapting federated averaging to offline policy learning with parametric function classes. The purpose of this portion was to present a feasible approach to offline policy learning in a federated setting that supplements our theoretical results. However, we do believe that there are additional questions at the intersection of federated learning and offline policy learning that merit further exploration.  For instance, the doubly robust offline policy learning approach employs a two-stage process of estimating nuisance functions before defining the optimization objective, and we simply employed a local nuisance function estimation procedure. Intriguing questions arise on aggregate nuisance parameter estimation among similar clients and how it could enhance performance. Furthermore, our FedAvg-based approach currently only applies to parametric policy classes. It would be valuable to develop a federated offline policy learning optimization procedure for non-parametric policy classes, such as decision trees, commonly used in real-world applications for their interpretability in decision-making.

Reviewer comment (Theory):
The regret analysis holds only when the algorithm converges to the optimal policy, which is not always guaranteed. In the paper, the authors claimed that it is still possible to achieve optimal policies via some additive term and appropriate choice of policy class, but it is vague and not convincing enough. It would be nice if you could provide further clarifications on the solution that enables the algorithm to achieve the optimal policy for general policy classes beyond linear policy classes.

Author response:
The reviewer astutely pointed out that the optimal policy in the optimization procedure is not always guaranteed, with exceptions for specific cases like linear policy classes as examined in our experiments. However, our analysis can be readily extended to encompass scenarios where the optimal policy is not necessarily attainable, introducing the possibility of approximation errors. In the concluding paragraph of Section 6, we explain how this approximation error simply becomes an additional additive term in the regret bounds, that is, $\text{Regret bound}+\epsilon$, where $\epsilon$ represents the policy value optimality gap—indicating how far the policy value of the learned policy deviates from that of the optimal policy. While this approximation error may not necessarily diminish with increased data, with a judicious choice of policy class and the corresponding optimization procedure, the optimality gap can be rendered insignificant or of a similar order of magnitude as other terms in the regret bounds. This holds particularly true in moderately to highly heterogeneous environments, as the local regret bound contains an additive distribution shift term that is also inherently irreducible and captures the extent of heterogeneity.

Reviewer comment (Theory):
It seems that Assumption 1-(c) requires a full exploration of all actions, which is quite strong given that there are some offline RL works suggesting that full coverage of state-action space is not necessary.

Author response:
We agree with the reviewer's observation that the overlap assumption (Assumption 1c) is strong. We note that this assumption is standard and commonly adopted in the offline policy learning literature to facilitate the identification of causal effects in observational studies. However, we agree with the reviewer's point that it may not be necessary as recent work has shown that, under a pessimism principle, overlap only under the optimal policy is sufficient. This insight is discussed in Section 3.4, following Assumption 1. The decision to adopt the more stringent standard assumption was made to simplify our analysis and maintain a focus on our contributions on the effects of data heterogeneity on policy learning.

    Nevertheless, we acknowledge the intriguing questions that arise when considering what the pessimism principle would grant in our setting. Specifically, would it be necessary to have coverage under the locally optimal policy for each data source, or is coverage under the globally optimal policy sufficient? Additionally, how do the mixture weights impact the satisfaction of this assumption? We believe our work provides an initial study that opens up numerous follow-up questions in this direction.

Reviewer comment (Experiments):
In the experiments, the algorithm was demonstrated in very limited settings. I wonder if this could be applied to more general and realistic settings. Also, it would be nice if you could compare the performance with other baseline algorithms.

Author response:
Given that our algorithm builds on FedAvg, extending its application to more general settings, such as more general policy classes or cross-device scenarios with thousands of devices, is relatively straightforward. As for baseline comparisons in our empirical evaluation, we compared our approach with the baseline of standard offline policy learning where all the data is assumed to be consolidated into a single source. As our work is pioneering in studying offline policy learning with bandit feedback data in a federated setting, identifying direct comparison baselines relevant to our results proved challenging. Moreover, the objective of our experiments was to contrast the policies learned in a federated setting with those in a non-federated setting. While we stress that our primary contribution is theoretical rather than empirical, we acknowledge the potential for further development in the empirical evaluations. Your feedback is valuable, and we appreciate your understanding of the balance we aimed to strike between the theoretical and empirical aspects of our work.

Reviewer comment (Theory):
The upper bounds were based on a few detailed data assumptions, such as local ignorability, unconfoundedness, and overlap. Although the paper mentioned that some of the assumptions can potentially be relaxed, there is a lack of details on the discussion, and which assumption may not be relaxed fundamentally.

Author response:
Due to space limitations, our discussion regarding the necessity of certain assumptions may appear concise. Another contributing factor to this condensed discussion is that these assumptions are relatively standard within the offline policy learning literature for ensuring consistent policy value estimation. The cited works in this section serve as references indicating how some of these assumptions have been relaxed in recent literature. We acknowledge the need for more detailed exploration of these assumptions in future revisions.

Reviewer comment (Novelty):
In the theoretical analysis part (section 4-5), the main innovation part for the algorithm / estimator design and regret analysis in comparison to prior works was not clearly highlighted.

Author response:
Our primary contribution lies in the adaptation of offline policy learning to novel problem settings, specifically in the context of learning from multiple heterogeneous datasets and federated learning. The theoretical results presented in our work almost directly correspond with established findings in offline policy learning within the standard setting involving a single data source. We made efforts to underscore this parallel throughout the theoretical results section and the paper as a whole. Acknowledging the constraints posed by limited space, we recognize that our discussion might not have fully elaborated on this aspect. In response to your feedback, we have incorporated additional citations and context to enhance clarity. However, the space constraints still limit our ability to delve extensively into this matter. We appreciate the reviewer's input and we are open to addressing any specific concerns regarding comparisons with prior work here. We also kindly encourage you to refer to our responses to other reviewers' comments, where we delve more deeply into the question of the novelty of our work.

Reviewer comment (Experiments):
The empirical evaluation did not include any comparisons with other baselines, or any real dataset, and the heterogeneous setting was fairly simply constructed.

Author response:
In our empirical evaluation, we compared our approach with the baseline of standard offline policy learning. This comparison involved assessing our method against the scenario where all the data is assumed to be consolidated in a single source, allowing for the application of standard offline policy learning methods. Given that our work is pioneering in the study of offline policy learning with bandit feedback data in a heterogeneous multiple-source and federated setting, we faced a challenge in identifying direct comparison baselines that we deemed relevant to our results. The objective of our experiments was to contrast the policies learned in a federated setting with those in a non-federated setting.

    It's crucial to emphasize that our primary contribution lies in the theoretical domain rather than the empirical one. While our focus is on advancing theoretical understanding, we acknowledge that there is potential for further development in the empirical evaluations. We appreciate your understanding of the balance we aimed to strike between the theoretical and empirical aspects of our work.

Reviewer comment (Theory):
Can the nuisance parameters be learnt jointly instead of being required to be separately known or estimated in the policy value estimates?

Author response:
Indeed, if the central server can group similar clients, it can learn the nuisance parameters more effectively. The doubly robust approach relies on a two-stage approach of estimating nuisance functions prior to forming the optimization objective. We relied on a simple local nuisance function estimation procedure. However, there are interesting questions regarding the aggregate nuisance parameter estimation among similar clients and how that may improve performance.

Reviewer comment (Novelty):
My main concern is the technical novelty, and I'd appreciate if the authors can provide more clarification on the contribution compared with the existing work mentioned below. Based on my understanding, the main difference of this paper, compared with existing off-policy learning method using doubly robust estimator, e.g., Zhou et al. (2022b), lie in the optimization oracle. i.e., this paper needs to solve the CSMC problem over multiple heterogeneous clients to update the policy, instead of in a centralized setting. However, I am not sure if this has led to any technical challenge in obtaining the global regret bound in Theorem 1.

Author response:
Our work differs from prior work in two main ways. First, we explore offline policy learning (Zhou et al. 2023) under a novel setting of learning from multiple heterogeneous data sources. Prior studies, such as those by He et al. (2019) and Kallus et al. (2021), have considered a related, yet more restricted, problem of offline policy learning from multiple logging policies, where multiple datasets are also leveraged, but the underlying environment is assumed to be the same and only the action sample strategy differs across datasets. The more general problem setting we consider in this work has not been investigated, possibly due to the challenge of understanding the tradeoffs of a policy learned from diverse provenance. In our study, we are able to exactly quantify the impact of such heterogeneity on downstream performance for any target client by characterizing finite-sample upper bounds on suitable notions of regret. Note that this main result is independent of whatever optimization oracle is used to solve the policy learning problem we consider.

    Second, we considered this policy learning problem in a setting where such heterogeneities may naturally arise in practice: in a federated setting. In federated settings, there tends to be less coordination on data collection procedures among clients (e.g., institutions, devices) compared to centralized learning settings. Consider a consortium of hospitals that conduct randomized control trials on a new treatment as an example; since these hospitals may not fully communicate their idiosyncratic treatment procedures, each hospital might conduct its unique trial on distinct populations, treatment formulations, and treatment assignment mechanisms. How should the central healthcare policymaker account for this heterogeneity across clients when training a treatment policy in a federated manner? With increasing privacy regulations, these coordination challenges may become more common, requiring federated policymakers to design policy learning procedures adaptable to such inherent heterogeneities for the intended purpose.

    Our work demonstrates a practical way of adapting the federated learning framework to policy learning while accounting for underlying heterogeneities. As a first approach to this novel problem setting, we relied on the most well-known federated learning procedure for parametric models: federated averaging. We are not claiming that we are introducing a novel algorithm to the federated learning literature. We are simply adapting a well-known federated learning procedure to solve our optimization problem in a setting of practical interest for the above reasons. We also emphasize that the optimization problem is somewhat orthogonal to the technical analysis we conducted to arrive at our regret bounds. The regret bounds hold for the solution to the optimization problem and our optimization procedure demonstrates the feasibility and simplicity of solving this problem in the federated setting. However, that being said, we believe that the adaptation of FedAvg to offline policy learning is not entirely trivial and there are interesting questions at the intersection of federated learning and offline policy learning that warrant further exploration. For instance, the doubly robust offline policy learning approach employs a two-stage process of estimating nuisance functions before defining the optimization objective, and we simply employed a local nuisance function estimation procedure. Intriguing questions arise on aggregate nuisance parameter estimation among similar clients and how it could enhance performance. Furthermore, our FedAvg-based approach currently only applies to parametric policy classes. It would be valuable to develop a federated offline policy learning optimization procedure for non-parametric policy classes, such as decision trees, commonly used in real-world applications for their interpretability in decision-making.

Reviewer comment (Theory):
With CSC, we typicaly have a non-concave non-convex objective function to optimize, which makes finding the policy that maximizes Eq 8 difficult. Therefore, I expected the regret analysis to cover the situation where the FedAVG-CSMC procedure can only provide policy with certain approximation error. Moreover, even if the objective function is easy to optimize, it still seems to be unrealistic to assume we can obtain the exact maximizer of Eq 8 under federated setting, as this requires infinite number of iteration/communication rounds. I'd appreciate it if the authors can provide more insights on how the current analysis can be extended to allow for approximation error of Eq 8.

Author response:
The reviewer astutely pointed out that the maximization objective is not necessarily concave, with exceptions for specific cases like linear policy classes as examined in our experiments. Therefore, determining the optimal policy in Eq. 8 may not be so straightforward. However, our analysis can be readily extended to encompass scenarios where the optimal policy is not necessarily attainable, introducing the possibility of approximation errors. In the concluding paragraph of Section 6, we explain how this approximation error simply becomes an additional additive term in the regret bounds, that is, $\text{Regret bound}+\epsilon$, where $\epsilon$ represents the policy value optimality gap—indicating how far the policy value of the learned policy deviates from that of the optimal policy. While this approximation error may not necessarily diminish with increased data, with a judicious choice of policy class and the corresponding optimization procedure, the optimality gap can be rendered insignificant or of a similar order of magnitude as other terms in the regret bounds. This holds particularly true in moderately to highly heterogeneous environments, as the local regret bound contains an additive distribution shift term that is also inherently irreducible and captures the extent of heterogeneity.

Reviewer comment (Writing):
In Section 3.2, notations like $X^{C}, Y^{C}$ are not formaly defined.

Author response:
In definition 1, we defined $(X^c, Y^c(a_1), \dots, Y^c(a_d))$ to be a sample from the local data-generating distribution $\mathcal{D}_c$ over the decision space $\mathcal{X}\times\mathcal{Y}^d$. However, for further clarity, in the updated draft we introduce this notation earlier in Section 3.1.

Reviewer comment (Theory):
The author mentioned that Assumption 3 can be easilly satisfied under regularity assumption. Can the authors provide a more formal description of the regularity assumption?

Author response:
Assumption 3 is a requirement on the estimation error rates of the two separately estimated nuisance functions. For simplicity, note that this assumption is sufficiently satisfied under any conditions that warrant $o(1/\sqrt{n})$ error rates for each nuisance function. Therefore, the vaguely mentioned regularity conditions essentially refer to whatever sufficient conditions enable efficient estimation of these two functions at the desired rates. For example, if the nuisance functions are parametric, then the least squares estimators under conditions like boundedness achieve O(1/n) rates. For non-parametric function classes conditions like belonging to a reproducing kernel Hilbert space with smooth kernels or having Lipschitz smoothness (along with their appropriate estimation procedures) suffice. The efficient estimation of these nuisance parameters is, although necessary, somewhat orthogonal to the policy learning. This is, in fact, the benefit of double machine learning.