FedKL: Tackling Data Heterogeneity in Federated Reinforcement Learning by Penalizing KL Divergence

A typical FL system consists of one central server and $N$ agents (clients/devices), where the data of the agents are jointly utilized to train a global model without violating their privacy. In particular, the global optimization problem can be formulated as [18]

where $f_{n}(\theta)$ represents the local objective function, e.g. the loss function, of the $n$-th agent with respect to the model parameters $\theta$. $q_{n}$ is a weighting factor often set to $q_{n}=l_{n}/L$, with $l_{n}$ denoting the number of data points at the $n$-th agent and $L=\sum_{n=1}^{N}l_{n}$. In each communication round, $K$ out of the total $N$ agents are selected and the central server broadcasts the global model to the selected agents. Then, each agent updates its local model utilizing the local data and uploads the trained model to the server for model aggregation. This process will repeat for some rounds until a preset learning objective is achieved.

The system setup of FRL is similar, i.e., one central server federates the learning of $N$ distributed agents. In the $t$-th training round, the central server broadcasts the current global policy $\pi^{t}$ to $K$ selected agents which will perform $I$ iterations of local training. In each iteration, the agent interacts with its environment to collect $T$ time-steps of data, and optimizes its local policy according to the local objective with gradient ascent methods. At the end of each round, the training results will be uploaded to the server for aggregation. An example of FRL systems for autonomous driving is illustrated in Fig. 1.

We model the local training problem of each agent as a finite Markov Decision Process (MDP). Accordingly, the FRL system consists of $N$ finite MDPs $\{(\mathcal{S},\mu_{n},\mathcal{A},P_{n},\gamma)|n\in[1,N]\}$, where $\mathcal{S}$ denotes a finite set of states, $\mu_{n}$ represents the initial state distribution of the n-th MDP, $\mathcal{A}$ is a finite set of actions and $\gamma\in(0,1)$ is the discount factor. The dynamics function $P_{n}(s^{\prime},r|s,a):\mathcal{S}\times\mathbb{R}\times\mathcal{S}\times\mathcal{A}\to\mathbb{R}$ represents the probability that the n-th MDP transits from state $s$ to $s^{\prime}$ and obtain a reward of $r\in\mathbb{R}$ after taking action $a$, which characterizes the environment dynamics of the n-th MDP [11]. We can define the transition probability $P_{n}(s^{\prime}|s,a):\mathcal{S}\times\mathcal{S}\times\mathcal{A}\to\mathbb{R}$ as $P_{n}(s^{\prime}|s,a)=\sum_{r\in\mathbb{R}}P_{n}(s^{\prime},r|s,a)$. Furthermore, based on the dynamics function, we can define the reward function $\mathcal{R}_{n}(s,a):\mathcal{S}\times\mathcal{A}\to\mathbb{R}$ as $\mathcal{R}_{n}(s,a)=\sum_{r\in\mathbb{R}}r\sum_{s^{\prime}\in\mathcal{S}}P_{n}(s^{\prime},r|s,a)$, which gives the expected reward of the n-th MDP for taking action $a$ in state $s$. As a result, the $n$-th MDP $\mathcal{M}_{n}$ can be represented by a 5-tuple $(\mathcal{S},\mu_{n},\mathcal{A},P_{n},\gamma)$ sharing the same state and action space with other agents, but with possibly different initial state distributions and environment dynamics. An agent’s behavior is controlled by a stochastic policy $\pi_{n}:\mathcal{S}\times\mathcal{A}\rightarrow[0,1]$ which outputs the probability of taking an action $a$ in a given state $s$. Specifically, each agent $n$ aims to train a local policy to maximize its expected discounted reward over the long run

where the notation $\mathbb{E}_{s_{0}\thicksim\mu_{n},a_{t}\thicksim\pi,s_{t+1}\thicksim P_{n}}$ indicates that the reward is averaged over all states and actions according to the initial state distribution, the transition probability and the policy. Accordingly, the optimization problem for FRL can be formulated as

Note that $\mu_{n}$ is added to the formulation in (2) to represent the different initial state distributions for different agents and will be discussed in details later. The above formulation covers both IID and non-IID cases. In particular, the different MDPs, i.e., different initial state distributions and environment dynamics, represent the heterogeneous environments experienced by different agents (non-IID). All MDPs will be identical for the IID case [35].

