Paper: Federated Q-Learning: Linear Regret Speedup with Low Communication Cost (ICLR-2024)

Reviewer comment (Experiments):
It would be great if the authors could provide some empirical validations for their algorithm. I understand this is a theoretical work, but it is always helpful to corroborate the theoretical results with some experiments.

Author response:
**Response:** Thank you for the suggestion. We have performed some experiments and updated our draft to include the experimental results of FedQ-Hoeffding and the comparison with the single-agent Q-learning algorithm UCB-Hoeffding in Appendix F. The source code has been uploaded as a supplementary file. The results clearly exhibit the linear speedup of the learning regret and logarithmic communication cost, which corroborate our theoretical results. We will perform more experiments to thoroughly validate the performances of FedQ-Hoeffding and FedQ-Bernstein, and include the results in the next version of this paper.

Reviewer comment (Writing):
Why define $\tilde{C}$ as it is used only once (above Eq. (2)) in the main paper?

Author response:
**Response:** Thank you for the careful reading.  $\tilde{C}$ has been used frequently in the proofs in the appendix. We have followed your suggestion to only keep this notation in the statements of the theorems in the main paper.

Reviewer comment (Writing):
Case 1 at Page 5, it should be $\eqqcolon i_0$ instead of $\coloneqq i_0$.

Author response:
**Response:** Thank you for your careful reading. We have revised it accordingly in the updated version of this paper.

Reviewer comment (Theory):
Although the algorithm considers a setting that agents can collaboratively explore by changing their policies, the algorithm requires all agents to use the same fixed policy during local iterations, which seems to be quite restrictive. It would be nice if you could elaborate on the necessity of these restrictions.

Author response:
**Response:** The main reason that our algorithm adopts the same policy for all agents is two-fold: *First, such a fixed policy design facilitates tractable global information aggregation in each round*. In federated RL, the distribution of the collected samples in each episode depends on the executed policy. Different exploration policy leads to different distribution of the data samples. If we allow local clients to adopt different policies, or allow the exploration policies to update locally in each episode, the data samples will no longer be IID across clients or episodes in the same round, leading to significant challenges when aggregating them for an analytically tractable global estimate.

*Second, the fixed policy design naturally fits the homogeneous MDP setting and the low regret and communication cost requirements.* Given that the MDP environments encountered by the local clients are homogeneous, in order to achieve low learning regret across all clients, it is expected that the policies adopted by the clients will eventually converge to the same optimal policy. Meanwhile, in order to save the communication cost, it is desirable to keep the local policy fixed, such that the required information exchange is minimized.

We note that such a fixed exploration policy design is widely adopted in the study of federated RL and federated bandits, such as [3]-[6].

[3] Yiding Chen, Xuezhou Zhang, Kaiqing Zhang, Mengdi Wang, and Xiaojin Zhu. Byzantine-robust online and offline distributed reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pp. 3230–3269. PMLR, 2023.

[4] Flint Xiaofeng Fan, Yining Ma, Zhongxiang Dai, Wei Jing, Cheston Tan, and Bryan Kian Hsiang Low. Fault-tolerant federated reinforcement learning with theoretical guarantee. Advances in Neural Information Processing Systems, 34:1007–1021, 2021.

[5] Chengshuai Shi and Cong Shen. Federated multi-armed bandits. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 9603–9611, 2021.

[6] Ruiquan Huang, Weiqiang Wu, Jing Yang, and Cong Shen. Federated linear contextual bandits. Advances in neural information processing systems 34 (2021): 27057-27068.

Reviewer comment (Evaluation):
Although the paper claims that the event-triggered synchronization method is a key to reducing communication costs, the order of communication costs they showed in this paper seems to be larger than the one shown in [1], which uses just a fixed communication period. The communication cost shown in [1] not only logarithmically scales with T, but also is more efficient in terms of other factors (M: number of agents, $(1-\gamma)^{-1} (\approx H)$: length of horizon). I understand that direct comparison might be difficult given the settings are different, but I’m still not convinced that the communication cost shown in this paper is especially low. It would be nice if you could provide more detailed comparisons with recent literature to help better understand on the communication efficiency.

Author response:
**Response:** As pointed out by the reviewer, Woo et al (2023) considers a very different setting (infinite-horizon MDP with time-invariant transition kernels) with different learning objective (near-optimal Q-estimate at the central server). It also assumes the accessibility to a simulator for each agent in the synchronous setting, which generates a new sample
for *every state-action pair* independently at every iteration. Thus, it does not require any exploration policy, which is in stark contrast to our problem and cannot be compared fairly. 

