Paper: Accelerated Policy Gradient: On the Nesterov Momentum for Reinforcement Learning (ICLR-2024)

Reviewer comment (Theory):
In Theorem 1, the authors mentioned that the softmax parameterized policy is tabular. What does it mean by tabular here? Would it mean the policy acts like a lookup table? Do the conclusions still apply if removing tabular? It is confusing in the paper.

Author response:
### **Q. The softmax tabular parameterization.**
The softmax tabular parameterization, extensively used and discussed in the RL literature (e.g., [Szepesvári et al., 2022]), serves as a method to express a policy in a well-defined tabular format (and hence we follow the convention by using the term "tabular"). This representation involves storing and manipulating the policy explicitly in a table, where each row is associated with a state, and each column is associated with each possible action.

Moreover, it is important to note that tabular parameterization is a commonly used scheme in the convergence rate analysis of policy gradient methods (e.g., [Agarwal et al., 2020; Mei et al., 2020; Xiao 2022].

Reviewer comment (Experiments):
The experimental results are not promising. The benchmark models are quite simple. The authors should present more complex benchmark models. Continuous environment should be utilized to showcase APG's superiority.

Author response:
### **Q. Empirical performance of APG in more complex benchmark environments**
To begin with, we would like to first highlight that the main focus of this paper is the theoretical convergence rate of APG, and the experimental results in small-scale MDPs in Figures 1-2 are included mainly to validate the $\tilde{O}(1/t^2)$ convergence rate. 

On the other hand, as suggested by the reviewer, we further evaluate the performance of APG beyond the true gradient setting in more diverse RL environments. Specifically, we extend APG to the stochastic setting (please see Algorithm 7 in Appendix H for the pseudo code) and test the stochastic APG in LunarLander-v2 and BipedalWalker-v3 using the estimated gradients. The detailed experimental results and the learning curves can be found in Appendix H in the revised manuscript. At a glance, the following table highlights that APG consistently outperforms HBPG and PG in terms of convergence behavior in both LunarLander-v2 and BipedalWalker-v3. These results showcase the empirical superiority of the APG method.

| Total Environment Interactions / Algorithms | LunarLander-v2 (1.5e6 steps) | LunarLander-v2 (2.5e6 steps) | LunarLander-v2 (4e6 steps) | BipedalWalker-v3 (6e6 steps) | BipedalWalker-v3 (1.2e7 steps) | BipedalWalker-v3 (2e7 steps) |
|---------------------------------------------|:----------------------------:|:----------------------------:|:--------------------------:|:----------------------------:|:------------------------------:|------------------------------|
|           **APG (mean / std dev)**          |    **59.41709 / 28.38159**   |   **177.36543 / 82.97050**   |   **235.86339 / 3.75312**  |   **115.29717 / 51.89256**   |    **254.83298 / 19.24260**    |   **249.83565 / 28.13334**   |
|            HBPG (mean / std dev)            |      27.76393 / 20.25072     |     180.82551 / 56.37636     |    151.22692 / 83.93418    |      55.92438 / 113.5947     |      216.47174 / 70.40871      |     232.24193 / 59.57783     |
|             PG (mean / std dev)             |     -144.25341 / 25.66134    |     -128.69148 / 13.92372    |    -107.66071 / 3.56015    |      1.47716 / 23.27089      |      191.22486 / 20.54212      |     249.27057 / 12.74333     |

**(The results reported above are the averaged over 3 random seeds.)**

Reviewer comment (Theory):
Definitions 1 and 2 are a bit ah-hoc for the convergence proof in this work. If they are existing, the authors should cite references. Otherwise, the authors should justify why they are needed.

Author response:
### **Q. Elucidating the need for introducing Definitions 1 and 2**
We thank the reviewer for the suggestion. Here, we highlight the motivation for introducing Definitions 1 and 2, underscoring the innovative nature of Definition 1. In comparing our $C$-near concavity with traditional concavity, we accentuate that $C$-near concavity entails less stringent conditions than classic concavity. The differentiation emerges due to the inclusion of the tolerance constant $C$ in our definition of concavity. This signifies that there exists a $C$ times larger tolerance than the original concavity, rendering $C$-near concavity a more flexible and adaptable concept. For a more detailed discussion, please refer to our paper, particularly on page 7, where we have provided additional explanations on these definitions.

Regarding Definition 2, this definition is to characterize a set of directions that are effective in enhancing policy performance under APG. Specifically, this definition is meant to characterize the update directions that maximize the increase in the probability of the optimal action. Under this definition, we show that the RL objective is $C$-nearly concave in these directions.

In summary, both Definitions 1 and 2 constitute novel contributions in our work. Beyond their effectiveness in our specific settings and proofs, Definition 1 introduces a less stringent property for the  RL objective function, paving the way for further research. Moreover, Definition 2 serves as a characterization of update directions for policy improvement under APG.

Reviewer comment (Theory):
Section 6.2 did not really present anything new on the lower bound. How to relate Theorem 3 to APG? The authors only said due to Theorem 3, there was no tighter lower bound for APG. This seems to me quite simple. They should show the contradiction if there existed a tighter lower bound for APG.

Author response:
### **Q. Elucidating the motivation for introducing Section 6.2**
Section 6.2 aims to underscore the absence of contradiction with the lower bound established in [Mei et al., 2020]. This is attributed to APG's ability to traverse a larger region of the domain than what PG can reach, elucidating why no contradiction exists.

Reviewer comment (Novelty):
The major concern about this work is that the authors did not present a detailed discussion of the work [1]. The work [1] has shown that a linear convergence rate, which is faster than $O(1/t^2)$, can be achieved by the policy gradient with an exponentially increasing step size. Moreover, the result in [1] is also based on the non-regularized MDP, which is the same setting as in this submission. The authors need to provide a detailed comparison of the theoretical results in [1] and this submission and also discuss the significance of the result in this submission, given the linear convergence rate in [1].

Author response:
## Rebuttal 1 (ID: WsjhYXWgQ0)
**We thank the reviewer for the detailed comments and helpful suggestions.**

### **Q. A detailed comparison between APG and the results in [Xiao, 2022]**
We thank the reviewer for the helpful suggestion on offering a comparison between our work and [Xiao, 2022]. Notably, one major difference lies in the algorithm used for acceleration: [Xiao, 2022] focuses on the policy mirror descent methods, which include natural policy gradient (NPG) as a special case (under direct tabular parameterization), while we focus on Nesterov accelerated gradient for RL (under softmax parameterization), which is a momentum-based approach. Accordingly, we provide a detailed comparison of these two family of approaches:

**(1) Momentum:**

**Empirical Aspect:** In practice, momentum is one of the most commonly used approaches for acceleration in policy optimization for RL, mainly due to its simplicity. Specifically, a variety of classic momentum methods (e.g., Nesterov momentum, heavy-ball method, and AdaGrad) have been already implemented in optimization solvers and deep learning frameworks (e.g., PyTorch) as a basic machinery and shown to improve the empirical convergence results than the standard policy gradient methods like A2C and PPO (e.g., see [Henderson et al., 2018]).

**Theoretical Aspect:** However, despite the popularity of these approaches in empirical RL, it has remained largely unknown whether momentum could indeed improve the theoretical global convergence rates of RL algorithms. To address this fundamental question, we establish the very first global convergence rate of Nesterov momentum in RL (termed APG in our paper) and show that APG could improve the rate of PG from $O(1/t)$ to $\tilde{O}(1/t^2)$ (in the constant step-size regime). Given the wide application of momentum, our analytical framework and the insights from our convergence result (e.g., local near-concavity) serve as an important first step towards better understanding the momentum approach in RL.

**Theoretical Advantage:** In addition, we would like to highlight that in the constant step-size regime (which is one of the most widely adopted setting in practice), our proposed APG indeed achieves a superior global convergence rate of $\tilde{O}(1/t^2)$, which is better than the $O(1/t)$ rate of vanilla PG and the $O(1/t)$ rate of NPG under constant step sizes.

**(2) Natural Gradients (or more generally, Policy Mirror Descent):**
Theoretical Aspect: Natural policy gradient (NPG), another RL acceleration technique, borrows the idea from the natural gradient, which uses the inverse of Fisher information matrix as the preconditioner of the gradient and can be viewed as achieving approximate steepest descent in the distribution space [Amari, 1998; Kakade, 2001]. Moreover, under direct policy parameterization, NPG is also known as a special case of policy mirror descent (PMD) with KL divergence as the proximal term [Shani et al., 2020; Xiao et al., 2022]. Regarding the global convergence rates, NPG has been shown to achieve:
(i) $O(1/t)$ rate under constant step sizes [Agarwal et al., 2020; Xiao, 2022] 
(ii) linear convergence either under adaptively increasing step sizes [Khodadadian et al., 2021] or non-adaptive geometrically increasing step sizes [Xiao, 2022]. 
Notably, under direct policy parameterization and geometrically increasing step sizes, similar linear convergence results can also be established for the more general PMD method [Xiao, 2022].

**Empirical Issues:**
Despite the above convergence rates, as also pointed out by [Mei et al., 2021a], it is known that NPG in general requires solving a costly optimization problem in each iteration (cf. [Kakade, 2001; Agarwal et al., 2020]) even with the help of compatible approximation theorem [Kakade, 2001] (unless the policy parametrization adopts some special forms, e.g., log-linear policies [Yuan et al., 2023]). Notably, the TRPO algorithm, a variant of NPG, is known to suffer from this additional computational overhead. Similarly, the PMD in general also requires solving a constrained optimization problem in each policy update [Xiao, 2022]. This computational complexity could be a critical factor for RL in practice.

**(We further add a new section in Appendix I for a detailed comparison of the acceleration methods in RL.)**

Reviewer comment (Theory):
The upper bound in this paper has a dependence on the factor $||\frac{1}{\mu}||\_\infty$. The recent work on policy gradient, e.g., [1] [2], has a convergence rate dependent on a tighter factor $||\frac{d\_{\rho}^{\pi\_*}}{\mu}||\_\infty$. Is it possible to sharpen such a factor in the result of this submission?

Author response:
### **Q. Sharpening the factor in the convergence rate of APG**

We agree with the reviewer that it is indeed feasible to sharpen the bound by replacing $\lVert\frac{1}{\mu}\rVert_\infty$ with $\lVert\frac{d_{\rho}^{\pi_*}}{\mu}\rVert_\infty$. Specifically, in Equation (60) of the original manuscript, we opt for simplicity by directly taking 1 as an upper bound of $d_{\rho}^{\pi_*}$, resulting in an ostensibly looser bound $\lVert\frac{1}{\mu}\rVert_\infty$, albeit with a more lucid expression. Accordingly, we have updated the paper to incorporate the tighter bound with the term $\lVert\frac{d_{\rho}^{\pi_*}}{\mu}\rVert_\infty$ in Lemma 11 on page 18, Theorem 2 on pages 7 & 35, and Theorem 7 on pages 45 & 46.

Reviewer comment (Experiments):
How do you compare the acceleration provided by momentum with faster acceleration methods for policy gradient, such as natural policy gradient and normalization, as well as regularization?

Author response:
### **Q. Comparison of the acceleration methods for policy gradient in RL (continued)**

**(3) Natural Gradients (or more generally, Policy Mirror Descent):**

**Theoretical Aspect:** Natural policy gradient (NPG), another RL acceleration technique, borrows the idea from the natural gradient, which uses the inverse of Fisher information matrix as the preconditioner of the gradient and can be viewed as achieving approximate steepest descent in the distribution space [Amari, 1998; Kakade, 2001]. Moreover, under direct policy parameterization, NPG is also known as a special case of policy mirror descent (PMD) with KL divergence as the proximal term [Shani et al., 2020; Xiao et al., 2022]. Regarding the global convergence rates, NPG has been shown to achieve:

(i) $O(1/t)$ rate under constant step sizes [Agarwal et al., 2020; Xiao, 2022];

(ii) linear convergence either under adaptively increasing step sizes [Khodadadian et al., 2021] or non-adaptive geometrically increasing step sizes [Xiao, 2022]. 
Notably, under direct policy parameterization and geometrically increasing step sizes, similar linear convergence results can also be established for the more general PMD method [Xiao, 2022].

**Empirical Issues:**
Despite the above convergence rates, as also pointed out by [Mei et al., 2021a], it is known that NPG in general requires solving a costly optimization problem in each iteration (cf. [Kakade, 2001; Agarwal et al., 2020]) even with the help of compatible approximation theorem [Kakade, 2001] (unless the policy parametrization adopts some special forms, e.g., log-linear policies [Yuan et al., 2023]). Notably, the TRPO algorithm, a variant of NPG, is known to suffer from this additional computational overhead. Similarly, the PMD in general also requires solving a constrained optimization problem in each policy update [Xiao, 2022]. This computational complexity could be a critical factor for RL in practice. 

**(4) Normalization:**

**Theoretical Aspect:** In the exact gradient setting, the normalization approach, which exploits the non-uniform smoothness by normalizing the true policy gradient by its L2 gradient norm, has also been shown to achieve acceleration in RL. Specifically, the geometry-aware normalized (GNPG) has been shown to exhibit linear convergence rate under softmax policies [Mei et al., 2021a]. This rate could be largely attributed to the effective increasing step size induced by the normalization technique. 

**Empirical Issues:** Despite the linear convergence property, it has been already shown that GNPG could suffer from divergence in the standard on-policy stochastic setting, even in simple 1-state MDPs. This phenomenon is mainly due to the committal behavior induced by the increasing effective step size. This behavior could substantially hinder the use of GNPG in practice.

**We have also included the above detailed comparison in the updated manuscript (cf. Appendix I).**

Reviewer comment (Theory):
Any idea of using generalizing the methods to stochastic settings, where the policy gradient has to be estimated from samples?

Author response:
### **Q. Generalizing the APG method to stochastic settings, where the policy gradient has to be estimated from samples?**

**Algorithmic Aspect:** In the stochastic setting, we can extend APG (called Stochastic APG, or SAPG) by following a procedure similar to the classic stochastic policy gradient method by sampling a batch of data at each iteration (cf. Algorithm 7 in Appendix H in the updated manuscript). Subsequently, we estimate the stochastic gradients based on this batch and proceed to update our policy parameters. As for momentum, we directly update by incorporating the direction calculated as the difference between the current parameters and the previous parameters, scaled by the momentum step size. The above design resembles the typical Nesterov acceleration in the stochastic optimization literature [Ghadimi and Lan, 2016].

**Empirical Aspect:** We conduct additional experiments to evaluate the Stochastic APG on two types of environments:

(i) A 5-state, 5-action Markov Decision Process (MDP): In this part, we would like to have a glance at whether the improved convergence rate of APG over PG could be preserved in the stochastic setting. Please refer to Figure 4 in Appendix H.1 in our updated paper for the learning curves. The results show that in this example, stochastic PG has a rate of about $O(1/t)$ and SAPG exhibits a rate faster than $O(1/t)$. This shows the good potential of APG in the stochastic setting. 

(ii) More complex benchmark environments, including LunarLander-v2 and Bipedal Walker: 
In this part, we compare SAPG with the stochastic PG and the stochastic Heavy Ball Policy Gradient (HBPG). As shown in the table below, we can see that SAPG still exhibits faster convergence over PG in both LunarLander and BipedalWalker. Furthermore, SAPG also enjoys better stability than HBPG in LunarLander. Please refer to Appendix H in our updated paper for the full learning curves.

| Total Environment Interactions / Algorithms | LunarLander-v2 (1.5e6 steps) | LunarLander-v2 (2.5e6 steps) | LunarLander-v2 (4e6 steps) | BipedalWalker-v3 (6e6 steps) | BipedalWalker-v3 (1.2e7 steps) | BipedalWalker-v3 (2e7 steps) |
|---------------------------------------------|:----------------------------:|:----------------------------:|:--------------------------:|:----------------------------:|:------------------------------:|------------------------------:|
|           **APG (mean / std dev)**          |    **59.41709 / 28.38159**   |   **177.36543 / 82.97050**   |   **235.86339 / 3.75312**  |   **115.29717 / 51.89256**   |    **254.83298 / 19.24260**    |   **249.83565 / 28.13334**   |
|            HBPG (mean / std dev)            |      27.76393 / 20.25072     |     180.82551 / 56.37636     |    151.22692 / 83.93418    |      55.92438 / 113.5947     |      216.47174 / 70.40871      |     232.24193 / 59.57783     |
|             PG (mean / std dev)             |     -144.25341 / 25.66134    |     -128.69148 / 13.92372    |    -107.66071 / 3.56015    |      1.47716 / 23.27089      |      191.22486 / 20.54212      |     249.27057 / 12.74333     |

**(The results reported above are the averaged over 3 random seeds.)**

**Theoretical Aspect:**
To extend the convergence analysis of APG to the stochastic setting, we expect to build on the novel local near-concavity property and borrow ideas from the stochastic Nesterov's Accelerated Gradient  literature (e.g., [Ghadimi and Lan, 2016]). Specifically, we need to establish the following two results in the stochastic setting: (i) We need to demonstrate that SAPG will reach the nearly locally concave region within a finite number of time steps with high probability, regardless of the initialization. (ii) We also need to relax the convergence result of stochastic NAG typically for concave functions to that with only near concavity.