In this section, we will introduce the policy gradient (PG) method [15], which is taken by the agents to update their local policies in this paper. While action-value methods learn the values of actions and then select actions based on their estimated values, PG methods directly learn the decision-making policy.

To update the policy by gradient ascent methods, we need to estimate the policy gradient. For that purpose, we first define the state-value function and the action-value function. Specifically, the state-value function $V_{\pi,P_{n}}(s)$ of a state $s$ is the expected return when the agent starts from state $s$ and follows policy $\pi$ thereafter:

where the expectation is performed over actions sampled from policy $\pi$ and the next state is sampled based on the transition probability $P_{n}$. Similarly, the action-value function $Q_{\pi,P_{n}}(s,a)$ of a state-action pair is the expected return when the agent starts from state $s$, takes the action $a$, and follows policy $\pi$ thereafter:

By subtracting the state-value function from the action-value function, we can define the advantage function $A_{\pi,P_{n}}(s,a)=Q_{\pi,P_{n}}(s,a)-V_{\pi,P_{n}}(s),$ which measures the value of action $a$ in state $s$. The author of [15] proved that the parameters $\theta$ of a differentiable policy $\pi(\theta)$ can be updated using the policy gradient

where $\nabla_{\theta}$ denotes the derivative with respect to $\theta$.

However, the high variance of the gradient estimation often leads to large policy updates and causes instability. As a result, the step size must be limited. For that purpose, [36, 37] investigated the performance gap for taking two different polices $\pi$ and $\pi^{\prime}$,

where

denotes the unnormalized discounted visitation frequency.

Equation (7) implies that the reward can be increased by any policy update that improves the expected advantage. However, it is difficult to optimize the right-hand side (RHS) of (7) since $\rho_{\pi^{\prime},\mu_{n},P_{n}}$ is dependent on $\pi^{\prime}$. Therefore, the following surrogate objective was proposed

where

denotes the policy advantage that measures the advantage of policy $\pi^{\prime}$ over $\pi$, starting from state $s\thicksim\mu_{n}$ and following $\pi^{\prime}$. $L_{\pi,\mu_{n},P_{n}}$ is an approximation to $\eta_{n}(\pi^{\prime})$ but easier to optimize because it replaces the dependency of $\rho_{\pi^{\prime},\mu_{n},P_{n}}$ on $\pi^{\prime}$ with that of $\rho_{\pi,\mu_{n},P_{n}}$ on $\pi$.

Although the above step involves an approximation, it is shown that a sufficiently small step that improves $L_{\pi,\mu_{n},P_{n}}(\pi^{\prime})$ will also improve $\eta_{n}(\pi^{\prime})$. Specifically, the following policy improvement bound was derived for any policies $\pi$ and $\pi^{\prime}$

or equivalently

where $c=\frac{4\max_{s,a}|A_{\pi,P_{k}}(s,a)|\gamma}{(1-\gamma)^{2}}$ is a constant and $D_{KL}^{max}(\pi,\pi^{\prime})=\max_{s}D_{KL}\left(\pi(\cdot|s)\|\pi^{\prime}(\cdot|s)\right)$ denotes the maximum KL divergence between two policies $\pi$ and $\pi^{\prime}$ among all states. As a result, one can optimize Equation (2) by optimizing the surrogate objective. A tighter bound for the policy improvement was later derived [38]

where $c^{\text{CPO}}=\frac{2\max_{s}\left|\mathbb{E}_{a\thicksim\pi^{\prime}(\cdot|s)}\left[A_{\pi,P_{k}}(s,a)\right]\right|\gamma}{(1-\gamma)^{2}}$ is upper bounded by $c$ and $D_{TV}(\pi(\cdot|s)\|\pi^{\prime}(\cdot|s))=\frac{1}{2}\sum_{a}\left|\pi(a|s)-\pi^{\prime}(a|s)\right|$ denotes the TV distance between two discrete probability distributions $\pi(\cdot|s)$ and $\pi^{\prime}(\cdot|s)$. For the continuous case, it is straightforward to replace the summation by integration. In this work, (13) is useful to derive the policy improvement bound of FRL.

Given (8), the expected discounted reward for the $n$-th MDP $\mathcal{M}_{n}$, defined in (2), can be written as

