Global Convergence of Policy Gradient in Average Reward MDPs

Thank you for your helpful comments.

Response to Weaknesses:

1. We will rewrite the related work section with these inputs in our revised version.

2. Since the focus of the paper was to provide the first comprehensive convergence analysis of average reward policy gradient, we focused on why the current state of the art in discounted reward policy gradient could not be leveraged to obtain non trivial bounds in the average reward counterpart. Hence, for sake of clarity we excluded all $\epsilon$ dependences and focused on the role of the discount factor alone in our bounds. However, we will now include the $\epsilon$ dependence in our revised version.

3. The convergence rate with constant step sizes, is $O(\epsilon^{-1})$ in (Xiao, 2022b), and linear convergence rate of $O(\log\epsilon^{-1})$ with increasing stp sizes. In our work, we used constant step size, hence we compared our result with its counterpart result in Xiao 2022b. Moreover, we believe that the our result can be extended to attain linear convergence rate with aggressive step sizes as well. But we leave that for future work. However, we note that, while we did not explicitly study the learning problem, increasing step sizes will not work when the $Q$ function has to be estimated and hence, we did not consider the increasing step-size case. However, we have linear convergence in Theorem 1, with constant step sizes for simple MDPs.

4. Thank you for bringing this paper to our notice. This paper considers the natural policy gradient algorithm whereas we are dealing with projected policy gradient. Besides, they also work under the assumption that the average reward is smooth. Nonetheless, we shall include this paper in our related work section.

5. We only included regret because such a notion would be useful if one were to extend our result to the case where learning is needed to estimate the policy gradient. For this, we have not discussed regret significantly in the paper. We only mention it in the abstract; we will remove it from the abstract and make a small comment in the paper clarifying when the notion of regret would be useful.

Response to Questions:

1. (a) The linear rate can be achieved for Projected Mirror Descent (also Policy Gradient) but it requires aggressively increasing step sizes. This aggressive step sizes makes the algorithm very similar to policy iteration (take step sizes to $\infty$). However, this aggressive step sizes are not suitable when noise is present (or model is not known), and this is where policy gradient shines. Hence, we have limited our study to constant non-aggressive step sizes.

   (b) Moreover, the linear convergence in (Lin Xiao, 2022) requires an additional assumption on the mismatch coefficient. Under this assumption and aggressive step size, it is likely that their analysis applies to our case as well since our sub-optimality recursions are similar.

   (c) We have linear convergence in our Theorem 1 under non-aggressive step sizes for simple MDPs.

2. We will mention in the table that the definition of $C_e$ and $\lambda$ can be found in Assumption 1.

3. The proof techniques in this paper depend on the ergodicity of the policy class. Whether non-ergodic MDPs admit a smooth average reward remains an open question. We leave this for future work, anticipating that if true, it will require a fundamentally different proof approach. However, most MDPs can be converted to satisfy our assumption at an $O(\epsilon)$ loss of optimality using the following trick: suppose the original MDP has a probability transition kernel $P(s'|s,a),$ define a new MDP with transition kernel $\frac{\epsilon}{|A|} \sum_{a'}P(s|s,a')+(1-\epsilon)P(s'|s,a),$ where $A$ is the action space. If choosing an action uniformly at random in each state makes the resulting Markov chain irreducible (which is usually the case), then our assumption is automatically satisfied with an $O(\epsilon)$ loss of optimality.

Thank you for your helpful comments.

Response to Weaknesses:

1. We thank the reviewer for pointing it out, we will add it this discussion in the final version.

   (a) The work [1, 2] considers discounted reward setting, and our core contribution is in the average reward setting.

   (b) Furthermore, [1,2] achieves a linear rate for PMD with increasing step sizes. PMD with aggressive step sizes effectively reduces to policy iteration (under suitable condition), which is less suitable for scenarios with noisy gradients (though it is an important yet distinct algorithm).

   (c) We have a linear convergence for simple MDPs with constant step sizes (in average reward case, which can be trivially extended to discounted case.)

   (d) We think, PMD in average case, can have linear convergence too, using the similar techniques in [1,2] and our analysis, with aggressive step sizes. However, it is a different algorithm (with aggressive step sizes PMD becomes closer to PI rather than PG) and hence deserves its own study.

2. We have fixed typos and improved upon notations in the revised version.

3. $d^\pi(s)$ is defined and elaborated upon in the paragraph below Equation 2 in the revised version (to be uploaded soon)

4. The step-size is chosen as $\eta<\frac{1}{L^\Pi_2}$, where $L^\Pi_2$ is the restricted smoothness constant. We will include this in the main theorem statement.

5. We have added a reference for this result in the revised version.

6. We have fixed the formatting errors in the revised version.

