Title: Improved Communication Efficiency in Federated Natural Policy Gradient via ADMM-based Gradient Updates
Venue: NeurIPS 2023
Accepted: yes
Keywords: reinforcement learning

federated learning
Abstract: Federated reinforcement learning (FedRL) enables agents to collaboratively train a global policy without sharing their individual data. However, high communication overhead remains a critical bottleneck, particularly for natural policy gradient (NPG) methods, which are second-order. To address this issue, we propose the FedNPG-ADMM framework, which leverages the alternating direction method of multipliers (ADMM) to approximate global NPG directions efficiently. We theoretically demonstrate that using ADMM-based gradient updates reduces communication complexity from $\mathcal{O}({d^{2}})$ to $\mathcal{O}({d})$ at each iteration, where $d$ is the number of model parameters. Furthermore, we show that achieving an $\epsilon$-error stationary convergence requires $\mathcal{O}(\frac{1}{(1-\gamma)^{2}{\epsilon}})$ iterations for discount factor $\gamma$, demonstrating that FedNPG-ADMM maintains the same convergence rate as standard FedNPG. Through evaluation of the proposed algorithms in MuJoCo environments, we demonstrate that FedNPG-ADMM maintains the reward performance of standard FedNPG, and that its convergence rate improves when the number of federated agents increases.

Improved Communication Efficiency in Federated Natural Policy Gradient via ADMM-based Gradient Updates

Abstract

Federated reinforcement learning (FedRL) enables agents to collaboratively train a global policy without sharing their individual data. However, high communication overhead remains a critical bottleneck, particularly for natural policy gradient (NPG) methods, which are second-order. To address this issue, we propose the FedNPG-ADMM framework, which leverages the alternating direction method of multipliers (ADMM) to approximate global NPG directions efficiently. We theoretically demonstrate that using ADMM-based gradient updates reduces communication complexity from O ( d 2 ) to O ( d ) at each iteration, where d is the number of model parameters. Furthermore, we show that achieving an ϵ -error stationary convergence requires O ( 1 (1 -γ ) 2 ϵ ) iterations for discount factor γ , demonstrating that FedNPG-ADMM maintains the same convergence rate as the standard FedNPG. Through evaluation of the proposed algorithms in MuJoCo environments, we demonstrate that FedNPG-ADMM maintains the reward performance of standard FedNPG, and that its convergence rate improves when the number of federated agents increases.

1 Introduction

Policy gradient methods are commonly used to solve reinforcement learning (RL) problems in various applications, with recent popular examples being InstructGPT [30] and ChatGPT (GPT-4 [29]). Although first-order methods, such as Policy Gradient (PG) and Proximal Policy Optimization (PPO) [40], are favored for their simplicity, second-order methods, such as Natural Policy Gradient (NPG) [16, 36] and its practical version Trust Region Policy Optimization (TRPO) [38], are often seen to exhibit superior convergence behavior, e.g., on the Swimmer [40] and Humanoid [10] tasks in MuJoCo environments [44]. Furthermore, recent works have demonstrated that the convergence guarantees of NPG with Kullback-Leibler (KL) divergence constraints (second-order methods) are superior to those of PG (first-order methods) [21, 23], motivating closer inspection for practical use.

Many contemporary RL applications are large-scale and rely on high volumes of data for model training [20, 22, 33]. A conventional approach to enabling training on massive data is transmitting data collected locally by different agents to a central server, which can then be used for policy learning. However, this is not always feasible in real-world systems where communication bandwidth

is limited and long delays are not acceptable [19, 28], such as in edge devices [19], Internet of Things [5, 27], autonomous driving [41], and vehicle transportation [3, 11]. Moreover, sharing individual data collected by agents raises privacy and legal issues [15, 26].

Federated learning (FL) [12, 25] offers a promising solution to the challenges posed by data centralization, where agents communicate locally trained models rather than raw datasets. However, FL is typically applied to supervised learning problems. Recent work has expanded the scope of FL to federated reinforcement learning (FedRL), where N agents collaboratively learn a global strategy without sharing the trajectories they collected during agent-environment interaction [14, 17, 50]. Federated reinforcement learning has been studied in tabular RL [2, 14, 17], control tasks [47, 48] and for value-function based algorithms [46, 50], where linear speedup has been demonstrated. For policy-gradient based algorithms, we note that linear speedup is easy to see as the trajectories collected by each agent could be parallelized.

Efficient state-of-the-art guarantees for federated policy gradient based approaches can be achieved by Federated NPG (FedNPG), which is a second-order method. However, sharing second-order information increases the communication complexity, which is one of the fundamental challenges in FL [7, 42, 37]. In supervised FL, works including FedNL [37], BL [34], Newton-Star/Learn [13] and FedNew [7] have been recently proposed to reduce the communication complexity of second-order methods, typically by approximating Hessian matrices for convex or strongly convex problems. However, federated second-order reinforcement learning has not been investigated, which provides the key motivation for our work. The key question that this paper aims to address is:

Can we reduce the communication complexity for second-order federated natural policy gradient approach while maintaining performance guarantees?