It can be observed that the performance of a policy on $\mathcal{M}_{n}$ is jointly determined by its initial state distribution $\mu_{n}$ and the dynamics function $P_{n}$, both of which can be either stochastic or deterministic. As a result, agents with different initial state distributions and environment dynamics will end up with different policies, causing the heterogeneity issue. Therefore, we classify the data heterogeneity in FRL into two types:

1. 1.

   Heterogeneity due to initial state distribution;

2. 2.

   Heterogeneity due to environment dynamics.

For ease of illustration, we consider two extreme cases where only one type of heterogeneity exists. Consider two finite MDPs, $\mathcal{M}_{1}$ and $\mathcal{M}_{2}$, with the same environment dynamics, $P_{1}=P_{2}$ but different initial state distributions, $\mu_{1}\neq\mu_{2}$. This corresponds to the situation where two self-driving vehicles have the same car model and training routine, but tend to enter the same road from different entrances. Now, consider the other extreme case, where $\mu_{1}=\mu_{2}$ but $P_{1}\neq P_{2}$, i.e., the two MDPs have the same initial state distribution but different environment dynamics. This can be utilized to model the situation where two different cars always enter the same road at the same entrance. The two types of heterogeneity can occur simultaneously. But, the heterogeneity due to different initial state distributions is of more concern for episodic tasks whose on-policy distribution is different from the stationary distribution [11] and heavily relies on the initial state distribution.

In the following, we define a measurement regarding the level (impact) of heterogeneity which will be utilized to investigate the connection between the global and local policy advantages. For ease of illustration, we define $\mathbf{D}_{\pi,\mu_{n},P_{n}}$ as a $\left|\mathcal{S}\right|\times\left|\mathcal{S}\right|$ diagonal matrix with $\rho_{\pi,\mu_{n},P_{n}}$ being the $i$-th diagonal entry, where $\left|\mathcal{S}\right|$ denotes the cardinality of $\mathcal{S}$. Similarly, we denote $\mathbf{\Pi}_{\pi}$ as a $\left|\mathcal{S}\right|\times\left|\mathcal{A}\right|$ matrix whose $(i,j)$th entry is $\pi(a_{j}|s_{i})$, and $\mathbf{A}_{\pi,P_{n}}$ as a $\left|\mathcal{S}\right|\times\left|\mathcal{A}\right|$ matrix whose $(i,j)$th entry is $A_{\pi,P_{n}}(s_{i},a_{j})$. With the above notation, we can rewrite the policy advantage in (10) as

where ${\rm Tr}(\cdot)$ denotes the trace operator. It can be observed that the advantage of taking a new policy $\pi^{\prime}$ on a specific agent $n$ is realized through the matrix product $\mathbf{D}_{\pi,\mu_{n},P_{n}}\mathbf{A}_{\pi,P_{n}}$. Thus, we propose the following definition regarding the level of heterogeneity. Without loss of generality, we assume that all the states are reachable, i.e., $\rho_{\pi,\mu_{n},P_{n}}(s)>0$ for all $s$ and $\pi$.

Note that $\mathbf{B}_{\pi,\mu_{n},P_{n}}$ measures the deviation of the policy advantage of the $n$-th agent from the average policy advantage of all agents. Given the local policy advantage directly affects the policy gradient, its deviation from the average will lead to potential gradient divergence. For the IID case, $\mathbf{B}_{\pi,\mu_{n},P_{n}}=\mathbf{0}$ and thus $\left\|\mathbf{B}_{\pi,\mu_{n},P_{n}}\right\|_{F}=0$ for arbitrary policy $\pi$. With non-IID data, $\left\|\mathbf{B}_{\pi,\mu_{n},P_{n}}\right\|_{F}$ will be positive and its magnitude indicates the level of heterogeneity of the $n$-th agent. Intuitively, if the level of heterogeneity of one agent is too high, i.e., the policy advantage calculated over its MDP is too different from that of the others, improving the global policy according to the local environment may not be beneficial for the whole system. In Corollary I, we will prove a necessary condition, with respect to $\left\|\mathbf{B}_{\pi,\mu_{n},P_{n}}\right\|_{F}$, for the global policy to be learn-able from the local policy.

