Federated Reinforcement Learning with Constraint Heterogeneity

TL;DR

This paper introduces federated primal-dual policy optimization methods for reinforcement learning with multiple constraints across distributed agents, ensuring collaborative policy learning while respecting local constraint signals.

Contribution

It proposes novel federated primal-dual algorithms based on policy gradient methods, with convergence guarantees and practical implementations using NPG and PPO.

Findings

FedNPG achieves global convergence with rate (1/\u221A T)

FedPPO effectively handles complex tasks with deep neural networks

The methods address constraint heterogeneity in federated RL scenarios

Abstract

We study a Federated Reinforcement Learning (FedRL) problem with constraint heterogeneity. In our setting, we aim to solve a reinforcement learning problem with multiple constraints while $N$ training agents are located in $N$ different environments with limited access to the constraint signals and they are expected to collaboratively learn a policy satisfying all constraint signals. Such learning problems are prevalent in scenarios of Large Language Model (LLM) fine-tuning and healthcare applications. To solve the problem, we propose federated primal-dual policy optimization methods based on traditional policy gradient methods. Specifically, we introduce $N$ local Lagrange functions for agents to perform local policy updates, and these agents are then scheduled to periodically communicate on their local policies. Taking natural policy gradient (NPG) and proximal policy optimization (PPO) as policy optimization methods, we mainly focus on two instances of our algorithms, ie, {FedNPG} and {FedPPO}. We show that FedNPG achieves global convergence with an $\tilde{O}(1/\sqrt{T})$ rate, and FedPPO efficiently solves complicated learning tasks with the use of deep neural networks.