We answer this question in the affirmative by introducing FedNPG-ADMM, an algorithm that estimates global NPG directions using alternating direction method of multipliers (ADMM) [8, 9, 45]. This estimation reduces the communication complexity from O ( d 2 ) to O ( d ) , where d is the number of model parameters. However, it is non-trivial to see whether FedNPG-ADMM will maintain similar convergence guarantees as FedNPG. We show in this work that FedNPG-ADMM indeed does maintain these guarantees, and provides a speedup with the number of agents. The key contributions that we make are summarized as follows:

We propose a novel federated NPG algorithm, FedNPG-ADMM, where the global NPG directions are estimated through ADMM.

Using the ADMM-based global direction estimation, we demonstrated that the communication complexity reduces by O ( d ) as compared to transmitting the second-order information (standard FedNPG).

We prove the FedNPG-ADMM method achieves an ϵ -error stationary convergence with O ( 1 (1 -γ ) 2 ϵ ) iterations for discount factor γ . Thus, it achieves the same convergence rate as the standard FedNPG.

Experimental evaluations in MuJoCo environments demonstrate that FedNPG-ADMM maintains the convergence performance of FedNPG. We also show improved performance as more federated agents engage in collecting trajectories.

2 Background

Markov Decision Process: We consider the Markov decision process (MDP) as a tuple ⟨S , A , P , R , γ ⟩ , where S is the state space, A is a finite action space, P : S × A × S → R is a Markov kernel that determines transition probabilities, R : S × A → R is a reward function, and γ ∈ (0 , 1) is a discount factor. At each time step t , the agent executes an action a t ∈ A from the current state s t ∈ S , following a stochastic policy π , i.e., a t ∼ π ( ·| s t ) . For on policy π , a state value function is defined as

Similarly, a state-action value function (Q-function) is defined as

An advantage function is then define as A π ( s, a ) = Q π ( s, a ) -V π ( s ) . With continuous states, the policy is parametrized by θ ∈ R d , and then the policy is referred as π θ (Deep RL parametrizes π θ by deep neural networks). A state-action visitation measure induced by π θ is given as

where starting state s 0 is drawn from a distribution ρ . The goal of an agent is to maximize the expected discounted return defined as

The gradient of J ( θ ) can be written as [39]:

where τ = ( s 0 , a 0 , s 1 , a 1 , · · · ) is a trajectory induced by policy π θ . We denote the policy gradient by g for short. In practice, we can sample ( s, a ) ∼ ν π θ k and obtain the unbiased estimate ̂ A π θ k ( s, a ) using Algorithm 3 in [1].

Natural Policy Gradient (NPG): At the k -th iteration, natural policy methods with a trust region [38] update policy parameters as follows

D ( · ) is the KL-divergence operation, and δ > 0 is the radius of the trust region. Practically, using the first-order Taylor expansion for the target value and the second-order Taylor expansion for the divergence constraint, (6) is expanded as follows

where H = ∇ 2 θ D ( θ ∥ θ k ) ∈ R d × d , and the individual elements are given by H ij = ∂ ∂θ i ∂ ∂θ j E s [ D ( π θ ( ·| s ) ∥ π θ k ( ·| s ) )] ∣ ∣ ∣ θ = θ k . Using Lagrangian duality, the iterates of NPG ascent are expressed as

Federated NPG: FedNPG is a paradigm in that N agents collaboratively train a common global policy with parameters θ as illustrated in Figure 1 (a). During the training process, each agent computes gradients (and second-order matrices) using its local data ; then, the gradients of all agents are transmitted to a central server. In particular, one FedNPG training iteration consists of the following three steps:

where

Downlink Transmission: The server broadcasts the current global policy parameters θ ∈ R d to all N agents.

Uplink Transmission: Collecting its own local data D i based on the common policy π θ , each agent i computes its local gradient g i ∈ R d and second-order matrix H i ∈ R d × d . Then, it sends g i and H i back to the server.

Global Update: The server averages local gradients and second-order matrices to get the global gradient and matrix as follows:

The server then updates global policy parameters as

We use |D i | to denote the size of collected data set D i . Without loss of generality, we take |D 1 | = · · · = |D N | for simplicity. As proven in [49], gradient update methods are immune to whether collected data is i.i.d., or not.

1

N

Server

1

local compute

N

Server

local compute

(a) FedNPG

(b) FedNPG-ADMM

Bottleneck of Federated NPG: As shown in Figure 1, at each iteration, the server collects { H i ∈ R d × d , g i ∈ R d } N i =1 from N agents and updates global policy parameters as follows

We call this method a standard average FedNPG. The uplink communication complexity from each agent is O ( d 2 ) in each iteration round. As the uplink is highly limited [43], applying (12) is not practical when d is large.

In the next section, we introduce our ADMM-based approach as shown in Figure 1 (b), which reduces the communication complexity to O ( d ) at each iteration and meanwhile keeps convergence performances.

3 FedNPG via ADMM

To minimize the communication overhead in each round of communication, we begin by formulating a quadratic problem. The solution to this problem provides the updating direction in (12), as follows:

This minimization problem is equivalent to

where ρ > 0 is a penalty constant and ∥ · ∥ denotes the Euclidean norm. This equivalence transforms the original quadratic problem into a distributed manner and is the key process of efficient FedNPG.

