Momentum for the Win: Collaborative Federated Reinforcement Learning across Heterogeneous Environments

TL;DR

This paper introduces two momentum-based algorithms for federated reinforcement learning that effectively handle highly heterogeneous environments, ensuring convergence to optimal policies with linear speedup proportional to the number of agents.

Contribution

The paper proposes FedSVRPG-M and FedHAPG-M algorithms that achieve convergence in heterogeneous environments and demonstrate linear speedup with respect to agent count.

Findings

Algorithms converge to stationary points regardless of environment heterogeneity.

Achieve state-of-the-art sample complexity of O(ε^{-3/2}/N).

Exhibit linear convergence speedup with more agents.

Abstract

We explore a Federated Reinforcement Learning (FRL) problem where $N$ agents collaboratively learn a common policy without sharing their trajectory data. To date, existing FRL work has primarily focused on agents operating in the same or ``similar" environments. In contrast, our problem setup allows for arbitrarily large levels of environment heterogeneity. To obtain the optimal policy which maximizes the average performance across all potentially completely different environments, we propose two algorithms: FedSVRPG-M and FedHAPG-M. In contrast to existing results, we demonstrate that both FedSVRPG-M and FedHAPG-M, both of which leverage momentum mechanisms, can exactly converge to a stationary point of the average performance function, regardless of the magnitude of environment heterogeneity. Furthermore, by incorporating the benefits of variance-reduction techniques or Hessian approximation, both algorithms achieve state-of-the-art convergence results, characterized by a sample complexity of $\mathcal{O}\left(\epsilon^{-\frac{3}{2}}/N\right)$. Notably, our algorithms enjoy linear convergence speedups with respect to the number of agents, highlighting the benefit of collaboration among agents in finding a common policy.