# Heterogeneous Federated Reinforcement Learning Using Wasserstein Barycenters

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

## Key facts

- Cited by 4 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 9 of its 28 references that are in the map.
- First posted on arXiv 18 Jun 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, 9 of 9)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Federated Learning with Differential Privacy: Algorithms and Performance Analysis](https://mlsanity.com/paper/1911.00222/) (2019)
- [Release Strategies and the Social Impacts of Language Models](https://mlsanity.com/paper/1908.09203/) (2019)
- [Model Fusion via Optimal Transport](https://mlsanity.com/paper/1910.05653/) (2019)
- [Federated Reinforcement Learning: Techniques, Applications, and Open Challenges](https://mlsanity.com/paper/2108.11887/) (2021)
- [Federated Reinforcement Learning with Environment Heterogeneity](https://mlsanity.com/paper/2204.02634/) (2022)
- [Multi-task Learning with Attention for End-to-end Autonomous Driving](https://mlsanity.com/paper/2104.10753/) (2021)
- [Wasserstein Barycenter-based Model Fusion and Linear Mode Connectivity of Neural Networks](https://mlsanity.com/paper/2210.06671/) (2022)

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

- [Can Optimal Transport Improve Federated Inverse Reinforcement Learning?](https://mlsanity.com/paper/2601.00309/) (2026)
- [Federated Distributional Reinforcement Learning with Distributional Critic Regularization](https://mlsanity.com/paper/2603.17820/) (2026)
- [FedGuide: Diffusion Prior Alignment and Value Baseline Guidance for Heterogeneous Federated Reinforcement Learning](https://mlsanity.com/paper/2609.18964/) (2026)
- [Optimal Transport Meets Reinforcement Learning: A Survey](https://mlsanity.com/paper/2610.01413/) (2026)

## Abstract (arXiv)

In this paper, we first propose a novel algorithm for model fusion that leverages Wasserstein barycenters in training a global Deep Neural Network (DNN) in a distributed architecture. To this end, we divide the dataset into equal parts that are fed to "agents" who have identical deep neural networks and train only over the dataset fed to them (known as the local dataset). After some training iterations, we perform an aggregation step where we combine the weight parameters of all neural networks using Wasserstein barycenters. These steps form the proposed algorithm referred to as FedWB. Moreover, we leverage the processes created in the first part of the paper to develop an algorithm to tackle Heterogeneous Federated Reinforcement Learning (HFRL). Our test experiment is the CartPole toy problem, where we vary the lengths of the poles to create heterogeneous environments. We train a deep Q-Network (DQN) in each environment to learn to control each cart, while occasionally performing a global aggregation step to generalize the local models; the end outcome is a global DQN that functions across all environments.