Remark : This approach does not aim to approximate the Hessian from each agent, but approximates the global direction ( ∑ N i =1 H i ) -1 ∑ N i =1 g i directly. Note that this differentiates our approach from the work mentioned in the introduction.

The Lagrangian function associated to optimization problem (14) is

where { λ i ∈ R d } N i =1 are dual variables. Next, we solve the distributed optimization problem through the alternating direction method of multipliers (ADMM) [4]. The policy update of FedNPG via one-step ADMM is given in Algo. 1.

Algorithm 1 FedNPG-ADMM

Input: MDP ⟨S , A , P , R , γ ⟩ ; Number of timesteps T ; Penalty constant ρ ; Step size η ; Initial θ 0 ∈ R d , y 0 ∈ R d , { y 0 i = y 0 } N i =1 , { λ i ∈ R d } N i =1 . 1: for k = 1 , · · · , K do 2: ▷ Server broadcast 3: Broadcast y k -1 and θ k -1 to N agents. 4: ▷ Agent update 5: for each agent i ∈ { N } do in parallel 6: λ i ← λ i + ρ ( y k -1 i -y k -1 ) 7: g k i ← 1 |D i | ∑ τ ∈D i ∑ T t =0 ( ∇ θ k -1 log π θ k -1 ( a t | s t ) ) ̂ A π θ k -1 ( s t , a t ) 8: y k i ← ( H k i + ρ I ) -1 ( g k i -λ i + ρ y k -1 ) 9: Transmit y k i ∈ R d and g k i ∈ R d to the server. 10: end for 11: ▷ Server update 12: y k ← 1 N ∑ N i =1 y k i 13: θ k ← θ k -1 + η √ 2 Nδ ( ∑ N i =1 g k i ) ⊤ y k · y k 14: end for Output: θ K

Agent i computes H i and g i based on locally collected data. At each step of ADMM, agent i updates y i in line 8 as follows

where I ∈ R d × d is the identity matrix. In practical implementation, conjugate gradient methods can be used to compute y i for efficiency (Appendix C in [38]).

In line 12, after receiving y i and g i from all N agents, the server updates the global search direction as follows

Dual variables in line 6 are updated by each agent as follows

Combining (17) and (18), we have ∑ N i =1 λ i = 0 .

In line 13 of Algo. 1, after the ADMM process, the server updates global policy parameters, where η ∈ (0 , 1) is the step size. The server then broadcasts the updated parameters to all N agents at the next iteration.

In every communication round, agent i only transmits y i and g i , with a communication complexity of O ( d ) . In contrast, the standard average approach in (12) requires transmitting H i and g i , with a communication complexity of O ( d 2 ) . This efficient communication approach allows second-order methods scalable to large-scale systems.

4 Convergence Analysis

In this section, we derive the convergence rate of FedNPG based on ADMM. In order to derive the guarantees, we make the following standard assumptions [1, 6, 23, 31, 51] on policy gradients, second-order matrices, and rewards.

Assumption 4.1.

The score function is bounded as ∥∇ θ log π θ ( a | s ) ∥ ≤ G , for all θ ∈ R d , s ∈ S , and a ∈ A .

Policy gradient is M -Lipschitz continuous. In other words, ∀ θ i , θ j ∈ R d , s ∈ S , and a ∈ A , we have

The reward function is bounded as r ( s, a ) ∈ [0 , R ] , for all s ∈ S , and a ∈ A .

Assumption 4.2. For all θ ∈ R d , the Fisher information matrix induced by policy π θ and initial state distribution ρ is positive definite as

for some constant µ F > 0 . For any two symmetric matrices with the same dimension, A ≽ B denotes the eigenvalues of A -B are greater or equal to zero.

Theorem 4.3. For a target error ϵ of stationary-point convergence, each agent samples O ( 1 (1 -γ ) 4 Nϵ ) trajectories, and the server obtains the update direction y k at each iteration. Choose η = µ 2 F 4 G 2 (56 G 2 + L J ) and

We have:

Proof sketch. The main idea in our proof is to show that the approximation error between the updating direction given by FedNPG-ADMM and NPG geometrically decreases up to some additional term, which depends on the statistical error. This term appears since we do not have access to the exact gradient. However, it decays at a rate proportional to sample size. As a result, FedNPG-ADMM achieves the same convergence rate as NPG as shown by constructing an appropriate Lyapunov function. The complete proof of Theorem 4.3 is given in Appendix A.

By the definition of stationary points, we need to find a parameter θ such that E ∥∇ J ( θ ) ∥ 2 ≤ ϵ , for all ϵ > 0 . The result in (22) achieves the stationary-point convergence for policy gradient methods as provided in [51]. Our approach keeps the sample complexity as that in the NPG method [23] and thanks to the federated scenario we consider, enjoys a much lower communication complexity.

We summarize the complexity improvement in Table 1. Recall in (5) that the estimated gradient g is an average over collected trajectories. The total trajectories ∑ N i =1 |D i | are collected by N agents equally using the common global policy π θ at each iteration. Each agent i samples |D i | = O ( 1 (1 -γ ) 4 Nϵ ) at each iteration and enjoys a federated sampling benefit compared to a single agent with |D i | = O ( 1 (1 -γ ) 4 ϵ ) . During the whole training process, each agent i has K |D i | = O ( 1 (1 -γ ) 6 Nϵ 2 ) sample complexity and K · 2 d = O ( d (1 -γ ) 2 ϵ ) communication complexity to achieve a stationary convergence.