[1] Johnson, E., Pike-Burke, C., & Rebeschini, P. (2023). Optimal convergence rate for exact policy mirror descent in discounted markov decision processes. [2] Xiao, L. (2022). On the convergence rates of policy gradient methods.

Response to Questions:

1. Since the focus of the paper was to provide the first comprehensive convergence analysis of average reward policy gradient, we focussed on why the current state of the art in discounted reward policy gradient could not be leveraged to obtain non trivial bounds in the average reward counterpart. Hence, for sake of clarity we excluded all $\epsilon$ dependences and focussed on the role of the discount factor in our bounds. However, we will now include the $\epsilon$ dependence in our revised version.

2. The existing bounds are of the form $\frac{C}{k(1-\gamma)}$, where $C \gg 1$ is very large constant compared to the worst sub-optimality of $\frac{2}{1-\gamma}$. Now, let's say for $k =10$, we have $\frac{C}{10(1-\gamma)} \gg \frac{2}{1-\gamma}$, which is the largest possible difference in the discounted rewards between two policies. And this is true for all $k \leq C$, which is a very big number. Hence, the bound $\frac{C}{(1-\gamma)k}$, yields a meaningful bound only after a large $k$ which may not be preferable. On the other hand, our bound is of the form $\frac{1}{\frac{1-\gamma}{2} + \frac{k}{C}}$, and this bound is meaningful for all $k\geq 0$. Observe that our bound aysmptotically becomes $\approx \frac{C}{k}$ which improves upon the existing result [2] by a factor $\frac{1}{1-\gamma}$. Regarding softmax policies: The work [1] improves upon [2], but also uses the recursion $a_k-a_{k-1}\geq a_k^2$ in their work, yielding $\frac{CA}{(1-\gamma)k}$. If our methodology is applied for solving the recursion, it yields the rate of $\frac{1}{\frac{1-\gamma}{2}+\frac{k}{CA}}$. This is a more meaningful bound for small values of $k$ and, asymptotically for large $k,$ provides an improvement by a factor $\frac{1}{1-\gamma}$ asymptotically. We would like to remind, that the work [1,2] is for discounted case, and our core contribution of the work is the analysis on average reward case. The resulting improvement on discounted case is a welcome bonus of our analysis.

3. This is a good question. A lower bound of $O(\frac{1}{\epsilon})$ is known for policy gradient in the discounted reward case with softmax parametrization. It may be possible to use a similar approach in the average reward case, we will certainly look into this.

[1] Liu, J., Li, W., & Wei, K. (2024). Elementary analysis of policy gradient methods. [2] Mei, J., Xiao, C., Szepesvari, C., & Schuurmans, D. (2020, November). On the global convergence rates of softmax policy gradient methods.

Thank you for your helpful comments.

Response to Key Comments and Questions:

1. Theorem 1 states that the suboptimality at iteration $k$ is of $O(\frac{1}{k})$. Hence the total regret accumulated after $T$ iterations of the algorithm is $\sum_{k=1}^T O(\frac{1}{k}) = O(\log T)$. We will include a line on this in the main result.

2. It is true that Assumption 1 is restrictive when compared to MDP classes such as weakly communicative ones. However, most MDPs can be converted to satisfy our assumption at an $O(\epsilon)$ loss of optimality using the following trick: suppose the original MDP has a probability transition kernel $P(s'|s,a),$ define a new MDP with transition kernel $\frac{\epsilon}{|A|} \sum_{a'}P(s|s,a')+(1-\epsilon)P(s'|s,a),$ where $A$ is the action space. If choosing an action uniformly at random in each state makes the resulting Markov chain irreducible (which is usually the case), then our assumption is automatically satisfied with an $O(\epsilon)$ loss of optimality. Additionally, even if we are able to relax this assumption for the planning problem, often model-free learning requires something like our assumption for algorithms such as TD learning to provide good estimates. This is due to mixing time conditions needed in the analysis of stochastic approximation algorithms (TD learning is an example of one).

3. We agree with the reviewer that $C_{PL}$ can be quite large. However, it is important to note that $C_{PL}$ is non-trivial in our analysis. In contrast, for discounted reward MDPs, the equivalent constant is $\frac{1}{(1-\gamma)\min_{s}\mu(s)}$, where $\mu(s)$ represents the initial state distribution. As $\gamma$ approaches 1, this constant diverges to infinity, thus providing a vacuous upper bound on the PL constant. In our case, we manage to keep $C_{PL}$ non-trivial, even though it is admittedly large. Moreover, to the best of our knowledge, ours is the first complete proof of the global convergence of policy gradient methods for average-cost problems. Thus, we believe that the fact that $C_{PL}$ is large does not detract from the significance of our contribution; however, we agree with the reviewer that tightening this bound is an important future direction.

