SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Guangchen Lan
lan44@purdue.edu
Han Wang
hw2786@columbia.edu
anderson@ee.columbia.edu
Christopher Brinton
cgb@purdue.edu
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}\left(\frac{1}{(1 - \gamma)^2\epsilon}\right)$  iterations for discount factor  $\gamma$ , 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.
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  $\mathcal{O}(d^2)$  to  $\mathcal{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:
1. We propose a novel federated NPG algorithm, FedNPG-ADMM, where the global NPG directions are estimated through ADMM.
2. Using the ADMM-based global direction estimation, we demonstrated that the communication complexity reduces by  $\mathcal{O}(d)$  as compared to transmitting the second-order information (standard FedNPG).
3. We prove the FedNPG-ADMM method achieves an  $\epsilon$ -error stationary convergence with  $\mathcal{O}\left(\frac{1}{(1 - \gamma)^2\epsilon}\right)$  iterations for discount factor  $\gamma$ . Thus, it achieves the same convergence rate as the standard FedNPG.
4. 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.
Markov Decision Process: We consider the Markov decision process (MDP) as a tuple  $\langle S, \mathcal{A}, \mathcal{P}, \mathcal{R}, \gamma \rangle$ , where  $\mathcal{S}$  is the state space,  $\mathcal{A}$  is a finite action space,  $\mathcal{P}: \mathcal{S} \times \mathcal{A} \times \mathcal{S} \to \mathbb{R}$  is a Markov kernel that determines transition probabilities,  $\mathcal{R}: \mathcal{S} \times \mathcal{A} \to \mathbb{R}$  is a reward function, and  $\gamma \in (0,1)$  is a discount factor. At each time step  $t$ , the agent executes an action  $a_t \in \mathcal{A}$  from the current state  $s_t \in S$ , following a stochastic policy  $\pi$ , i.e.,  $a_t \sim \pi(\cdot | s_t)$ . For on policy  $\pi$ , a state value function is defined as
$$
V _ {\pi} (s) = \underset { \begin{array}{c} a _ {t} \sim \pi (\cdot | s _ {t}), \\ s _ {t + 1} \sim P (\cdot | s _ {t}, a _ {t}) \end{array} } {\mathbb {E}} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r \left(s _ {t}, a _ {t}\right) | s _ {0} = s \right]. \tag {1}
$$
Similarly, a state-action value function (Q-function) is defined as
$$
Q _ {\pi} (s, a) = \underset { \begin{array}{c} a _ {t} \sim \pi (\cdot | s _ {t}), \\ s _ {t + 1} \sim P (\cdot | s _ {t}, a _ {t}) \end{array} } {\mathbb {E}} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r (s _ {t}, a _ {t}) | s _ {0} = s, a _ {0} = a \right]. \tag {2}
$$
An advantage function is then defined as  $A_{\pi}(s,a) = Q_{\pi}(s,a) - V_{\pi}(s)$ . With continuous states, the policy is parametrized by  $\theta \in \mathbb{R}^d$ , and then the policy is referred as  $\pi_{\theta}$  (Deep RL parametrizes  $\pi_{\theta}$  by deep neural networks). A state-action visitation measure induced by  $\pi_{\theta}$  is given as
$$
\nu_ {\pi_ {\theta}} (s, a) = (1 - \gamma) \underset {s _ {0} \sim \rho} {\mathbb {E}} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} P \left(s _ {t} = s, a _ {t} = a \mid s _ {0}, \pi_ {\theta}\right) \right], \tag {3}
$$
where starting state  $s_0$  is drawn from a distribution  $\rho$ . The goal of an agent is to maximize the expected discounted return defined as
$$
J (\theta) = \underset {s \sim \rho} {\mathbb {E}} \left[ V _ {\pi_ {\theta}} (s) \right]. \tag {4}
$$
The gradient of  $J(\theta)$  can be written as [39]:
$$
\nabla_ {\theta} J (\theta) = \underset {\tau} {\mathbb {E}} \left[ \sum_ {t = 0} ^ {\infty} \left(\nabla_ {\theta} \log \pi_ {\theta} \left(a _ {t} \mid s _ {t}\right)\right) A _ {\pi_ {\theta}} \left(s _ {t}, a _ {t}\right) \right], \tag {5}
$$
where  $\tau = (s_0, a_0, s_1, a_1, \dots)$  is a trajectory induced by policy  $\pi_{\theta}$ . We denote the policy gradient by  $\mathbf{g}$  for short. In practice, we can sample  $(s, a) \sim \nu^{\pi_{\theta}k}$  and obtain the unbiased estimate  $\widehat{A}_{\pi_{\theta}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
$$
\theta^ {k + 1} = \arg \max  _ {\theta} \mathbb {E} _ {s, a} \left[ \frac {\pi_ {\theta} (a | s)}{\pi_ {\theta k} (a | s)} A _ {\pi_ {\theta k}} (s, a) \right] \tag {6}
$$
$$
\mathrm {s . t .} \bar {D} (\theta \| \theta^ {k}) \leq \delta .
$$
$$
\bar {D} (\theta \| \theta^ {k}) = \underset {s} {\mathbb {E}} \left[ D \left(\pi_ {\theta} (\cdot | s) \| \pi_ {\theta^ {k}} (\cdot | s)\right) \right], \tag {7}
$$
$D(\cdot)$  is the KL-divergence operation, and  $\delta > 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
$$
\theta^ {k + 1} = \underset {\theta} {\arg \max } \mathbf {g} ^ {\top} (\theta - \theta^ {k})
$$
$$
s. t. \frac {1}{2} \left(\theta - \theta^ {k}\right) ^ {\top} \mathbf {H} \left(\theta - \theta^ {k}\right) \leq \delta , \tag {8}
$$
where  $\mathbf{H} = \nabla_{\theta}^{2}\overline{D} (\theta \| \theta^{k})\in \mathbb{R}^{d\times d}$ , and the individual elements are given by  $\mathbf{H}_{ij} = \frac{\partial}{\partial\theta_i}\frac{\partial}{\partial\theta_j}\mathbb{E}_s\left[D\big(\pi_\theta (\cdot |s)\| \pi_{\theta^k}(\cdot |s)\big)\right]\Big|_{\theta = \theta^k}$ . Using Lagrangian duality, the iterates of NPG ascent are expressed as
$$
\theta^ {k + 1} = \theta^ {k} + \sqrt {\frac {2 \delta}{\mathbf {g} ^ {\top} \mathbf {H} ^ {- 1} \mathbf {g}}} \mathbf {H} ^ {- 1} \mathbf {g}. \tag {9}
$$
Federated NPG: FedNPG is a paradigm in that  $N$  agents collaboratively train a common global policy with parameters  $\theta$  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:
- Downlink Transmission: The server broadcasts the current global policy parameters  $\theta \in \mathbb{R}^d$  to all  $N$  agents.
- Uplink Transmission: Collecting its own local data  $\mathcal{D}_i$  based on the common policy  $\pi_{\theta}$ , each agent  $i$  computes its local gradient  $\mathbf{g}_i \in \mathbb{R}^d$  and second-order matrix  $\mathbf{H}_i \in \mathbb{R}^{d \times d}$ . Then, it sends  $\mathbf{g}_i$  and  $\mathbf{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:
$$
\mathbf {H} \leftarrow \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {H} _ {i}, \quad \mathbf {g} \leftarrow \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {g} _ {i}. \tag {10}
$$
The server then updates global policy parameters as
$$
\theta \leftarrow \theta + \sqrt {\frac {2 \delta}{\mathbf {g} ^ {\top} \mathbf {H} ^ {- 1} \mathbf {g}}} \mathbf {H} ^ {- 1} \mathbf {g}. \tag {11}
$$
We use  $|\mathcal{D}_i|$  to denote the size of collected data set  $\mathcal{D}_i$ . Without loss of generality, we take  $|\mathcal{D}_1| = \dots = |\mathcal{D}_N|$  for simplicity. As proven in [49], gradient update methods are immune to whether collected data is i.i.d., or not.
Figure 1: An illustration of federated learning based on second-order methods with  $N$  agents. (a) FedNPG via standard average. In the uplink, transmitting the matrix  $\mathbf{H}_i$  brings  $\mathcal{O}(d^2)$  communication complexity. (b) FedNPG-ADMM in this paper with only  $\mathcal{O}(d)$  communication complexity.
Bottleneck of Federated NPG: As shown in Figure 1, at each iteration, the server collects  $\{\mathbf{H}_i\in \mathbb{R}^{d\times d}$ $\mathbf{g}_i\in \mathbb{R}^d\}_{i = 1}^N$  from  $N$  agents and updates global policy parameters as follows
$$
\theta \leftarrow \theta + \sqrt {\frac {2 N \delta}{\left(\sum_ {i = 1} ^ {N} \mathbf {g} _ {i}\right) ^ {\top} \left(\sum_ {i = 1} ^ {N} \mathbf {H} _ {i}\right) ^ {- 1} \sum_ {i = 1} ^ {N} \mathbf {g} _ {i}}} \left(\sum_ {i = 1} ^ {N} \mathbf {H} _ {i}\right) ^ {- 1} \sum_ {i = 1} ^ {N} \mathbf {g} _ {i}. \tag {12}
$$
We call this method a standard average FedNPG. The uplink communication complexity from each agent is  $\mathcal{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  $\mathcal{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:
$$
\left(\sum_ {i = 1} ^ {N} \mathbf {H} _ {i}\right) ^ {- 1} \sum_ {i = 1} ^ {N} \mathbf {g} _ {i} = \underset {\mathbf {y}} {\arg \min } \frac {1}{2} \mathbf {y} ^ {\top} \left(\sum_ {i = 1} ^ {N} \mathbf {H} _ {i}\right) \mathbf {y} - \mathbf {y} ^ {\top} \sum_ {i = 1} ^ {N} \mathbf {g} _ {i}. \tag {13}
$$
This minimization problem is equivalent to
$$
\min  _ {\mathbf {y}, \left\{\mathbf {y} _ {i} \right\} _ {i = 1} ^ {N}} \sum_ {i = 1} ^ {N} \left(\frac {1}{2} \mathbf {y} _ {i} ^ {\top} \mathbf {H} _ {i} \mathbf {y} _ {i} - \mathbf {y} _ {i} ^ {\top} \mathbf {g} _ {i} + \frac {\rho}{2} \| \mathbf {y} _ {i} - \mathbf {y} \| ^ {2}\right) \tag {14}
$$
$$
\begin{array}{l} \text {s . t .} \mathbf {y} = \mathbf {y} _ {i}, \quad i = 1, \dots , N, \end{array}
$$
where  $\rho > 0$  is a penalty constant and  $\|\cdot\|$  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  $(\sum_{i=1}^{N}\mathbf{H}_{i})^{-1}\sum_{i=1}^{N}\mathbf{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
$$
\mathcal {L} (\mathbf {y}, \left\{\mathbf {y} _ {i} \right\} _ {i = 1} ^ {N}, \left\{\lambda_ {i} \right\} _ {i = 1} ^ {N}) = \sum_ {i = 1} ^ {N} \left(\frac {1}{2} \mathbf {y} _ {i} ^ {\top} \mathbf {H} _ {i} \mathbf {y} _ {i} - \mathbf {y} _ {i} ^ {\top} \mathbf {g} _ {i} + \frac {\rho}{2} \| \mathbf {y} _ {i} - \mathbf {y} \| ^ {2} + \langle \lambda_ {i}, \mathbf {y} _ {i} - \mathbf {y} \rangle\right), \tag {15}
$$
where  $\{\lambda_i\in \mathbb{R}^d\}_{i = 1}^N$  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  $\langle S, \mathcal{A}, \mathcal{P}, \mathcal{R}, \gamma \rangle$ ; Number of timesteps  $T$ ; Penalty constant  $\rho$ ; Step size  $\eta$ ; Initial  $\theta_0 \in \mathbb{R}^d$ ,  $\mathbf{y}^0 \in \mathbb{R}^d$ ,  $\{\mathbf{y}_i^0 = \mathbf{y}^0\}_{i=1}^N$ ,  $\{\lambda_i \in \mathbb{R}^d\}_{i=1}^N$ .  
1: for  $k = 1, \dots, K$  do  
2: ▷ Server broadcast  
3: Broadcast  $\mathbf{y}^{k-1}$  and  $\theta^{k-1}$  to  $N$  agents.  
4: ▷ Agent update  
5: for each agent  $i \in \{N\}$  do in parallel  
6:  $\lambda_i \gets \lambda_i + \rho (\mathbf{y}_i^{k-1} - \mathbf{y}^{k-1})$   
7:  $\mathbf{g}_i^k \gets \frac{1}{|\mathcal{D}_i|} \sum_{\tau \in \mathcal{D}_i} \sum_{t=0}^T (\nabla_{\theta^{k-1}} \log \pi_{\theta^{k-1}}(a_t | s_t)) \widehat{A}_{\pi_{\theta^{k-1}}} (s_t, a_t)$   
8:  $\mathbf{y}_i^k \gets (\mathbf{H}_i^k + \rho \mathbf{I})^{-1} (\mathbf{g}_i^k - \lambda_i + \rho \mathbf{y}^{k-1})$   
9: Transmit  $\mathbf{y}_i^k \in \mathbb{R}^d$  and  $\mathbf{g}_i^k \in \mathbb{R}^d$  to the server.  
10: end for  
11: ▷ Server update  
12:  $\mathbf{y}^k \gets \frac{1}{N} \sum_{i=1}^N \mathbf{y}_i^k$   
13:  $\theta^k \gets \theta^{k-1} + \eta \sqrt{\frac{2N\delta}{(\sum_{i=1}^N \mathbf{g}_i^k)^{\top} \mathbf{y}^k}} \cdot \mathbf{y}^k$   
14: end for  
Output:  $\theta^K$
Agent  $i$  computes  $\mathbf{H}_i$  and  $\mathbf{g}_i$  based on locally collected data. At each step of ADMM, agent  $i$  updates  $\mathbf{y}_i$  in line 8 as follows
$$
\mathbf {y} _ {i} = \underset {\mathbf {y} _ {i}} {\arg \min } \left(\frac {1}{2} \mathbf {y} _ {i} ^ {\top} \mathbf {H} _ {i} \mathbf {y} _ {i} - \mathbf {y} _ {i} ^ {\top} \mathbf {g} _ {i} + \frac {\rho}{2} \| \mathbf {y} _ {i} - \mathbf {y} + \frac {\lambda_ {i}}{\rho} \| ^ {2}\right) \tag {16}
$$
$$
\stackrel {8:} {=} \left(\mathbf {H} _ {i} + \rho \mathbf {I}\right) ^ {- 1} (\mathbf {g} _ {i} - \lambda_ {i} + \rho \mathbf {y}),
$$
where  $\mathbf{I} \in \mathbb{R}^{d \times d}$  is the identity matrix. In practical implementation, conjugate gradient methods can be used to compute  $\mathbf{y}_i$  for efficiency (Appendix C in [38]).
In line 12, after receiving  $\mathbf{y}_i$  and  $\mathbf{g}_i$  from all  $N$  agents, the server updates the global search direction as follows
$$
\begin{array}{l} \mathbf {y} = \underset {\mathbf {y}} {\arg \min } \sum_ {i = 1} ^ {N} \left(\frac {\rho}{2} \| \mathbf {y} _ {i} - \mathbf {y} \| ^ {2} + \langle \lambda_ {i}, \mathbf {y} _ {i} - \mathbf {y} \rangle\right) \tag {17} \\ = \frac {1}{N} \sum_ {i = 1} ^ {N} \left(\mathbf {y} _ {i} + \frac {\lambda_ {i}}{\rho}\right) \stackrel {1 2:} {=} \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {y} _ {i}. \\ \end{array}
$$
Dual variables in line 6 are updated by each agent as follows
$$
\lambda_ {i} \leftarrow \lambda_ {i} + \rho (\mathbf {y} _ {i} - \mathbf {y}), \quad i = 1, \dots , N. \tag {18}
$$
Combining (17) and (18), we have  $\sum_{i=1}^{N} \lambda_i = 0$ .
In line 13 of Algo. 1, after the ADMM process, the server updates global policy parameters, where  $\eta \in (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  $\mathbf{y}_i$  and  $\mathbf{g}_i$ , with a communication complexity of  $\mathcal{O}(d)$ . In contrast, the standard average approach in (12) requires transmitting  $\mathbf{H}_i$  and  $\mathbf{g}_i$ , with a communication complexity of  $\mathcal{O}(d^2)$ . This efficient communication approach allows second-order methods scalable to large-scale systems.
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.
1. The score function is bounded as  $\| \nabla_{\theta}\log \pi_{\theta}(a\mid s)\| \leq G$  , for all  $\theta \in \mathbb{R}^d$ $s\in S$  , and  $a\in A$
2. Policy gradient is  $M$ -Lipschitz continuous. In other words,  $\forall \theta_{i}, \theta_{j} \in \mathbb{R}^{d}$ ,  $s \in S$ , and  $a \in \mathcal{A}$ , we have
$$
\left\| \nabla_ {\theta_ {i}} \log \pi_ {\theta_ {i}} (a \mid s) - \nabla_ {\theta_ {j}} \log \pi_ {\theta_ {j}} (a \mid s) \right\| \leq M \| \theta_ {i} - \theta_ {j} \|. \tag {19}
$$
3. The reward function is bounded as  $r(s, a) \in [0, R]$ , for all  $s \in S$ , and  $a \in \mathcal{A}$ .
Assumption 4.2. For all  $\theta \in \mathbb{R}^{\mathrm{d}}$ , the Fisher information matrix induced by policy  $\pi_{\theta}$  and initial state distribution  $\rho$  is positive definite as
$$
F (\theta) = \underset {(s, a) \sim \nu_ {\pi_ {\theta}}} {\mathbb {E}} \left[ \nabla_ {\theta} \log \pi_ {\theta} (a \mid s) \nabla_ {\theta} \log \pi_ {\theta} (a \mid s) ^ {\top} \right] \succcurlyeq \mu_ {F} \cdot \mathbf {I} \tag {20}
$$
for some constant  $\mu_F > 0$ . For any two symmetric matrices with the same dimension,  $A \succcurlyeq B$  denotes the eigenvalues of  $A - B$  are greater or equal to zero.
Theorem 4.3. For a target error  $\epsilon$  of stationary-point convergence, each agent samples  $\mathcal{O}(\frac{1}{(1 - \gamma)^4N\epsilon})$  trajectories, and the server obtains the update direction  $\mathbf{y}^k$  at each iteration. Choose  $\eta = \frac{\mu_F^2}{4G^2(56G^2 + L_J)}$  and
$$
K = \frac {\left(J ^ {\star} - J \left(\theta^ {1}\right)\right) \left(5 6 G ^ {2} + L _ {J}\right) ^ {2} 1 6 G ^ {2} + 2 8 G ^ {2} \mu_ {F} ^ {3}}{\left(5 6 G ^ {2} + L _ {J} - 5 6 G ^ {2} \mu_ {F}\right) \mu_ {F} ^ {2} \epsilon} = \mathcal {O} \left(\frac {1}{\left(1 - \gamma\right) ^ {2} \epsilon}\right). \tag {21}
$$
$$
\frac {1}{K} \sum_ {k = 1} ^ {K} \mathbb {E} \left[ \left\| \nabla J \left(\theta^ {k}\right) \right\| ^ {2} \right] \leq \epsilon . \tag {22}
$$
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  $\theta$  such that  $\mathbb{E}\left\| \nabla J(\theta)\right\|^2 \leq \epsilon$ , for all  $\epsilon > 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.
Table 1: Complexity comparison in each agent.
NPG [23]	FedNPG	FedNPG-ADMM
Sample complexity	O(1/(1-γ)6ε2)	O(1/(1-γ)6Nε2)	O(1/(1-γ)6Nε2)
Communication complexity	-	O(d2/(1-γ)2ε)	O(d/(1-γ)2ε)
We summarize the complexity improvement in Table 1. Recall in (5) that the estimated gradient  $\mathbf{g}$  is an average over collected trajectories. The total trajectories  $\sum_{i=1}^{N}|\mathcal{D}_i|$  are collected by  $N$  agents equally using the common global policy  $\pi_\theta$  at each iteration. Each agent  $i$  samples  $|\mathcal{D}_i| = \mathcal{O}\left(\frac{1}{(1-\gamma)^4N\epsilon}\right)$  at each iteration and enjoys a federated sampling benefit compared to a single agent with  $|\mathcal{D}_i| = \mathcal{O}\left(\frac{1}{(1-\gamma)^4\epsilon}\right)$ . During the whole training process, each agent  $i$  has  $K|\mathcal{D}_i| = \mathcal{O}\left(\frac{1}{(1-\gamma)^6N\epsilon^2}\right)$  sample complexity and  $K \cdot 2d = \mathcal{O}\left(\frac{d}{(1-\gamma)^2\epsilon}\right)$  communication complexity to achieve a stationary convergence.