Policy Gradient–Based Reinforcement Weighted Aggregation for Efficient Federated Learning on Heterogeneous Data
摘要
Federated Learning (FL) is a decentralized paradigm that enables collaborative model training without sharing raw data. However, traditional FL methods often struggle with slow and unstable convergence, especially under non-Independent and non-Identically Distributed (non-IID) data, which degrades model performance and training stability. Most existing aggregation strategies rely on static, uniform, or heuristic-based weighting schemes that fail to adapt dynamic client behavior and overlook variations in local training quality. Although some recent approaches explore client selection or value-based reinforcement processes (e.g., Q-learning), they fall short of directly optimizing the aggregation process in an end-to-end fashion. To address these limitations, we propose Federated Reinforcement Weighted Aggregation (FedRWA), a policy gradient-based aggregation approach that dynamically assigns client-specific weights based on local training characteristics on the server side. Each client generates a feature vector consisting of training loss and accuracy, which is fed into a central policy network. This network outputs aggregation weights and is trained via a reinforcement process, using the global model’s validation performance as the reward. Extensive experiments on the UNSW-NB15 and NSL-KDD datasets, under both IID and non-IID conditions, show that FedRWA consistently outperforms state-of-the-art baselines in terms of classification metrics and communication efficiency. In addition, the learning curves demonstrate the superior convergence speed and training stability of FedRWA, validating its effectiveness as a robust solution for adaptive aggregation in heterogeneous FL environments.