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

TL;DR

This work addresses data heterogeneity in Federated Reinforcement Learning by showing that local policy updates can aid the global objective only when appropriately constrained by distribution-space divergences. It introduces FedKL, a KL-divergence based penalty with adaptive coefficients to regulate local updates and enable larger, stable steps, backed by a convergence analysis. The authors classify FRL heterogeneity into initial-state and dynamics differences, derive a necessary condition for global learnability, and demonstrate that FedKL accelerates convergence and stabilizes training under non-IID conditions. The approach reduces communication rounds and improves robustness in distributed RL, with implications for scalable, heterogeneous FRL deployments.

Abstract

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.

Figures (10)

• Figure 1: An example of FRL systems: autonomous driving. Many vehicles run on different roads. In each training round, selected vehicles (white cars) will participate in the training. With the PG method, policies are represented by ML models, which will be communicated between agents and the server for model aggregation and synchronization. The model aggregation is the same as that of FedAvg.
• Figure 2: Reacher's field splitting. Sub-fields in grey are reachable by the robotic arm. Sub-environments are created based on these sub-fields.
• Figure 3: Vehicle Placement Schema. Red cars are human-driven vehicles. White cars are RL-controlled vehicles.
• Figure 4: Visualization of $\left\Vert\mathbf{D}_{\pi,\mu_{n},P_{n}} \mathbf{A}_{\pi,P_{n}}\right\Vert_{F}-\left\Vert\mathbf{D}_{\pi,\mu_{n},P_{n}}\mathbf{B}_{\pi,\mu_{n},P_{n}}\right\Vert_{F}$ on modified Reacher-v2 environments. Results are averaged across three runs.
• Figure 5: The effect of negative $G$ on modified Reacher-V2 environments with different environment dynamics (action noise $\thicksim\mathcal{N}(0, 0.4)$). Red lines indicate the rounds where $G<0$ happens.

Theorems & Definitions (29)

• Definition I
• Theorem I
• Corollary I
• Remark I
• Remark II
• Corollary II
• Remark III
• Remark IV
• Theorem II
• Theorem III