A Computationally Efficient Algorithm for Infinite-Horizon Average-Reward Linear MDPs

Thank you very much for your thorough review and suggestions for improving the paper. Here are our responses.

Minor error in proof. Thank you very much for catching this. We will replace $m_{t + 1}$ with $m_t$ for the equality to hold.

Relation to previous work. Thank you for the suggestion. We fully agree that detailed comparison with Hong 2025 v1 will benefit the community. A short explanation is that Hong 2025v1 that uses "max-pooling" of Q functions, which had an error, happened to allow computationally efficient clipping without additional algorithmic modification. Hong 2025v2, which fixes the error using double loop structure could no longer use the computational efficient clipping. Our submission introduces an algorithmic modification that allows for computationally efficient clipping under the double loop structure. We will mention the attempt made by Hong 2025v1 in Section 2.4, and add detailed comparison in the Appendix.

Knowledge of T, H. We will make it clear to the reader that $T$ is required by adding $T$ in the list of input parameters in the algorithm box, and mentioning it in the analysis when we set $\gamma = 1 - 1 / \sqrt{T}$. We will add a remark following the Theorem statement that using doubling-trick, we can get away with the knowledge of $T$ at the cost of incurring additional factor of $\sqrt{\log T}$ in the regret. As for the knowledge of $H$, we will remind the reader after the Optimism lemma (Lemma 3.4) and after the main theorem statement that knowledge of an upper bound of span is required. As mentioned in our response to Reviewer CyKa, we will also mention that there is a recent work that relaxes this assumption for the tabular setting, and that we leave relaxing this assumption to future work.

Discussion on Bellman optimality equation. Thank you for the suggestion. We will add discussion on the relation to Communicating/Weakly Communicating MDPs and Ergodic assumption. Bellman optimality equation assumption is strictly weaker than Communicating/Weakly Communicating MDP which is strictly weaker than ergodic assumption. We will mention this fact with brief reasoning and citation.

On boundedness of reward vector. Note that the paper by Wei et al. 2021 "Learning infinite-horizon average-reward mdps with linear function approximation." discusses this in detail in Appendix A. (Based on your comment, we think you were referring to Wei et al. 2020. "Model-free reinforcement learning in infinite-horizon average-reward markov decision processes"). Their reasoning is as follows. Given any feature mapping $\phi$ with $\Vert \phi(\cdot, \cdot) \Vert_2 \leq 1$ and a reward function in the form $r(\cdot, \cdot) = \phi(\cdot, \cdot)^\top \theta$ such that $\vert r(\cdot, \cdot) \vert \leq 1$, there always exists an invertible matrix $A \in \mathbb{R}^{d \times d}$ such that $\Vert A \phi(\cdot, \cdot) \Vert_2 \leq 1$ and $\Vert A^{-1} \theta \Vert_2 \leq \sqrt{d}$, which enables us to work with $\widetilde\phi = A \phi$ and $\widetilde\theta = A^{-1} \theta$ instead that satisfy $\Vert \widetilde\phi \Vert_2 \leq 1$ and $\Vert \widetilde\theta \Vert_2 \leq \sqrt{d}$. They argue in their proof that such a transformation $A$ is the transformation that transforms the MVEE of the feature mappings into a unit ball. When citing their work, we will add which section of their paper to refer to (Appendix A), and briefly explain their argument.

Thank you for your review. Here is our response.

Simulation in toy setting. Thank you for the suggestion for running a simulation in a toy setting. We have confirmed that our algorithm runs in a tabular setting (i.e. linear MDP setting with feature vectors orthogonal) with average regret reasonably low when treating the bonus factor beta as a hyperparameter, as is usually done in UCB-type algorithms.

To our knowledge, no study to date has conducted a simulation under the linear MDP setting. Nevertheless, we agree that providing a simple simulation would be valuable to the community, and we will aim to include a simulation. One reason why simulation under linear MDP is absent in the literature is due to the complication of specifying the vector of measures $\mu$ that ensures $P(s' | s, a) = \langle \mathbf\varphi(s, a), \mathbf\mu(s') \rangle$ is a valid probability measure. However, we believe we can get around the problem by restricting the feature vector $\varphi(s, a)$ to have L1-norm equal to 1, and each of the $d$ measures in $\mathbf\mu$ are probability measures. We will aim to conduct experiments under the restricted linear MDP setting.

Thank you for your time for the review and for the suggestion for improving the paper.

