Particle Based Stochastic Policy Optimization

1 INTRODUCTION . Deep reinforcement learning ( DRL ) leverages the power of neural-network function approximators and has shown great promise in a diverse field , such as control ( Lillicrap et al. , 2015 ; Andrychowicz et al. , 2020 ) , robotics ( Haarnoja et al. , 2017 ; Gu et al. , 2017 ; Plappert et al. , 2018 ) . Classic views on the notion of optimality in reinforcement learning state that an optimal policy can be deterministic ( Sutton & Barto , 2011 ) . However , recent studies by Daniel et al . ( 2012 ) ; Ziebart ( 2010 ) demonstrate that deterministic policy may suffer from several drawbacks , e.g. , the multi-modal behaviors have significant applications in real robotic control tasks ( Daniel et al. , 2012 ) and inflexibility on exploration ( Ostrovski et al. , 2017 ) . Particularly in offline RL , the deterministic policy lacks the structure to manage risk towards uncertainty dynamics . The branch of probabilistic RL is then introduced where the optimal solution is stochastic . Probabilistic RL not only enables the stochasticity of the learned policy but also provides a different perspective to reveal the connection between reinforcement learning and probabilistic inference ( Todorov , 2008 ) , where the optimization procedure of the policy could be considered a statistical divergence minimization procedure towards distribution in the state-action or trajectory space , which is implied by the corresponding optimal RL policy . Besides the different definitions of optimality and above mentioned priorities , another appealing property of probabilistic reinforcement learning lies in that advanced techniques in approximate inference can then be applied in the optimization for RL . There are two widely applied frameworks , i.e. , the pseudo-likelihood framework and maximum entropy ( MERL ) framework . The pseudo-likelihood defines the optimality by considering a specified positive and bounded function of reward as the density function . By regressing to the reweighted samples , the pseudo-likelihood based methods not only is not an explicit probabilistic model but also suffers from high variance and poor risk handling ( Levine , 2018 ) . MERL is a prominent choice for policy improvement for stochastic policy which tends to maximize the expected reward and the expected conditional entropy . It avoids the drawbacks of pseudo-likelihood that it suffers from high variance and poor risk handling ability ( Levine , 2018 ) . In practice , the MERL is generally implemented with the reverse KL minimization which promotes the mode-seeking behavior and could fade the merits of a stochastic policy . For example , in offline setting , the mode-seeking policy will result in an overly-optimistic estimation of Q-function ( Kumar et al. , 2020 ) which could be problematic when generalizing in the test scenario . In this work , we explore the possibility of using a broad class of distribution discrepancies for stochastic policy improvement , and find it is possible by leveraging particle-based methods . By providing better distribution matching methods , our proposed framework , ParPI , enjoys the power of probabilistic RL . Moreover , ParPI is naturally compatible with the offline RL setting . And we show that with specific discrepancy measure selected , ParPI could conjoin the benefit of different policy constraint methods . We firstly demonstrate that using a broader class of distribution discrepancies achieves significant improvements in stabilizing the policy optimization and achieving the state-ofthe-art expected returns on high-dimensional and complex continuous control tasks ( Todorov et al. , 2012 ) . Besides , we conduct extensive experiments in the offline setup ( Yu et al. , 2020 ) to show that the ParPI could be more defensive for optimization thus results in more generalizable optimal policy . 2 PRELIMINARIES . 2.1 PROBABILISTIC INFERENCE FRAMEWORK FOR RL . The reinforcement learning ( RL ) problem can be modeled as the following infinite-horizon Markov Decision Process ( MDP ) , which is implied by the tuple 〈S , A , r , p , p0 , γ〉 . We consider the common cases where the state space S and action space A are assumed to be continuous . p ( st+1|st , at ) refers to the transition probability , i.e . the probability density of next state st+1 given the current state st and action at . when an agent chooses an action a in state s , the environment emits a reward r ( s , a ) : S ×A → [ rmin , rmax ] during the transition . The policy distribution π ( a|s ) is presented as a distribution over the action space conditioned on any state . The distribution of state-action pairs when navigating the environment with policy π could be defined as ρπ ( s , a ) , which is also referred to as occupancy measure ( Ho & Ermon , 2016 ) , and the marginal distribution on the state space is ρπ ( s ) . With a predefined MDP , the notation trajectory τπ refers to the following Markov sequence τπ = { ( s0 , a0 ) , · · · ( sT , aT ) } generated by interacting with environment under policy π . The discount factor is denoted as γ . Standard reinforcement learning seeks to find the optimal policy that could maximize the expected return , and the optimal solution could be deterministic ( Sutton & Barto , 2011 ) . However , considering the exploration ( Schulman et al. , 2015 ) , robustness ( Ziebart , 2010 ) and multimodal objectives ( Daniel et al. , 2012 ) in real-world application scenarios , a stochastic/multi-modal policy is desired . To address this desideratum , optimal maximum entropy policy ( Ziebart et al. , 2008 ; Haarnoja et al. , 2017 ) is defined as : π∗ : = argmax π ∑ t Eρπ ( st , at ) [ r ( st , at ) + αH [ π ( ·|st ) ] ] . and the soft Bellman optimality equation ( Ziebart , 2010 ) is defined as : Q∗ ( st , at ) = rt + γEst+1 [ V ∗ ( st+1 ) ] , ( 1 ) which gives the soft Bellman operator T πQπ ( · ) = Qπ ( · ) and T π· : = r ( h ) + γEh′∼p ( s′|h ) π ( a′|s′ ) [ · ] and also referred to as Maximum Entropy RL ( MERL ) . Ziebart ( 2010 ) ; Haarnoja et al . ( 2017 ) and Haarnoja et al . ( 2018 ) show that the target of the policy update is given by the Q-function in an energybased form , π ( a|s ) ∝ exp ( −Q ( s , a ) /α ) . Due to the intractability in estimating the normalizing constant Z ( s ) = ∫ exp ( −Q ( s , a ) /α ) da for continuous or large discrete action space , the energybased form poses difficulties in utilizing the policy such as sampling or density estimation , so various methods use a standalone parametric policy model πφ ( a|s ) to distill from the energy function , and the MERL task is turned into iteratively updating the value model Qθ and the policy model πφ ( a|s ) . 3 PARTICLE-BASED POLICY IMPROVEMENT . 3.1 DISTRIBUTION MATCHING FOR POLICY IMPROVEMENT . We aim to update a policy model to match a better target p ( a|s ) , which is demonstrated by particles . We now formalize it as a distribution matching task under a general distribution metric/discrepancy D [ ·‖· ] . The first problem is that the policy model parameterizes a conditional distribution πφ ( a|s ) , so we need to match the target p ( a|s ) by minimizing D ( πφ ( a|s ) , p ( a|s ) ) , however , it is intractable for any given s. To make the minimization practical , we note that in almost all modern policy update tasks , we also have an associated state trajectory , e.g. , experiences in a replay buffer in off-policy RL , or simulated trajectories in inverse RL ( Finn et al. , 2016 ) and demonstrative trajectories in imitation learning ( Ho & Ermon , 2016 ) . This state trajectory demonstrates a marginal distribution on the state , ρ ( s ) , which is widely assumed to be stationary among different policies within the same MDP ( Tsitsiklis & Van Roy , 1996 ; Haarnoja et al. , 2017 ; 2018 ) . Hence we instead minimize D ( ρ ( s ) πφ ( a|s ) , ρ ( s ) p ( a|s ) ) . When the joint distributions match , the conditionals also match . This idea has been exploited in other fields , such as expectation propagation algorithms ( Minka , 2001 ) for Bayesian inference . And we further provide discussions in the Appendix F. Along with the state trajectory , we also have samples/particles from the augmented target , ρ ( s ) p ( a|s ) . Particularly for existing MERL methods , they minimize the reverse KL divergence DKL ( πφ ( ·|s ) , p ( ·|s ) ) averaged over a state distribution ρ ( s ) which can be implemented by conducting sampling from the replay buffer : Eρ ( s ) [ DKL ( πφ ( ·|s ) , p ( ·|s ) ] . This objective can be reformulated as DKL ( ρ ( s ) πφ ( a|s ) , ρ ( s ) p ( a|s ) ) , corresponding to our proposed method with D instantiated as the reverse KL . As stressed , this framework allows matching the distribution by minimizing more powerful metrics/discrepancies that do not have a particular propensity , like the reverse KL . To balance utility and implementability , we consider the f -divergences DF and the Wasserstein distance ( WD ) DW as two instances of D [ ·‖· ] of the proposed framework . Minimizing f-divergences for policy updates . The f -divergence is defined as DF ( q , p ) : = Ep [ F ( q/p ) ] , where F is a real-valued convex function with F ( 1 ) = 0 . Although the forward and reverse KL are also f-divergence instances , we do not consider them here since they would introduce the respective propensity of the learned model . Nowozin et al . ( 2016 ) show a useful lower bound of an f-divergence that benefits practical optimization : DF ( q , p ) > sup f ( D̃F ( q , p ; f ) : = Eq [ f ( · ) ] − Ep [ F ∗ ( f ( · ) ) ] ) ( 2 ) where F ∗ ( t ) : = supx { xt− F ( x ) } is the Fenchel conjugate of F . If the parametric model is defined via a reparameterization h ∼ qφ ( h ) : h = gφ ( ) with following a fixed and easy-to-sample distribution q ( ) like the standard Gaussian , the first term can also be estimated as Eq ( ) [ f ( gφ ( ) ) ] . By also parameterizing f in some function class , the supremum can be estimated after optimizing over f , which also tightens the lower bound . Thus minimizing D̃F well serves as minimizing DF . Minimizing WD for policy updates . The WD is defined as the minimal cost of transferring from one distribution to the other by a probabilistic transferring plan . It is shown to have an optimization utility even when the two distributions do not have overlapping support ( Arjovsky et al. , 2017 ) . Its formulation under Kantorovich-Rubinstein duality makes it convenient to optimize : DW ( q , p ) = sup ‖f‖L61 Eq [ f ] − Ep [ f ] ( 3 ) For distributions q , p with a bounded expectation on a Polish metric space ( Villani , 2008 , Particular Case 5.16 ) ( e.g. , the common Euclidean space ) , where ‖·‖L denotes the Lipschitz constant of a function . To enforce the Lipschitz constant constraint , various implementations are proposed , e.g . parameter clipping ( Arjovsky et al. , 2017 ) , gradient penalty ( Gulrajani et al. , 2017 ) and spectral norm / hinge loss regularizations ( Miyato et al. , 2018 ) . Framework 1 . The Particle-based Policy Improvement ( ParPI ) framework , minφD ( ρ ( s ) πφ ( a|s ) , ρ ( s ) p ( a|s ) ) can be used for the sub-task of updating policy in RL , as long as the improved policy can be demonstrated by particles . The process is shown in Fig . 1 . Remark 1 . The ParPI framework can be applied to policy update in SQL ( Ziebart , 2010 ; Haarnoja et al. , 2017 ) and SAC ( Haarnoja et al. , 2018 ) . In SQL , once the optimal Q-function Q∗ is achieved , the optimal policy is shown to be π∗ ( a|s ) = exp ( Q∗ ( s , a ) /α ) /Z∗ ( s ) where Z∗ ( s ) is a normalizing constant . In SAC , it is shown that π′ ( a|s ) : = exp ( Qπ ( s , a ) /α ) /Z ( s ) is a better policy than π where Qπ is the Q-function under policy π . In both cases , the policy update target is given in an energy-based form , whose particles can be drawn using MCMC algorithms . 3.2 PARPI FOR MERL The particle-based policy improvement ( ParPI ) framework can be applied to any RL task where a state-action trajectory demonstrates a better policy . We now show in detail how it helps to update policy in the MERL framework . As introduced , MERL ( Ziebart et al. , 2008 ; Ziebart , 2010 ; Haarnoja et al. , 2017 ; Liu et al. , 2017 ; Haarnoja et al. , 2018 ; Zhang et al. , 2018 ) defines the optimal policy which also considers the entropy to cover multiple optimalities and encourage exploration . Based on the definition , a policy update target is given in the form p ( a|s ) ∝ exp ( Q ( s , a ) /α ) , where Q ( s , a ) is either the optimal Q-function or the Q-function of the current policy , which can be estimated from simulated trajectories . These methods update the policy model by minimizing the reverse KL DKL ( ρ ( s ) πφ ( a|s ) , ρ ( s ) p ( a|s ) ) . The merit of it is that estimating the troublesome normalizing constant is not required in optimization , which is also widely exploited in variational inference . However , it has been pointed out ( Huszár , 2015 ; Theis et al. , 2016 ) that minimizing the reverse KL promotes the mode-seeking behavior in finite optimization steps , making ρ ( s ) πφ ( a|s ) concentrate to one mode of the target ρ ( s ) p ( a|s ) . This is because a substantial penalty can be incurred when ρ ( s ) πφ ( a|s ) puts more mass to where ρ ( s ) p ( a|s ) is dilute , but a regular loss in the other way . The behavior is also observed in practice , and there are attempts to ameliorate it in other fields , like variational inference ( Hernández-Lobato et al. , 2016 ; Li & Gal , 2017 ) As the goal of MERL is to capture the diversity and multi-modality of policy , such behavior would weaken the spirit of MERL . The ParPI framework allows MERL to benefit from more powerful metrics/discrepancies , as long as the trajectory samples are available . To draw from the target policy , we employ the Langevin dynamics , which only requires an unnormalized density function of the target distribution . Sampling with Langevin Dynamics . To minimize the general metrics/discrepancies above , samples from the target distribution ρπ ( s ) p ( a|s ) are required . Fortunately , this can be done by subsequently sampling from ρπ ( s ) and p ( a|s ) , and both can be implemented . Samples from ρπ ( s ) can be drawn by simulating the ( s , a ) trajectory following the current policy π ( a|s ) and environment transition . Sampling from p ( a|s ) can be done by running dynamics-based MCMC algorithms , which only require the unnormalized density , i.e . exp ( Q ( s , a ) /α ) . From various particle-based inference algorithms , we employ the Langevin dynamics ( Roberts et al. , 1996 ; Roberts & Stramer , 2002 ; Welling & Teh , 2011 ) due to the low computational cost , whose transition is given by : a ( i+1 ) = a ( i ) + ε∇ log p ( a ( i ) |s ) +N ( 0 , 2εI ) = a ( i ) + ( ε/α ) ∇aQ ( a ( i ) , s ) +N ( 0 , 2εI ) ( 4 ) It has been shown to achieve δ precision in total variance ( Dalalyan , 2017 ) , Wasserstein distance ( Durmus & Moulines , 2016 ) and KL divergence ( Cheng & Bartlett , 2017 ) in O ( 1/δ2 ) steps for strongly log-concave densities . Note that although we can use MCMC to draw samples from the policy target , the sampling is slow in deploying the policy , where we need to run Markov chains for every state along the trajectory , so explicitly modeling policy πφ ( a|s ) is still necessary . Value Function Update . The proposed policy improvement rule described above only requires a model for the state-action value function Qθ ( s , a ) , so it can be applied with any value function update method . The soft-Q learning method ( Haarnoja et al. , 2017 ) aims to update the value function towards optimal by a temporal difference implementation based on the soft Bellman optimality equation 1 : JQ ( θ ) = E ( st , at ) ∼D [ 1 2 ( Qθ ( st , at ) − r ( st , at ) + logEst+1∼p ( st+1|st , at ) [ exp ( V ( st+1 ) ) ] ) 2 ] ( 5 ) The soft actor-critic ( SAC ) method ( Haarnoja et al. , 2018 ) instead updates the state-value function model towards the state-action value of the current policy which gives a better policy target as following : JV ( ψ ) = Est∼ρ ( s ) [ 1 2 ( Vψ ( st ) − Eat∼πφ [ Qθ ( st , at ) − log πφ ( at|st ) ] ) 2 ] JQ ( θ ) = E ( st , at ) ∼D [ 1 2 ( Qθ ( st , at ) − r ( st , at ) + γEst+1∼p ( st+1|st , at ) [ Vψ̄ ( st+1 ) ] ) 2 ] ( 6 ) And we provide convergence analysis of ParPI in Appendix E .

Summary: This work extends the popular RL as inference framework to accommodate for more general divergence measures such as the f-divergence and Wasserstein distance. The author proposed a particle-based optimization framework for learning stochastic policies where the policy is learned using samples generated via Langevin dynamics. The framework can be applied to both online and offline settings.