5 Simulations

Setup : We consider three MuJoCo tasks [44] with the MIT License, which have continuous state spaces. Specifically, a Swimmer-v4 task with small state and action spaces, a Hopper-v4 task with middle state and action spaces, and a Humanoid-v4 task with large state and action spaces are considered as described in Table 3. Policies are parameterized by fully connected multi-layer perceptions (MLPs) with settings in Table 4. We follow the practical settings in TRPO with line search [38], and in stable-baselines3 [35] with generalized advantage estimation ( 0 . 95 ) [39] and the Adam optimizer [18] in our implementation. Convergence performances are measured over 10 runs with random seeds from 0 to 9 . The solid lines in Figure 2 and 3 are averaged results, and the shadowed areas are confidence intervals with 95% confidence level. We use PyTorch [32] to implement deep neural networks and RL algorithms. The tasks are trained on NVIDIA RTX 3080 GPU with 10 GB of memory.

Performance metrics : We consider performance metrics as follows:

Communication overhead: the data size transmitted from each agent;

Rewards: the average trajectory rewards across the batch collected at each iteration;

Convergence: rewards versus iterations during the training process.

We first evaluate the influence of the number of federated agents. In Figure 2, with different numbers of agents, we test the convergence performances of standard average FedNPG in (12) with O ( d 2 ) communication complexity and FedNPG-ADMM in Algo. 1 with O ( d ) communication complexity at each iteration. The x-axis denotes the number of iterations in federated learning. Both algorithms converge faster, have lower variance, and achieve higher final rewards when more agents are engaged to collect trajectories. Compared to the standard average FedNPG, FedNPG-ADMM does not only reduce communication complexity, but also keeps convergence performances. The hyperparameter and MLP settings are described in Table 4. We summarize the final reward results with standard deviations in Table 2. FedNPG-ADMM achieves similar final rewards compared to the standard average FedNPG. It also works with slightly high variance when only one agent is engaged.

In Figure 3, we compare the performances of the standard average FedNPG, FedNPG-ADMM, and first-order FedPPO algorithms. The number of federated agents N is fixed to 8 . PPO clipping parameter is set as 0 . 2 . FedNPG methods outperform FedPPO in these tasks with faster convergence and higher final rewards. It is noticeable that FedNPG-ADMM has similar convergence rates and achieves slightly higher final rewards than the standard average FedNPG for the Swimmer-v4 task, while it becomes slightly lower for the Humanoid-v4 task. Generally, there is no significant difference in convergence rates after ADMM approximation. The communication overhead is measured by the number of transmitted parameters with double precision in each agent. FedNPG-ADMM keeps the communication overhead as the first-order methods, while FedNPG has much higher costs. The ADMMmethod reduces the cost by 4 orders of magnitude in the Swimmer-v4 task and about 6 orders of magnitude in the Humanoid-v4 task.

(a) FedNPG (Swimmer-v4)

(b) FedNPG-ADMM (Swimmer-v4)

(c) FedNPG (Hopper-v4)

(d) FedNPG-ADMM (Hopper-v4)

120

120

100

100

Rewards

S

80

Rewards

S

80

60

agent

60

agent

N =1

N =1

40

40

N =2

N =2

20

N = 4

20

N = 4

N=8

N=8

0

0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

#Iterations

1e3

#Iterations

1e3

3000

3000

agent

agent

2500

N=1

2500

N = 1

N=2

N =2

2000

N=4

2000

N = 4

vards

ards

N=8

N=8

1500

1500

ew:

R

R

1000

1000

500

500

0

0

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

#Iterations

1e3

#Iterations

1e3

(a) Swimmer-v4 Rewards

(b) Humanoid-v4 Rewards

(c) Swimmer-v4 Overhead

(d) Humanoid-v4 Overhead

120

1000

100

800

Rewards

80

Rewards

600

60

400

40

FedPPO

FedPPO

20

FedNPG

200

FedNPG

FedNPG-ADMM

FedNPG-ADMM

0

0

0.0

0.2

0.4

0.6

0.8

1.0

0

1

2

3

#Iterations

1e3

#Iterations

1e3

10

16

overhead

overhead

10

10

10

14

ication

uor

10

12

10°

8

icatic

10

10

10

6

FedPPO

FedPPO

FedNPG

10

8

FedNPG

C

FedNPG-ADMM

FedNPG-ADMM

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.5

1.0

1.5

2.0

2.5

3.0

#Iterations

1e3

#Iterations

1e3

In Figure 4, we test performances with agent selection. In the Swimmer task, we randomly select 75% and 50% of agents in each iteration, and the performances only drop slightly (final rewards drop less than 6% ). Thus, our proposed method is robust for agents with disconnection in practice.

120

100

S

80

Rew:

60

R

selection

40

50%

20

75%

100%

0

0.0

0.2

0.4

0.6

0.8

1.0

#Iterations

1e3

6 Conclusion & Discussion

