Federated Q-Learning: Linear Regret Speedup with Low Communication Cost

We thank the reviewer for the careful reading and thoughtful comments. We address the reviewer's questions in the following and have revised the paper accordingly. The changes are marked in blue in our revision. We hope the responses below and changes in the paper address the reviewer's concerns.

Weakness: 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.

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.

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

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.

Question 2: Case 1 at Page 5, it should be $=:i _0$ instead of $:=i _0$.

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

We thank the reviewer again for the helpful comments and suggestions for our work. If our response resolves your concerns to a satisfactory level, we kindly ask the reviewer to consider raising the rating of our work. Certainly, we are more than happy to address any further questions that you may have.

We thank the reviewer for the prompt reply. As acknowledged by the reviewer, this is a theoretical work, and we followed the reviewer's suggestion to conduct some experiments within the short rebuttal window in order to verify the linear speedup and logarithmic communication cost shown in our theoretical results. We have worked tirelessly to implement UCB-H and FedQ-Hoeffding within the short time window and obtained those experimental results, which indeed corroborate our theoretical results as we expected.

We agree with the reviewer that the experimental results are not comprehensive at this stage. This is because those works listed in Table 1 are all theoretical works, and they do not provide any experimental results or source code. It is thus impossible for us to implement those algorithms and perform the comparison given the short deadline we have. We are currently working on the implementation of FedQ-Beinstein, and we will include more comprehensive comparison with other works (if possible) in the next version of this work.

We thank the reviewer for the careful reading and thoughtful comments. We address the reviewer's questions in the following and have revised the paper accordingly. The changes are marked in blue in our revision. We hope the responses below and changes in the paper address the reviewer's concerns.

Q1: The main weakness is that the experimental evaluation is missing. The paper would be strengthened if an experimental section were added.

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.

We thank the reviewer again for the helpful comments and suggestions for our work. If our response resolves your concerns to a satisfactory level, we kindly ask the reviewer to consider raising the rating of our work. Certainly, we are more than happy to address any further questions that you may have.

Weakness 2: Communication cost comparison with Woo et al (2023).

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.

Question 1: The previous federated Q-learning literature Woo et al (2023) 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?

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.

We thank the reviewer for the careful reading and thoughtful comments. We address the reviewer's questions in the following and have revised the paper accordingly. The changes are marked in blue in our revision. We hope the responses below and changes in the paper address the reviewer's concerns. Due to character limits of openreview, our reply is splitted into three parts.

Weakness 1: 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.

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.

Question 1: 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?

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.

Question 2: Can the authors provide either some lower bound or some additional justification for the overhead terms?

Response: Thank you for the question. We explain why those overhead terms linear in $M$ appear in the regret upper bound as follows.

In our paper, $N _h^k(x,a)$ is the total visiting number of $(x,a,h)$ before starting round $k$ and $n _h^{m,k}(x,a)$ is the visiting number of $(x,a,h)$ of agent $m$ in round $k$. Our design for both the Hoeffding-type and the Bernstein-type algorithms guarantees that $n _h^{m,k}(x,a)\leq \max \left(1,\left\lfloor\frac{N _h^k(x,a)}{MH(H+1)}\right\rfloor\right).$ During the early stage of the algorithms when $N _h^k(x,a)$ is small, we have $n _h^{m,k}(x,a)\leq 1$. Therefore, the total number of visits during round $k$ satisfies that $n _h^k(x,a) = \sum _{m \in [M]} n _h^{m,k}(x,a)\leq M .$ Only when $N _h^k(x,a)$ becomes sufficiently large, the algorithms enter the second stage and the statistics start to kick in.

In other words, the overhead terms is contributed by the $O(M)$ samples collected in the first stage, i.e., the burn-in cost. Such a burn-in cost is arguably inevitable in the federated setting. Intuitively, the learning agent needs to collect at least one sample for each $(x,a,h)$ tuple for a proper initialization. In the federated setting where all local clients collect their samples in parallel, it is hard to coordinate their exploration to reduce the dependency on $M$. We note that such linear overhead terms exist in [1] above as well.

Question 3: For theorem 4.1, why can we ignore the overhead terms that are on the order of in the general FL setting?

Response: Please refer to our response to Weakness 1.

We thank the reviewer again for the helpful comments and suggestions for our work. If our response resolves your concerns to a satisfactory level, we kindly ask the reviewer to consider raising the rating of our work. Certainly, we are more than happy to address any further questions that you may have.

We thank the reviewer for the careful reading and thoughtful comments. We address the reviewer's questions in the following and have revised the paper accordingly. The changes are marked in blue in our revision. We hope the responses below and the changes in the paper address the reviewer's concerns. Due to character limits of openreview, our reply is splitted into five parts.

Weakness 1: 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 Theorem 4.1 and 5.1. Immediately after Theorem 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.

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.

I am not convinced by the comment on the "linear speedup".

• The overhead no longer depends on $M$ because an average is taken over all the machines. On a per machine basis, that means there is at least a part of the regret that does not benefit from the distributed setup whatsoever.

• Moreover, in the discussion for Theorem 4.1 in [2], an example provided here, the authors of [2] explicitly mentions that we see linear speedup only for certain problem settings, but not overall. The referenced work clearly discusses when linear speedup is possible and when it is not by deriving a condition for $d$ under which the overhead term is dominated. Similar discussions are missing in the submission.

