Stochastic Second-Order Methods Improve Best-Known Sample Complexity of SGD for Gradient-Dominated Functions
Abstract
We study the performance of Stochastic Cubic Regularized Newton (SCRN) on a class of functions satisfying gradient dominance property with 1  ↵  2 which holds in a wide range of applications in machine learning and signal processing.
This condition ensures that any ﬁrst-order stationary point is a global optimum.
We prove that the total sample complexity of SCRN in achieving ✏ -global optimum is O ( ✏  7 / (2 ↵ )+1 ) for 1  ↵ < 3 / 2 and ˜ O
( ✏  2 / ( ↵ ) ) for 3 / 2  ↵  2 .
SCRN improves the best-known sample complexity of stochastic gradient descent.
Even under a weak version of gradient dominance property, which is applicable to policy-based reinforcement learning (RL), SCRN achieves the same improvement over stochastic policy gradient methods.
Additionally, we show that the average sample complexity of SCRN can be reduced to O ( ✏  2 ) for ↵ = 1 using a variance reduction method with time-varying batch sizes.
Experimental results in various RL settings showcase the remarkable performance of SCRN compared to ﬁrst-order methods.
Introduction
Consider the following unconstrained stochastic non-convex optimization problem:
[Equation]
where the random variable ⇠ is sampled from an underlying distribution P ⇠ .
In order to optimize the objective function F ( x ) , we have access to the ﬁrst and second derivatives of stochastic function f
( x , ⇠ ) .
The above optimization problem covers a wide range of problems, from the ofﬂine setting where the objective function is minimized over a ﬁxed number of samples, to the online setting where the samples are drawn sequentially.
In the deterministic case (where we have access to the derivatives of F ( x ) ), the gradient descent (GD) algorithm in the non-convex setting only guarantees convergence to a ﬁrst-order stationary point (FOSP) (i.e., a point x such that kr F ( x ) k = 0 ), which can be a local minimum, a local maximum, or a saddle point.
In contrast, second-order methods accessing the Hessian of F (or Hessian of f ( x , ⇠ ) in the stochastic setting) can exploit the curvature information to effectively escape saddle points and converge to a second-order stationary point (SOSP)
(i.e., such that kr F ( x ) k = 0 , r 2 F ( x ) ⌫ 0 ).
In their seminal work, Nesterov and Polyak [25] proposed the so-called cubic-regularized Newton (CRN) algorithm which exploits Hessian information and globally converges to an SOSP at a sub-linear rate of O (1 /k 2 / 3 ) , where k is the number of iterations.
In recent years, the performance of stochastic CRN (SCRN) for general non-convex functions has been the focus of several studies (for more details, see the related work in Section 1.1).
A variancereduced version of SCRN [3] can ﬁnd ( ✏ ,  ) -SOSP
(i.e., a point x such that kr F ( x ) k  ✏ and r 2 F ( x ) ⌫   I ) with the sample complexity of ˜ O ( ✏  3 ) .
Moreover, this rate is optimal for achieving ✏ -approximate FOSP (i.e., a point x such that kr F ( x ) k  ✏ ) and it cannot be improved using any stochastic p -th order methods for p  2 [3].
Nesterov and Polyak [25] studied CRN under the gradient dominance property (See Assumption 2).
Under Assumption 2 with ↵ = 2 , they showed that iterates of F ( x t )  min x F ( x ) converges to zero super-linearly.
The case of ↵ = 2 (commonly called Polyak- Ł ojasiewicz (PL) condition) was originally introduced by Polyak in [28], who showed that GD achieves linear convergence rate.
The gradient dominance property or its weak variations are satisﬁed in a quite wide range of machine learning applications such as neural networks with one hidden layer [18] or ResNet with linear activation [12] (for more details, see Section 1.1).
A particularly important application of the weak version of gradient dominance property with ↵ = [-
-]1 (Assumption 6) is in policy-based reinforcement learning (RL) [41].
Khaled et al.
[16] showed that under PL condition, the stochastic GD (SGD) with time-varying step-size returns a point ˆ x with a sample complexity of O (1 / ✏ ) , to reach E
[ F (ˆ x )]  min x F ( x )  ✏ .
Furthermore, the dependency of the sample complexity of SGD on ✏ is optimal [26].
Recently,
Fontaine et al.
[10] obtained a sample complexity of O ( ✏  4 ↵ +1 ) for SGD under gradient dominance property with 1  ↵  2 .
This shows that the worst sample complexity occurs for ↵ = 1 , which is precisely the value of ↵ that ﬁnds important applications in policy-based RL.
Indeed, under a weak version of gradient dominance property with ↵ = 1 , it has been shown that stochastic policy gradient (SPG) converges to the optimum point with a sample complexity of ˜ O ( ✏  3 ) [41, 9].
We know that in the deterministic case, CRN outperforms GD under gradient dominance property for all ↵ 2 [1 , 2] [25, 50] 2 .
Therefore, a natural question that arises is whether this holds true in the stochastic setting as well?
That is, does SCRN improve upon SGD under the gradient dominance property?
Herein, we address this question.
Speciﬁcally, our main contributions are as follows:
• We analyze the sample complexity of SCRN under gradient dominance property for 1  ↵  2 in order to return an ✏ -global stationary point ˆ x satisfying F (ˆ x )  min x F ( x )  ✏ (in expectation or with high probability).
As stated in Table 1, SCRN improves upon the best-known sample complexity of SGD for all 1  ↵ < 2 .
The largest improvement is for ↵ = 1 , and is O ( ✏  0 . 5 ) .
• In the setting of policy-based RL, under the weak version of gradient dominance property with ↵ = 1 (Assumption 6), we show that SCRN achieves a sample complexity of ˜ O ( ✏  2 . 5 ) , improving over the best-known sample complexity of SPG by a factor of ˜ O ( ✏  0 . 5 ) .
• We show that an adaptation of a variance-reduced SCRN [47] with time-varying batch sizes further improves the sample complexity of SCRN.
For ↵ = 1 , the average sample complexity is reduced to O ( ✏  2 ) .
[Table]
1.1 Related work
Gradient dominance property and its applications: The gradient dominance property with ↵ =(commonly called PL condition) was originally introduced by Polyak in [28].
It was shown by
Karimi et al.
[15] to be weaker than the most recent global optimality conditions that appeared in the literature of machine learning [19, 23, 42].
The gradient dominance property is also satisﬁed (sometimes locally rather than globally, and also under distributional assumptions) for the population risk in some learning models including neural networks with one hidden layer [18], ResNet with linear activation [12]
[Table]
and generalized linear model and robust regression [11].
Moreover, in policybased reinforcement learning (RL), a weak version of gradient dominance property with ↵ = 1 (see
Assumption 6) holds for some classes of policies (such as Gaussian policy and log-linear policy).
Variants of cubic regularized Newton method: For non-convex optimization, Nesterov and Polyak [25] proposed the CRN algorithm, which converges to a SOSP with the convergence rate of O (1 /k 2 / 3 ) (where k is the number of iterations) by solving a cubic sub-problem in each iteration.
Cartis et al.
[5] presented an adaptive framework for cubic regularization method.
In [17, 39], sub-sampled versions of gradient and Hessian were used in CRN to overcome the computational burden of Hessian matrix evaluations in high dimensional settings.
In the context of stochastic optimization, Tripuraneni et al.
[33] proposed a stochastic cubic regularization algorithm that obtains ✏ -SOSP with sample complexity of ˜ O ( ✏  3 . 5 ) .
Arjavani et al.
[3] improved the sample complexity to ˜ O ( ✏  3 ) using variance reduction.
In the convex setting, Song et al.
[30] presented a proximal CRN and its accelerated version, and proved a sample complexity of ˜ O ( ✏  2 ) to reach ✏ -global stationary point as long as the approximated Hessian in each iteration satisﬁes certain properties.
In ﬁnite-sum non-convex setting,
Zhou et al.
[48] proposed an adaptive sub-sampled CRN method that requires ˜ O ( N + N 4 / 5 ✏  3 / 2 ) to ﬁnd ✏ -SOSP, where N is the total number of samples.
Sample complexity was further reduced to ˜ O
( N + N 2 / 3 ✏  3 / 2 ) using various variance reduction methods [36, 45, 49].
To the best of our knowledge, no previous work on analyzing SCRN for gradient-dominant functions exists.
1.2 Notations
We adopt the following notation in the sequel.
Calligraphic letters (e.g., S ) denote spaces.
Upper-case bold letters (e.g., A ) denote matrices, and the lower-case bold letters (e.g., x ) denote vectors.
k · k denote the ` 2 -norm for vectors and the operator norm for matrices ( k A k :=
max ( A T A ) where  max ( X ) is the maximum eigenvalue of matrix X ), respectively.
A ⌫ B indicates that A  B is positive semi-deﬁnite.
We use the notation O to hide constants, and the notation ˜ O to hide both constants and logarithmic factors.
X  1   Y denotes that random variable X is less than or equal to random variable Y with the probability at least 1   .
Setup
Recall the stochastic non-convex optimization problem in (1), where the goal is to minimize the objective function F ( x ) having access to stochastic gradients
r f ( x , ⇠ ) and stochastic Hessian matrices r 2 f ( x , ⇠ ) .
We make the following assumption about the objective function in (1).
Assumption 1.
The Hessian of F is Lipschitz continuous with constant L 2 , i.e.,
[Equation]
Consider the empirical estimators g t := 1 n
P n
i
=1 r f ( x t , ⇠ i ) , and H t := 1 n
P n 2 i =1 r 2 f ( x t , ⇠ i ) where n 1 and n 2 are the numbers of samples used for estimating the gradient vector and Hessian matrix, respectively.
Input: Batch sizes n 1 , n 2 , initial point x 0 , accuracy ✏ , cubic penalty parameter M , maximum number of iterations T 1: t 2: k  0 k = 1 3: while k  t  1 k  2 ↵ p ✏ or t  T do 4: g t 1 n
P n
i
=1 r f ( x t , ⇠ i ) 5:
H t 1 n
P n 2 i =1 r 2 f ( x t , ⇠ i ) 6:  t arg min  2 R d h g t ,  i + 12 h  , H t  i + M 6 k  k 3 7: x t + x t +  t 8: t t +9: end while 10: return x t
Deﬁnition 1 (Total sample complexity) .
Given ✏ ,  > 0 , the total sample complexity is the number of calls (queries) of stochastic gradient and stochastic Hessian along the iterations until reaching a point x that satisﬁes one of the following: (1) for high probability analysis: F ( x )  F ( x ⇤ )  ✏ with probability at least 1   ; or (2) for analysis in expectation: E [ F ( x )]  F ( x ⇤ )  ✏ .
Our goal is to study the performance of stochastic cubic regularized Newton (SCRN) for objective functions that satisfy the gradient dominance property.
Algorithm 1 describes the steps in SCRN.
At each iteration t , we take batches of stochastic gradient vectors and Hessian matrices (lines 4 and 5) and then solve the following sub-problem to obtain  t (line 6):
[Equation]
We assume that there is an oracle that returns a global solution for this sub-problem (This assumption will be relaxed subsequently, see Remark 6).
Finally, we update x t in line 7.
SCRN under gradient dominance property
We shall study the performance of SCRN for functions satisfying gradient dominance property, deﬁned as follows.
Assumption 2.
Function F ( x ) satisﬁes gradient dominance property when for every x 2 R d ,
[Equation]
where x ⇤ 2 arg min x F ( x ) , ⌧ F > 0 , and ↵ 2 [1 , 2] are two constants.
The case ↵ = 2 is often referred to as PL condition [28, 15].
In this paper, we consider all ↵ ’s in the interval [1 , 2] .
The gradient dominance property holds for a large class of functions including sub-analytic functions, logarithm, and exponential functions, and semi-algebraic functions.
These function classes cover some of the most common non-convex objectives used in practice (see related work in Section 1.1).
In the following lemma, we present a recursion inequality that captures the behaviour of the function
F ( x t )  F ( x ⇤ )
at each iteration t for SCRN under gradient dominance property.
Lemma 1.
Assume that function F satisﬁes Assumption 1 (Lipschitz Hessian) and Assumption(gradient dominance property) for ↵  1 .
Then the resulting update x t +1 in Algorithm 1 (line 7) after plugging in  t , the solution of sub-problem in (3) , satisﬁes the following:
[Equation]
where C, C g , C H > 0 are constants depending on M, L 2 , and ⌧ F , and deﬁned in (34) .
Due to space limitations, all proofs are moved to the appendix.
In the following, we ﬁrst provide an analysis in expectation of SCRN under gradient dominance property with ↵ 2 [1 , 3 / 2] .
Next, we study the same algorithm for ↵ 2 (3 / 2 , 2] using a high probability analysis.
3.1 SCRN under gradient dominance property with ↵ 2 [1 , 3 / 2]
We make the following assumption on the stochastic gradients and Hessians.
Assumption 3.
For a given ↵ 2 [1 , 3 / 2] and for each query point x 2 R d :
[Equation]
where  1 and  2 , ↵ are two constants.
Remark 1.
The assumption E [ kr 2 f ( x , ⇠ )  r 2 F ( x ) k 2 ↵ ]   22 , ↵ for 1  ↵  3 / 2 is slightly stronger than the usual assumption E [ kr 2 f ( x , ⇠ )  r 2 F ( x ) k 2 ]
  22 .
We need this assumption because of the speciﬁc form of the error of Hessian estimator in recursion inequality in (5) .
As a result of the assumption in (7) , using a version of matrix moment inequality (See Lemma 4 in Appendix), we show that the dependency of the Hessian sample complexity in Theorem 1 on dimension d is in the order of log d .
Theorem 1.
Let F ( x ) satisfy Assumptions 1 and 2 for a given ↵ and the stochastic gradient and
Hessian satisfy Assumption 3 for the same ↵ .
Moreover, assume that an exact solver for sub-problem (3) exists.
Then Algorithm 1 outputs a point x T such that E [ F ( x T )]
F ( x ⇤ )  ✏ after T iterations, where
(i) if ↵ 2 [1 , 3 / 2) , T = O
( ✏
2 ↵ 2 ↵ ) , with
access to the
following numbers of samples of the stochastic gradient
and
Hessian per iteration:
[Equation]
where C 0 H is deﬁned in (38) and depends on log( d ) .
(ii) if ↵ = 3 / 2 , T
= O (log(1 / ✏ )) with the same numbers of samples per iteration as in (35) .
3.2 SCRN under gradient dominance property with ↵ 2 (3 / 2 , 2]
Deﬁnition 2 (Bernstein’s condition for matrices) .
A zero-mean symmetric random matrix X satisﬁes the Bernstein condition with parameter b > 0 if
[Equation]
where Var( X ) :
= E [ X 2 ]  ( E [ X ]) 2 .
Assumption 4.
We assume that the symmetric version of each centered gradient estimator G ( x , ⇠ ) :=  0 1 ⇥g ( x , ⇠ ) T g ( x , ⇠ )
d ⇥ d  where g ( x , ⇠ ) :=
r f ( x , ⇠ )  r F ( x ) and each centered Hessian estimator
H ( x , ⇠ ) :
= r 2 f ( x , ⇠ )  r 2 F ( x ) satisfy Bernstein’s condition (9) with parameters M 1 and M 2 , respectively.
Remark 2.
It is noteworthy that most previous work analyzing SCRN [33, 49, 36] assumed that centered gradient and centered Hessian estimators are bounded, i.e., kr f ( x , ⇠ )  r F ( x ) k 2 a.s. 
M 1 , kr 2 f ( x , ⇠ )  r 2 F ( x ) k a.s.  M 2 .
This is a stronger assumption than Assumption 4 as it implies
Bernstein’s condition (9) for g ( x , ⇠ ) and H ( x , ⇠ ) .
Theorem 2.
Suppose that F ( x ) satisﬁes Assumptions 1, 2, the stochastic gradient and Hessian satisfy Assumption 3 (with ↵ = 1 ) and Assumption 4, and there exists an exact solver for sub-problem (3) .
Then, Algorithm 1, with probability 1   , outputs a solution x T such that F ( x T )  F ( x ⇤ )
 ✏ after T
= O (log (log(1 / ✏ ))) iterations with the following numbers of samples for the stochastic gradient and Hessian, respectively per iteration:
[Equation]
where ˜ C
= 1 + ⌧
[Equation]
[Equation]
↵ .
Remark 3.
It is noteworthy that the assumptions in Theorem 2 sufﬁce to obtain the sample complexities in Theorem 1 for 1  ↵  3 / 2 up to a logarithmic factor.
However, the assumptions of Theorem 2 are stronger than those in Theorem 1.
Remark 4.
Theorem 2 implies that the sample complexity of Algorithm 1 for ↵ = 2 is
O (log log log(1 / ✏ ) · log(log(1 / ✏ )) / ✏ ) .
We summarized the sample complexity of SCRN in Theorems 1 and 2 and the best-known sample complexity of SGD under gradient dominance property in Table 1.
The optimal sample complexity of SGD under gradient dominance property with ↵ = 2 is discussed in detail in [16, Section 5.2].
For the general case of ↵ 2 [1 , 2] , the sample complexity of SGD under gradient dominance property was derived in [10, Theorem 10].
This rate is achieved with a time-varying step-size.
In the fourth column of Table 1, we provide the improvement of the sample complexity of SCRN with respect to that of SGD.
The improvement is O  ✏ [Equation] ↵
[Equation]
for ↵ 2 [1 , 3 / 2) and is ˜ O
( ✏  2 / ↵ +1 ) for ↵ 2 [3 / 2 , 2] , which are decreasing functions of ↵ .
The largest improvement is for ↵ = 1 and is O ( ✏  0 . 5 ) .
Remark 5.
The special case ↵ = 2 generalizes strong convexity.
For strongly convex objectives, the dependence of sample complexity of SGD on ⌧ F is O ( ⌧ 2 F / ✏ ) [16, Corollary 2] while the sample complexity of SCRN is O ( ⌧ 7 / 4 F / ✏ ) .
This is an improvement by a factor of ⌧ 1 / 4 F .
In Appendix A.2.3, we provide simulation results comparing the performance of SCRN and SGD over synthetic functions by varying ↵ and ⌧ F .
Remark 6.
In our analysis, we assumed that we have access to the exact solution of sub-problem in (3) .
Although no closed-form solution nor exact solver exists for this sub-problem, there are algorithms that approximate the exact solution with high probability [2, 4].
We emphasize that solving the sub-problem in (3) requires extra computation, but extra gradient and Hessian queries are not needed.
In particular, Carmon and Duchi [4] proposed a perturbed GD-based algorithm that returns an approximate solution ˜
t such that m t ( ˜  t )  1   0 m t (  t ) + ✏ 0 with O (log(1 /  0 ) / ✏ 0 ) iterations, for any given  0 , ✏ 0 > 0 .
In Appendix A.2.2, we prove the following lemma: Lemma 2. Theorems 1 and 2 are still true if we were to use an inexact sub-solver which returned an approximate solution ˜
t such that kr m t ( ˜  t )
k  ✏ 1 / ↵ ( ✏ 1 / ↵ -stationary point).
Moreover, under some mild assumptions (same as those in [4]), a GD-based algorithm indeed returns such a solution in O ( ✏  2 / ↵ ) iterations.
Remark 7. In the iterations of GD-based algorithm, instead of directly computing the Hessian matrix (which could be computationally expensive in high dimensions), we could compute Hessian-vector products by running Pearlmutter’s algorithm [27].
Thus, the total computational complexity of evaluating gradients and Hessian vector products of SCRN is in the order of O ( d ✏  3 / ↵ ) while the one of SGD is O ( d ✏  4 / ↵ +1 ) , where d is the dimension of x .
For ↵ 2 [1 , 3 / 2) , it can be seen that the computational complexity of SGD is less than the one of SCRN by a factor of O ( ✏  1 / (2 ↵ ) ) .
For ↵ 2 [3 / 2 , 2] , this factor is ˜ O ( ✏  (1  1 / ↵ ) ) .
3.3 Further improvements
In our analysis of SCRN, we set the batch sizes such that the stochastic error terms C g kr F ( x t )  g t k ↵ and C H kr 2 F ( x t )
H t k 2 ↵ in (5) are in the order of ✏ (either in expectation or with high probability).
However, for 1  ↵ < 3 / 2 , it is just needed to make sure that the error terms at iteration t are
O ( t  (2 ↵ ) / (3  2 ↵ ) ) , which equals the convergence rate of the function values F ( x t )  F ( x ⇤ )
(see
Lemma 11 in Appendix A.3).
Moreover, as stated in the following theorem, incorporating timevarying batch sizes in conjunction with variance reduction improves sample complexity results.
Assumption 5.
We assume that f ( x , ⇠ ) satisﬁes L {+
+}
1 -average smoothness and L 0 2 -average Hessian Lipschitz continuity, i.e., E [ kr f ( x , ⇠ )  r f[-( y , ⇠ ) k 2 ]  L 0 21 k x  y k 
-]
( y , ⇠ ) k 2 ]  L 0 21 k x  y k 2 and E [ kr 2 f ( x , ⇠ )  r 2 f ( y , ⇠ ) k 2 ]  L 0 22 k x  y k 2 for all x , y 2 R d .
Theorem 3.
Suppose that F ( x ) satisﬁes the gradient dominance property with ↵ = 1 .
Assume that Assumptions 1, 3, and 5 for ↵ = 1 hold.
Then variance-reduced SCRN (See Algorithm 2 in Appendix
A.3) achieves ✏ -global stationary point in expectation by making O ( ✏  2 ) stochastic gradients and O ( ✏  1 ) stochastic Hessian on average, and the iteration complexity is O ( ✏  1 / 2 ) .
We will see in the next section that there are a class of objective functions in RL setting that satisﬁes a weak version of gradient dominance with ↵ = 1 and Lipschitz gradient property and then we can take advantage of variance-reduced SCRN (See Remark 11).
SCRN under weak gradient dominance property in RL Setting
In this section, we showcase practical relevance of our result in Section 3.1 by applying SCRN to model-free RL.
Speciﬁcally, in Theorem 4, we prove that as long as the expected return satisﬁes the weak version of gradient dominance property (Assumption 6), SCRN improves upon the best-known sample complexity of stochastic policy gradient (SPG)
[41] by a factor of O (1 / p ✏ ) (see Appendix
A.4.1 for the related work).
4.1 RL Setup
Consider a discrete Markov decision process (MDP) M
= ( S , A , P, R, ⇢ ,  ) , where S is the state space and A is the action space.
P ( s 0 | s, a ) denotes the probability of state transition from s to safter taking action a and R ( · , · ) : S ⇥ A ! [  R
[Equation] , R max ] is a bounded reward function, where R max is a positive scalar.
⇢ represents the initial distribution on state space S and  2 (0 , 1) is the discount factor.
The parametric policy ⇡ ✓ is a probability distribution over S ⇥
A with parameter ✓ 2 R d , and ⇡ ✓ ( a | s ) denotes the probability of taking action a at a given state s .
Let ⌧ = { s t , a t } t  0 ⇠ p ( ⌧ | ⇡ ✓ ) be a trajectory generated by the policy ⇡ ✓ , where p ( ⌧ | ⇡ ✓ ) :
= ⇢ ( s 0 )
Q
[Equation]
0 ⇡ ✓
[Equation]
s t ) P ( s t +1 | s t , a t ) .
The expected return of ⇡ is deﬁned as J ( ⇡ ✓ ) :
= E ⌧ ⇠ p ( ·| ⇡ ✓ ) [
P 1 t =
t R ( s t , a t )] .
In the sequel, we consider a set of parameterized policies { ⇡ ✓ : ✓ 2 R d } , with the assumption that ⇡ ✓ is differentiable with respect to ✓ .
For ease of presentation, we denote J ( ⇡ ✓ ) by J ( ✓ ) .
The goal of policy-based RL is to ﬁnd ✓
⇤ = arg max ✓ J ( ✓ ) .
However, in many cases, J ( ✓ ) is a non-concave function and instead we settle for obtaining an ✏ -FOSP, ˆ ✓ , such that kr J (ˆ ✓ ) k  ✏ .
It can be shown that: r J ( ✓ ) =
E [
P 1 h =0 ( P 1 t = h  t R ( s t , a t )) r log ⇡ ✓ ( a h | s h )] .
In practice, the full gradient cannot be computed due to the inﬁnite horizon length.
Instead, it is commonly truncated to a length H horizon as follows.
r J H ( ✓ ) =
E [
P
H  1 h = h ( ⌧ ) r log ⇡ ✓ ( a h | s h )] , where h ( ⌧ ) =
P
H
[Equation]
= h  t R ( s t
[Equation]
a t ) .
Assume that we sample m trajectories ⌧ i =
{ s it , a it } t  0 , 1  i  m , and then compute ˆ r m J ( ✓ ) =m
P mi =
P H  1 h = h ( ⌧ i ) r log ⇡ ✓ ( a ih | s ih ) , which is an unbiased estimator for r J H ( ✓ ) .
The vanilla
SPG method is based on the following update: ✓
[Equation]
where ⌘ is the learning rate.
It can be shown that the Hessian matrix of J H ( ✓ ) can be obtained as follows [29, Appendix 7.2]:
[Equation]
✓ [Equation]
⌧
[Equation]
] , where  ( ✓ ; ⌧ ) = [Equation] ⇡ ✓ ( a h | s h ) .
As a result, for trajectories ⌧ i = { s it , a it } t  0 , 1  i  m , ˆ r [Equation]
m
[Equation] is an unbiased estimator of Hessian matrix r 2 J H ( ✓ ) .
4.2 Sample Complexity of SCRN
In our analysis, we consider the recently introduced relaxed weak gradient dominance property with ↵ = 1 [41].
Assumption 6 (Weak gradient dominance property with ↵ = 1 ) .
J satisﬁes the weak gradient dominance property if for all ✓ 2 R d , there exist ⌧ J > 0 and ✏ 0 > 0 such that
[Equation]
where J ⇤ : = max ✓ J ( ✓ ) .
Remark 8.
Two commonly made assumptions in the literature: non-degenerate Fisher matrix [20, 9] and transferred compatible function approximation error [35, 1] imply Assumption 6.
See Appendix
A.4.1 for more details.
Further, we also make Lipschitz and smooth policy (LS) assumptions which are widely adopted in the analysis of vanilla policy gradient (PG) [46] as well as variance-reduced PG methods, e.g. in [29].
Assumption 7 (LS) .
There exist constants G 1 , G 2 > 0 such that for every state s 2 S , the gradient and Hessian of log ⇡ ✓ ( ·| s ) satisfy kr ✓ log ⇡ ✓ ( a | s ) k  G 1 and kr 2 ✓ log ⇡ ✓ ( a | s ) k  G 2 .
Lemma 3.
Under Assumption 7, we have kr J ( ✓ ) r
J H ( ✓ ) k  D g  H and kr 2 J ( ✓ ) r 2 J H ( ✓ [Equation] k 
D H  H , where D g =
G 1 R max 1
[Equation]
Assumption 8 (Lipschitz Hessian) .
There exists a constant ¯ L 2 such that the Hessian of log ⇡ ✓ ( a | s ) satisﬁes
[Equation]
The Lipschitz Hessian assumption is commonly used to ﬁnd SOSP in policy gradient algorithms [40].
For the Gaussian policy (136), r 2 log ⇡ ✓ ( a | s ) reduces to the matrix   ( s )  ( s ) T /  2 , which is a constant function of ✓ and thus satisﬁes condition (13).
Soft-max policy also satisﬁes this assumption (See appendix A.4.4).
Theorem 4.
For a policy ⇡ ✓ satisfying Assumptions 7, 8, and the corresponding objective function
J ( ✓ ) satisfying Assumption 6, SCRN outputs the solution ✓ T such that J ⇤  E
[ J ( ✓ T )]  ✏ + ✏ 0 and the sample complexity (the number of observed state-action pairs) is: T ⇥ m ⇥ H =
˜ O ( ✏  2 . 5 ) for ✏ 0 = 0 and T ⇥ m ⇥ H =
˜ O ( ✏  0 .
5 ✏ 0 2 ) for ✏ 0 > 0 .
Remark 9.
Under weak gradient dominance property with ↵ = 1 (Assumption 6), it has been shown that the sample complexity of SPG is ˜ O ( ✏  3 ) in case of ✏ 0 = 0 and ˜ O
( ✏  1 ✏ 0 2 ) in case of ✏ 0 > 0 [41,
Theorem C.1].
Therefore, SCRN improves upon the best-known sample complexity of SPG in both cases ✏ 0 = 0 and ✏ 0 > 0 by a factor of O ( ✏  0 . 5 ) .
Remark 10.
Having access to exact gradient and Hessian (the deterministic case), under Assumption 6, PG algorithm achieves global convergence (i.e., J ⇤  J ( ✓ T )  ✏ ) with ˜ O
( ✏  1 ) iterations [41] while CRN requires ˜ O ( ✏  0 . 5 ) iterations.
Remark 11.
Under the same assumptions as in Theorem 4 and Assumption 5 for ↵ = 1 and bounded variance of importance sampling weights (See Assumption 15), a variance-reduced version of SCRN (See Algorithm 3) achieves global convergence (i.e., J ⇤  E
[ J ( ✓ T )]  ✏ ) with a sample complexity of ˜ O ( ✏  2 ) .
See Appendix A.4.6 for a proof.
Experiments
In this section, we evaluate the performance of SCRN in the two following RL settings.
First, we consider grid world environments with ﬁnite state and action spaces, and next some robotic control tasks with continuous state and action spaces.
The details of the implementations for all methods and some additional experiments appear in the appendix and the codes are available in the supplementary material.
Environments with ﬁnite state and action spaces: We consider two grid world environments in our experiments: cliff walking [31, Example 6.6], and random mazes [51].
In cliff walking, the agent’s aim is to reach a goal state from a start state, avoiding a region of cells called “cliff”.
The episode is terminated if the agent enters the cliff region, or the number of steps exceeds 100 without reaching the goal.
Moreover, we consider a soft-max tabular policy in the experiments of this part.
Indeed, it has been shown that a variant of gradient dominance property with ↵ = 1 holds for soft-max tabular policy in environments with ﬁnite state and action spaces [22].
We compare SCRN with two existing ﬁrst-order methods: vanilla SPG and REINFORCE [37].
For the ﬁrst-order methods, we use a time-varying learning rate and tune the parameters.
To improve the performance of the ﬁrst-order methods, we also add an entropy regularization term to the reward function.
In Fig.
1, the average length of paths traversed by the agent and the average episode return are depicted against the number of episodes for each method.
The results are averaged over 64 instances [Figure] [Figure] and the shaded region shows the 90% conﬁdence interval.
As can be seen in Fig.
1 (a), all the algorithms have a phase at the beginning during which the agent falls off the cliff in most episodes and the average length of paths are small.
Then, the agent learns to avoid the cliff but could still not reach the goal in most cases.
Finally, it ﬁnds a path to the goal and tries to reduce the path length.
Note that SCRN ﬁnds the path very quickly and signiﬁcantly outperforms the other algorithms.
In fact, SPG and REINFORCE get stuck in the start state for some period of time, while SCRN easily escapes the ﬂat plateau (for more details, please see the demonstrations in the supplementary ﬁle).
The performance of SPG and REINFORCE improves with entropy regularization, but SCRN still outperforms them.
Moreover, SPG and REINFORCE fail to reach the goal in some instances while
SCRN is successful in almost all instances.
In the captions of ﬁgures, for each algorithm, we also provide the percentage of instances in which the agent reached the goal.
Please note that these percentages are obtained based on the parameters from the last update of each algorithm.
Additionally, we studied the performance of the three aforementioned algorithms on a random maze and a random shape maze [51].
In the random shape maze, random shape blocks are placed on a grid and the agent tries to reach the goal state ﬁnding the shortest path.
As shown in Fig.
2, SCRN again [Figure] outperforms the ﬁrst-order methods in both environments.
In Appendix A.5, we also provided results for SPG and REINFROCE with entropy regularization which have slightly better performance.
Environments with continuous state and action spaces: We consider the following control tasks in MuJoCo simulator [32]: Walker, Humanoid, Reacher, and HalfCheetah.
We compare SCRN with ﬁrst-order methods such as REINFORCE, and two state-of-the-art representatives of variance-reduced PG methods, HAPG [29] and MBPG [13], both with guaranteed convergence to ✏ -FOSP in general non-convex settings.
HAPG uses second order information (Hessian vector products) for variance reduction and MBPG is a recent work based on STORM, a batch-free state-of-the-art variance reduction approach [8].
We report average episode return against system probes as our performance measure.
That is, the number of observed state-action pairs (see Fig. 3).
At each point, we run the trained policy 10 times and compute the empirical estimate of the mean and the 90% conﬁdence interval of the episode return.
As seen in Fig.
3, SCRN outperforms the other methods, especially in more complex environments such as HalfCheetah and Humanoid.
In Appendix A.5, we also provided results for variance-reduced
SCRN which improves upon SCRN in Humanoid and Reacher environments.
Conclusion
We studied the performance of SCRN for objectives satisfying the gradient dominance property for 1  ↵  2 , which holds in various machine learning applications.
We showed that SCRN improves the best-known sample complexity of SGD.
The largest improvement is in the case of ↵ = 1 .
Moreover, for ↵ = 1 , the average sample complexity of SCRN can be reduced to O ( ✏  2 ) by utilizing a variance reduction method with time-varying batch sizes.
A weak version of gradient dominance for ↵ = 1 is satisﬁed in some policy-based RL settings.
In the RL setting, we showed that SCRN achieves the same improvement over SPG under the weak version of gradient dominance property for ↵ = 1 .
Acknowledgement
The authors would like to thank Mohammadsadegh Khorasani for conducting the experiments for continuous state-action environments.
This research is partially supported by National Centre of
Competence in Research (NCCR), grant agreement no. 51NF40_180545.