In this paper, we proposed a novel federated natural policy gradient (NPG) algorithm that estimates global directions using ADMM (FedNPG-ADMM). Our ADMM-based gradient updates significantly reduce communication complexity from O ( d 2 ) to O ( d ) , where d is the number of model parameters. It thus enables second-order policy gradient methods to be used for large-scale federated reinforcement learning problems. We proved that our FedNPG-ADMM algorithm achieves stationary convergence without requiring more samples as compared to standard FedNPG. Furthermore, we empirically showed the improved performances of FedNPG-ADMM in the MuJoCo environments as compared to FedNPG and the first order methods.

Overall, our proposed FedNPG-ADMM algorithm provides a scalable and efficient solution for large-scale federated reinforcement learning problems. Our contributions include a novel direction on policy-based FedRL, a new algorithm, reduced communication complexity, and convergence analysis. We believe that our approach can have a significant impact on a wide range of real-world applications, where large-scale distributed reinforcement learning with communication and privacy constraints is critical.

Limitations: (1) While the communication complexity improvement is loosely tied to the privacy advantage, exploring deep connections with differential privacy improvement is open. (2) Partial agent participation is not studied under this framework, in which communication complexity can be further reduced. (3) In experiments, it would be more persuasive to extend the number of federated agents to a larger scale.

Acknowledgments

We thank all reviewers for their invaluable feedback that helped us significantly strengthen the work. C. Brinton acknowledges financial support from NSF CNS-2146171 and DARPA D22AP00168-00. Han Wang and James Anderson acknowledge financial support from NSF grants ECCS-2144634 and ECCS-2231350.

References

Agarwal, A., Kakade, S.M., Lee, J.D., Mahajan, G.: On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research 22 (1), 4431-4506 (2021)

Agarwal, M., Ganguly, B., Aggarwal, V.: Communication efficient parallel reinforcement learning. In: Uncertainty in Artificial Intelligence (2021)

Al-Abbasi, A.O., Ghosh, A., Aggarwal, V.: DeepPool: Distributed model-free algorithm for ride-sharing using deep reinforcement learning. IEEE Transactions on Intelligent Transportation Systems 20 (12), 4714-4727 (2019)

Boyd, S., Parikh, N., Chu, E., Peleato, B., Eckstein, J., et al.: Distributed optimization and statistical learning via the alternating direction method of multipliers. Foundations and Trends® in Machine learning 3 (1), 1-122 (2011)

Chen, T., Giannakis, G.B.: Bandit convex optimization for scalable and dynamic iot management. IEEE Internet of Things Journal 6 (1), 1276-1286 (2018)

Ding, D., Zhang, K., Basar, T., Jovanovic, M.: Natural policy gradient primal-dual method for constrained Markov decision processes. In: Advances in Neural Information Processing Systems (2020)

Elgabli, A., Issaid, C.B., Bedi, A.S., Rajawat, K., Bennis, M., Aggarwal, V.: FedNew: A communication-efficient and privacy-preserving Newton-type method for federated learning. In: International Conference on Machine Learning (2022)

Elgabli, A., Park, J., Bedi, A.S., Bennis, M., Aggarwal, V.: GADMM: Fast and communication efficient framework for distributed machine learning. Journal of Machine Learning Research 21 (76), 1-39 (2020)

Elgabli, A., Park, J., Bedi, A.S., Issaid, C.B., Bennis, M., Aggarwal, V .: Q-GADMM: Quantized group ADMM for communication efficient decentralized machine learning. IEEE Transactions on Communications 69 (1), 164-181 (2020)

Engstrom, L., Ilyas, A., Santurkar, S., Tsipras, D., Janoos, F., Rudolph, L., Madry, A.: Implementation matters in deep RL: A case study on PPO and TRPO. In: International Conference on Learning Representations (2020)

Haliem, M., Mani, G., Aggarwal, V., Bhargava, B.: A distributed model-free ride-sharing approach for joint matching, pricing, and dispatching using deep reinforcement learning. IEEE Transactions on Intelligent Transportation Systems 22 (12), 7931-7942 (2021)

Hosseinalipour, S., Brinton, C.G., Aggarwal, V., Dai, H., Chiang, M.: From federated to fog learning: Distributed machine learning over heterogeneous wireless networks. IEEE Communications Magazine 58 (12), 41-47 (2020)

Islamov, R., Qian, X., Richtarik, P.: Distributed second order methods with fast rates and compressed communication. In: Proceedings of the 38th International Conference on Machine Learning (2021)

Jin, H., Peng, Y., Yang, W., Wang, S., Zhang, Z.: Federated reinforcement learning with environment heterogeneity. In: International Conference on Artificial Intelligence and Statistics (2022)

