FedPolicy: An RL-Guided Redistribution Policy for Synergizing Local-Global Optimization in Federated Learning

Abstract: Statistical heterogeneity remains a central challenge in federated learning. Existing methods primarily address this problem through improved local objectives, aggregation strategies, or personalization mechanisms, while the post-aggregation redistribution step is typically applied uniformly and receives little explicit treatment. This design becomes problematic under heterogeneous client distributions, where repeatedly overwriting local models with the same aggregated parameters can disrupt client-specific adaptation and induce negative transfer. We propose \textsf{FedPolicy}, an RL-guided post-aggregation redistribution framework that treats the return path of the aggregated model as a client-specific decision problem. Rather than broadcasting the same update to every client, FedPolicy learns which part of the aggregated model should be transferred back to each client by selecting among full-model, backbone-only, and head-only parameter blocks. This formulation identifies post-aggregation redistribution as a previously underexplored control axis in federated optimization, improving the balance between global transfer and local specialization. Extensive experiments under heterogeneous federated settings show that FedPolicy consistently outperforms strong baselines across FMNIST, CIFAR-10, and CIFAR-100, with the clearest gains appearing in the more challenging heterogeneous regimes. Across all datasets and heterogeneity settings, FedPolicy achieves an average relative gain of $3.93\%$ over the strongest baseline, with the largest improvement reaching $8.40\%$ on CIFAR-100 under severe heterogeneity, while converging faster and delivering a more favorable cost-to-accuracy trade-off with negligible overhead. These results highlight client-specific post-aggregation redistribution as an underexplored yet impactful design dimension in heterogeneous federated learning.

Submission Type: Regular submission (no more than 12 pages of main content)

Assigned Action Editor: ~Han_Zhao1

Submission Number: 8251