In FRL, the local objective $\eta_{n}(\pi)$ serves as a proxy of the global objective $\eta(\pi)$. Specifically, each agent tries to optimize its local objective, but the global objective is the actual performance measure. Given the global policy $\pi^{t}$ in the $t$-th training round, the $n$-th agent interacts with its environment and obtains a new local policy $\pi_{n}^{t+1}$. Without access to any information of other agents, the $n$-th agent can only update the local policy by optimizing the local objective $\eta_{n}(\pi)$. However, a good local performance does not necessarily lead to a good global performance, especially with data heterogeneity. The question we want to answer is how we can optimize the local policy to best improve the global performance. For that purpose, we start from the local improvement.

It follows from (13) that

which indicates that the improvement brought by the local policy $\pi_{n}^{t+1}$ to the $n$-th agent over the current global policy $\pi^{t}$ is lower bounded by the policy advantage. The next question is whether we can bound the global improvement similarly. For that purpose, we first investigate the relation between the global and local policy advantage, and the result is given in the following theorem.

The proof is given in Appendix A. Note that the left-hand-side (LHS) of (19) represents the average policy advantages of $\pi_{n}^{t+1}$ over $\pi^{t}$ when applying $\pi_{n}^{t+1}$ on all agents and the RHS represents the penalized policy advantage at the n-th agent. Thus, Theorem I indicates that the policy advantage of $\pi_{n}^{t+1}$ on all agents is lower bounded by its local policy advantage with a penalty term. As a result, improving the penalized local policy advantage of a given agent may improve its local policy’s performance on all agents.

Based on Theorem I, we have the following corollary.

However, the average policy advantage is not a direct measure of the final performance, as it is a proxy for the improvement of the global objective. A more useful result should directly measure the improvement in terms of the global expected discounted reward. For that purpose, we have the following corollary.

The proof is given in Appendix C. The LHS of (II) represents the average of expected discounted reward of the n-th local policy $\pi_{n}^{t+1}$ over all agents. The first term on the RHS of (II) is a constant, representing the return of the current policy, and the second term represents the policy advantage of $\pi_{n}^{t+1}$ over $\pi_{n}^{t}$ in the $n$-th agent. Corollary II indicates that it is possible to improve a local policy’s average performance on all agents by improving the local policy advantage. In particular, by penalizing the local update with the TV distance, weighted by a coefficient ($\alpha$) proportional to the heterogeneity level, we may improve the performance of a local policy over all environments. In fact, it can be proved that optimizing the RHS of (II) will effectively improve the expected discounted reward on all agents. Please refer to Appendix D for more details.

The last two terms in (II) are penalty terms representing the expected TV distance and the maximum TV divergence between $\pi_{n}^{t+1}$ and $\pi_{n}^{t}$. The two penalties indicate that even if it is possible to improve the global performance by optimizing the local objective, the improvement only happens when the local policy $\pi_{n}^{t+1}$ is close to the global policy $\pi^{t}$. Given the first term on the RHS of (II) is a constant, Corollary II suggests the following objective for the local training

Next, we show the convergence of policy updates according to (22). The theoretical analysis is motivated by the convergence analysis of Heterogeneous-Agent Trust Region Policy Optimization (HATPRO) algorithm [39]. Similar to [33], we focus our discussion on the tabular case, i.e. we assume the aggregation step can linearly combine local policies to obtain the global policy. A summary of this learning procedure is presented in Algorithm 1.

The proof is given in Appendix E. With Theorem II, the convergence of Algorithm 1 is given by the following theorem.

The proof is given in Appendix F. When policies are parameterized with tabular methods, Theorem II and III are exact. For non-linear parameterization, the algorithm can be similarly implemented, but the error term introduced by the aggregation step should be considered.

To handle the first issue, we replace the maximum of TV divergence over all states $D_{TV}^{\max}(\pi^{t},\pi_{n}^{t+1})$ by its average value $\mathbb{E}_{s\thicksim\rho_{\pi^{t},\mu_{n},P_{n}}}\left[D_{TV}(\pi^{t}(\cdot|s)\lVert\pi_{n}^{t+1}(\cdot|s))\right]$. It has been shown that this heuristic approximation has similar empirical performance as the one with the maximum divergence [36]. Furthermore, we replace the TV distance in (22) by the KL divergence based on the Pinsker’s inequality [40], i.e. $D_{TV}(p\lVert q)\leq\sqrt{\frac{1}{2}D_{KL}(p\lVert q)}$, and obtain

