Bregman Gradient Policy Optimization

1 INTRODUCTION . Policy Gradient ( PG ) methods are a class of popular policy optimization methods for Reinforcement Learning ( RL ) , and have achieved significant successes in many challenging applications ( Li , 2017 ) such as robot manipulation ( Deisenroth et al. , 2013 ) , the Go game ( Silver et al. , 2017 ) and autonomous driving ( Shalev-Shwartz et al. , 2016 ) . In general , PG methods directly search for the optimal policy by maximizing the expected total reward of Markov Decision Processes ( MDPs ) involved in RL , where an agent takes action dictated by a policy in an unknown dynamic environment over a sequence of time steps . Since the PGs are generally estimated by Monte-Carlo sampling , such vanilla PG methods usually suffer from very high variances resulted in slow convergence rate and destabilization . Thus , recently many fast PG methods have been proposed to reduce variances of vanilla stochastic PG . For example , Sutton et al . ( 2000 ) introduced a baseline to reduce variances of the stochastic PG . Konda & Tsitsiklis ( 2000 ) proposed an efficient actor-critic algorithm by estimating the value function to reduce effects of large variances . ( Schulman et al. , 2015b ) proposed the generalized advantage estimation ( GAE ) to control both the bias and variance in policy gradient . More recently , some faster variance-reduced PG methods ( Papini et al. , 2018 ; Xu et al. , 2019a ; Shen et al. , 2019 ; Liu et al. , 2020 ) have been developed based on the variance-reduction techniques in stochastic optimization . Alternatively , some successful PG algorithms ( Schulman et al. , 2015a ; 2017 ) improve convergence rate and robustness of vanilla PG methods by using some penalties such as Kullback-Leibler ( KL ) divergence penalty . For example , trust-region policy optimization ( TRPO ) ( Schulman et al. , 2015a ) ensures that the new selected policy is near to the old one by using KL-divergence constraint , while proximal policy optimization ( PPO ) ( Schulman et al. , 2017 ) clips the weighted likelihood ratio to implicitly reach this goal . Subsequently , Shani et al . ( 2020 ) have analyzed the global convergence properties of TRPO in tabular RL based on the convex mirror descent algorithm . Liu et al . ( 2019 ) have also studied the global convergence properties of PPO and TRPO equipped with overparametrized neural networks based on mirror descent iterations . At the same time , Yang et al . ( 2019 ) tried to propose the PG methods based on the mirror descent algorithm . More recently , mirror descent policy optimization ( MDPO ) ( Tomar et al. , 2020 ) iteratively updates the policy beyond the tabular RL by approximately solving a trust region problem based on convex mirror descent algorithm . In addition , Agarwal et al . ( 2019 ) ; Cen et al . ( 2020 ) have studied the natural PG methods for regularized RL . However , Agarwal et al . ( 2019 ) mainly focuses on tabular policy and log-linear , neural policy classes . Cen et al . ( 2020 ) mainly focuses on softmax policy class . Although these specific PG methods based on mirror descent iteration have been recently studied , which are scattered in empirical and theoretical aspects respectively , it lacks a universal framework for these PG methods without relying on some specific RL tasks . In particular , there still does not exist the convergence analysis of PG methods based on the mirror descent algorithm under the nonconvex setting . Since mirror descent iteration adjusts gradient updates to fit problem geometry , and is useful in regularized RL ( Geist et al. , 2019 ) , there exists an important problem to be addressed : Could we design a universal policy optimization framework based on the mirror descent algorithm , and provide its convergence guarantee under the non-convex setting ? In the paper , we firmly answer the above challenging question with positive solutions and propose an efficient Bregman gradient policy optimization framework based on Bregman divergences and momentum techniques . In particular , we provide a convergence analysis framework of the PG methods based on mirror descent iteration under the nonconvex setting . In summary , our main contributions are provided as follows : a ) We propose an effective Bregman gradient policy optimization ( BGPO ) algorithm based on the basic momentum technique , which achieves the sample complexity of O ( −4 ) for finding -stationary policy only requiring one trajectory at each iteration . b ) We propose an accelerated Bregman gradient policy optimization ( VR-BGPO ) algorithm based on the variance-reduced momentum technique . Moreover , we prove that the VRBGPO reaches the best known sample complexity of O ( −3 ) under the nonconvex setting . c ) We design a unified policy optimization framework based on mirror descent iteration and momentum techniques , and provide its convergence analysis under nonconvex setting . In Table 1 shows thet sample complexities of the representative PG algorithms based on mirror descent algorithm . Shani et al . ( 2020 ) ; Liu et al . ( 2019 ) have established global convergence of a mirror descent variant of PG under some pre-specified setting such as over-parameterized networks ( Liu et al. , 2019 ) by exploiting these specific problems ’ hidden convex nature . Without these special structures , global convergence of these methods can not be achieved . However , our framework does not rely on any specific policy classes , and our convergence analysis only builds on the general nonconvex setting . Thus , we only prove that our methods convergence to stationary points . Geist et al . ( 2019 ) ; Jin & Sidford ( 2020 ) ; Lan ( 2021 ) ; Zhan et al . ( 2021 ) studied a general theory of regularized MDPs based on policy function space that generally is discontinuous . Since both the state and action spaces S and A generally are very large in practice , the policy function space is large . While our methods build on policy parameter space that is generally continuous Euclidean space and relatively small . Clearly , our methods and theoretical results are more practical than the results in ( Geist et al. , 2019 ; Jin & Sidford , 2020 ; Lan , 2021 ; Zhan et al. , 2021 ) . ( Tomar et al. , 2020 ) also proposes mirror descent PG framework based on policy parameter space , but it does not provide any theoretical results and only focuses on Bregman divergence taking form of KL divergence . While our framework can collaborate with any Bregman divergence forms . 2 RELATED WORKS . In this section , we review some related work about mirror descent algorithm in RL and variancereduced PG methods , respectively . 2.1 MIRROR DESCENT ALGORITHM IN RL . Due to easily deal with the regularization terms , mirror descent ( a.k.a. , Bregman gradient ) algorithm ( Censor & Zenios , 1992 ; Beck & Teboulle , 2003 ) has shown significant successes in regularized RL , which is first proposed in ( Censor & Zenios , 1992 ) based on the Bregman distances ( Bregman , 1967 ; Censor & Lent , 1981 ) . For example , Neu et al . ( 2017 ) have shown both the dynamic policy programming ( Azar et al. , 2012 ) and TRPO ( Schulman et al. , 2015a ) algorithms are approximate variants of mirror descent algorithm . Subsequently , Geist et al . ( 2019 ) have introduced a general theory of regularized MDPs based on the convex mirror descent algorithm . More recently , Liu et al . ( 2019 ) have studied the global convergence properties of PPO and TRPO equipped with overparametrized neural networks based on mirror descent iterations . At the same time , Shani et al . ( 2020 ) have analyzed the global convergence properties of TRPO in tabular policy based on the convex mirror descent algorithm . Wang et al . ( 2019 ) have proposed divergence augmented policy optimization for off-policy learning based on mirror descent algorithm . MDPO ( Tomar et al. , 2020 ) iteratively updates the policy beyond the tabular RL by approximately solving a trust region problem based on convex mirror descent algorithm . 2.2 ( VARIANCE-REDUCED ) PG METHODS . PG methods have been widely studied due to their stability and incremental nature in policy optimization . For example , the global convergence properties of vanilla policy gradient method in infinite-horizon MDPs have been recently studied in ( Zhang et al. , 2019 ) . Subsequently , Zhang et al . ( 2020 ) have studied asymptotically global convergence properties of the REINFORCE ( Williams , 1992 ) , whose policy gradient is approximated by using a single trajectory or a fixed size mini-batch of trajectories under soft-max parametrization and log-barrier regularization . To accelerate these vanilla PG methods , some faster variance-reduced PG methods have been proposed based on the variance-reduction techniques of SVRG ( Johnson & Zhang , 2013 ) , SPIDER ( Fang et al. , 2018 ) and STORM ( Cutkosky & Orabona , 2019 ) in stochastic convex and non-convex optimization . For example , fast SVRPG ( Papini et al. , 2018 ; Xu et al. , 2019a ) algorithm have been proposed based on SVRG . Fast HAPG ( Shen et al. , 2019 ) and SRVR-PG ( Xu et al. , 2019a ) algorithms have been presented by using SPIDER technique . Subsequently , the momentum-based PG methods , i.e. , ProxHSPGA ( Pham et al. , 2020 ) and IS-MBPG ( Huang et al. , 2020 ) , have been developed based on variance-reduced technique of STORM/Hybrid-SGD ( Cutkosky & Orabona , 2019 ; Tran-Dinh et al. , 2019 ) . More recently , ( Ding et al. , 2021 ) studied the global convergence of momentum-based policy gradient methods . ( Zhang et al. , 2021 ) proposed a truncated stochastic incremental variance-reduced policy gradient ( TSIVR-PG ) method to relieve the uncheckable importance weight assumption in above variance-reduced PG methods and provided the global convergence of the TSIVR-PG under overparameterizaiton of policy assumption . 3 PRELIMINARIES . In the section , we will review some preliminaries of Markov decision process and policy gradients . 3.1 NOTATIONS . Let [ n ] = { 1 , 2 , · · · , n } for all n ∈ N+ . For a vector x ∈ Rd , let ‖x‖ denote the ` 2 norm of x , and ‖x‖p = ( ∑d i=1 |xi|p ) 1/p ( p ≥ 1 ) denotes the p-norm of x . For two sequences { ak } and { bk } , we denote ak = O ( bk ) if ak ≤ Cbk for some constant C > 0 . E [ X ] and V [ X ] denote the expectation and variance of random variable X , respectively . 3.2 MARKOV DECISION PROCESS . Reinforcement learning generally involves a discrete time discounted Markov Decision Process ( MDP ) defined by a tuple { S , A , P , r , γ , ρ0 } . S and A denote the state and action spaces of the agent , respectively . P ( s′|s , a ) : S × A → 4 ( S ) is the Markov kernel that determines the transition probability from the state s to s′ under taking an action a ∈ A. r ( s , a ) : S ×A → [ −R , R ] ( R > 0 ) is the reward function of s and a , and ρ0 = p ( s0 ) denotes the initial state distribution . γ ∈ ( 0 , 1 ) is the discount factor . Let π : S → 4 ( A ) be a stationary policy , where4 ( A ) is the set of probability distributions on A . Given the current state st ∈ S , the agent executes an action at ∈ A following a conditional probability distribution π ( at|st ) , and then the agent obtains a reward rt = r ( st , at ) . At each time t , we can define the state-action value function Qπ ( st , at ) and state value function V π ( st ) as follows : Qπ ( st , at ) = Est+1 , at+1 , ... [ ∞∑ l=0 γlrt+l ] , V π ( st ) = Eat , st+1 , ... [ ∞∑ l=0 γlrt+l ] . ( 1 ) We also define the advantage function Aπ ( st , at ) = Qπ ( st , at ) − V π ( st ) . The goal of the agent is to find the optimal policy by maximizing the expected discounted reward max π J ( π ) : = Es0∼ρ0 [ V π ( s0 ) ] . ( 2 ) Given a time horizonH , the agent collects a trajectory τ = { st , at } H−1t=0 under any stationary policy . Then the agent obtains a cumulative discounted reward r ( τ ) = ∑H−1 t=0 γ tr ( st , at ) . Since the state and action spaces S and A are generally very large , directly solving the problem ( 2 ) is difficult . Thus , we let the policy π be parametrized as πθ for the parameter θ ∈ Θ ⊆ Rd . Given the initial distribution ρ0 = p ( s0 ) , the probability distribution over trajectory τ can be obtained p ( τ |θ ) = p ( s0 ) H−1∏ t=0 P ( st+1|st , at ) πθ ( at|st ) . ( 3 ) Thus , the problem ( 2 ) will be equivalent to maximize the expected discounted trajectory reward : max θ∈Θ J ( θ ) : = Eτ∼p ( τ |θ ) [ r ( τ ) ] . ( 4 ) In fact , the above objective function J ( θ ) has a truncation error of O ( γ H 1−γ ) compared to the original infinite-horizon MDP .

Summary: The authors consider the optimization problem of an MDP. They designed two policy gradient algorithms based on the mirror descent method, which are named BGPO and VR-BGPO. The BGPO algorithm is a momentum mirror descent method that finds an $\epsilon$-stationary policy with $O(\epsilon^{-4})$ samples. The VR-BGPO algorithm is a STORM-type variance reduced mirror descent method that finds an $\epsilon$-stationary policy with $O(\epsilon^{-3})$ samples. The analysis is nicely organized and the authors also provide a couple of experiments to verify their theoretical findings.