For the asynchronous setting, since the objective is to learn a near-optimal Q-estimate *at the central server*, the local clients adopt *fixed exploration policies* through the learning process to collect trajectories only. 
As a result, it is able to achieve a communication round upper bound in $O(\frac{ M}{\mu(1-\gamma)}\log T)$ with a fixed communication period design. Here $\mu = \mu _{min}$ with equal weight assignment, and $\mu = \mu _{avg}$ with unequal weight assignment, where $\mu _{min} = \min _{x,a,m} \mu^m(x,a)$ and $\mu _{avg} = \min _{x,a}\frac{1}{M}\sum _{m=1}^M\mu^m(x,a)$, and $\mu^m(x,a)$ is the visitation frequency for $(x,a)$ under the exploration policy at agent $m$.

We note that the total number of communication rounds in both our work and Woo et al (2023) scales linearly in $M$. As pointed out in our response to **Question 1 of reviewer tdCn**, the dependency on $M$ under our algorithm can be further reduced or even removed without affecting the regret if we relax the synchronization requirement.

In terms of the dependency $S,A,H$, we first note that $\mu _{min} \leq \mu _{avg} \leq \frac{1}{SA}$. Besides, if we use the approximation $(1-\gamma)^{-1} \approx H$ as suggested by the reviewer, the bound in Woo et al (2023) becomes $O(MHSA\log T)$. Compared with our bound $O(MH^3SA\log T)$, we note that the only difference lies in the order of $H$. We emphasize that in the episodic MDP we consider, we have to handle *$H$ different transition kernels* for the $H$ steps in each episode. Roughly speaking, learning such MDPs is at least $H$ times harder than stationary infinite-horizon MDPs. Thus, it is no surprising that our dependency on $H$ has a higher order. Whether we can reduce the dependency on $H$ from $H^3$ to $H$ by assuming a constant transition kernel in each step requires some significant modification of the algorithm design and analysis, and we leave it to our future work.

Reviewer comment (Novelty):
The previous federated Q-learning literature [1] already showed that communication cost logarithmically scaling with T is achievable without using event-triggered synchronization (with fixed communication period). Is there any reason to introduce event-triggered synchronization especially in this setting?

Author response:
**Response:** As pointed out in our response to **Weakness 2** above, Woo et al (2023) consider a very different setting (infinite-horizon MDP with time-invariant transition kernels) with different learning objective (near-optimal Q-estimate at the central server), which enables them to achieve logarithmic communication cost with fixed communication period. However, in our setting, we consider episodic MDPs with different transition kernels $P _h$ in each step $h$, and we aim to achieve low learning regret across all clients. Thus, in contrast to the fixed policy adopted by the clients throughout the learning process in Woo et al (2023), in our setting, it is necessary for the clients to adaptively change their local exploration policy and eventually converge to the optimal policy. On the other hand, in order to reduce the communication cost, it is also desirable to have the local exploration policy stays fixed for a period of time. It thus requires a sophisticated design to achieve low regret and low communication cost at the same time, **which in general cannot be achieved by a fixed communication period design**. 

Motivated by the event-triggered policy switching in the single-agent setting, which has been shown to achieve low regret with logarithmic policy switching cost, we adopt the event-triggered synchronization in our setting. As we have explained in our response to **Question 1 of reviewer tdCn**, such event-triggered synchronization essentially guarantees **exponentially increasing round lengths**, which is critical to guarantee the logarithmic communication cost. Besides, when combining this technique with the new equal weight assignment in global aggregation, we are able to bound the estimation error in each round using the novel round-wise approximation technique, and obtain low learning regret at the same time.

Reviewer comment (Experiments):
The algorithm seems to fix the Q-values and behavior policies to be the same for all agents during local iterations. However, a setting that agents can flexibly change their policy based on their local observations before the next synchronization seems more natural to me, especially in the federated setting. Would it hurt the performance if they can change their policies and Q-estimates locally during local updates? I wonder if letting agents to change their policies can introduce some diversity in their exploration, which might be an advantage in learning.

Author response:
**Response:** Please refer to our response to **Weakness 1** regarding why a fixed exploration policy is desirable in our setting.

Reviewer comment (Experiments):
The main weakness is that the experimental evaluation is missing. The paper would be strengthened if an experimental section were added.

Author response:
**A1:** Thank you for the suggestion. We have performed some experiments and updated our draft to include the experimental results of FedQ-Hoeffding and the comparison with the single-agent Q-learning algorithm UCB-Hoeffding in Appendix F. The source code has been uploaded as a supplementary file. The results clearly exhibit the linear speedup of the learning regret and logarithmic communication cost, which corroborate our theoretical results. We will perform more experiments to thoroughly validate the performances of FedQ-Hoeffding and FedQ-Bernstein, and include the results in the next version of this paper.

