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 paper proposes a particle-based sampling method for stochastic policy optimization in reinforcement learning. The main observation is that many policy optimization algorithms use reverse KL as the objective to optimize policies, which is mode-seeking and could be problematic in practice. To mitigate this problem, the authors propose to use particle-based sampling such that one can directly use KL for policy optimization.