To optimize (23), we need to calculate the policy advantage and the penalty terms. To this end, we assume all agents adopt the training algorithm proposed by [26]. Next, we show how we can approximate (23) by Monte Carlo simulation. In each training round, the agent will perform $I$ iterations of local training. Let $\pi^{t+1}_{n,i}$ denote the local policy of the $n$-th agent after $i$ iterations of local training with $i=0,...,I$, where $\pi^{t+1}_{n,0}=\pi^{t}$ and $\pi^{t+1}_{n,I}=\pi^{t+1}_{n}$. Thus, the objective function for the $i$-th iteration of the $n$-th agent to obtain $\pi^{t+1}_{n,i}$ from $\pi^{t+1}_{n,i-1}$ can be given by

where $c_{1}$ is used to approximate $(\alpha+\beta+\delta)$, the last term is the KL penalty introduced by the local training, and $c_{2}$ is used to approximate $c$ defined in (11). Here, $\omega(s,a)=\frac{\rho_{\pi^{t},\mu_{n},P_{n}}(s)}{\rho_{\pi^{t+1}_{n,i-1},\mu_{n},P_{n}}(s)}\frac{\pi_{n,i}^{t+1}(a|s)}{\pi_{n,i-1}^{t+1}(a|s)}$ is the importance sampling ratio where we estimate $\mathbb{A}_{\pi^{t},\mu_{n},P_{n}}(\pi_{n}^{t+1})$ by states and actions sampled from $\rho_{\pi^{t+1}_{n,i-1},\mu_{n},P_{n}}$ and $\pi^{t+1}_{n,i-1}$. The expected value of the first term in (24) is an unbiased estimation of the policy advantage $\mathbb{A}_{\pi^{t},\mu_{n},P_{n}}(\pi_{n}^{t+1})$. However, it is difficult and inefficient to compute $\omega(s,a)$ directly. Therefore, we ignore the difference in the discounted visitation frequency and use $\omega(s,a)=\frac{\pi_{n,i}^{t+1}(a|s)}{\pi_{n,i-1}^{t+1}(a|s)},$ which yields a slightly biased estimation of the policy advantage but with smaller variance.

In the following discussion, we will refer to the second and third terms of (24) as the global and local penalties, respectively. In particular, the second term penalizes the divergence of one agent’s policy from the global policy and the third term penalizes the policy outputs in two consecutive iterations. Given both the two penalties are related to KL divergence, we will refer to the proposed algorithm as FedKL.

In FRL, the initial state distribution and dynamics function of one agent is not known to others. Thus, the penalty coefficients $c_{1}$ and $c_{2}$ cannot be known locally. In this paper, we adopt the adaptive method in [26] to determine the coefficients. Specifically, we introduce two hyper-parameters, $d_{local}$ and $d_{global}$, as the target local and global KL divergence. Then, $c_{1}$ and $c_{2}$ will be adjusted in an adaptive manner so that the target $d_{local}$ and $d_{global}$ are achieved in each policy update. After initializing $c_{1},c_{2}$, we perform a three-phases update in each policy update as follows:

• •

  Perform policy update using E epochs of minibatch gradient ascent to optimize $h_{n}(\pi_{n,i}^{t+1};\pi^{t},\pi_{n,i-1}^{t+1})$ in (24).

• •

  Compute
  $d=\mathbb{E}_{s\thicksim\rho_{\pi^{t+1}_{n,i-1},\mu_{n},P_{n}}}\left[D_{KL}\left(\pi_{n,i-1}^{t+1}(\cdot|s)\lVert\pi_{n,i}^{t+1}(\cdot|s)\right)\right]$

  • –

    $\text{If }d<d_{local}/1.1,c_{2}\leftarrow c_{2}/2$

  • –

    $\text{If }d>d_{local}\times 1.1,c_{2}\leftarrow c_{2}\times 2$

• •

  Compute $d=\sqrt{\frac{1}{2}D_{KL}(\pi^{t}(\cdot|s)\lVert\pi_{n,i}^{t+1}(\cdot|s))}$

  • –

    $\text{If }d<d_{global}/1.1,c_{1}\leftarrow c_{1}/2$

  • –

    $\text{If }d>d_{global}\times 1.1,c_{1}\leftarrow c_{1}\times 2$

Specifically, in each training round, FedKL performs several iterations of SGD update. After each iteration, FedKL updates its estimation of $c_{1}$ and $c_{2}$ for the next iteration so that the target local and global KL divergence are achieved.