Reviewer comment (Presentation):
Are there any simple experimental figures for the proposed algorithms?

Author response:
We would like to let the reviewer know that we have successfully implemented FedQ-Beinstein and UCB-B. The new experimental results have been included in the latest version of this paper. We are striving to implement other algorithms from related works in order to have some comparison between those works and our algorithms. We will update the draft when such results become available.

Reviewer comment (Theory):
The paper slightly over-claims its results. The regret speedup is nearly linear, but not exactly linear, due to the overhead terms that are linear in $M$, the number of machines, that appears in both Thm 4.1 and 5.1. Immediately after Thm 4.1, the paper also states that the algorithm enjoys a linear speedup in terms of $M$ in the general FL setting, despite the presence of the overhead terms.

Author response:
**Response**:  Thank you for your comment. In Theorem 4.1, we claim that with high probability, the cumulative regret over all clients can be upper bounded as
    $$\mbox{Regret}(T)\leq O\left(\sqrt{H^4\iota MTSA} + M\mbox{poly}(H,S,A)\sqrt{\iota}\right),$$
where $H, S, A$ are episode length, total number of states, and total number of actions, respectively, $T$ is the time horizon, and $\iota$ is a $\log$ factor.

Correspondingly, the per-client regret under our algorithm scales in 
$$ O\left(\sqrt{\frac{H^4SAT\iota}{M}} + \mbox{poly}(H,S,A)\sqrt{\iota}\right),$$ in which **the overhead term no longer depends on $M$**.
    
Thus, when we consider the regret speedup, we ignore the overhead and just focus on the impact of $M$ on the dominating term $\sqrt{\frac{H^4SAT\iota}{M}} $. This is why we claim "linear" regret speedup when $T$ is sufficiently large.

We note that "linear speedup" is a commonly adopted terminology to describe the impact of multiple agents on various performance metrics in the federated setting, even though similar overhead terms exist and linear speedup is only observed in the dominant term. For example, Theorem 2 of [1] shows the sample complexity for FedAsynQ-EqAvg required to find an $\varepsilon$-optimal Q function scales in  $ \tilde{O}\left(\frac{1}{M\varepsilon^2} + M\right)$, and they claim it achieves a linear speedup in the sample complexity with respect to $M$, even though an overhead term $M$ exists.  Similar examples can be found in Theorem 4.1 of [2].

[1] Jiin Woo, Gauri Joshi, and Yuejie Chi. The blessing of heterogeneity in federated Q-learning: Linear speedup and beyond. In International Conference on Machine Learning, pp. 37157–37216, 2023.

[2] Kumar Kshitij Patel, Lingxiao Wang, Aadirupa Saha, and Nathan Srebro. Federated online and bandit convex optimization. In International Conference on Machine Learning, pp. 27439–27460, 2023.

Reviewer comment (Theory):
While the typical martingale-style concentration analysis in the single agent RL setting cannot be directly applied, bounding each local term's "drift" from some "averaged parameter update path" is a commonly used technique in federated learning.

Author response:
**Response**: We agree with with reviewer that bounding each local term's "drift" from some "averaged parameter update path" is a commonly used technique in federated learning (FL). However, such bounding techniques in FL rely on certain assumptions (e.g., smoothness, convexity, bounded variation in gradients, uniform data sampling), which do not exist in federated RL. In particular, due to the online exploration nature of RL, the observations are collected as actions are taken sequentially at local agents. Due to the inherent randomness in the environment, even if the same policy is taken at all local agents, it may result in very different trajectories. Thus, it is hard to obtain an easy-to-track "averaged parameter update path" in federated RL, or a tight bound on the local terms' drifts from such averaged parameter update path. 

