Momentum-Based Contextual Federated Reinforcement Learning

Federated Reinforcement Learning (FRL) is an attractive edge learning paradigm for decision-making applications, which has garnered significant interest recently. However, owing to the inherent spatio-temporal non-stationarity of local state-action distributions, current FRL approaches typically suffer from high interaction and communication costs. In this paper, we introduce a new FRL method, which incorporates momentum, importance sampling, and server-side adjustments, capable of controlling the gradient shifts induced by the non-stationary data. We prove that by proper selection of momentum parameters and interaction frequency, it can achieve <inline-formula xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink"> <tex-math notation="LaTeX">$\tilde {\mathcal {O}}(H N^{-1}\epsilon ^{-3/2})$ </tex-math></inline-formula> and <inline-formula xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink"> <tex-math notation="LaTeX">$\tilde {\mathcal {O}}(\epsilon ^{-1})$ </tex-math></inline-formula> interaction and communication complexities (N represents the agent number), where the interaction complexity achieves linear speedup with the number of agents, and the communication complexity aligns with the best achievable among existing first-order FL algorithms. Further, we leverage attention-based contextual representation extraction to enable the learning policy to adapt to heterogeneous tasks and environments. Extensive experiments demonstrate that our proposed method significantly outperforms existing baselines on a range of complex, high-dimensional single-task and multi-task benchmarks.