Reinforcement Learning Picks the Right Devices to Speed Up Federated Learning

Federated learning has become one of the most important ideas in modern machine learning because it allows many devices to train a shared model without ever sending their raw data to a central server. Instead of collecting millions of private records in one place, the server distributes the current model, each participant improves it locally on its own data, and only the resulting parameter updates are sent back and combined. This architecture promises to protect privacy while still exploiting the enormous amount of information that lives on smartphones, sensors, vehicles and other connected devices. Yet the very feature that makes federated learning attractive, namely its reliance on thousands of geographically and technically diverse participants, also creates one of its hardest unsolved problems: the participants are not equal, and treating them as if they were can badly damage training.

Researchers at Guizhou University in China, led by Haiming Luo and corresponding author Fu Yan, together with Yuling Chen and Changgen Peng, have now proposed a new way to decide which devices should participate in each round of federated training. Writing in the journal Complex & Intelligent Systems, they describe a client selection framework called FLASH-DRM, built on a reinforcement learning technique known as Double Deep Q-Learning. The central insight of their work is that the choice of which clients to involve in a given training round is itself a decision problem that can be learned, and that a well-trained selection agent can simultaneously improve accuracy, reduce latency and cut the energy consumed by the whole system.

The problem the authors attack is usually described as heterogeneity, and it comes in two distinct forms. Statistical heterogeneity refers to the fact that different clients hold different data distributions. A keyboard prediction model trained on the phones of ten thousand users will see wildly different vocabularies, languages and typing habits from one device to the next, so the locally computed updates are not drawn from the same distribution and can pull the shared model in conflicting directions. System heterogeneity, by contrast, is about the hardware and network conditions of the participants themselves: some devices are powerful and connected over fast Wi-Fi, while others are battery-constrained, intermittently online or reachable only through slow links.

In very large Internet of Things deployments, where the number of available devices can be enormous, it is not practical for every client to take part in every round. The communication overhead alone would be crushing, because each round requires the server to broadcast the model and each selected client to upload its update. Selecting a subset of clients per round has therefore become a common requirement, but the selection strategy matters enormously. If the server systematically picks fast, powerful devices, the rounds complete quickly but the data held by slow or unpopular devices is underrepresented, and the global model becomes biased. If it picks devices purely to balance the data distribution, it may repeatedly wait on stragglers that delay convergence. Balancing this trade-off between system heterogeneity and statistical heterogeneity is the key challenge that FLASH-DRM is designed to address, because the choice directly determines both the efficiency and the stability of global training.

The framework treats client selection as a sequential decision process. An intelligent agent observes the state of the federated system and chooses which subset of clients to activate in each round. Rather than a hand-crafted rule, the agent learns a selection policy through Double Deep Q-Learning, an algorithm that improves on classical Q-learning by using two separate neural networks to decouple the selection of the best action from the evaluation of that action’s value. This decoupling is well known to reduce the overestimation bias that can destabilise value-based reinforcement learning, which matters here because the agent must estimate the long-term contribution of each client choice across many training rounds.

Two design elements distinguish FLASH-DRM from earlier attempts. The first is a utility-based reputation function that quantifies the contribution of individual clients to the global model. Each client accumulates a reputation score based on how useful its updates actually turn out to be, which gives the agent a principled signal for rewarding devices that reliably improve the shared model rather than devices that merely appear attractive on paper. The second element is a reward path that directly links the agent’s client selection to the global accuracy of the federated model. This means the agent does not optimise a proxy; it is explicitly trained to consider how both the individual utility of each selected client and their collective utility as a group affect the quality of the final global model in each round.

Perhaps the most appealing practical property of the approach is what it does not cost. The combined utility measure integrated into FLASH-DRM incurs virtually no additional communication overhead, which is a crucial consideration in settings where bandwidth is the dominant bottleneck. Many sophisticated federated learning techniques have failed to move from paper to practice precisely because the bookkeeping they require consumes the very resources they are meant to save. By keeping the information flow between server and clients essentially unchanged, the framework ensures that its intelligence is applied at the point of decision, the server-side selection step, without inflating the per-round data transfer.

The authors evaluated FLASH-DRM in a variety of heterogeneous dataset environments, training it on three widely used benchmark test suites: CIFAR-10, MNIST and EMNIST. These datasets span natural images, handwritten digits and handwritten letters, and by partitioning them in non-identical ways across simulated clients the experiments recreate the unbalanced, non-independent data distributions that characterise real federated deployments. The framework was compared against a strong field of state-of-the-art baselines, including FedAvg, the foundational averaging algorithm of federated learning, FAVOR, FedGCS, FLASH-RL, FedProx and FedMSC, each of which represents a different strategy for coping with heterogeneity or for steering client participation.

The experimental results show that FLASH-DRM generally outperforms all of these competing methods across three metrics that matter in deployed systems: accuracy, latency and energy consumption. Achieving gains on all three at once is notable because these objectives frequently conflict. A selection policy that maximises accuracy by repeatedly involving slow but data-rich clients will typically inflate both latency and energy use, while a policy that optimises for speed may starve parts of the data distribution and depress final accuracy. The reinforcement learning agent, guided by its accuracy-linked reward and its reputation scores, appears to find selections that respect both sides of the trade-off, producing a global model that converges to higher quality in less time and with lower total energy expenditure across the participating devices.

The implications reach well beyond the benchmarks. As federated learning moves into healthcare, automotive systems, industrial IoT and edge computing, the diversity of participating hardware and data will only grow, and naive selection strategies will become an increasingly expensive liability. The Guizhou team’s work suggests that the tools of deep reinforcement learning, already proven in domains from game playing to robotics, can be repurposed as orchestration engines for distributed training itself. By learning to answer the deceptively simple question of which devices should contribute in the next round, FLASH-DRM offers a template for making collaborative machine learning faster, greener and more accurate without asking any participant to give up a single byte of their private data. The study is published open access, allowing researchers worldwide to build on the framework as the field continues to wrestle with heterogeneity in the federated setting.

Subject of Research: Reinforcement learning-based client selection to mitigate system and statistical heterogeneity in federated learning

Article Title: Mitigating client heterogeneity in federated learning: a reinforcement learning approach for efficient client selection

Article References: Luo, H., Yan, F., Chen, Y., & Peng, C. (2026). Mitigating client heterogeneity in federated learning: a reinforcement learning approach for efficient client selection. Complex & Intelligent Systems. https://doi.org/10.1007/s40747-026-02519-6

Image Credits: AI Generated

DOI: 10.1007/s40747-026-02519-6

Keywords: federated learning, reinforcement learning, client selection, Double Deep Q-Learning, system heterogeneity, statistical heterogeneity, Internet of Things, DDQL, distributed machine learning, energy efficiency, CIFAR-10, MNIST