Kairouz, P., McMahan, H.B., Avent, B., Bellet, A., Bennis, M., Bhagoji, A.N., Bonawitz, K., Charles, Z., Cormode, G., Cummings, R., D'Oliveira, R.G.L., Eichner, H., Rouayheb, S.E., Evans, D., Gardner, J., Garrett, Z., Gascón, A., Ghazi, B., Gibbons, P.B., Gruteser, M., Harchaoui, Z., He, C., He, L., Huo, Z., Hutchinson, B., Hsu, J., Jaggi, M., Javidi, T., Joshi, G., Khodak, M., Konecný, J., Korolova, A., Koushanfar, F., Koyejo, S., Lepoint, T., Liu, Y., Mittal, P., Mohri, M., Nock, R., Özgür, A., Pagh, R., Qi, H., Ramage, D., Raskar, R., Raykova, M., Song, D., Song, W., Stich, S.U., Sun, Z., Suresh, A.T., Tramèr, F., Vepakomma, P., Wang, J., Xiong, L., Xu, Z., Yang, Q., Yu, F.X., Yu, H., Zhao, S.: Advances and open problems in federated learning. Foundations and Trends® in Machine Learning 14 (1-2), 1-210 (2021)

Kakade, S.M.: A natural policy gradient. In: Advances in Neural Information Processing Systems (2001)

Khodadadian, S., Sharma, P., Joshi, G., Maguluri, S.T.: Federated reinforcement learning: Linear speedup under Markovian sampling. In: International Conference on Machine Learning (2022)

Kingma, D.P., Ba, J.: Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 (2014)

Lan, G., Liu, X.Y., Zhang, Y., Wang, X.: Communication-efficient federated learning for resource-constrained edge devices. IEEE Transactions on Machine Learning in Communications and Networking 1 , 210-224 (2023)

Li, X., Li, Y ., Zhan, Y ., Liu, X.Y .: Optimistic bull or pessimistic bear: Adaptive deep reinforcement learning for stock portfolio allocation. In: International Conference on Machine Learning Workshop on Applications and Infrastructure for Multi-Agent Learning (2019)

Liu, B., Cai, Q., Yang, Z., Wang, Z.: Neural trust region/proximal policy optimization attains globally optimal policy. In: Advances in Neural Information Processing Systems (2019)

Liu, X.Y., Li, Z., Yang, Z., Zheng, J., Wang, Z., Walid, A., Guo, J., Jordan, M.I.: ElegantRLPodracer: Scalable and elastic library for cloud-native deep reinforcement learning. In: Advances in Neural Information Processing Systems Workshop on Deep Reinforcement Learning (2021)

Liu, Y., Zhang, K., Basar, T., Yin, W.: An improved analysis of (variance-reduced) policy gradient and natural policy gradient methods. In: Advances in Neural Information Processing Systems (2020)

Makhdoumi, A., Ozdaglar, A.: Convergence rate of distributed ADMM over networks. IEEE Transactions on Automatic Control 62 (10), 5082-5095 (2017)

McMahan, B., Moore, E., Ramage, D., Hampson, S., y Arcas, B.A.: Communication-efficient learning of deep networks from decentralized data. In: International Conference on Artificial Intelligence and Statistics (2017)

Mothukuri, V., Parizi, R.M., Pouriyeh, S., Huang, Y., Dehghantanha, A., Srivastava, G.: A survey on security and privacy of federated learning. Future Generation Computer Systems 115 , 619-640 (2021)

Nguyen, D.C., Ding, M., Pathirana, P.N., Seneviratne, A., Li, J., Poor, H.V.: Federated learning for internet of things: A comprehensive survey. IEEE Communications Surveys & Tutorials 23 (3), 1622-1658 (2021)

Niknam, S., Dhillon, H.S., Reed, J.H.: Federated learning for wireless communications: Motivation, opportunities, and challenges. IEEE Communications Magazine 58 (6), 46-51 (2020)

OpenAI: GPT-4 technical report. arXiv preprint arXiv:2303.08774 (2023)

Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Gray, A., Schulman, J., Hilton, J., Kelton, F., Miller, L., Simens, M., Askell, A., Welinder, P., Christiano, P., Leike, J., Lowe, R.: Training language models to follow instructions with human feedback. In: Advances in Neural Information Processing Systems (2022)

Papini, M., Binaghi, D., Canonaco, G., Pirotta, M., Restelli, M.: Stochastic variance-reduced policy gradient. In: Proceedings of the 35th International Conference on Machine Learning (2018)

Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al.: PyTorch: An imperative style, high-performance deep learning library. In: Advances in Neural Information Processing Systems (2019)

Provost, F., Fawcett, T.: Data science and its relationship to big data and data-driven decision making. Big Data 1 (1), 51-59 (2013)

Qian, X., Islamov, R., Safaryan, M., Richtarik, P.: Basis matters: Better communicationefficient second order methods for federated learning. In: Proceedings of The 25th International Conference on Artificial Intelligence and Statistics (2022)

Raffin, A., Hill, A., Gleave, A., Kanervisto, A., Ernestus, M., Dormann, N.: Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research 22 (268), 1-8 (2021)

Rajeswaran, A., Lowrey, K., Todorov, E.V., Kakade, S.M.: Towards generalization and simplicity in continuous control. Advances in Neural Information Processing Systems (2017)

Safaryan, M., Islamov, R., Qian, X., Richtarik, P.: FedNL: Making Newton-type methods applicable to federated learning. In: International Conference on Machine Learning (2022)

Schulman, J., Levine, S., Moritz, P., Jordan, M.I., Abbeel, P.: Trust region policy optimization. arXiv preprint arXiv:1502.05477 (2017)

