# Optimized Local Updates in Federated Learning via Reinforcement Learning

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in RL for energy, networks and operations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 26 of its 48 references that are in the map.
- First posted on arXiv 31 May 2025.

## Placed in

- Deep reinforcement learning methods / [RL for energy, networks and operations](https://mlsanity.com/deeprl/rl-for-energy-networks-and-operations/)

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification](https://mlsanity.com/paper/1909.06335/) (2019)
- [Byzantine-Robust Distributed Learning: Towards Optimal Statistical Rates](https://mlsanity.com/paper/1803.01498/) (2018)
- [Flower: A Friendly Federated Learning Research Framework](https://mlsanity.com/paper/2007.14390/) (2020)
- [Bayesian Nonparametric Federated Learning of Neural Networks](https://mlsanity.com/paper/1905.12022/) (2019)
- [A Survey on Federated Learning Systems: Vision, Hype and Reality for Data Privacy and Protection](https://mlsanity.com/paper/1907.09693/) (2019)
- [Personalized Cross-Silo Federated Learning on Non-IID Data](https://mlsanity.com/paper/2007.03797/) (2020)
- [FedDC: Federated Learning with Non-IID Data via Local Drift Decoupling and Correction](https://mlsanity.com/paper/2203.11751/) (2022)
- [FedALA: Adaptive Local Aggregation for Personalized Federated Learning](https://mlsanity.com/paper/2212.01197/) (2022)
- [Stochastic Controlled Averaging for Federated Learning with Communication Compression](https://mlsanity.com/paper/2308.08165/) (2023)
- [Personalized Federated Learning for Intelligent IoT Applications: A Cloud-Edge based Framework](https://mlsanity.com/paper/2002.10671/) (2020)
- [Momentum Benefits Non-IID Federated Learning Simply and Provably](https://mlsanity.com/paper/2306.16504/) (2023)
- [On the Privacy-Robustness-Utility Trilemma in Distributed Learning](https://mlsanity.com/paper/2302.04787/) (2023)
- [PeFLL: Personalized Federated Learning by Learning to Learn](https://mlsanity.com/paper/2306.05515/) (2023)
- [FedImpro: Measuring and Improving Client Update in Federated Learning](https://mlsanity.com/paper/2402.07011/) (2024)
- [FedCompass: Efficient Cross-Silo Federated Learning on Heterogeneous Client Devices using a Computing Power Aware Scheduler](https://mlsanity.com/paper/2309.14675/) (2023)
- [FedHyper: A Universal and Robust Learning Rate Scheduler for Federated Learning with Hypergradient Descent](https://mlsanity.com/paper/2310.03156/) (2023)
- [Deep Reinforcement Learning for Modelling Protein Complexes](https://mlsanity.com/paper/2405.02299/) (2024)
- [Calibre: Towards Fair and Accurate Personalized Federated Learning with Self-Supervised Learning](https://mlsanity.com/paper/2412.20020/) (2024)
- [Trustworthy Personalized Bayesian Federated Learning via Posterior Fine-Tune](https://mlsanity.com/paper/2402.16911/) (2024)
- [DBA-DFL: Towards Distributed Backdoor Attacks with Network Detection in Decentralized Federated Learning](https://mlsanity.com/paper/2501.15005/) (2025)

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

- [On the Fragility of Data Attribution When Learning Is Distributed](https://mlsanity.com/paper/2605.15520/) (2026)

## Abstract (arXiv)

Federated Learning (FL) is a distributed framework for collaborative model training over large-scale distributed data, enabling higher performance while maintaining client data privacy. However, the nature of model aggregation at the centralized server can result in a performance drop in the presence of non-IID data across different clients. We remark that training a client locally on more data than necessary does not benefit the overall performance of all clients. In this paper, we devise a novel framework that leverages a Deep Reinforcement Learning (DRL) agent to select an optimized amount of data necessary to train a client model without oversharing information with the server. Starting without awareness of the client's performance, the DRL agent utilizes the change in training loss as a reward signal and learns to optimize the amount of training data necessary for improving the client's performance. Specifically, after each aggregation round, the DRL algorithm considers the local performance as the current state and outputs the optimized weights for each class, in the training data, to be used during the next round of local training. In doing so, the agent learns a policy that creates an optimized partition of the local training dataset during the FL rounds. After FL, the client utilizes the entire local training dataset to further enhance its performance on its own data distribution, mitigating the non-IID effects of aggregation. Through extensive experiments, we demonstrate that training FL clients through our algorithm results in superior performance on multiple benchmark datasets and FL frameworks. Our code is available at https://github.com/amuraddd/optimized_client_training.git.