# FedKL: Tackling Data Heterogeneity in Federated Reinforcement Learning by Penalizing KL Divergence

arXiv: https://arxiv.org/abs/2204.08125 · HTML: https://mlsanity.com/paper/2204.08125/

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 48 references that are in the map.
- First posted on arXiv 18 Apr 2022.

## Builds on (most cited first, 28 of 28)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environments](https://mlsanity.com/paper/1706.02275/) (2017)
- [Federated Learning: Challenges, Methods, and Future Directions](https://mlsanity.com/paper/1908.07873/) (2019)
- [Federated Learning: Strategies for Improving Communication Efficiency](https://mlsanity.com/paper/1610.05492/) (2016)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [Federated Learning with Non-IID Data](https://mlsanity.com/paper/1806.00582/) (2018)
- [Constrained Policy Optimization](https://mlsanity.com/paper/1705.10528/) (2017)
- [Federated Learning on Non-IID Data Silos: An Experimental Study](https://mlsanity.com/paper/2102.02079/) (2021)
- [Federated Machine Learning: Concept and Applications](https://mlsanity.com/paper/1902.04885/) (2019)
- [Federated Learning on Non-IID Data: A Survey](https://mlsanity.com/paper/2106.06843/) (2021)
- [Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning](https://mlsanity.com/paper/1807.06629/) (2018)
- [Client-Edge-Cloud Hierarchical Federated Learning](https://mlsanity.com/paper/1905.06641/) (2019)
- [Trust Region Policy Optimisation in Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2109.11251/) (2021)
- [QSGD: Communication-Efficient SGD via Gradient Quantization and Encoding](https://mlsanity.com/paper/1610.02132/) (2016)
- [Local SGD with Periodic Averaging: Tighter Analysis and Adaptive Synchronization](https://mlsanity.com/paper/1910.13598/) (2019)
- [Federated Reinforcement Learning: Techniques, Applications, and Open Challenges](https://mlsanity.com/paper/2108.11887/) (2021)
- [Lifelong Federated Reinforcement Learning: A Learning Architecture for Navigation in Cloud Robotic Systems](https://mlsanity.com/paper/1901.06455/) (2019)
- [Federated Reinforcement Learning with Environment Heterogeneity](https://mlsanity.com/paper/2204.02634/) (2022)
- [Fault-Tolerant Federated Reinforcement Learning with Theoretical Guarantee](https://mlsanity.com/paper/2110.14074/) (2021)
- [Federated Transfer Reinforcement Learning for Autonomous Driving](https://mlsanity.com/paper/1910.06001/) (2019)
- [Expected Policy Gradients for Reinforcement Learning](https://mlsanity.com/paper/1801.03326/) (2018)
- [Communication-Efficient Policy Gradient Methods for Distributed Reinforcement Learning](https://mlsanity.com/paper/1812.03239/) (2018)
- [The Gradient Convergence Bound of Federated Multi-Agent Reinforcement Learning with Efficient Communication](https://mlsanity.com/paper/2103.13026/) (2021)

## Built on by (most cited first, 29 of 29)

- [Finite-Time Analysis of On-Policy Heterogeneous Federated Reinforcement Learning](https://mlsanity.com/paper/2401.15273/) (2024)
- [Improved Communication Efficiency in Federated Natural Policy Gradient via ADMM-based Gradient Updates](https://mlsanity.com/paper/2310.19807/) (2023)
- [Asynchronous Federated Reinforcement Learning with Policy Gradient Updates: Algorithm Design and Convergence Analysis](https://mlsanity.com/paper/2404.08003/) (2024)
- [Federated Offline Reinforcement Learning](https://mlsanity.com/paper/2206.05581/) (2022)
- [Stateful active facilitator: Coordination and Environmental Heterogeneity in Cooperative Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2210.03022/) (2022)
- [Federated Learning Across Decentralized and Unshared Archives for Remote Sensing Image Classification](https://mlsanity.com/paper/2311.06141/) (2023)
- [The Sample-Communication Complexity Trade-off in Federated Q-Learning](https://mlsanity.com/paper/2408.16981/) (2024)
- [Momentum for the Win: Collaborative Federated Reinforcement Learning across Heterogeneous Environments](https://mlsanity.com/paper/2405.19499/) (2024)
- [SCAFFLSA: Taming Heterogeneity in Federated Linear Stochastic Approximation and TD Learning](https://mlsanity.com/paper/2402.04114/) (2024)
- [Federated Ensemble-Directed Offline Reinforcement Learning](https://mlsanity.com/paper/2305.03097/) (2023)
- [Global Convergence Guarantees for Federated Policy Gradient Methods with Adversaries](https://mlsanity.com/paper/2403.09940/) (2024)
- [FOVA: Offline Federated Reinforcement Learning with Mixed-Quality Data](https://mlsanity.com/paper/2512.02350/) (2025)
- [Federated Learning Resilient to Byzantine Attacks and Data Heterogeneity](https://mlsanity.com/paper/2403.13374/) (2024)
- [Towards Fast Rates for Federated and Multi-Task Reinforcement Learning](https://mlsanity.com/paper/2409.05291/) (2024)
- [FedCFA: Alleviating Simpson's Paradox in Model Aggregation with Counterfactual Federated Learning](https://mlsanity.com/paper/2412.18904/) (2024)
- [Client Selection for Federated Policy Optimization with Environment Heterogeneity](https://mlsanity.com/paper/2305.10978/) (2023)
- [Analysis and Optimization of Wireless Federated Learning with Data Heterogeneity](https://mlsanity.com/paper/2308.03521/) (2023)
- [Energy-Efficient Channel Decoding for Wireless Federated Learning: Convergence Analysis and Adaptive Design](https://mlsanity.com/paper/2407.13703/) (2024)
- [Efficient Federated Learning against Byzantine Attacks and Data Heterogeneity via Aggregating Normalized Gradients](https://mlsanity.com/paper/2408.09539/) (2024)
- [On the Convergence Rates of Federated Q-Learning across Heterogeneous Environments](https://mlsanity.com/paper/2409.03897/) (2024)
- [Collaborative Value Function Estimation Under Model Mismatch: A Federated Temporal Difference Analysis](https://mlsanity.com/paper/2503.17454/) (2025)
- [Achieving Tighter Finite-Time Rates for Heterogeneous Federated Stochastic Approximation under Markovian Sampling](https://mlsanity.com/paper/2504.11645/) (2025)
- [The Actor-Critic Update Order Matters for PPO in Federated Reinforcement Learning](https://mlsanity.com/paper/2506.01261/) (2025)
- [Federated Hierarchical Reinforcement Learning for Adaptive Traffic Signal Control](https://mlsanity.com/paper/2504.05553/) (2025)
- [Epidemic Decision-making System Based Federated Reinforcement Learning](https://mlsanity.com/paper/2311.01749/) (2023)
- [Single-Loop Federated Actor-Critic across Heterogeneous Environments](https://mlsanity.com/paper/2412.14555/) (2024)
- [Federated Reinforcement Learning in Heterogeneous Environments](https://mlsanity.com/paper/2507.14487/) (2025)
- [On the Effectiveness of Adaptation Strategies for VLM-Based Federated Learning in Remote Sensing](https://mlsanity.com/paper/2608.04791/) (2026)
- [FedGuide: Diffusion Prior Alignment and Value Baseline Guidance for Heterogeneous Federated Reinforcement Learning](https://mlsanity.com/paper/2609.18964/) (2026)

## Abstract (arXiv)

As a distributed learning paradigm, Federated Learning (FL) faces the communication bottleneck issue due to many rounds of model synchronization and aggregation. Heterogeneous data further deteriorates the situation by causing slow convergence. Although the impact of data heterogeneity on supervised FL has been widely studied, the related investigation for Federated Reinforcement Learning (FRL) is still in its infancy. In this paper, we first define the type and level of data heterogeneity for policy gradient based FRL systems. By inspecting the connection between the global and local objective functions, we prove that local training can benefit the global objective, if the local update is properly penalized by the total variation (TV) distance between the local and global policies. A necessary condition for the global policy to be learn-able from the local policy is also derived, which is directly related to the heterogeneity level. Based on the theoretical result, a Kullback-Leibler (KL) divergence based penalty is proposed, which, different from the conventional method that penalizes the model divergence in the parameter space, directly constrains the model outputs in the distribution space. Convergence proof of the proposed algorithm is also provided. By jointly penalizing the divergence of the local policy from the global policy with a global penalty and constraining each iteration of the local training with a local penalty, the proposed method achieves a better trade-off between training speed (step size) and convergence. Experiment results on two popular Reinforcement Learning (RL) experiment platforms demonstrate the advantage of the proposed algorithm over existing methods in accelerating and stabilizing the training process with heterogeneous data.