Sample Complexity of Average-Reward Q-Learning: From Single-agent to Federated Reinforcement Learning
Yuchen Jiao, Jiin Woo, Gen Li, Gauri Joshi, Yuejie Chi
TL;DR
This paper addresses learning in infinite-horizon average-reward MDPs with finite state/action spaces under the weakly communicating assumption. It introduces a simple, model-free Q-learning framework that operates in both single-agent and federated settings, using epoch-wise parameter scheduling to balance horizon bias and discounted value estimation without requiring problem-dependent constants. The key results establish a single-agent sample complexity of $\widetilde{O}\left(\frac{| S|| A|\|h^{\star}\|_{\mathsf{sp}}^3}{\varepsilon^3}\right)$ and a federated per-agent complexity of $\widetilde{O}\left(\frac{| S|| A|\|h^{\star}\|_{\mathsf{sp}}^3}{M\varepsilon^3}\right)$, with communication rounds bounded by $\widetilde{O}\left(\frac{\|h^{\star}\|_{\mathsf{sp}}}{\varepsilon}\right)$. The federated formulation yields a linear speedup in sample efficiency across $M$ agents and features two communication-scheduling schemes that keep rounds independent of $M$. Additionally, the work shows policy-learning guarantees with per-agent sample complexity $\widetilde{O}\left(\frac{SA\|h^{\star}\|_{\mathsf{sp}}^5}{M\varepsilon^5}\right)$, maintaining the same communication efficiency. This is the first federated Q-learning treatment for average-reward MDPs, offering scalable, practical efficiency for long-horizon RL problems. The results bridge theory from discounted to average-reward RL and demonstrate the viability of federated approaches in settings requiring long-run average optimization. All results are stated with explicit dependence on the span norm $\|h^{\star}\|_{\mathsf{sp}}$ and other problem dimensions, and the analysis relies on a generative-model, synchronous sampling framework.
Abstract
Average-reward reinforcement learning offers a principled framework for long-term decision-making by maximizing the mean reward per time step. Although Q-learning is a widely used model-free algorithm with established sample complexity in discounted and finite-horizon Markov decision processes (MDPs), its theoretical guarantees for average-reward settings remain limited. This work studies a simple but effective Q-learning algorithm for average-reward MDPs with finite state and action spaces under the weakly communicating assumption, covering both single-agent and federated scenarios. For the single-agent case, we show that Q-learning with carefully chosen parameters achieves sample complexity $\widetilde{O}\left(\frac{|\mathcal{S}||\mathcal{A}|\|h^{\star}\|_{\mathsf{sp}}^3}{\varepsilon^3}\right)$, where $\|h^{\star}\|_{\mathsf{sp}}$ is the span norm of the bias function, improving previous results by at least a factor of $\frac{\|h^{\star}\|_{\mathsf{sp}}^2}{\varepsilon^2}$. In the federated setting with $M$ agents, we prove that collaboration reduces the per-agent sample complexity to $\widetilde{O}\left(\frac{|\mathcal{S}||\mathcal{A}|\|h^{\star}\|_{\mathsf{sp}}^3}{M\varepsilon^3}\right)$, with only $\widetilde{O}\left(\frac{\|h^{\star}\|_{\mathsf{sp}}}{\varepsilon}\right)$ communication rounds required. These results establish the first federated Q-learning algorithm for average-reward MDPs, with provable efficiency in both sample and communication complexity.