Our technique to bound the global estimation error in $Q$ for each round is significantly different from the "drift"-based analysis in FL. In particular, as we outlined in the Proof Sketch of Theorem 4.1, the critical step in our analysis is to bound $\left|\sum _{i=1}^{t^k} \tilde{\theta} _{t^k}^iX _i\right|$ where $X _i = V _{h+1}^\star(x _{h+1,i}) - \mathbb{E}\left[V _{h+1}^\star(x _{h+1})|(x _h,a _h) = (x,a)\right].$ Since $\tilde{\theta} _{t^k}^i=\sum _{i\in\mathcal{I}^{k'}} \theta _{t^k}^i / n _h^{k'}(x,a)$ if $N _h^k(x,a)\geq i _0$ (See Sec. 3.1, Case 2), which requires non-causal information $n _h^{k'}(x,a)$ to determine the weight for the $i$-th visit in the current round, $\tilde{\theta} _{t^k}^i,i=1,2...$ is a non-martingale sequence and the analysis of $\left|\sum _{i=1}^{t^k} \tilde{\theta} _{t^k}^iX _i\right|$ thus becomes challenging. 

To overcome this challenge, in the proof for Lemma C.3, we introduce a novel technique for **round-wise approximation and concentration**. In particular, we decompose the summation as
$$\sum _{i=1}^{t^k}\tilde{\theta} _{t^k}^iX _i=\sum _{i=1}^{t^k} \theta _{t^k}^iX _i +\sum _{i=1}^{t^k}(\tilde{\theta} _{t^k}^i - \theta _{t^k}^i)X _i,$$
and bound each term separately. While the first term preserves the martingale structure and can be bounded using Azuma-Hoeffding inequality, bounding the second term is still not straightforward.  

If we try to **bound each local term's drift**, we can only show that
$$\left|\sum _{i=1}^{t^k} (\tilde{\theta} _{t^k}^i - \theta _{t^k}^i)X _i\right|\leq H\sum _{i=1}^{t^k}\left|\tilde{\theta} _{t^k}^i - \theta _{t^k}^i\right| \leq O(1)$$
under our current algorithm design. This is because $\tilde{\theta} _{t^k}^i/\theta _{t^k}^i\in [\exp\left(-O(1/H)\right),\exp\left(O(1/H)\right)]$ as shown in (d) of Lemma B.3 and $\sum _{i = 1}^{t^k}\theta _{t^k}^i = 1$ due to the definition of $\theta _t^i$. If we stick to this technique, we have to further reduce the sample size generated in one round in order to have a tighter bound on $\tilde{\theta} _{t^k}^i/\theta _{t^k}^i$ and reduce the "drift" $\left|\tilde{\theta} _{t^k}^i - \theta _{t^k}^i\right|$. This would in turn increase the communication cost to $O(\sqrt{T})$, which is much worse compared to our current communication cost $O(\log T)$.

To overcome this difficulty, instead of bounding the local drift in each time step, we choose to **group the "drift" terms based on the corresponding rounds** as
 $$\left|\sum _{i=1}^{t^k} (\tilde{\theta} _{t^k}^i - \theta _{t^k}^i)X_i\right|\leq \sum _{k'\in [k]}\left|\sum _{i\in\mathcal{I}^{k'}}(\tilde{\theta} _{t^k}^i - \theta _{t^k}^i)X _i\right|$$ 
and then bound the total deviation in each round $\left|\sum _{i\in \mathcal{I}^{k'}}(\tilde{\theta} _{t^k}^i - \theta _{t^k}^i)X _i\right|$ separately.

In particular, within round $k'$, $X _i,i\in [N _h^{k'}(x,a)+1,N _h^{k'+1}(x,a)]$ is a martingale and $\tilde{\theta} _{t^k}^i-\theta _{t^k}^i=\alpha^c(t^{k'}+1,t^k)(\tilde{\theta} _{t^{k'}}^i - \theta _{t^{k'}}^i)$ thanks to the specially crafted weight assignment during global aggregation. Note that $\tilde{\theta} _{t^{k'}}^i-\theta _{t^{k'}}^i$ only depends on the value of $t^{k'}$, and does not depend on the realization after round $k'$, which enables us to obtain **a tighter bound compared with bounding each local drifts separately** without increasing the communication cost as follows:
 $$\left|\sum _{i=1}^{t^k} (\tilde{\theta} _{t^k}^i - \theta _{t^k}^i)X _i\right|\leq \sum _{k'\in [k]}\alpha^c(t^{k'}+1,t^k)\left|\sum _{i\in \mathcal{I}^{k'}}(\tilde{\theta} _{t^{k'}}^i - \theta _{t^{k'}}^i)X _i\right|\leq \tilde{O}(\sqrt{H^2/t^k}).$$
The detailed analysis can be found between Eqn (17) and Eqn (19) in Appendix C.2. To the best of our knowledge, this technique is new and may be useful to handle the randomness of visits in other federated RL algorithms.

Reviewer comment (Novelty):
Compared with contemporary or prior works on federated RL, the paper focuses on a more "vanilla" setting, where updates are allowed to be adversarial or asynchronous, the setting considered in this paper is a fairly simplified, bare bones version of federated RL.

Author response:
**Response**: First, as pointed out in Appendix C.1, our algorithms naturally accommodate the **asynchronous setting** where each agent may generate different number of episodes during each round, which may be caused by different processing speed or latency in communications, and our theoretical analysis in the appendix applies to this general setting as well. 

We would like to emphasize that our paper is the first federated model-free RL algorithm that enjoys both (near) linear regret speedup and low communication cost. Even though it does not cover settings such as adversarial MDP, we believe that the results presented in this work fill important research gaps in the existing literature, and serve as a first step to the exploration of model-free federated RL in other more complicated settings.

Reviewer comment (Writing):
The term $\beta$ in eq (3) and (4) should be (at least) indexed to reflect the fact that it changes over time. Preferably, it should also indicate that the term changes with $(x, a, h)$. Otherwise, as of now, it appears that a constant is added to the Q-function estimate at every single round and is misleading.

Author response:
**Response**: Thank you for your suggestion. We have revised our paper accordingly to make the notation clearer.

Reviewer comment (Theory):
It is a bit surprising that the number of communication rounds is linear in $M$. Can the authors intuitively explain why this must be the case, or if this term could be removed by more involved technical analysis?

Author response:
**Response**: 
 The current linear dependency of the number of communication rounds on $M$ is due to the event-triggered synchronization design in our algorithm. Specifically, if in round $k$, the visitation to any $(x,a,h)$ tuple at any agent $m$ equals $\max\left(1,\left\lfloor\frac{N _h^k(x,a)}{MH(H+1)}\right\rfloor\right)$, the current round ends and communications happen.
Intuitively, as more agents are involved in this procedure, such an event happens more frequently, triggering more communication rounds.

Mathematically,  under this synchronization protocol, for each round $k$, we can only claim that there exists one $(x,a,h,m)$ tuple such that the equality holds: 
$$n _h^{m,k}(x,a)= \max\left(1,\left\lfloor\frac{N _h^k(x,a)}{MH(H+1)}\right\rfloor\right).$$
For the $(x,a,h,m)$ that meets this equality, after the early sequential updating stage, $n_h^{m,k}(x,a)$ is at least $\frac{N _h^k(x,a)}{MH(H+1)}$, which implies that $N _h^{k+1}(x,a)$ is at least $(1+\frac{1}{MH(H+1)}) N _h^{k}(x,a)$, i.e., increasing exponentially from a round to the next round.
 
On the other hand, by the pigeonhole principle, if in total there are $K$ communication rounds over time $T$, there must exist one $(x,a,h)$ tuple such that the above equality holds for at least $\left\lceil\frac{K}{HSA}\right\rceil$ rounds. Combining with the exponential increase in the visiting number $n_h^{m,k}(x,a)$ during those rounds, we have that
 $$MT\geq \Omega\left(\left[1+\Omega\left(\frac{1}{MH(H+1)}\right)\right]^{\left\lceil K/(HSA)\right\rceil}\right).$$
 We note that $M$ appears in the denominator of the base on the RHS but is absent from the exponent, leading to the linear dependency on $M$ in the communication cost.

 We point out that such a linear dependency on $M$ in the communication round upper bound commonly exists in federated RL algorithms, such as [1], [3].


**The linear dependency in $M$ can be improved by removing the synchronization design.** Specifically, we can let each agent $m$ continue its current exploration round until the above equation is met for an $(x,a,h)$ tuple locally. Thus, out of $K$ rounds, the equality must hold for at least $KM$ times. By the pigeonhole principle, there exists one $(x,a,h)$ such that the above equality holds for at least $\left\lceil\frac{KM}{HSA}\right\rceil$ times. Intuitively, since now the length of each round is determined locally, the total number of communication rounds does not depend on $M$. A detailed analysis can be performed to show that the communication round scales in $O(\log T)$. We note that the implicit price for the improved dependency in $M$ is that the clients who finish their exploration in the current round earlier need to wait until all clients are done before proceeding to the next round. 

 
[1] Jiin Woo, Gauri Joshi, and Yuejie Chi. The blessing of heterogeneity in federated Q-learning: Linear speedup and beyond. In International Conference on Machine Learning, pp. 37157–37216, 2023.


[3] Yiding Chen, Xuezhou Zhang, Kaiqing Zhang, Mengdi Wang, and Xiaojin Zhu. Byzantine-robust online and offline distributed reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pp. 3230–3269. PMLR, 2023.

Reviewer comment (Theory):
For theorem 4.1, why can we ignore the overhead terms that are on the order of $O(HSAM\sqrt{H^3\iota} + H^4SAM)$ in the general FL setting?

Author response:
**Response**: Please refer to our response to **Weakness 1**.