[论文解读] FedKL: Tackling Data Heterogeneity in Federated Reinforcement Learning by Penalizing KL Divergence

Zhijie Xie, Shenghui Song|arXiv (Cornell University)|Apr 18, 2022

Privacy-Preserving Technologies in Data参考文献 46被引用 4

一句话总结

FedKL 提出了一种新颖的联邦强化学习框架，通过惩罚本地策略与全局策略之间的 Kullback-Leibler（KL）散度来缓解数据异质性问题，从而提升收敛速度与稳定性。通过联合应用全局和本地 KL 惩罚，该方法在非独立同分布（non-IID）数据上实现了训练效率与策略性能之间的更好权衡。

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.

研究动机与目标

• 解决联邦强化学习（FRL）中因数据异质性导致的收敛性与性能下降问题。
• 通过定义异质性的类型与程度，形式化数据异质性对基于策略梯度的 FRL 系统的影响。
• 建立理论条件，确保在存在数据异质性时，本地策略更新能对全局目标产生积极贡献，并将可学习性与异质性水平相关联。
• 设计一种基于 KL 散度的惩罚机制，直接在输出空间对策略分布进行约束，而非在参数空间中。
• 通过结合全局与本地 KL 正则化的双重惩罚策略，实现更快、更稳定的训练。

提出的方法

• 定义一个全局目标函数，引入本地与全局策略之间的总变差（TV）距离，以确保策略对齐。
• 提出一种基于 KL 散度的惩罚项，直接约束本地与全局策略之间的分布差异。
• 引入双重正则化方案：对策略差异施加全局惩罚，并在每次训练迭代中施加本地惩罚，以平衡步长与收敛性。
• 将本地目标函数表述为一个最大化问题，其中包含优势项、TV 距离项以及基于 KL 的正则化项。
• 通过证明策略序列的极限点为 FedKL-平稳点（即满足惩罚目标下的最优性条件），证明所提算法的收敛性。
• 利用价值函数与策略分布的连续性，确保在非独立同分布（non-IID）数据下理论上的稳定性和收敛性。

实验结果

研究问题

• RQ1 数据异质性如何具体影响基于策略梯度的联邦强化学习的收敛性与性能？
• RQ2 在存在数据异质性的情况下，何种理论条件可确保本地策略更新对全局目标产生正向贡献？
• RQ3 在策略输出空间中使用 KL 散度是否可作为有效正则化器，以在非独立同分布数据下稳定训练？
• RQ4 所提出的双重惩罚机制（全局与本地 KL 惩罚）如何改善训练速度与收敛稳定性的权衡？
• RQ5 在非独立同分布数据分布下，所提出的 FedKL 算法的收敛行为如何？

主要发现

• 在 MuJoCo 和 Procgen 环境中，与现有 FRL 方法相比，所提出的 FedKL 算法在非独立同分布（non-IID）数据上实现了更快的收敛速度与更高的训练稳定性。
• 理论分析证明，若异质性水平满足基于 TV 距离推导出的必要条件，则全局策略可从本地策略中学习得到。
• 算法收敛性得到保证，策略序列的极限点为 FedKL-平稳点，即满足惩罚目标下的最优性条件。
• 双重 KL 惩罚机制——同时应用全局与本地正则化——实现了对步长与收敛性的更好权衡，优于仅使用参数空间正则化的基线方法。
• 实验结果表明，FedKL 显著降低了训练方差并加速了学习过程，尤其在高异质性设置下表现突出。
• 该方法在最终回报与训练稳定性方面均优于 FedProx 与 FedAvg，在极端数据偏移条件下展现出更强的鲁棒性。