Finite Sample Analysis of Average-Reward TD Learning and $Q$-Learning

Abstract: Sheng Zhang, Zhe Zhang, Siva Theja Maguluri

Reviews and author responses:

 We thank the reviewer for the valuable feedback. Please find our responses below:

**(Main Review)[Comment 1.]**  "For the first part, I think the TD($\lambda$) policy evaluation with linear function approximation is not that hard even if it's under the average-reward setting. Essentially it is centered around the positive definiteness of $ \Phi^\top D \left(I - P\right)\Phi$ and the rest is just regular routine. For the positive definiteness, it calls the lemma from the MDP book. The potential non-uniqueness of the optimal solution is usually assumed away in previous works but handled in this paper, which I believe is an advantage."

**[Response]** A main contribution of our work in the case of average-reward TD is handling the non-uniqueness of the TD fixed point. None of the prior work in the literature deal with the non-uniqueness of the projected Bellman equation. Even the classical work  [[Tsitsiklis \& Van Roy, 1999](https://web.stanford.edu/~bvr/pubs/average-td.pdf)] that establishes asymptotic convergence makes an additional assumption to guarantee uniqueness. More recent work [[Qiu et al., 2021](https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=9435807&casa_token=1Zm72H39ztIAAAAA:yaq8teAJzlcrstprDKily4bSg8xzDpl_Gd16idCZXOx6SXXScLIakmw8G4aUk-Hvvx0EvSpk7g&tag=1); [Wu et al., 2020](https://arxiv.org/pdf/2005.01350.pdf)] on finite-sample bounds also make similar assumptions.
These assumptions are in general not satisfied, which is illustrated by the fact that they are not satisfied even in the simpler tabular setting. Since the projected Bellman equation is linear, we indeed exploit the positive definiteness of $\Phi^\top D \left(I - P^{(\lambda)}\right)\Phi$, but more importantly, exploiting the linearity, we introduce an additional (easy to implement) projection step (i.e. using the projection $\Pi_{2,E}$) to handle non-uniqueness.

**(Main Review)[Question 1)]**  "Compared with [17] Chen et al. what is the new technique in this paper, except taking the non-uniqueness into account?"

**[Response]** It is well-known in MDP theory [[Puterman, 2014](https://scholar.google.com/citations?view_op=view_citation&hl=en&user=uCCMYzIAAAAJ&citation_for_view=uCCMYzIAAAAJ:u5HHmVD_uO8C)] that the average reward problems are a lot more challenging to study than their discounted-reward counterparts. This is mainly due to the lack of contraction property in the average-reward setting, and non-uniqueness is a symptom of that. All these issues make the study of average-reward $Q$-learning an especially challenging problem, due to which there is no prior work on finite-time bounds on average reward $Q$-learning. This illustrates the importance of our contribution on finite-time bounds on the average-reward $Q$-learning. Even though many problems in practice are naturally modeled as average-reward RL problems, due to the lack of understanding of average-reward RL algorithms, discounted-reward variants are used. In particular, a non-discounted performance measure is used to evaluate the trained policies in many benchmarks for reinforcement learning algorithms such as Atari games [[Mnih et al., 2013](https://arxiv.org/pdf/1312.5602.pdf)] and MuJoCo [[Todorov et al., 2012](https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=6386109&casa_token=5lMgsNSA4S8AAAAA:AuDF_FnDRBMHcyPFfiCE06RkB-mrMozzKSKKMM_54YMhMfW-YoMNJQA0-LS_IKsF2jrx1Ep1sQ&tag=1)]. Most modern Deep RL algorithms introduce a discount factor during training even when the natural objective of interest is undiscounted. This leads to a discrepancy between the evaluation and training objective.

In the case of $Q$-learning, the idea of using a simple projection step to handle non-uniqueness does not work, due to the underlying nonlinearity, so we develop a novel approach. Chen et al. [17] studies the discounted setting, where the Bellman operator has a contraction property, which is exploited in their analysis. In contrast, we don't have such a contraction property with respect to some norm. We instead have contraction with respect to the span seminorm. This naturally leads to the non-uniqueness, because constant shifted differential $Q$ functions are essentially the same. Therefore, we introduce a new Lyaponov function to measure the distance between $Q$ equivalent classes, rather than $Q$ functions, to show $Q$ equivalent classes converge. This Lyapunov function is also designed to take advantage of the span seminorm contraction property.

**(Main Review)[Question 2)]**  "Assumption 3 seems non-trivial to me. Line 272 says it's not restrictive since the aperiodic transformation can be applied. That transformation is used to guarantee aperiodicity if I understand correctly. How that can lead to Assumption 3? I feel this claim needs more reasoning and explanation. Also, what is the relation between $J$ and $|S|$ ?"

**[Response]** Assumption 3 is discussed in Theorem 8.5.2 of [[Puterman, 2014](https://scholar.google.com/citations?view_op=view_citation&hl=en&user=uCCMYzIAAAAJ&citation_for_view=uCCMYzIAAAAJ:u5HHmVD_uO8C)]. Under the unichain assumption (i.e. Assumption 2, line 241-243), if all deterministic stationary policies are aperiodic (which can be made through the aperiodicity transformation in section 8.5.4 of [[Puterman, 2014](https://scholar.google.com/citations?view_op=view_citation&hl=en&user=uCCMYzIAAAAJ&citation_for_view=uCCMYzIAAAAJ:u5HHmVD_uO8C)]), then Assumption 3 holds.  We will include more details in the camera-ready version. In particular, Theorem 8.5.2 of [[Puterman, 2014](https://scholar.google.com/citations?view_op=view_citation&hl=en&user=uCCMYzIAAAAJ&citation_for_view=uCCMYzIAAAAJ:u5HHmVD_uO8C)] shows that the contraction factor $\gamma$ is bounded by one minus the minimum probability of any two deterministic policies from any two states ending up at the same state, denoted by $\eta$. Thus $\gamma$ is dependent on the smallest entry of the $J$-step transition matrix $P_{\pi}^J$ for any deterministic policy $\pi$. In addition, the aperiodic transformation in section 8.5.4 of [[Puterman, 2014](https://scholar.google.com/citations?view_op=view_citation&hl=en&user=uCCMYzIAAAAJ&citation_for_view=uCCMYzIAAAAJ:u5HHmVD_uO8C)] ensures a positive probability of returning to the same state for any action (by essentially adding self-loops). Since the distance between any two states is at most $|S|$, the transformation ensures that starting from any state, there is a non-zero probability being any state in the system. Thus $\eta>0$ and $\gamma<1$ after at most $|S|$ steps and thus $J \leq |S|$.

**(Main Review)[Typo)]** "Some typo I found: Line 65: an additional period sign Line 553: fellowing -> following"

**[Response]** Good catch on typos. We will fix them in the camera-ready version.

 **(Main Review)[Concern 1.]** "The equivalence in Line 453 does not seem trivial to me...Can the author provide a detailed computation to show this equivalence?"
  
**[Response]** Suppose $\theta^*$ is a solution of the equation $\Phi \theta = \Pi_{D,W^*} T^{(\lambda)} \Phi \theta$. Then we know that $\theta^*$ must lie in the subspace $E$. Thus, we have $\Phi \theta^* = \Phi \Pi_{2,E} \theta^*$. By the definition of the projection operator $\Pi_{D,W^*}$, we have
$$\Pi_{D,W^*} V = \underset{\bar{V} \in \\{ \Phi \theta | \theta \in E \\} }{\operatorname{arg min}} \Vert V - \bar{V} \Vert_D =  \underset{\bar{V} \in \\{\Phi \Pi_{2,E} \theta | \theta \in R^d\\}}{\operatorname{arg min}} \Vert V - \bar{V} \Vert_D.$$

Therefore, using $\Phi \theta^* = \Phi \Pi_{2,E} \theta^*$, the equation $\Phi \theta^* = \Pi_{D,W^*} T^{(\lambda)} \Phi \theta^*$ is equivalent to
$$\Phi \Pi_{2,E} \theta^* \in \underset{\bar{V} \in \\{\Phi \Pi_{2,E} \theta | \theta \in R^d\\}}{\operatorname{arg min}} \Vert T^{(\lambda)} \Phi \Pi_{2,E} \theta^* - \bar{V} \Vert_D.$$

Thus, by the first-order optimality condition and the definition of $T^{(\lambda)}$, we have
$$\Pi_{2,E} \Phi^\top D \left[\mathcal{R}^{(\lambda)} - \frac{r(\mu)}{1-\lambda} e + P^{(\lambda)} \Phi \Pi_{2,E} \theta^* - \Phi \Pi_{2,E} \theta^* \right] = 0$$

Using $\Phi \theta^* = \Phi \Pi_{2,E} \theta^*$ again and rearranging terms, we have
$$\Pi_{2,E}\Phi^\top D \left(P^{(\lambda)} - I \right)\Phi \theta = \Pi_{2,E}\Phi^\top D \left[\frac{r(\mu)}{1-\lambda} e - \mathcal{R}^{(\lambda)} \right].$$

On the other hand, suppose $\theta^*$ is in the subspace $E$ and satisfies 

$$\Pi_{2,E}\Phi^\top D \left(P^{(\lambda)} - I \right)\Phi \theta = \Pi_{2,E}\Phi^\top D \left[\frac{r(\mu)}{1-\lambda} e - \mathcal{R}^{(\lambda)} \right].$$

Then, following the same arguments above reversely, we can show that $\theta^*$ is a solution of the equation $\Phi \theta = \Pi_{D,W^*} T^{(\lambda)} \Phi \theta.$

We will show the equivalence in the camera-ready version.

**(Main Review)[Concern 2.]** "In Line 508, the bound of (A1), why does the 4th inequality hold?"

**[Response]** The 4th inequality holds because for any $0 \leq t_1 < t_2$ satisfying $\beta_{t_1,t_2-1} \leq \frac{1}{4\eta}$, the following inequality (see lemma 2.3 of [[Chen et al., 2019]]((https://arxiv.org/pdf/1905.11425.pdf) for the proof) hold:
$$\Vert \Theta_{t_2} - \Theta_{t_1} \Vert_2 \leq 4\eta \beta_{t_1,t_2-1} \left(\Vert \Theta_{t_2} \Vert_2 + 1\right).$$

Since we have assumed that $\beta_{t-\tau(\beta_t),t-1} \leq \frac{1}{4\eta}$, then we have 
$$\Vert \Theta_t - \Theta_{t-\tau(\beta_t)} \Vert_2 \leq 4\eta \beta_{t-\tau(\beta_t),t-1} \left(\Vert \Theta_{t} \Vert_2 + 1\right),$$
which implies the 4th inequality in line 508.

We will clarify the inequality in the camera-ready version.

**(Main Review)[Concern 3.]** "In Line 510, the bound of (A2,1), why does the 4th inequality hold?"

**[Response]** The 4th inequality holds because for any $0 \leq t_1 < t_2$ satisfying $\beta_{t_1,t_2-1} \leq \frac{1}{4\eta}$, the following inequality (implied by the inequality in the previous response) hold:
$$\Vert \Theta_{t_2} - \Theta_{t_1} \Vert_2 \leq \Vert \Theta_{t_2} \Vert_2 + 1.$$

Since we have assumed that $\beta_{t-\tau(\beta_t),t-1} \leq \frac{1}{4\eta}$, then we have $\Vert\Theta_t- \Theta_{t-\tau(\beta_t)} \Vert_2 \leq \Vert\Theta_t \Vert_2 + 1$, which implies the 4th inequality in line 510.

We will clarify the inequality in the camera-ready version.

**(Main Review)[Suggestion]** "Anyway, I suggest the authors to briefly describe [17] so the readers not familiar with [17] can have some context."

**[Response]** Thank you for the suggestion. We will include a brief discussion of [17] in the camera-ready version to provide readers with more background information and to highlight our contributions.

 We thank the reviewer for the valuable feedback. Please find our responses below:

**(Main Review)[Comment]** "For the second part, it seems the proposed algorithm requires the sample starting from each state-action pair, which implies either a simulator of the environment or extra interactions with the environment are required. For either case, this makes the derived results deviate from the purpose of analyzing the sample efficiency, but only the computational efficiency (under a restrictive setting). Also, I believe Assumption 3 deserves more discussion, e.g., the intuitive reason why it holds and the magnitude of  $\gamma$ with regard to J, |S|, or other quantities of the transition kernel matrix." 

**[Response]** Indeed, a limitation of our current result on $Q$-learning is that we consider synchronous $Q$-learning which cannot be implemented using a single trajectory. However, there was no known prior result on average-reward $Q$-learning even in the synchronous setting due to the lack of the contraction property. We are the first one to obtain the finite-sample bounds, which we do by exploiting the span seminorm contraction property. Since it is well-known that the average-reward problems are more challenging than their discounted-reward counterparts, we believe that this itself is a significant contribution. We believe that building upon our analysis, future work will consider the asynchronous setting. Results in the discounted setting were obtained in a similar step-by-step procedure. Finite-sample bounds for the synchronous version were first obtained in [[Chen et al., 2020](https://arxiv.org/pdf/2002.00874.pdf)] and [[Wainwright, 2019](https://arxiv.org/pdf/1905.06265.pdf)], and these breakthrough results later led to similar bounds for asynchronous version [[Qu & Wierman, 2020](http://proceedings.mlr.press/v125/qu20a/qu20a.pdf); [Chen et al., 2021](https://arxiv.org/pdf/2102.01567.pdf)].

We now present an approach one can take to study the asynchronous version, and highlight the key technical challenge. Because of the $J$-step Bellman operator $H^J$ (introduced in section 4.2, line 256), one can use the following asynchronous sampling scheme. Given a behavior policy $\pi$, an arbitrary starting point $s^0$ and an initial value function $Q_0$, the following trajectory can be drawn:
$$(s^0,\pi(s^0)) \rightarrow (s^1, \mu_{Q_0}(s^1)) \rightarrow \ldots \rightarrow (s^{J-1}, \mu_{Q_0}(s^{J-1})) \rightarrow (s^{J}, \pi(s^{J})) \rightarrow (s^{J+1}, \mu_{Q_1}(s^{J+1})) \rightarrow \ldots \rightarrow (s^{2J-1}, \mu_{Q_1}(s^{2J-1})) \rightarrow (s^{2J}, \pi(s^{2J})) \rightarrow (s^{2J+1}, \mu_{Q_2}(s^{2J+1})) \rightarrow (s^{2J+2}, \mu_{Q_2}(s^{2J+2})) \rightarrow \ldots \rightarrow (s^{3J - 1}, \mu_{Q_2}(s^{3J-1})) \rightarrow (s^{3J}, \pi(s^{3J})) \ldots,$$
  where $\mu_{Q}$ denotes the greedy policy with respect to the value function $Q$.
  Since $\mu_{Q_0}$ is fixed during the first $J$ steps, the sample estimate $\mathcal{R}(s^0,\pi(s^0)) + \sum_{j=1}^{J-1} \mathcal{R}(s^j, \mu_{Q_0}(s^j))$, being an unbiased estimator of $H^{J}(Q_0)(s^0,\pi(s^0)),$ can be used to update $Q_1(s^0,\pi(s^0))$ in the same way as Algorithm 2. Analyzing such an algorithm is especially challenging because the underlying stochastic process is a time varying Markov Chain; the $\mu_Q$ policy is changed every $J$-step. 
This challenge can be overcome by using recent techniques that are used to study algorithms with time-varying Markov chains [[Khodadadian et al., 2021](https://arxiv.org/pdf/2101.10506.pdf); [Wu et al., 2020](https://arxiv.org/pdf/2005.01350.pdf)]. We believe that such a highly technical analysis is beyond the scope of the current paper, and so is deferred to future work.

Assumption 3 is discussed in Theorem 8.5.2 of [[Puterman, 2014](https://scholar.google.com/citations?view_op=view_citation&hl=en&user=uCCMYzIAAAAJ&citation_for_view=uCCMYzIAAAAJ:u5HHmVD_uO8C)]. We will include more details in the camera-ready version. In particular, Theorem 8.5.2 of [[Puterman, 2014](https://scholar.google.com/citations?view_op=view_citation&hl=en&user=uCCMYzIAAAAJ&citation_for_view=uCCMYzIAAAAJ:u5HHmVD_uO8C)] shows that the contraction factor $\gamma$ is bounded by one minus the minimum probability of any two deterministic policies from any two states ending up at the same state, denoted by $\eta$. Thus $\gamma$ is dependent on the smallest entry of the $J$-step transition matrix $P_{\pi}^J$ for any deterministic policy $\pi$. In addition, the aperiodicity transformation in section 8.5.4 [[Puterman, 2014](https://scholar.google.com/citations?view_op=view_citation&hl=en&user=uCCMYzIAAAAJ&citation_for_view=uCCMYzIAAAAJ:u5HHmVD_uO8C)] ensures a positive probability of returning to the same state for any action (by essentially adding self-loops). Since the distance between any two states is at most $|S|$, the aperiodicity transformation ensures that starting from any state, there is a non-zero probability being any state in the system. Thus $\eta>0$ and $\gamma<1$ after at most $|S|$ steps.

This paper provides finite-sample convergence analysis of TD algorithms in both policy evaluation and optimization, under the average reward setting.   This paper provides finite-sample convergence analysis of TD algorithms in both policy evaluation and optimization, under the average reward setting. Such a setting is basically more challenging than the finite horizon and discount reward settings. The paper can be divided into two quite separate parts. The first part studies policy evaluation with linear function approximations and the second part studies policy optimization (Q-learning) with tabular MDPs. 

For the first part, the paper provides a solid analysis of the convergence rate/sample efficiency by showing that the underlying ODE of the solution path has a negative drift and characterizing the error from Markovian noise. The paper shows that the equilibrium of the ODE is the fixed point of the projected Bellman equation and also discussed the case when it is not unique. I did not check the complete proof but believe the analysis route is correct. The properties of the ODE of TD($\lambda$) under the average-reward setting with linear function approximations (Lemma 2) can be a solid contribution to the theoretical foundation of reinforcement learning.

For the second part, it seems the proposed algorithm requires the sample starting from each state-action pair, which implies either a simulator of the environment or extra interactions with the environment are required. For either case, this makes the derived results deviate from the purpose of analyzing the sample efficiency, but only the computational efficiency (under a restrictive setting). Also, I believe Assumption 3 deserves more discussion, e.g., the intuitive reason why it holds and the magnitude of $\gamma$ with regard to J, |S|, or other quantities of the transition kernel matrix.

Overall, the reviewer believes the results on policy evaluation with linear approximation are nontrivial contributions while the results on Q-learning are not very relevant (to be put in the same paper) and with too restrictive environments. Although both algorithms use TD updates, the parametrizations, proof techniques, and sampling schemes are completely different.

 yes

This paper provides finite-sample analysis for the average-reward TD($\lambda$) policy evaluation with linear function approximation and the average-reward tabular Q-learning algorithm.    The paper is clearly written and easy to follow. It presents two separate results. One is the TD($\lambda$) policy evaluation with linear approximation and the other is tabular Q-learning. Both are under the average-reward setting, which is more difficult compared with the discounted infinite-horizon setting.

For the first part, I think the TD($\lambda$) policy evaluation with linear function approximation is not that hard even if it's under the average-reward setting. Essentially it is centered around the positive definiteness of $\Phi^{\top} D (I - P) \Phi$ and the rest is just regular routine. For the positive definiteness, it calls the lemma from the MDP book. The potential non-uniqueness of the optimal solution is usually assumed away in previous works but handled in this paper, which I believe is an advantage.

For the second part, since I am not quite familiar with relevant papers I am not going to judge the technical originality and significance. Some questions are 1) Compared with [17] Chen et al. what is the new technique in this paper, except taking the non-uniqueness into account? 2) Assumption 3 seems non-trivial to me. Line 272 says it's not restrictive since the aperiodic transformation can be applied. That transformation is used to guarantee aperiodicity if I understand correctly. How that can lead to Assumption 3? I feel this claim needs more reasoning and explanation. Also, what is the relation between $J$ and $|S|$?

Some typo I found:
Line 65: an additional period sign
Line 553: fellowing -> following Yes.

The authors provide a finite sample analysis for average reward TD with linear function approximation. Importantly, the authors succeeded in eliminating a commonly used assumption. The authors also provide a J-step synchronous average reward Q-learning and a corresponding finite sample analysis. One central idea to the success of the paper is to work on equivalent classes directly to deal with the challenge that differential Bellman equation has multiple solutions.  Overall I like the idea to work on equivalence classes directly and it appears novel to me.
For the TD side, it is great to see that the assumption X\theta \neq ce is finally eliminated. And I believe the study of the TD fixed point in Lemma 1 can be VERY useful for many future work. I briefly checked the proofs of the TD part, it looks correct for me. But I have the following concerns:

1. The equivalence in Line 453 does not seem trivial to me. It is easy to compute \Pi_{2, E} but computing \Pi_{D, W^*} looks quite involving. Can the author provide a detailed computation to show this equivalence?

2. In Line 508, the bound of (A1), why does the 4th inequality hold? Shouldn't it contain ||\Theta_{t-1}||_2 - 1, \dots, ||\Theta_{t-\tau(\beta_t)}||_2 - 1?

3. In Line 510, the bound of (A2,1), why does the 4th inequality hold? It looks some clarification of the property of \Theta_{t - \tau(\beta_t)} is missing.

For the Q-learning part, overall the story makes sense to me and I like the idea. However, it looks checking the technical correctness in details requires a good understanding of [17] which I don't have. So I can't comment on this.
Anyway, I suggest the authors to briefly describe [17] so the readers not familiar with [17] can have some context. 
 Yes

Meta-review: This paper considers RL in average-reward setting. In particular, it provides sample complexity guarantees for (1) TD(lambda) with linear function approximation for policy evaluation; (2) Q-learning in the tabular setting for policy optimization. Reviewers believe there is a clear disconnection of the two parts of the paper (since algorithm/setting/proof techniques are all different). Furthermore, reviewers are concerned with the assumptions and technical contribution of the second part. Although reviewers agree the first part does have some interesting technical contribution, but these contribution of the first part alone may not be sufficient for NeurIPS.

Decision: accept