[Paper Review] Federated Offline Reinforcement Learning

Doudou Zhou, Yufeng Zhang+4 more|arXiv (Cornell University)|Jun 11, 2022

Blockchain Technology Applications and Security4 citations

TL;DR

This paper proposes Federated Offline Dynamic Treatment Regime (FDTR), a communication-efficient algorithm for multi-site offline reinforcement learning that jointly learns site-specific and shared policy effects using a multi-site linear MDP model. The method achieves suboptimality comparable to centralized learning with only one round of communication, demonstrating superior performance on a sepsis dataset across nine ICUs.

ABSTRACT

Evidence-based or data-driven dynamic treatment regimes are essential for personalized medicine, which can benefit from offline reinforcement learning (RL). Although massive healthcare data are available across medical institutions, they are prohibited from sharing due to privacy constraints. Besides, heterogeneity exists in different sites. As a result, federated offline RL algorithms are necessary and promising to deal with the problems. In this paper, we propose a multi-site Markov decision process model that allows for both homogeneous and heterogeneous effects across sites. The proposed model makes the analysis of the site-level features possible. We design the first federated policy optimization algorithm for offline RL with sample complexity. The proposed algorithm is communication-efficient, which requires only a single round of communication interaction by exchanging summary statistics. We give a theoretical guarantee for the proposed algorithm, where the suboptimality for the learned policies is comparable to the rate as if data is not distributed. Extensive simulations demonstrate the effectiveness of the proposed algorithm. The method is applied to a sepsis dataset in multiple sites to illustrate its use in clinical settings.

Motivation & Objective

• Address the challenge of learning optimal dynamic treatment regimes (DTRs) from distributed, non-sharable electronic health records (EHR) across multiple healthcare sites.
• Overcome data privacy constraints that prevent direct aggregation of multi-site EHR data while accounting for site-level heterogeneity in treatment protocols and patient populations.
• Develop a communication-efficient algorithm that enables collaborative policy learning across sites without sharing raw data, using only summary statistics.
• Enable the analysis of both shared (common) and site-specific (heterogeneous) effects in sequential clinical decision-making using a structured multi-site MDP model.
• Provide theoretical guarantees on suboptimality that match centralized offline RL performance under the same model assumptions.

Proposed method

• Formulate a multi-site linear Markov decision process (MDP) where transition kernels and reward functions are linear in known feature maps, with shared effects (ϕ₀) and site-specific effects (ϕ₁).
• Decompose the value function into common and site-specific components using a bilinear model structure, allowing joint estimation of global and local policy parameters.
• Design a federated policy optimization algorithm that requires only a single round of communication, exchanging only summary statistics (e.g., gradients or moment estimators) across sites.
• Use a step-importance sampling estimator for off-policy evaluation of the value function, enabling unbiased comparison of learned policies across sites.
• Incorporate site-level covariates (e.g., ICU size, staffing ratios) as features in the shared component to model institutional heterogeneity.
• Apply shrinkage or regularization techniques implicitly through the model structure to improve generalization and reduce overfitting in low-sample regimes.

Experimental results

Research questions

• RQ1 Can a federated offline RL framework effectively learn optimal dynamic treatment regimes across multiple healthcare sites without sharing raw patient data?
• RQ2 How does the inclusion of site-specific and shared effects in the MDP model improve policy performance compared to homogeneous or site-specific baselines?
• RQ3 To what extent does the proposed communication-efficient, single-round algorithm maintain performance comparable to centralized offline RL?
• RQ4 How do site-level covariates influence treatment decisions across time in a clinical setting like sepsis management?
• RQ5 Does the model maintain robustness when the linear MDP assumption is slightly violated, especially under limited data per site?

Key findings

• FDTR significantly outperformed all baseline methods in value function estimation on a multi-site sepsis dataset, with a mean value function estimate of 0.87 on the test set across nine ICUs.
• The estimated coefficients for site-level covariates (e.g., patient flow) were consistently positive, indicating that higher patient volume in care units was associated with better outcomes under the same treatment.
• Site-level covariate effects on treatment decisions decreased over time, suggesting that later-stage treatment decisions are more driven by patient-specific factors than institutional characteristics.
• Incorporating nonlinear basis functions did not improve performance, likely due to a bias-variance trade-off in the presence of limited sample sizes per site.
• The algorithm achieved suboptimality rates comparable to centralized offline RL, demonstrating that federated learning with a single communication round can preserve learning efficiency.
• The method successfully captured heterogeneous treatment effects across sites, enabling personalized policy learning even when individual sites had limited data or biased action distributions.