4. In prior average reward literature, most of the work in gradient methods have considered the learning problem rather than planning problem where they make unproven assumptions on the nature of the average reward (such as smoothness). The planning problem without these assumptions remained unsolved but the learning problem has been studied under the assumption that policy gradient converges for the planning problem and hence, we cited those papers. In fact, one of the contributions of our work is in proving smoothness (which was an assumption in previous papers) previously assumed in many learning-based papers.

Response to Minor Comments:

1. $\pi_k$ is the policy obtained at $k$-th iteration of the projected policy gradient algorithm in discounted MDPs. We will include this description in the revised manuscript.

2. We will make these changes to the citation style in the revised manuscript.

3. It's the same version. We will remove the redundancy in the revised version.

4. We will change the notation to $\nabla(\mathcal{A})$ in the revised version.

5. We will change Eq. 8 to include $d_{\mu,\gamma}^{\pi^*}$ to maintain consistency .

6. We will remove the citation to Boyd and Vandenberghe.

7. We have updated the draft, it should be reflected in the final version. $C_p, C_m$ are defined using operator norm w.r.t. $L_\infty$ norm. Precisely, $C_m = \max_{\pi} \max_{||v||\_\infty\leq 1} ||(I- \Phi P^\pi)\^{-1}v||\_\infty$, and $C_p = \max_{\pi,\pi'\in\Pi}\max_{||v||\_\infty\leq 1}\frac{||(P^{\pi'}- P^\pi)v||\_\infty}{||\pi'-\pi||\_2}$.

Regarding Typos:

Thank you for pointing out these typos. We have fixed it in the revised version. Yes, by $L$ we do mean $L_2^\Pi$. We denote $L_2^\Pi$ to specify second derivative continuity on the restricted space $\Pi$.

Thank you for your helpful comments.

Response to Weaknesses:

1. Creating MDPs with a fixed $C_p$ is infeasible because modifying the transition kernel inevitably changes $C_m$, preventing the isolation of a single variable's effect. However, for $C_r$, this is precisely what is examined in Figure 1b. In this experiment, all MDPs share the same transition kernel and identical minimal and maximal reward values, ensuring $C_p$, $C_m$, and $\kappa_r$ remain constant. We start by freezing the transition kernel, ensuring that $C_p$ and $C_m$ remain identical across all MDPs. To keep $\kappa_r$ consistent, we fix all rewards to $\pm1$. Then, we vary $C_r$ by adjusting the proportion of actions yielding a reward of $-1$. For the maximal $C_r$, half of the actions return $+1$ and the other half return $-1$. To achieve faster convergence, we increase the proportion of actions that return $+1$. More details on this can be found in the appendix.

2. Thank you for the formatting suggestions. We will upload a revised manuscript with these changes soon.

Response to the Questions:

1. This is very interesting direction to explore, thanks for the question.

   A) Parametric classes of policies: The core idea would be to use the chain rule in our result.

   B) Infinite state-action space: This is very interesting direction, we thank the reviewer for this insightful question. The traditional bounds are $O(\frac{|S| |A| }{\epsilon})$ which yields vacuous bounds for infinite state-action spaces.

   i) Whereas (for starters), the current version of our result, can deal with infinite action space (possibly with some structure) as our bound is $O(\frac{SL^\Pi_2}{\epsilon})$ and $L^\Pi_2$ depends on the hardness coefficients such as $C_m,C_r,C_p$ which can be finite for infinite action space, hence yielding some meaningful bounds.

   ii) For infinite state space, getting rid of $|S|$ in the bound is challenging. It is obtained from the bound on diameter of the policy class $diam(\Pi)^2 \leq |S|$ (see eq. 160 of the draft). However, it is possible to bound $diam(\Pi)^2$, for policy class with some structure such as low rank policy class for infinite-state space. We thank the reviewer for this intriguing question, we will definitely add this discussion in the main text of the final version. Besides, in order to characterize suboptimality associated with a policy, we need a finite concentrability coefficient $C_{PL}$ (See Lemma 7). This constant will be infinite when the state space is also infinite. Current approach requires $C_{PL}$ to be finite. Overcoming this state space constraint might require alternate ways to bound the suboptimality.

2. The extension of this result to the discounted case is as follows. In the smoothness analysis of the discounted MDP case, $\Phi P^\pi, \Phi R^\pi,$ need to be replaced with the $\gamma P^\pi$ and $R^\pi$ respectively. Only significant change is $C_m = \max_{\pi} \max_{||v||_{\infty} \leq 1} ||(I - \gamma P^{\pi})^{-1} v || \_{\infty} \leq \frac{1}{1-\gamma}$. We will add a subsection in the final version in the appendix, outlining this proof.