Policy Optimization with Stochastic Mirror Descent

1 INTRODUCTION . Reinforcement learning ( RL ) is one of the most wonderful fields of artificial intelligence , and it has achieved great progress recently ( Mnih et al. , 2015 ; Silver et al. , 2017 ) . To learn the optimal policy from the delayed reward decision system is the fundamental goal of RL . Policy gradient methods ( Williams , 1992 ; Sutton et al. , 2000 ) are powerful algorithms to learn the optimal policy . Despite the successes of policy gradient method , suffering from high sample complexity is still a critical challenge for RL . Many existing popular methods require more samples to be collected for each step to update the parameters ( Silver et al. , 2014 ; Lillicrap et al. , 2016 ; Schulman et al. , 2015 ; Mnih et al. , 2016 ; Haarnoja et al. , 2018 ) , which partially reduces the effectiveness of the sample . Although all the above existing methods claim it improves sample efficiency , they are all empirical results which lack a strong theory analysis of sample complexity . To improve sample efficiency , in this paper , we explore how to design an efficient and stable algorithm with stochastic mirror descent ( SMD ) . Due to its advantage of the simplicity of implementation , low memory requirement , and low computational complexity ( Nemirovsky & Yudin , 1983 ; Beck & Teboulle , 2003 ; Lei & Tang , 2018 ) , SMD is one of the most widely used methods in machine learning . However , it is not sound to apply SMD to policy optimization directly , and the challenges are two-fold : ( I ) The objective of policy-based RL is a typical non-convex function , but Ghadimi et al . ( 2016 ) show that it may cause instability and even divergence when updating the parameter of a non-convex objective function by SMD via a single batch sample . ( II ) Besides , the large variance of gradient estimator is the other bottleneck of applying SMD to policy optimization for improving sample efficiency . In fact , in reinforcement learning , the non-stationary sampling process with the environment leads to the large variance of existing methods on the estimate of policy gradient , which results in poor sample efficiency ( Papini et al. , 2018 ; Liu et al. , 2018 ) . Contributions To address the above two problems correspondingly , in this paper ( I ) We analyze the theoretical dilemma of applying SMD to policy optimization . Our analysis shows that under the common Assumption 1 , for policy-based RL , designing the algorithm via SMD directly can not guarantee the convergence . Hence , we propose the MPO algorithm with a provable convergence guarantee . Designing an efficiently computable , and unbiased gradient estimator by averaging its historical policy gradient is the key to MPO . ( II ) We propose the VRMPO : a sample efficient policy optimization algorithm via constructing a variance reduced policy gradient estimator . Specifically , we propose an efficiently computable policy gradient estimator , utilizing fresh information and yielding a more accurate estimation of the gradient w.r.t the objective , which is the key to improve sample efficiency . We prove VRMPO needs O ( −3 ) sample trajectories to achieve an -approximate firstorder stationary point ( -FOSP ) ( Nesterov , 2004 ) . To our best knowledge , our VRMPO matches the best-known sample complexity among the existing literature . Besides , we conduct extensive experiments , which further show that our algorithm outperforms state-of-the-art bandit algorithms in various settings . 2 BACKGROUND AND NOTATIONS . 2.1 POLICY-BASED REINFORCEMENT LEARNING . We consider the Markov decision processes M = ( S , A , P , R , ρ0 , γ ) , where S is state space , A is action space ; At time t , the agent is in a state St ∈ S and takes an action At ∈ A , then it receives a feedback Rt+1 ; P ass′ = P ( s ′ |s , a ) ∈ P is the probability of the state transition from s to s′ under taking a ∈ A ; The bounded reward function R : S × A → [ −R , R ] , Ras 7→ E [ Rt+1|St = s , At = a ] ; ρ0 : S → [ 0 , 1 ] is the initial state distribution and γ ∈ ( 0 , 1 ) is discounted factor . Policy πθ ( a|s ) is a probability distribution on S × A with the parameter θ ∈ Rp . Let τ = { st , at , rt+1 } Hτt=0 be a trajectory , where s0 ∼ ρ0 ( s0 ) , at ∼ πθ ( ·|st ) , rt+1 = R ( st , at ) , st+1 ∼ P ( ·|st , at ) , and Hτ is the finite horizon of τ . The expected return J ( πθ ) is defined as : J ( θ ) def = J ( πθ ) = ∫ τ P ( τ |θ ) R ( τ ) dτ = Eτ∼πθ [ R ( τ ) ] , ( 1 ) where P ( τ |θ ) = ρ0 ( s0 ) ∏Hτ t=0 P ( st+1|st , at ) πθ ( at|st ) is the probability of generating τ , R ( τ ) =∑Hτ t=0 γ trt+1 is the accumulated discounted return . Let J ( θ ) = −J ( θ ) , the central problem of policy-based RL is to solve the problem : θ∗ = arg max θ J ( θ ) ⇐⇒ θ∗ = arg min θ J ( θ ) . ( 2 ) Computing the∇J ( θ ) analytically , we have ∇J ( θ ) = Eτ∼πθ [ Hτ∑ t=0 ∇θ log πθ ( at|st ) R ( τ ) ] . ( 3 ) For any trajectory τ , let g ( τ |θ ) = ∑Hτ t=0∇θ log πθ ( at|st ) R ( τ ) , which is an unbiased estimator of ∇J ( θ ) . Vanilla policy gradient ( VPG ) is a straightforward way to solve problem ( 2 ) : θ ← θ + αg ( τ |θ ) , where α is step size . Assumption 1 . ( Sutton et al. , 2000 ; Papini et al. , 2018 ) For each pair ( s , a ) , any θ ∈ Rp , and all components i , j , there exists positive constants G , F s.t. , |∇θi log πθ ( a|s ) | ≤ G , | ∂2 ∂θi∂θj log πθ ( a|s ) | ≤ F. ( 4 ) According to the Lemma B.2 of ( Papini et al. , 2018 ) , Assumption 1 implies ∇J ( θ ) is L-Lipschiz , i.e. , ‖∇J ( θ1 ) −∇J ( θ2 ) ‖ ≤ L‖θ1 − θ2‖ , where L = RH ( HG2 + F ) / ( 1− γ ) , ( 5 ) Besides , Assumption 1 implies the following property of the policy gradient estimator . Lemma 1 ( Properties of stochastic differential estimators ( Shen et al. , 2019 ) ) . Under Assumption 1 , for any policy πθ and τ ∼ πθ , we have ‖g ( τ |θ ) −∇J ( θ ) ‖2 ≤ G 2R2 ( 1− γ ) 4 def = σ2 . ( 6 ) 2.2 STOCHASTIC MIRROR DESCENT . Now , we review some basic concepts of SMD ; in this section , the notation follows ( Nemirovski et al. , 2009 ) . Let ’ s consider the stochastic optimization problem , min θ∈Dθ { f ( θ ) = E [ F ( θ ; ξ ) ] } , ( 7 ) where Dθ ∈ Rn is a nonempty convex compact set , ξ is a random vector whose probability distribution , µ is supported on Ξ ∈ Rd and F : Dθ × Ξ → R. We assume that the expectation E [ F ( θ ; ξ ) ] = ∫ Ξ F ( θ ; ξ ) dµ ( ξ ) is well defined and finite-valued for every θ ∈ Dθ . Definition 1 ( Proximal Operator ( Moreau , 1965 ) ) . T is a function defined on a closed convex X , and α > 0.Mψα , T ( z ) is the proximal operator of T , which is defined as : Mψα , T ( z ) = arg min x∈X { T ( x ) + 1 α Dψ ( x , z ) } , ( 8 ) where ψ ( x ) is a continuously-differentiable , ζ-strictly convex function : 〈x− y , ∇ψ ( x ) −∇ψ ( y ) 〉 ≥ ζ‖x−y‖2 , ζ > 0 , Dψ is Bregman distance : Dψ ( x , y ) = ψ ( x ) −ψ ( y ) −〈∇ψ ( y ) , x−y〉 , ∀ x , y ∈ X . Stochastic Mirror Descent The SMD solves ( 7 ) by generating an iterative solution as follows , θt+1 =Mψαt , ` ( θ ) ( θt ) = arg minθ∈Dθ { 〈gt , θ〉+ 1 αt Dψ ( θ , θt ) } , ( 9 ) where αt > 0 is step-size , ` ( θ ) = 〈gt , θ〉 is the first-order approximation of f ( θ ) at θt , gt = g ( θt , ξt ) is stochastic subgradient such that g ( θt ) = E [ g ( θt , ξt ) ] ∈ ∂f ( θ ) |θ=θt , { ξt } t≥0 represents a draw form distribution µ , and ∂f ( θ ) = { g|f ( θ ) − f ( ω ) ≤ gT ( θ − ω ) , ∀ω ∈ dom ( f ) } . If we choose ψ ( x ) = 1 2 ‖x‖22 , which implies Dψ ( x , y ) = 1 2 ‖x − y‖22 , since then iteration ( 9 ) is the proximal gradient ( Rockafellar , 1976 ) view of SGD . Thus , SMD is a generalization of SGD . Convergence Criteria : Bregman Gradient Bregman gradient is a generation of projected gradient ( Ghadimi et al. , 2016 ) . Recently , Zhang & He ( 2018 ) ; Davis & Grimmer ( 2019 ) develop it to measure the convergence of an algorithm for the non-convex optimization problem . Evaluating the difference between each candidate solution x and its proximity is the critical idea of Bregman gradient to measure the stationarity of x . Specifically , let X be a closed convex set on Rn , α > 0 , T ( x ) is defined on X . The Bregman gradient of T at x ∈ X is : Gψα , T ( x ) = α −1 ( x−Mψα , T ( x ) ) , ( 10 ) whereMψα , T ( · ) is defined in Eq. ( 8 ) . If ψ ( x ) = 1 2 ‖x‖22 , then x∗ is a critical point of T if and only if Gψα , T ( x∗ ) = ∇T ( x∗ ) = 0 ( Bauschke et al . ( 2011 ) ; Theorem 27.1 ) . Thus , Bregman gradient ( 10 ) is a generalization of gradient . The following Remark 1 is helpful for us to understand the significance of Bregman gradient , and it gives us some insights to understand this convergence criterion . Remark 1 . Let T be a convex function , by the Proposition 5.4.7 of Bertsekas ( 2009 ) : x∗ is a stationarity point of T if and only if 0 ∈ ∂ ( T + δX ) ( x∗ ) , ( 11 ) where δX ( · ) is the indicator function on X . Furthermore , suppose ψ ( x ) is twice continuously differentiable , let x̃ =Mψα , T ( x ) , by the definition of proximal operatorM ψ α , T ( · ) , we have 0 ∈ ∂ ( T + δX ) ( x̃ ) + ( ∇ψ ( x̃ ) −∇ψ ( x ) ) ( ? ) ≈ ∂ ( T + δX ) ( x̃ ) + αGψα , T ( x ) ∇ 2ψ ( x ) , ( 12 ) Eq. ( ? ) holds due to the first-order Taylor expansion of∇ψ ( x ) . By the criteria of ( 11 ) , if Gψα , T ( x ) ≈ 0 , Eq . ( 12 ) implies the origin point 0 is near the set ∂ ( T + δX ) ( x̃ ) , i.e. , x̃ is close to a stationary point . In practice , we choose T ( θ ) = 〈−∇J ( θt ) , θ〉 , since then discriminant criterion ( 12 ) is suitable to RL problem ( 2 ) . For the non-convex problem ( 2 ) , we are satisfied with finding an -approximate First-Order Stationary Point ( -FOSP ) ( Nesterov , 2004 ) , denoted by θ , such that ‖Gψα , T ( θ ) ( θ ) ‖ ≤ . ( 13 )

Summary: This paper proposes MPO, a policy optimization method with convergence guarantees based on stochastic mirror descent that uses the average of previous gradients to update the policy parameters. A lower-variance method, VRMPO, is then proposed that matches the best known convergence rate in the in the literature. Experiments show that (1) MPO converges faster than basic policy optimization methods on a small task, and (2) VRMPO achieves a performance comparable to, and often better than, popular policy optimization methods (TD3, DDPG, PPO, and TRPO) on MuJoCo.