On the Linear Speedup of Personalized Federated Reinforcement Learning with Shared Representations

TL;DR

This paper introduces a personalized federated reinforcement learning framework that leverages shared representations among heterogeneous agents, proving linear convergence speedup and demonstrating improved learning and generalization in experiments.

Contribution

The paper proposes PFedRL, a novel personalized federated RL framework with shared representations, and proves the first linear convergence speedup with respect to the number of agents.

Findings

PFedTD-Rep achieves linear convergence speedup with more agents.

Experimental results show improved learning in heterogeneous environments.

PFedTD-Rep generalizes better to new environments.

Abstract

Federated reinforcement learning (FedRL) enables multiple agents to collaboratively learn a policy without sharing their local trajectories collected during agent-environment interactions. However, in practice, the environments faced by different agents are often heterogeneous, leading to poor performance by the single policy learned by existing FedRL algorithms on individual agents. In this paper, we take a further step and introduce a \emph{personalized} FedRL framework (PFedRL) by taking advantage of possibly shared common structure among agents in heterogeneous environments. Specifically, we develop a class of PFedRL algorithms named PFedRL-Rep that learns (1) a shared feature representation collaboratively among all agents, and (2) an agent-specific weight vector personalized to its local environment. We analyze the convergence of PFedTD-Rep, a particular instance of the framework with temporal difference (TD) learning and linear representations. To the best of our knowledge, we are the first to prove a linear convergence speedup with respect to the number of agents in the PFedRL setting. To achieve this, we show that PFedTD-Rep is an example of the federated two-timescale stochastic approximation with Markovian noise. Experimental results demonstrate that PFedTD-Rep, along with an extension to the control setting based on deep Q-networks (DQN), not only improve learning in heterogeneous settings, but also provide better generalization to new environments.