Q1: main technical novelties. We will integrate the following summary in the introduction section and in Section 3.2 when introducing the method.

1. The main technical novelty lies in designing a clipped value iteration algorithm that ensures each successively generated value function deviates from its predecessor by no more than the difference between their respective clipping thresholds.
2. We show that a naive adaptation of previous work cannot control this deviation in the linear MDP setting.
3. A natural workaround would be to clip each newly generated value function with a new threshold so that it deviates from the previously generated value functions by at most the difference in the thresholds. But doing so will make the value functions more and more complex successively, running into an issue when using covering argument for uniform concentration bound.
4. To address this, we design a novel way for controlling the deviation without running into covering issues. Specifically, we "pool" latest value functions in such a way that the difference in successive "pooled" value functions is bounded by the the difference in the thresholds. The pooled value function has low complexity so that the function class that captures the pooled value function has low covering number.

Thank you for your time for the detailed review. Your feedback will help us in improving our paper. Here are our responses to your questions.

Q1: Computational issue & Simulation. You are correct to note that, in the large state space regime, the number of unique states visited over $T$ time steps can be very large, potentially with no repeats. However, our algorithm’s computational complexity does not scale with the size of the entire state space, since we do not enumerate all states. Instead, in the worst case, where states never repeat, the complexity scales with $T$. Regarding simulation, due to the theoretical nature of the linear MDP assumption, designing a concrete simulation environment is challenging and, to our knowledge, no simulation has been performed in the linear MDP setting in the literature. we agree that providing a simple simulation would be valuable to the community, and we will aim to include a simulation under a restricted linear MDP setting described in our response to Reviewer w193.

Q2: Why assume reward to be bounded. Your intuition that the assumption can be removed is correct. And in fact, it is this intuition that It is standard in the literature to assume $r \in [0, 1]$ without loss of generality. When the actual reward lines in $[0, B]$, say, then we can scale it to $\widetilde{r} \in [0, 1]$, analyze the regret bound under the new reward, and then multiply the final regret bound by $B$ to account for the scaling. To avoid carrying the constant $B$ around in our derivations, we assumed $r \in [0, 1]$. We will add a clarifying note on this point in our final version of the paper.

Q3, Q5: Assumption on the knowledge of span. Thank you for bringing this up. Indeed, requiring prior knowledge of $\text{sp}(v^\ast)$ is a limitation of our approach. Whether sample-efficient learning is possible without knowing the span in advance has been an open question for a long time, and many existing works rely on this assumption. A recent result by Boone et al. [1] shows for the first time that it can be avoided in the tabular setting. However, extending their technique to the linear MDP setting is non-trivial and will likely require a major breakthrough. We leave this to future work. We will add a remark on this assumption and cite [1] in the final version of the paper.

Q4: Covering number. Thank you for the suggestion. In the paragraph following Lemma 3.3, we will describe the general covering argument that uses epslion-net for covering the function class for getting uniform concentration bound and briefly describe $\epsilon$-covering number of a function class and how the log covering number scales the concentration bound.

Q6: Setting discounting factor. The original problem the designer is trying to solve is the average-reward setting. We are proposing to approximate the average-reward setting by a discounted setting with a particular discounting factor. The designer will work with the discounted setting with discounting factor set to the proposed value. So, it is in the hand of the designer to set the discounting factor. We will make it clear in the beginning of Section 3 that the discounting factor is set by the designer to approximate the average-reward setting.

Q7: Size of the state space and the dimension. Thank you for the question. When we say "size of the state space", we mean the number of the states in the state space. In the tabular setting, we do not assume any structure in the MDP that allows generalizing to unseen states, and the computational and statistical complexity both scale with the number of states. However, in the linear MDP setting, we assume a structure in the MDP that allows for the generalization through the feature mapping $\varphi$ that maps state-action pairs to a $d$-dimensional feature vector. Due to the feature representation, conceptually, we no longer need to learn about all the states individually. Rather, we only need to learn about $d$ "directions", which allows for computational and statistical efficiency. Specifically, the computational complexity and the regret bound scales with $d$, not the size of the state space. With this intuition, we can see that if $d$ is large, then there is more "directions" to explore, making both computational complexity and regret large. We will add a summary of this discussion when defining the linear MDP setting in Section 2.3.

[1] Achieving Tractable Minimax Optimal Regret in Average Reward MDPs. NeurIPS 2024.