The heuristics for tuning $d_{local}$ and $d_{global}$ will be discussed in the experiments section. The training procedure of FedKL is illustrated in Algorithm 2.

To the best of authors’ knowledge, there are few works focusing on the data heterogeneity issue in FRL, making it difficult to find benchmarks for comparison purpose. Therefore, we made several light-weight modifications to popular RL simulators to accommodate the FRL setting. There are two widely used simulation environments for evaluating the performance of RL controllers, namely the MuJoCo simulator [41, 42] and the Flow simulator [43]. In the following, we introduce the two simulators and the settings for FRL experiments.

In this section, we will investigate the type and level of heterogeneity and how they affect the performance of FRL. For ease of illustration, we will use Experiment 1 to analyze the heterogeneity type and level. Detailed performance comparison between different algorithms and in-depth discussion regarding the key parameters will be given in the next section based on results from Experiment 2.

To illustrate the impact of data heterogeneity, we show in Fig. 4 the changes of $G=\left\|\mathbf{D}_{\pi,\mu_{n},P_{n}}\mathbf{A}_{\pi,P_{n}}\right\|_{F}-\left\|\mathbf{D}_{\pi,\mu_{n},P_{n}}\mathbf{B}_{\pi,\mu_{n},P_{n}}\right\|_{F}$. Each term in (20) is rescaled by $\mathbf{D}_{\pi,\mu_{n},P_{n}}$ to overcome the numerical instability introduced by the inverse operation $\mathbf{D}_{\pi,\mu_{n},P_{n}}^{-1}$. Note that according to Corollary I, the local training is useful only when $\left\|\mathbf{A}_{\pi,P_{n}}\right\|_{F}-\left\|\mathbf{B}_{\pi,\mu_{n},P_{n}}\right\|_{F}>0$. Thus, we may regard $G$ as the potential of one agent to contribute to the global policy. In particular, if we treat $\left\|\mathbf{D}_{\pi,\mu_{n},P_{n}}\mathbf{A}_{\pi,P_{n}}\right\|_{F}$ as the local training potential, then $G$ is the local potential penalized by the heterogeneity level, i.e., even if there is potential improvement from local training, it may not be helpful for the global policy when the heterogeneity level is high. We consider four different scenarios, including the IID case, the case with different initial state distributions ($Q=60$), and two cases with different level of environment dynamics, where the heterogeneity level is controlled by the variance of the action noises. It can be observed that $G$ decreases with training because training will reduce agents’ potential to further contribute to the global policy. On the other hand, comparing the two cases with different environment dynamics, we notice that higher level of heterogeneity will lead to smaller $G$, indicating that heterogeneity will limit the potential contribution of one agent. Note that the comparison between different types of heterogeneity may not be meaningful, as their local advantage function may have different ranges. Fluctuations of the curves are due to limited number of simulations and matrix estimation error.

To further illustrate the impact of the heterogeneity level, we show in Fig. 5 the training performance with high level of action noise. The red lines indicate the training rounds where $G<0$ happens in selected agents. It can be observed that, for those rounds, the global training will slow down, which agrees with Corollary I, i.e., $G<0$ indicates that the associate training is not helpful for the global policy. In Figs. 6 and 7, we show the performance comparison of FedAvg, FedProx, FMARL and FedKL with different initial state distributions and different transition probabilities, respectively. It can be observed that with the same learning step size, FedKL can converge to a stable performance, but both FedAvg and FedProx may diverge or fluctuate. We can certainly reduce the step size for FedAvg and FedProx, but this will lead to slow learning speed. See next section for more related discussions.

In the following, we show the results of Experiment 2 where both types of heterogeneity are present. The best performance of each algorithm is shown in Fig. 8, where the averaged total reward over three runs is reported. Each run uses the same seed for the SUMO simulator but different seeds for the parameter initialization of the neural network. It can be observed that FedKL converges much faster than FedAvg, FedProx and FMARL. The faster convergence of FedKL indicates a larger step size. Then, a natural question is whether we can increase the step size of FedAvg, FedProx and FMARL to speed up their training. In Fig. 9, we show the results where the $d_{local}$ for FedAvg, FedProx and FMARL are increased. As a result, all of them converge at a similar speed as FedKL, but they diverge after reaching the maximum reward. The above results provide empirical evidence that directly constraining the model output by KL divergence is more effective in handling the data heterogeneity than penalizing model divergence in the parameter space. Furthermore, applying a proper global KL penalty can prevent divergence from happening without loss of learning speed.