Schulman, J., Moritz, P., Levine, S., Jordan, M., Abbeel, P.: High-dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.02438 (2018)

Schulman, J., Wolski, F., Dhariwal, P., Radford, A., Klimov, O.: Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 (2017)

Shalev-Shwartz, S., Shammah, S., Shashua, A.: Safe, multi-agent, reinforcement learning for autonomous driving. arXiv preprint arXiv:1610.03295 (2016)

Shamir, O., Srebro, N., Zhang, T.: Communication-efficient distributed optimization using an approximate Newton-type method. In: International Conference on Machine Learning (2014)

speedtest.net: Speedtest market report in the United States (2023), http://www.speedtest. net/reports/united-states

Todorov, E., Erez, T., Tassa, Y.: MuJoCo: A physics engine for model-based control. In: IEEE/RSJ International Conference on Intelligent Robots and Systems (2012)

Wang, H., Marella, S., Anderson, J.: FedADMM: A federated primal-dual algorithm allowing partial participation. In: IEEE 61st Conference on Decision and Control (2022)

Wang, H., Mitra, A., Hassani, H., Pappas, G.J., Anderson, J.: Federated temporal difference learning with linear function approximation under environmental heterogeneity. arXiv preprint arXiv:2302.02212 (2023)

Wang, H., Toso, L.F., Mitra, A., Anderson, J.: Model-free learning with heterogeneous dynamical systems: A federated LQR approach. arXiv preprint arXiv:2308.11743 (2023)

Wang, H., Toso, L.F., Anderson, J.: FedSysID: A federated approach to sample-efficient system identification. In: Learning for Dynamics and Control Conference. pp. 1308-1320. PMLR (2023)

Woodworth, B.E., Patel, K.K., Srebro, N.: Minibatch vs local SGD for heterogeneous distributed learning. In: Advances in Neural Information Processing Systems (2020)

Xie, Z., Song, S.: FedKL: Tackling data heterogeneity in federated reinforcement learning by penalizing KL divergence. IEEE Journal on Selected Areas in Communications 41 (4), 1227-1242 (2023)

Xu, P., Gao, F., Gu, Q.: Sample efficient policy gradient methods with recursive variance reduction. In: International Conference on Learning Representations (2020)

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a communication-efficient algorithm, FedNPG-ADMM, for federated natural policy gradient by using a reformulation of quadratic problem. It reduces the communication complexity from $\\mathcal{O}(d^2)$ to $\\mathcal{O}(d)$. The convergence analysis is provided accordingly.\n\nSTRENGTHS:\n1. The paper is well-written and clearly presented.\n2. The reformulation of the original quadratic problem into a distributed manner (eq 14) is a smart move.\n3. Using ADMM-based global direction estimation to reduce the communication complexity by $\\mathcal{O}(d)$.\n4. Stationary convergence in non-convex case is provided, which is the same order as the standard FedNPG.\n\nWEAKNESSES:\nAs it is in the federated learning (FL) setting, I am curious about main difference between the FL setting considered in this paper and the classic distributed learning.\n1. data heterogeneity is one of the key features in FL, but how it impact the model performance does not clearly discussed in the paper, neither in theory or in experiments. \n2. To reduce the communication costs, local updates are often used within each agent for FL algorithms. But in the FedNPG-ADMM algorithm proposed in this paper, there is no local update step for each agent. Could local steps be directly applied in the FedNPG-ADMM and further reduce the communication cost while maintaining the same convergence?\nGiven above, to be precise, the setting for this paper is a distributed learning setting, rather than FL with data heterogeneity and local update.\n\nThe major contribution is the communication reduction as claimed in the paper, so I expect the experiments could have more comparison of direct communication costs.\n1. Directly comparison in terms of the communication cost should be given. In other words, the results with communication cost as the x-axis are also desired.\n2. Appreciate it if other baselines could be included other than FedNPG. I am curious about the actual communication costs among first-order methods, second-order methods (FedNPG), and approximated version of second-order methods (FedNPG-ADMM).\n\nQUESTIONS:\nsee above\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies how to train a global policy using distributed data in reinforcement learning. The authors propose a distributed natural policy gradient method by employing ADMM to approximately compute a natural policy gradient direction. The communication complexity is linear in the dimension of policy parameters. The authors also prove the sublinear iteration complexity for achieving a stationary point. The effectiveness of the proposed method is finally verified in MuJoCo experiments.\n\nSTRENGTHS:\n**originality**\n\nHigh communication complexity is a bottleneck for implementing many RL algorithms in federated training. I believe the novelty of this work lies in a new application of ADMM to a popular RL algorithm: natural policy gradient. This method achieves better communication complexity than one for the naive method. I am not aware of this method in the literature.  \n\n**quality**\n\n- The authors provide the stationary-point analysis of the proposed method in a distributed training setting.\n\n- The authors verify the use of the proposed method in experiments. \n\n**clarity**\n\nThe proposed method and theory are clearly explained.\n\n**significance**\n\n- The distributed training algorithm is important for applying RL in distributed systems. \n\n- The stationary-point convergence guarantee is provided.\n\nWEAKNESSES:\n- The naive method of fererated natural policy gradient only has quadratic communication complexity with respect to the dimension of policy parameter. This does not seem to be a real bottleneck.\n\n- Directly extending centralized RL algorithms is not very significant, since fault-tolerance and robustness issues are more important to distributed systems. \n\n- It is important to compare communication complexity with the literature, e.g., Fault-Tolerant Federated Reinforcement Learning with Theoretical Guarantee. \n\n- Some assumptions can be strong in practice, e.g., the invertibility of Hessian. It is useful to provide examples when Hessian is invertible. \n\n- The provided convergence property can be arbitrarily sub-optimal, and it is not clear how to check all assumptions. The stationary-point analysis has a gap with the convergence rate analysis of natural policy gradient in the literature. \n\n- The experimental setting is artificial since it is not originally designed for federated learning. There are no comparable baselines in experiments.\n\nQUESTIONS:\nSome questions are raised in Weaknesses. Here are some other questions. \n\n- Can the authors provide robustness or fault-tolerance guarantees?  \n\n- Can ADMM-based update be applied to other policy gradient methods? It is useful to discuss generalizability. \n\n- Since the global convergence exists in many policy gradient methods, can the authors strengthen the convergence guarantees? \n\n- Can the authors compare the proposed method with other baselines, e.g., Fault-Tolerant Federated Reinforcement Learning with Theoretical Guarantee?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-03"}

