SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jincheng Mei
jcmei@google.com
Bo Dai
bodai@google.com
Alekh Agarwal
alekhagarwal@google.com
Mohammad Ghavamzadeh
Amazon*
szepi@google.com
daes@ualberta.ca
We prove that, for finite-arm bandits with linear function approximation, the global convergence of policy gradient (PG) methods depends on inter-related properties between the policy update and the representation. First, we establish a few key observations that frame the study: (i) Global convergence can be achieved under linear function approximation without policy or reward realizability, both for the standard Softmax PG and natural policy gradient (NPG). (ii) Approximation error is not a key quantity for characterizing global convergence in either algorithm. (iii) The conditions on the representation that imply global convergence are different between these two algorithms. Overall, these observations call into question approximation error as an appropriate quantity for characterizing the global convergence of PG methods under linear function approximation. Second, motivated by these observations, we establish new general results: (i) NPG with linear function approximation achieves global convergence if and only if the projection of the reward onto the representable space preserves the optimal action's rank, a quantity that is not strongly related to approximation error. (ii) The global convergence of Softmax PG occurs if the representation satisfies a non-domination condition and can preserve the ranking of rewards, which goes well beyond policy or reward realizability. We provide experimental results to support these theoretical findings.
Policy gradient (PG) is a foundational concept in reinforcement learning (RL), centrally used in both policy-based and actor-critic methods [25]. Despite the non-convexity of the policy optimization objective [4], global convergence of PG methods has been recently established in the tabular case for standard configurations such as the softmax parameterization [4, 22] and stochastic on-policy sampling [20]. In practice, when an RL agent is faced with a problem with large state and/or action spaces, function approximation is needed to generalize across related states and actions. The behavior of PG methods in these settings is relatively under-explored. In this paper, we study this question for the case of linear function approximation, and establish a surprising result that
the classical Softmax PG method converges whenever there exists an adequate linear function that ranks actions in the same order as the ground-truth reward function.
*The work was done prior to joining Amazon, while the author was at Google Research. 37th Conference on Neural Information Processing Systems (NeurIPS 2023).
Understanding the behavior of PG methods under function approximation is crucial for describing the behavior of RL in practice, since one rarely faces domains small enough to explicitly enumerate over states and actions in parameterizing the policy. It is well known that, standard Softmax PG converges to stationary points if a "compatible" function approximation is used [25]; i.e., one that is able to exactly represent policy value functions. However, when exact policy values are non-realizable, "approximation error" is typically considered to be the key quantity for characterizing how well a function approximation captures relevant problem quantities, including transition dynamics, rewards and policy values. This paper shows that such an approximation error perspective is overly demanding when attempting to characterize the conditions that lead to global convergence of PG methods.
Using the concept of approximation error, global convergence results for PG methods have been recently established in an additive form,
$$
\text {s u b - o p t i m a l i t y} \leq \text {o p t i z i o n e r r o r} + \text {a p p r o x i m a t i o n e r r o r}, \tag {1}
$$
implying that if the approximation error is small, a diminishing optimization error implies a small sub-optimality gap. A representative result is the global convergence of natural policy gradient (NPG) [4, Table 2], where the optimization error will diminish as the algorithm updates. There have also been global convergence results for other PG variants under linear function approximation that follow a similar approximation error analysis [3, 8, 10, 28, 5, 1, 2]. However, an additive bound like Eq. (1) has the inherent weakness that the approximation error will never be zero if the function approximation is not able to perfectly represent the desired quantities. This prevents such a strategy from establishing global convergence in cases where the approximation error is non-zero but a PG method still reaches the best representable solution.
Therefore, in spite of this recent progress, using approximation error in PG global convergence with function approximations has left two major gaps in the literature. First, it has not been investigated whether small approximation error is necessary to achieve convergence to an optimal representable policy [4], diverting attention from feature designs that achieve useful properties beyond small approximation error. Second, it is not clear if standard Softmax PG (other than NPG) converges globally under small approximation errors. In particular, NPG contains a least squares regression step [4, Eq. (17)] that can be naturally characterized with an approximation error quantity. However, standard Softmax PG does not have such a projection step [25], and the results in [4] do not apply to this update. Whether standard Softmax PG can achieve global convergence with even linearly realizable rewards (zero approximation error) is still an open problem.
In this paper, we address the above questions and contribute the following results. First, we provide negative answers to questions on the role of approximation error in determining global convergence of PG methods:
(i) Global convergence can be achieved under linear function approximation with non-zero approximation error, for both the standard Softmax PG and natural policy gradient (NPG) updates.
(ii) Approximation error is not a key quantity for characterizing global convergence in either case.
(iii) The conditions that imply global convergence are different between these two algorithms.
Second, these results lead us to question whether approximation error is an appropriate quantity to consider the global convergence of PG methods under linear function approximation. We establish new general results that characterize the conditions for global convergence of PG methods:
(i) NPG with log-linear function approximation achieves global convergence if and only if the projection of the reward onto the representable space preserves the optimal action's rank. This result significantly extends previous results that use approximation error in the analysis [4, 3], since preserving the rank of the optimal action is not strongly related to approximation error (except in the realizable limit).
(ii) We show that the global convergence of Softmax PG follows if the representation satisfies a non-domination condition and can preserve the ranking of rewards, which goes well beyond policy or reward realizability. As a byproduct, we resolve an open question by showing that even for linearly realizable reward function, Softmax PG cannot always converge to globally optimal policies when the non-domination condition for representation is violated.
2 Settings
We study the policy optimization problem under one state with  $K$  actions. Given a reward vector  $r \in \mathbb{R}^K$ , the problem is to find a parametric policy  $\pi_{\theta}$  to maximize the expected reward,
$$
\sup  _ {\theta \in \mathbb {R} ^ {d}} \pi_ {\theta} ^ {\top} r, \tag {2}
$$
where  $\theta \in \mathbb{R}^d$  with  $d < K$  is the parameter, and  $\pi_{\theta} = \mathrm{softmax}(X\theta)$  is called a "log-linear policy" [4, 28] such that for all action  $a \in [K] \coloneqq \{1, 2, \dots, K\}$ ,
$$
\pi_ {\theta} (a) = \frac {\exp \left\{\left[ X \theta \right] (a) \right\}}{\sum_ {a ^ {\prime} \in [ K ]} \exp \left\{\left[ X \theta \right] \left(a ^ {\prime}\right) \right\}}, \tag {3}
$$
where  $X \in \mathbb{R}^{K \times d}$  is the feature matrix with full column rank  $d < K$ . There are two major difficulties with the policy optimization problem. First, Eq. (2) is a non-concave maximization w.r.t.  $\theta$ , due to the softmax transform [22, Proposition 1]. Second, the policy and reward can be unrealizable, in the sense that the parametric log-linear policy  $\pi_{\theta} = \mathrm{softmax}(X\theta)$  cannot well approximate every policy  $\pi$  in the  $K$ -dimensional probability simplex, and the score  $X\theta \in \mathbb{R}^K$  cannot well approximate the true mean reward  $r \in \mathbb{R}^K$ . Such limitations arise in the linear function approximation case because  $\pi_{\theta}$  and  $X\theta$  are restricted to low-dimensional manifolds via  $\theta \in \mathbb{R}^d$  for  $d < K$ .
To solve Eq. (2), we consider the standard Softmax PG [25] and NPG [13, 4] methods, shown in Algorithms 1 and 2. Softmax PG is an instance of gradient ascent, obtained by the chain rule,
$$
\frac {d \pi_ {\theta_ {t}} ^ {\top} r}{d \theta_ {t}} = \frac {d X \theta_ {t}}{d \theta_ {t}} \left(\frac {d \pi_ {\theta_ {t}}}{d X \theta_ {t}}\right) ^ {\top} \frac {d \pi_ {\theta_ {t}} ^ {\top} r}{d \pi_ {\theta_ {t}}} = X ^ {\top} \left(\operatorname {d i a g} \left(\pi_ {\theta_ {t}}\right) - \pi_ {\theta_ {t}} \pi_ {\theta_ {t}} ^ {\top}\right) r. \tag {4}
$$
On the other hand, NPG conducts updates using least squares regression (i.e., projection),
$$
\left(X ^ {\top} X\right) ^ {- 1} X ^ {\top} r = \underset {w \in \mathbb {R} ^ {d}} {\arg \min } \| X w - r \| _ {2} ^ {2}. \tag {5}
$$
As representative policy-based methods, in their general forms, Softmax PG and NPG lay the foundation for widely used RL methods, including REINFORCE [26], actor-critic [16, 7, 12], TRPO and PPO [23, 24]. The above Eqs. (4) and (5) are their updates applied to the one-state setting.
Algorithm 1 Softmax policy gradient (PG)
Input: Learning rate η > 0.
Output: Policies πθt = softmax(Xθt).
Initialize parameter θ1 ∈ Rd.
while t ≥ 1 do
θt+1 ← θt + η · X^T(diag(πθt) - πθt π^T) r
end while
Algorithm 2 Natural policy gradient (NPG)
Input: Learning rate η > 0.
Output: Policies πθt = softmax(Xθt).
Initialize parameter θ1 ∈ Rd.
while t ≥ 1 do
θt+1 ← θt + η · (X^T X)^{-1} X^T r.
end while
To understand the difficulty of the optimization problem in Eq. (2), it is helpful to consider previous work that has analyzed the convergence of PG methods.
In the tabular setting, where  $d = K$ ,  $X = \mathbf{Id}$ , and  $\pi_{\theta} = \mathrm{softmax}(\theta)$  with  $\theta \in \mathbb{R}^K$ , both the rewards and optimal policy can be arbitrarily well approximated. In this case it is known that NPG enjoys a  $O(1 / t)$  global convergence rate [4, Table 1], which has been recently improved to  $O(e^{-c\cdot t})$  [14, 20, 17, 27]. For the case of function approximation, such results have subsequently been extended to log-linear policies, where approximation error is used to characterize the projection step of Eq. (5) [4, 28]. In particular, NPG achieves the following sub-optimality gap for all  $t \geq 1$  [4, Table 2],
$$
\left(\pi^ {*} - \pi_ {\theta_ {t}}\right) ^ {\top} r \leq c _ {1} / \sqrt {t} + c _ {2} \cdot \epsilon_ {\text {a p p r o x}}, \quad \left(c _ {1} > 0, c _ {2} > 0\right) \tag {6}
$$
where  $c_{1}$  and  $c_{2}$  are problem specific constants,  $\pi^{*}$  is the globally optimal policy,  $\pi_{\theta_t}$  is produced by NPG, and  $\epsilon_{\mathrm{approx}}$  is the approximation error, i.e., the minimum error with which the policy values can be approximated using the features [4, Table 2]. The "optimization error" term  $c_{1} / \sqrt{t}$  in Eq. (6) has since been improved to  $O(e^{-c_3\cdot t})$  with  $c_{3} > 0$  in [28, 5]. Note that if  $\epsilon_{\mathrm{approx}} > 0$  then Eq. (6) is insufficient for establishing  $\pi_{\theta_t}^\top r\to r(a^*)\coloneqq \max_{a\in [K]}r(a)$  as  $t\to \infty$  even when such global convergence is achieved.
The understanding for the standard Softmax PG is even less clear. In the tabular case, it is known that Softmax PG achieves global convergence asymptotically, i.e.,  $\pi_{\theta_t}^\top r \to r(a^*)$  as  $t \to \infty$  [4], with an  $O(1/t)$  rate of convergence that exhibits undesirable problems and initialization-dependent constants
[21, 18]. Directly extending this global convergence result to the case of function approximation, i.e., log-linear policies, is impossible without any additional assumptions on the features, since there can be exponentially many sub-optimal local maxima in the worst case [9]. In fact, even with linearly realizable rewards (zero approximation error), whether standard Softmax PG achieves global convergence still remains unsolved [4]. One intuitive reason why this is a difficult result to establish is that standard Softmax PG uses the gradient Eq. (4) rather than projection (regression) to perform updates, which is less directly connected to the concept of approximation error.
3 The Limitations of Approximation Error in Characterizing Convergence
It is known that there exist representations  $X \in \mathbb{R}^{K \times d}$  with  $d < K$  and  $r \in \mathbb{R}^K$  that create exponentially many sub-optimal local maxima in Eq. (2) [9, Theorem 1], which makes it impossible to ensure global convergence of PG methods without imposing any structure on the function approximation. Before identifying specific conditions that ensure global convergence, we first explain how approximation error cannot be a useful structural measure for this purpose, by demonstrating that zero approximation error is not a necessary condition for global convergence, and illustrating problem instances with comparable approximation error that render starkly different convergence behaviors across different PG methods. Specifically, we illustrate these points with a set of concrete scenarios, each with 4 actions and 2-dimensional feature vectors describing each action. Since  $d < K$ , not every policy can be expressed in these representations, hence the problem instances are unrealizable.
3.1 Global Convergence is Achievable with Non-zero Approximation Error
The results of [9, Theorem 1] do not imply that sub-optimal local maxima always appear, as shown in the following.
Example 1.  $K = 4$ ,  $d = 2$ ,  $X^{\top} = \left[ \begin{array}{ccc}0 & -1 & 0 & 2\\ -2 & 0 & 1 & 0 \end{array} \right]$  and  $r = (9,8,7,6)^{\top}$ . The approximation error is  $\epsilon_{approx} = \min_{w\in \mathbb{R}^d}\| Xw - r\| _2 = \left\| X(X^\top X)^{-1}X^\top r - r\right\| _2 = \sqrt{202.6}\approx 14.2338.$
Note that the approximation error is larger than any sub-optimality gap, i.e., for any policy  $\pi$
$$
\left(\pi^ {*} - \pi\right) ^ {\top} r \leq 3 <   \epsilon_ {\text {a p p r o x}}, \tag {7}
$$
hence the bound in Eq. (6) does not imply global convergence for NPG in this example. Yet, despite the non-zero approximation error and the inability of existing results including Eq. (6) to establish global convergence on Example 1, both Algorithms 1 and 2 can be shown to reach a global maximum.
Proposition 1. Denote  $a^* \coloneqq \arg \max_{a \in [K]} r(a)$ . With constant  $\eta > 0$  and any initialization  $\theta_1 \in \mathbb{R}^d$ , both Algorithms 1 and 2 guarantee  $\pi_{\theta_t}^\top r \to r(a^*)$  as  $t \to \infty$  on Example 1.
All proofs can be found in the appendix due to space limits. The fact that Softmax PG achieves global convergence in Example 1 is much harder to establish than for NPG, since Eq. (4) involves a complex non-linearity given the presence of the softmax, unlike the linear least squares Eq. (5) used in NPG. To illustrate the intuition behind Proposition 1 we use a visualization of the optimization landscape.
Visualization. A visualization of the optimization landscape of Example 1 is shown in Figure 1(a). The bottom two-dimensional plane is the parameter space  $\mathbb{R}^d$  where  $d = 2$ . For each  $\theta \in \mathbb{R}^d$ , we calculate  $\pi_{\theta}$  using Eq. (3) and  $\pi_{\theta}^\top r$  using Eq. (2), and use  $\pi_{\theta}^\top r$  as the vertical axis value of  $\theta$ .
(b) Example 2.
(c) Example 3.
Figure 1: Visualizing the landscapes in the example problem instances.
To verify Proposition 1, we run Softmax PG and NPG on Example 1 with the same  $\theta_{1} = (6,8)^{\top}\in \mathbb{R}^{2}$ . In Figure 1(a), the optimization trajectories show 85 iterations of NPG and  $8.5\times 10^{6}$  iterations of
Softmax PG, both with learning rate  $\eta = 0.2$ . It can be clearly seen that both Softmax PG and NPG eventually achieve expected reward  $\pi_{\theta_t}^\top r \to 9 = r(a^*)$ , demonstrating global convergence (Figure 3(c) later shows that the sub-optimality gap  $(\pi^* - \pi_{\theta_t})^\top r$  approaches 0).
In summary, Example 1 shows that both Softmax PG and NPG are able to achieve global convergence on unrealizable problem instances with non-zero approximation error. This raises the question:
Is non-zero approximation error useful for characterizing global convergence?
3.2 Global Convergence is Irrelevant to Non-zero Approximation Error
We answer the above question negatively. By comparing alternative problem instances with similar approximation errors but different convergence behaviors, we illustrate how approximation error is not able to distinguish between scenarios where global versus local convergence is obtained.
Example 2.  $K = 4$ ,  $d = 2$ ,  $X^{\top} = \left[ \begin{array}{ccc}0 & 0 & -1 & 2\\ -2 & 1 & 0 & 0 \end{array} \right] \in \mathbb{R}^{d\times K}$ , and  $r = (9,8,7,6)^{\top} \in \mathbb{R}^{K}$ . The approximation error is  $\| X(X^{\top}X)^{-1}X^{\top}r - r\|_{2} = \sqrt{205} \approx 14.3178$ .
The only difference between Examples 1 and 2 is that the second and third columns of  $X^{\top}$  have been exchanged. The approximation error remains similar to that of Example 1. Using the upper bound of Eq. (6), one might therefore expect similar sub-optimality gaps  $(\pi^{*} - \pi_{\theta_{t}})^{\top}r$  to be demonstrated by the algorithms, since the r.h.s. contains similar approximation errors. However, as shown in Figure 1(b), using the same initialization and learning rate, Softmax PG obtains  $\pi_{\theta_t}^\top r\to 8 = r(2) < r(a^*)$  as it converges to a sub-optimal deterministic policy, while NPG continues to succeed.
Lest one believe that NPG is globally convergent, the following example, where the first and second columns of  $X^{\top}$  are swapped, illustrate an analogous failure for NPG but not Softmax PG.
Example 3.  $K = 4$ ,  $d = 2$ ,  $X^{\top} = \left[ \begin{array}{ccc} -1 & 0 & 0 & 2 \\ 0 & -2 & 1 & 0 \end{array} \right] \in \mathbb{R}^{d\times K}$ , and  $r = (9,8,7,6)^{\top} \in \mathbb{R}^{K}$ . The approximation error is  $\| X(X^{\top}X)^{-1}X^{\top}r - r\|_{2} = \sqrt{212} \approx 14.5602$ .
Here again the approximation error is close to that of Example 1. Yet, Figure 1(c) shows that NPG achieves  $\pi_{\theta_t}^\top r \to 8 < r(a^*)$  as it converges to a sub-optimal solution, while Softmax PG succeeds.
In summary, the Examples 1, 2 and 3 all have similar approximation errors, yet Softmax PG achieves global convergence on Example 1 but reaches a bad local maxima on Example 2, while NPG succeeds on Example 1 and fails on Example 3. Note that these examples can be re-scaled to have exactly the same approximation errors while demonstrating the same convergence behavior of the algorithms. From these findings we conclude that, if there is any quantity that can predict whether global versus local convergence is obtained by Softmax PG or NPG, that the quantity cannot be approximation error alone. This motivates to investigate the question: what is the right quantity to characterize global convergence for unrealizable problems?
3.3 Global Convergence Characterization is Algorithm Dependent
We make one more key point. From Figure 1(b) and Figure 1(c), NPG achieves global convergence on Example 2 but fails on Example 3, while, conversely, Softmax PG succeeds on Example 3 and fails on Example 2. This difference indicates that whatever condition characterizes global convergence, it must be algorithm dependent, even for the closely related algorithms Softmax PG and NPG. Therefore, one has to study the conditions for Softmax PG and NPG respectively (rather than one condition for both algorithms), which motivates the refined question:
What conditions characterize global convergence of Softmax PG and NPG in unrealizable problems?
4 New Characterizations of Global Convergence for PG Methods
From these observations, it is clear that whatever quantity characterizes the global convergence of PG methods, it cannot be based solely on approximation error and it must be algorithm dependent. Therefore, we study distinct global convergence conditions for Softmax PG and NPG respectively.
4.1 Reward Order Preservation with Adequate Features is Sufficient for PG Convergence
We now investigate a global convergence condition for Softmax PG under log-linear policies.
Intuition. Consider Example 1, where Softmax PG achieves global convergence. From the landscape shown in Figure 1(a), there appears to be a monotonic path from any initialization point that allows gradient ascent to reach the optimal plateau with reward  $r(a^{*}) = 9$ . Intuitively, this arises because the actions' rewards seem to be nicely "ordered". For example, starting from  $\theta_{1} = (6,8)^{\top} \in \mathbb{R}^{d}$  such that  $\pi_{\theta_1}^\top r \approx 6$ , Softmax PG is able to improve its expected reward eventually to  $\pi_{\theta_t}^\top r \approx 7$ , since there exists a sub-optimal plateau with a higher reward 7 right beside the lowest plateau with reward 6. Next, Softmax PG continues to improve its expected reward eventually to  $\pi_{\theta_t}^\top r \approx 8$  by "climbing" toward another neighboring plateau with a higher reward. Finally, this process ends with Softmax PG successfully arriving at the optimal plateau with reward  $r(a^{*}) = 9$ .
By contrast, in Example 2, as shown in Figure 1(b), Softmax PG gets stuck on a bad plateau with a local maximum reward of 8. Visually, Softmax PG stops improving its expected reward on this sub-optimal plateau, because it is "surrounded" by two lower plateaus with rewards 6 and 7, which breaks the nice "ordering" of the expected reward landscape and traps the gradient ascent trajectory on a sub-optimal plateau from which there is no monotonic ascent to global optimality.
Verifying reward order preservation. Based on the above intuition and observations, we conjecture that the ordering structure between the different rewards is a key property behind the global convergence of Softmax PG. We can verify this conjecture in each of the Examples 1 to 3 by determining whether the feature matrix  $X \in \mathbb{R}^{K \times d}$  allows the same action ordering as the reward vector  $r \in \mathbb{R}^K$  to be realized. For Example 1, note that with  $w = (-1, -1)^{\top} \in \mathbb{R}^{d}$ , we have
$$
r ^ {\prime} := X w = (2, 1, - 1, - 2) ^ {\top} \in \mathbb {R} ^ {K}, \tag {8}
$$
which preserves the ordering of  $r \in \mathbb{R}^K$ , such that for all  $i, j \in [K]$ ,  $r(i) > r(j)$  if and only if  $r'(i) > r'(j)$ . Similarly, for Example 3, if we let  $w = (-3, -1)^{\top}$  then we have  $r' := Xw = (3,2,-1,-6)^{\top}$ , which also preserves the order of  $r$  over actions. Softmax PG converges to a globally optimal reward in both of these examples.
By contrast, for Example 2, it is impossible to find any  $w \in \mathbb{R}^d$  such that  $Xw$  preserves the order of the rewards  $r$ . To see why, consider any  $w = (w(1), w(2))^{\top}$  and note that
$$
r ^ {\prime} := X w = (- 2 \cdot w (2), w (2), - w (1), 2 \cdot w (1)) ^ {\top}. \tag {9}
$$
To preserve the reward order, we require both  $-2 \cdot w(2) > w(2)$  (which would imply  $w(2) < 0$ ) and  $-w(1) > 2 \cdot w(1)$  (which would imply  $w(1) < 0$ ), but these two conditions imply  $w(2) < 0 < -w(1)$ , which must reverse the order of the second and third actions. This is an example where PG can fail to reach a global optimum.