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: In this paper, the authors propose a new policy gradient algorithm called TSIVR-PG for the policy optimization for a general utility function. They adopt a gradient truncation mechanism in TSIVR-PG to get rid of the uncheckable importance weights assumption which is frequently used in previous literatures. They establish convergence rate to stationary points, and global convergence results with additional assumptions. Moreover, they conduct experiments and compare with other previous methods.