comments: SUMMARY:
This paper applies the ADMM technique to the Fed-NPG algorithm in reinforcement learning and reduces the communication cost from $O(d^2)$ to $O(d)$ where $d$ is the number of parameters, which nearly maintains the convergence results of Fed-NPG. Empirical results verify the theoretical analysis.

STRENGTHS:
The communication cost reduction is impressive since in the federated learning setting, the communication cost is one of the bottlenecks, and the convergence results are nearly the same as Fed-NPG. The idea to combine natural policy gradient with ADMM style methods (or the penalty functions) is interesting (and at least novel to me).

WEAKNESSES:
I am not very familiar with reinforcement learning, and thus may not find potential weaknesses. However, I suggest comparing with the original policy gradient methods in the empirical section and even in the main content (show that Fed-NPG-ADMM converges much fast than the policy gradient variation in federated learning) and plotting all the communication cost for Fed-NPG, Fed-NPG-ADMM, and policy gradients.

QUESTIONS:
See the weakness section.

LIMITATIONS:
n/a
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe work proposed a new algorithm for federated policy optimization. The proposed work uses prima-dual update to replace the primal update of FedNPG, so that the communication cost reduces from $d^2$ to $d$. The proposed work enjoys same rate of convergence as FedNPG under certain assumptions, and numerical experiments were proposed to show the efficacy of the algorithm.\n\nSTRENGTHS:\nThe paper is well-rounded and the presentation is clear. The proposed method clearly enjoys the supremacy in terms of communication error. The numerical experiments were proposed to show the efficacy of the algorithm.\n\nWEAKNESSES:\n(Please reply to the Questions section directly) The motivation of the proposed algorithm seems a bit less significant. Moreover, the convergence theory requires policy gradient bounded, which seems quite strong. In the numerical experiments, more results could be included to further show the efficiency of the proposed algorithm.\n\nQUESTIONS:\nHere are the my questions and suggestions:\n\n1. Concerns on the motivations: the idea of combining reinforcement learning and federated learning is interesting, which is also largely un-explored in the literature. It is good to see the exploration on this direction. An intriguing question is that can we consider existing works on reformulating the TRPO, such as the KL-penalitzed objective in [1], to federated setting? Since if we get rid of the constraint, we don’t need to communicate the Hessian in the first place (and FedAvg-type algorithm may also enjoys similar rate of convergence in theory). In summary, the authors propose to study the policy optimization in federated setting, but eventually ended up at a less significant improvement over FedNPG (and this improvement comes from essentially the constrained federated learning, not from RL or policy optimization), which makes me feel less confident about the contribution in the RL field;\n\n2. In terms of the assumptions for convergence analysis, the authors proposed bounded policy gradient and Lipschitz continuous (Assumption 4.1) and claimed that they are fairly standard in the existing literature. I’m actually not very sure about the boundedness of the policy gradient. Could the authors just include the specific assumptions in the existing literatures? Same question for Assumption 4.2.\n\n3. The penalty parameter $\\rho$ is the penalty to make sure that each local $y_i$ consensus with the global $y$. In the experiment $\\rho$ is simply a small number as 0.1 or 0.01. My concern is that during the update, is $y_i - y\\rightarrow 0$ true or not? The authors could also plot the consensus errors.\n\n4. In the experiments, for example Figure 2 and 3, FedNPG-ADMM seems not better than FedNPG in terms of the rewards. It would be interesting to see why this is the case. I know that the main claim of FedNPG-ADMM is its commutation efficiency, and I think the authors could also plot the curves wrt the CPU times.\n\nReferences:\n\n[1] Schulman, John, et al. \"Proximal policy optimization algorithms.\" arXiv preprint arXiv:1707.06347 (2017).\n\nLIMITATIONS:\nThe limitation is well stated in weakness and question sections. The authors also include the discussions in limitations at the end of the work.\n\nI’m not aware of any potential negative social impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-23"}