Heterogeneous Federated Reinforcement Learning Using Wasserstein Barycenters
This addresses model aggregation in heterogeneous federated learning for reinforcement learning, but it is incremental as it applies an existing method (Wasserstein barycenters) to a new context.
The paper tackles the problem of heterogeneous federated reinforcement learning by proposing FedWB, an algorithm that uses Wasserstein barycenters for model fusion, and demonstrates it on the CartPole toy problem with varied pole lengths, resulting in a global DQN that works across all environments.
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.