Why KL Penalty is better? It can be observed from Fig. 9 (after Round 120) that the proximal term of FedProx can slow down the divergence, but failed in avoiding it. This is because the proximal term $\left\|\theta_{n}^{t+1}-\theta^{t}\right\|$ constrains the learning in the model parameter space by the $\mathcal{L}2$-norm. Although the change in the model parameter space will lead to changes in the output distribution space (actions) and eventually influence the training objective, it is more effective to directly constrain the training in the distribution space [44, 45], which justifies the performance advantage of the proposed KL-divergence based penalty.

Why Global Penalty Helps? We now investigate how $d_{global}$ helps speed up convergence and stabilize training. Assume there are $I$ iterations of local training in each round and let $d_{i}$ denote the policy update in the $i$-th iteration. We consider two cases, without and with the global penalty (target step size), respectively, as illustrated in Fig. 10. We use the red and green arrows to represent the optimal optimization direction of the global and local objectives, respectively. Note that they are not on the same direction due to data heterogeneity. For the first case, we only apply the local penalty which constrains the step size for each iteration with $|d_{i}|\leq d_{local}$, i.e., the radius of the green circle in part (a) of Fig. 10. Let $d_{max}$ denote the highest tolerable step size in one training round, without causing training divergence. Thus, we need $d_{local}\leq d_{max}/I$, which is determined by considering the worst case scenario where all updates are on the same direction. However, the worst case happens with a very low chance. As a result, $d_{local}$ is pessimistically small for most training. Performance of FedAvg in Fig. 8 corresponds to the case that $d_{local}$ is small enough to avoid divergence, but causes slow convergence. Performance of FedAvg in Fig. 9 corresponds to the case that $d_{local}$ is large enough for fast training, but causes divergence.

Now consider the second case with $d_{global}$, denoted by the radius of the red circle in part (b) of Fig. 10. Under such circumstances, we only need to guarantee $\left|\sum_{i}^{I}d_{i}\right|\leq d_{global}\leq d_{max}$ and $\left|d_{i}\right|\leq d_{local}$. As a result, $d_{local}$ can be much larger than $d_{max}/I$, leading to a much larger per-iteration step size, represented by the larger radius of the green circle in part (b) of Fig. 10. Thus, adding $d_{global}$ enables a larger $d_{local}$ and their joint efforts can achieve a better trade-off between training speed (step size) and convergence.

Heuristics for Tuning Target Step Size: The global KL divergence is a good indicator of model divergence and we usually observe a large global KL divergence followed by a drastic performance degradation. This observation suggests a heuristic way for tuning the global step-size target $d_{global}$. For example, one can first run the algorithm without activating the global KL penalty (set $c_{1}$ to 0). Then, one may find the optimal value of $d_{local}$ and observe the average value of the KL divergence between $\pi^{t}$ and $\pi^{t+1}_{n}$ for certain number of rounds. It was observed that using a $d_{global}$ slightly larger than the average, together with a $d_{local}$ larger than the optimal value found in previous runs, can significantly speed up learning in the early stage of training.

There are many benefits to solving the data heterogeneity issue, including faster convergence, stable performance, etc. Faster convergence will also lead to fewer rounds of communication and mitigate the communication bottleneck issue for FRL. For example, to achieve a total reward of 350 in Experiment 2, it needs around 110 rounds of communication for both FedAvg and FedProx, but only 50 rounds for FedKL, representing a saving of more than $50\%$ of the communication workload.

Local agents in FedKL use the Actor-Critic architecture [27, 15], where the actor refers to the policy $\pi$ and the critic refers to the value function or the advantage function. One possible way to reduce communication workload is to keep the critic local since it is not needed for inference. We will investigate this idea in future works.

Another issue for large-scale FRL applications like autonomous driving is the model aggregation over many agents moving in a large area. A possible solution is the hierarchical FL scheme [46], where aggregations happen in two layers. First, local aggregations among agents within a small region are performed on a local server. Then, the aggregation results from multiple local servers are sent to a global server for global aggregation. The local and global aggregations can happen with different frequencies, which may significantly reduce the communication workload. Unfortunately, the convergence analysis for the hierarchical scheme with data heterogeneity is not available, and it will be interesting to investigate the performance of FedKL under such a hierarchical framework.