• For [1], they claim linear speedup because in Theorem 1, the number of rounds required $T$ itself is shown to be $1/K$. In other words, they have achieved linear speedup because there are no additional terms that are not improved by having more machines, which is not the case for the regret bound in this paper. In later theorems, [1] also discusses the implications of $T_0$ and $\tilde{T}$ in the form of burn-in costs. For Theorems 2 and 3, [1] explicitly said that the burn-in costs need to be amortized over time in order to achieve some notion of linear speedup. Their burn-in costs are also slightly more appealing, as they are instance dependent and are not directly affected by $|S|$, $|A|$, and $H$.

I think the paper would benefit greatly if it could incorporate these discussions, including

1. The reasons behind the burn-in cost and a brief discussion on when is the regret bound dominated by the term that is $O(\sqrt{M})$ rather than the terms that are $O(M)$.
2. Challenges caused by the RL setting that are not found in federated convex optimization.

We thank the reviewer for carefully checking our responses and providing prompt feedback and insightful suggestions. We address the reviewer's comments in the following and have revised the paper accordingly. We hope the responses below and the changes in the paper address the reviewer's concerns. We would be more than happy to address any further concerns that you may have.

Q1: Condition for linear regret speedup in our work.

A1: We thank the reviewer for the helpful suggestion about the conditions under which the linear regret speedup holds for the dominating term. Rigorously speaking, the linear regret speedup can be guaranteed when the time steps $T$ is sufficiently large. Specifically, in Theorem 4.1, we show that with high probability, $\mbox{Regret}(T)\leq \tilde{O}(1)\times O(\sqrt{H^4MTSA} + M\mbox{poly}(H,S,A)).$ Thus, when $T\geq \Omega(M\mbox{poly}(H,S,A)),$ with high probability, FedQ-Hoeffding enjoys a linear speedup as follows $\mbox{Reget}(T)\leq \tilde{O}(\sqrt{H^4MTSA}).$

Similar arguments can be provided for FedQ-Bernstein. When $T\geq \tilde{\Omega}(M\mbox{poly}(H,S,A))$ where $\tilde{\Omega}$ hides a log factor that takes the form $\log^2 (MT\mbox{poly}(H,S,A))$), the term $\tilde{O}(\sqrt{H^3SAMT})$ in Theorem 5.1 becomes the dominating term, thus a linear regret speedup can be achieved.

We have added those conditions in the remarks after the theorems to make our statements more rigorous. We have also made according changes throughout the paper to clarify that the claimed linear speedup only holds when $T$ is sufficiently large.

Q2: Discussions on the burn-in costs.

A2: Thank you for your valuable suggestion! We have followed your advice to incorporate a brief discussion on the burn-in cost in the remark after Theorem 4.1.

We would like to clarify that the instance-dependent burn-in cost in [1] also depends on $H,S,A$ in the asynchronous setting, and thus is of the same flavor as in our work.

For asynchronous setting with equal weights (Theorem 2 in [1]), the burn in cost is $\tilde{O}\left(\frac{t _{min}^{max}\max (\frac{1}{1-\gamma}, M)}{\mu _{min}^2(1-\gamma)}\right).$ For asynchronous setting with unequal weights (Theorem 3 in [1]), it is $\tilde{O}\left(\frac{\max(t _{min}^{max},\frac{1}{1-\gamma}, M)}{\mu _{avg}(1-\gamma)}\right).$ Here $\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)$ in which $\mu^m(x,a)$ is the stationary visitation frequency for $(x,a)$ under the exploration policy at agent $m$, and constant $t _{min}^{max}\geq 1$ is the required time for reaching the stationary frequency.

We first note that $\mu _{min} \leq \mu _{avg} \leq \frac{1}{SA}$. Besides, if we use the approximation $(1-\gamma)^{-1} \approx H$ to convert the results from infinite-horizon MDPs to those under episodic MDPs, both of the bounds now explicitly depend on $M,S,A,H$.

Q3: Discussion on the challenges caused by the RL setting that are not found in federated convex optimization.

A3: Thank you for this great suggestion. We have included a brief discussion on the unique challenges in federated RL that are not found in FL in the Proof Sketch of Theorem 4.1. Specifically, we add the following paragragh:

We would like to emphasize that the techniques required to bound those non-martingale differences are fundamentally different from the commonly used techniques in federated learning (FL), which usually construct an "averaged parameter update path'' and then bound each local term's "drift'' from it. This is because such bounding techniques in FL rely on certain assumptions that do not exist in federated RL. 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. We overcome this challenge by relating $\{\tilde{\theta} _{t^k}^i\} _i$ with $\{{\theta} _{t^k}^i\} _i$. Instead of bounding the local drift $\tilde{\theta} _{t^k}^i-{\theta} _{t^k}^i$ in each time step, we choose to group the ``drift'' terms based on the corresponding rounds and then leverage the round-wise equal weight assignment adopted in our algorithm to obtain a tight bound.

We thank the reviewer again for your prompt reply and invaluable suggestions, which have helped us greatly to enhance the rigor and presentation of our work. If you have any further comments, questions or suggestions, we will be more than happy to address them. Meanwhile, if our responses have addressed your concerns satisfactorily, we kindly request that you re-evaluate our work based on the revision and raise the rating of our work.