On the Convergence and Sample Efficiency of Variance-Reduced Policy Gradient Method

1 Introduction . In this paper , we investigate the theoretical properties of Policy Gradient ( PG ) methods for Reinforcement Learning ( RL ) [ 43 ] . In view of RL as a policy optimization problem , the PG method 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . parameterizes the policy function and conduct gradient ascent search to improve the policy . In this paper , we consider the soft-max policy parameterization πθ ( a|s ) = exp { ψ ( s , a ; θ ) } ∑ a′ exp { ψ ( s , a′ ; θ ) } ( 1 ) where ( s , a ) is a state-action pair and ψ is some smooth function . Potentially , one can set the function ψ to be some deep neural network with weights θ and input ( s , a ) . The main problem considered in this paper is the policy optimization for a general utility function : max θ R ( πθ ) : = F ( λ πθ ) , ( 2 ) where F is a general smooth function , and λπθ denotes the unnormalized state-action occupancy measure ( also referred to as the visitation measure ) . For any policy π and initial state distribution ξ , λπ ( s , a ) : = +∞∑ t=0 γt · P ( st = s , at = a ∣∣π , s0 ∼ ξ ) , ( 3 ) where γ stands for the discount factor and P denotes the probability of a certain event . When F is linear , the problem reduces to the standard policy optimization problem where the objective is to maximize a cumulative sum of rewards . When F is nonlinear , problem ( 2 ) goes beyond standard Markov decision problems : examples include the max-entropy exploration [ 16 ] , risk-sensitive RL [ 50 ] , certain set constrained RL [ 27 ] , and so on . In the standard cumulative-return case ( i.e. , F is linear ) , numerous works have studied PG methods in various scenarios , see e.g . [ 47 , 5 , 58 , 22 , 21 , 37 , 38 , 23 ] . When directly optimizing over the policy space without any parameterization , the policy mirror descent ( PMD ) method [ 23 ] achieves an Õ ( −2 ) sample complexity to find an O ( ) -optimal solution . However , for the more practical parameterized policy optimization , to the authors ’ best knowledge , the most recent variant of PG methods , using the SARAH/Spider stochastic variance reduction technique [ 13 , 29 ] , find a local -stationary policy using O ( −3 ) samples [ 48 , 33 ] . This poses a contrast with the known Õ ( −2 ) sample complexity results that can be achieved by various value-based methods [ 4 , 42 , 41 ] and are provably matching information-theoretic lower bounds [ 12 , 3 , 4 ] . In this paper , we attempt to close this gap and prove an Õ ( −2 ) sample complexity bound for a PG method . Most importantly , when it comes to PG estimation , the application of the variance reduction technique typically relies on certain off-policy PG estimator , resulting in the difficulty of distribution shift . We notice that none of the existing variance-reduced PG methods attempt to address this challenge . Instead , they directly make an uncheckable assumption that the variance of the importance weight is bounded for every policy pair encountered in running the algorithm , see e.g . [ 31 , 48 , 49 , 33 ] . In this paper , we propose a simple gradient truncation mechanism to fix this issue . Next , let us go beyond cumulative return and consider policy optimization for a general utility where F may be nonlinear . However , much less is known in this setting . The nonlinearity of F invalidates the concept of Q-function and value function , leading to the failure of policy gradient theorem [ 44 ] . To overcome such difficulty , [ 51 ] showed that the policy gradient for the general utilities is the solution to a min-max problem . However , estimating a single PG is highly nontrivial in this case . It is still unclear how to make PG methods to use samples in a most efficient way . In this paper , we aim to investigate the convergence and sample efficiency of the PG method , using episodic sampling , for both linear F ( i.e. , cumulative rewards ) and nonlinear F ( i.e. , general utility ) . Observe that problem ( 2 ) is an instance of the Stochastic Composite Optimization ( SCO ) problem [ 45 , 46 ] : minx f ( Eν [ gν ( x ) ] ) , which involves an inner expectation that corresponds to the occupancy measure λπ . Motivated by this view point , we attempt to develop stochastic policy gradient method with provable finite-sample efficiency bounds . Main results . Our main results are summarized below . • We propose the TSIVR-PG algorithm to solve problem ( 2 ) via episodic sampling . It provides a conceptually simple stochastic gradient approach for solving general utility RL . • We provide a gradient truncation mechanism to address the distribution shift difficulty in variance-reduced PG methods . Such difficulty has never been addressed in previous works . • We show that TSIVR-PG finds an -stationary policy using Õ ( −3 ) samples if F and ψ are general smooth functions . When F is concave and ψ satisfies certain overparameterization condition , we show that TSIVR-PG obtains a gloal -optimal policy using Õ ( −2 ) samples . Technical contribution . Our analysis technique is also of independent interest in the relating areas . • For stochastic composite optimization ( SCO ) , most existing algorithms require estimating the Jacobian matrix of the inner mapping , which corresponds to ∇θλπθ in our setting . This is in practice prohibitive if the Jacobian matrix has high dimensions , which is exactly the case in our problem . Unlike SCO algorithms such as [ 24 , 54 , 53 , etc . ] , our analysis enables us to avoid the Jacobian matrix estimation . • For the stochastic variance-reduced gradient methods , our analysis implies a convergence of SARAH/Spider methods to global optimality and a new O ( −2 ) sample complexity for nonconvex problems with “ hidden convexity ” structure , which has not been studied in the optimization community yet . 2 Related Works . Policy gradient gives rises to a rich family of RL algorithms , such as REINFORCE and many of its variants [ 47 , 5 , 58 ] , as well as extensions such as the natural policy gradient methods [ 20 , 32 ] , the actor-critic methods [ 22 , 21 , 28 ] , the trust-region policy optimization [ 37 , 57 , 39 ] , and the proximal policy optimization method [ 38 , 25 ] , etc . In this paper we mainly focus on REINFORCE-type methods , where many of them need Õ ( −4 ) samples to find an -stationary solution , including the vanilla REINFORCE [ 47 ] , as well as its variants with baseline [ 58 , 43 ] and GPOMDP [ 5 ] , etc . By incorporating the stochastic variance reduction techniques , the sample efficiency of PG methods can be further improved . In [ 31 ] , the SVRG [ 19 ] variance reduction scheme is adopted and an O ( −4 ) sample complexity is achieved , which is later improved to O ( −10/3 ) by [ 48 ] . With additional Hessian information , [ 40 ] achieved an O ( −3 ) complexity . By utilizing a more efficient SARAH/Spider [ 29 , 13 ] variance reduction scheme , people are able to achieve O ( −3 ) sample complexity without second-order information [ 48 , 33 ] . We would like to comment that these results are only for finding -stationary ( rather than near-optimal ) solutions , and all of them requires an uncheckable condition on the importance weights in every iteration . Recently , for cumulative reward , a series of works have started to study the convergence of policy gradient method to global optimal solutions [ 1 , 14 , 56 , 6 , 26 , 7 , 9 , 55 ] . In particular , [ 51 ] exploited the hidden convexity property of the MDP problem and established the convergence to global optimality for general utility RL problem , as long as the policy gradient can be computed exactly . Our approach is related to the stochastic composite optimization ( SCO ) [ 45 , 46 ] . For the general composition problem , there have been numerous developments , including momentum-based and multi-time-scale algorithms [ 45 , 46 , 15 ] , and various composite stochastic variance-reduced algorithms [ 24 , 17 , 52 , 54 ] . Our approach is also inspired by variance reduction techniques that were initially used for stochastic convex optimization , see [ 19 , 36 , 11 , 29 ] ; and were later on extended to the stochastic nonconvex optimization problems [ 2 , 34 , 18 , 35 , 13 , 30 ] . In particular , we will utilize the SARAH/Spider scheme [ 13 , 30 ] . 3 Problem Formulation . Consider an MDP with a general utility function , denoted as MDP ( S , A , P , γ , F ) , where S is a finite state space , A is a finite action space , γ ∈ ( 0 , 1 ) is a discount factor , and F is some general utility function . For each state s ∈ S , a transition to state s′ ∈ S occurs when selecting an action a ∈ A following the distribution P ( ·|a , s ) . For each state s ∈ S , a policy π gives a distribution π ( ·|s ) over the action space A . Let ξ be the initial state distribution and let the unnormalized state-action occupancy measure λπ be defined by ( 3 ) , we define the general utility function F as a smooth function of the occupancy measure , and the goal of the general utility MDP is to maximize F ( λπ ) . With the policy πθ being parameterized by ( 1 ) , we propose to solve problem ( 2 ) , which is max θ R ( πθ ) : = F ( λ πθ ) . For notational convenience , we often write λ ( θ ) instead of λπθ . Such utility function is very general and includes many important problems in RL . We provide a few examples where F are concave . Example 3.1 ( Cumulative reward ) . When F ( λπθ ) = 〈r , λπθ 〉 , for some r ∈ R|S||A| . Then we recover the standard cumulative sum of rewards : R ( πθ ) = E [ ∑+∞ t=0 γ t · r ( st , at ) ∣∣πθ , s0 ∼ ξ ] . Example 3.2 ( Maximal entropy exploration ) . Let µπθ ( s ) = ( 1 − γ ) ∑ a λ πθ ( s , a ) , ∀s ∈ S be the state occupancy measure , which is the margin of λπθ over S. Let F ( · ) be the entropy function , then we recover the objective for maximal entropy exploration [ 16 ] : R ( πθ ) = − ∑ s∈S µ πθ ( s ) logµπθ ( s ) . Example 3.3 ( RL with Set Constraint ) . Let z ( st , at ) ∈ Rd be a vector feedback received in each step . The cumulative feedback is u ( πθ ) : = E [ ∑+∞ t=0 γ t · z ( st , at ) |s0 ∼ ξ , πθ ] = Mλπθ for some matrix M ∈ Rd×|S||A| . [ 27 ] proposed a set-constrained RL problem which aims to find a policy π s.t . u ( π ) ∈ U for some convex set U . This problem can be formulated as an instance of ( 2 ) by letting F ( · ) be the negative squared distance : R ( πθ ) = −minu′∈U ‖u′ − u ( πθ ) ‖2 .

Summary: There is a lot of interest in studying policy gradient type of algorithms for solving the RL problem recently.  This paper develops a variance reduced policy gradient method, where the objective function can be some general utility function instead of just being the cumulative reward. A sample complexity of $\mathcal{O}(\epsilon^{-3})$ was established under some mild conditions, and it was further improved to $\mathcal{O}(\epsilon^{-2})$ under a convexity assumption and a so-called over-parametrization assumption.