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

Federated reinforcement learning (FedRL) enables multiple agents to collaboratively learn a policy without needing to share the local trajectories collected during agent-environment interactions. However, in practice, the environments faced by different agents are often heterogeneous, but since existing FedRL algorithms learn a single policy across all agents, this may lead to poor performance. In this paper, we introduce a 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 federated twotimescale 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.