Beyond Prioritized Replay: Sampling States in Model-Based Reinforcement Learning via Simulated Priorities

Prioritized Experience Replay ( ER ) has been empirically shown to improve sample efficiency across many domains and attracted great attention ; however , there is little theoretical understanding of why such prioritized sampling helps and its limitations . In this work , we take a deep look at the prioritized ER . In a supervised learning setting , we show the equivalence between the error-based prioritized sampling method for mean squared error and uniform sampling for cubic power loss . We then provide theoretical insight into why it improves convergence rate upon uniform sampling during early learning . Based on the insight , we further point out two limitations of the prioritized ER method : 1 ) outdated priorities and 2 ) insufficient coverage of the sample space . To mitigate the limitations , we propose our modelbased stochastic gradient Langevin dynamics sampling method . We show that our method does provide states distributed close to an ideal prioritized sampling distribution estimated by the brute-force method , which does not suffer from the two limitations . We conduct experiments on both discrete and continuous control problems to show our approach ’ s efficacy and examine the practical implication of our method in an autonomous driving application . 1 INTRODUCTION . Experience Replay ( ER ) ( Lin , 1992 ) has been a popular method for training large-scale modern Reinforcement Learning ( RL ) systems ( Degris et al. , 2012 ; Adam & Busoniu , 2012 ; Mnih et al. , 2015a ; Hessel et al. , 2018 ; François-Lavet et al. , 2018 ) . In ER , visited experiences are stored in a buffer , and at each time step , a mini-batch of experiences is uniformly sampled to update the training parameters in the value or policy function . Such a method is empirically shown to effectively stabilize the training and improve the sample efficiency of deep RL algorithms . Several follow-up works propose different variants to improve upon it ( Schaul et al. , 2016 ; Andrychowicz et al. , 2017 ; Oh et al. , 2018 ; de Bruin et al. , 2018 ; Horgan et al. , 2018 ; Zha et al. , 2019 ; Novati & Koumoutsakos , 2019 ; Sun et al. , 2020 ) . The most relevant one to our work is prioritized ER ( Schaul et al. , 2016 ) , which attempts to improve the vanilla ER method by sampling those visited experiences proportional to their absolute Temporal Difference ( TD ) errors . Empirically , it can significantly improve sample efficiency upon vanilla ER on many tested domains . To gain an intuitive understanding of why the prioritized ER method works , one may recall ModelBased RL ( MBRL ) methods ( Kaelbling et al. , 1996 ; Bertsekas , 2009 ; Sutton & Barto , 2018 ) . ER can be thought of as an instance of a classical model-based RL architecture—Dyna ( Sutton , 1991 ) , using a ( limited ) non-parametric model given by the buffer ( van Seijen & Sutton , 2015 ; van Hasselt et al. , 2019 ) . A Dyna agent uses real experience to update its policy as well as its reward and dynamics model . In-between taking actions , the agent can get hypothetical experiences from the model to further improve the policy . Existing works show that smart ways of sampling those experiences can further improve sample efficiency of a MBRL agent ( Sutton et al. , 2008 ; Gu et al. , 2016 ; Goyal et al. , 2019 ; Holland et al. , 2018 ; Pan et al. , 2018 ; Corneil et al. , 2018 ; Janner et al. , 2019 ; Chelu et al. , 2020 ) . Particularly , prioritized sweeping ( Moore & Atkeson , 1993 ) improves upon vanilla Dyna . The idea behind prioritized sweeping is quite intuitive : we should give high priority to states whose absolute TD errors are large because they are likely to cause the most change in value estimates . Hence , applying TD error-based prioritized sampling to ER is a natural idea in a model-free RL setting . This work provides a theoretical insight into the prioritized ER ’ s advantage and points out its two drawbacks : outdated priorities and insufficient sample space coverage , which may significantly weaken its efficacy . To mitigate the two issues , we propose to leverage the flexibility of using an environment model to acquire hypothetical experiences by simulating priorities . Specifically , we bring in the Stochastic Gradient Langevin Dynamics ( SGLD ) sampling method to acquire states . We demonstrate that the hypothetical experiences generated by our method are distributed closer to the desired TD error-based sampling distribution , which does not suffer from the two drawbacks . Finally , we demonstrate the utility of our method on various benchmark discrete and continuous control domains and an autonomous driving application . 2 BACKGROUND . In this section , we firstly review basic concepts in RL . Then we briefly introduce the prioritized ER method , which will be examined in-depth in the next section . We conclude this section by discussing a classic MBRL architecture called Dyna ( Sutton , 1991 ) and its recent variants , which are most relevant to our work . Basic notations . We consider a discounted Markov Decision Process ( MDP ) framework ( Szepesvári , 2010 ) . A MDP can be denoted as a tuple ( S , A , P , R , γ ) including state space S , action space A , probability transition kernel P , reward function R , and discount rate γ ∈ [ 0 , 1 ] . At each environment time step t , an RL agent observes a state st ∈ S , takes an action at ∈ A , and moves to the next state st+1 ∼ P ( ·|st , at ) , and receives a scalar reward signal rt+1 = R ( st , at , st+1 ) . A policy is a mapping π : S ×A → [ 0 , 1 ] that determines the probability of choosing an action at a given state . A popular algorithm to find an optimal policy is Q-learning ( Watkins & Dayan , 1992 ) . With function approximation , parameterized action-values Qθ are updated using θ = θ + αδt∇θQθ ( st , at ) for stepsize α > 0 with TD-error δt def = rt+1 + γmaxa′∈AQθ ( st+1 , a ′ ) − Qθ ( st , at ) . The policy is defined by acting greedily w.r.t . these action-values . ER methods . ER is critical when using neural networks to estimate Qθ , as used in DQN ( Mnih et al. , 2015b ) , both to stabilize and speed up learning . ER method uniformly samples a mini-batch of experiences from those visited ones in the form of ( st , at , st+1 , rt+1 ) to update neural network parameters . Prioritized ER ( Schaul et al. , 2016 ) improves upon it by prioritized sampling experiences , where the probability of sampling a certain experience is proportional to its TD error magnitude , i.e. , p ( st , at , st+1 , rt+1 ) ∝ |δt| . However , the underlying theoretical mechanism behind this method is still not well understood . MBRL . With a model , an agent has more flexibility to sample hypothetical experiences . We consider a one-step model which maps a state-action pair to its possible next state and reward : P : S × A 7→ S × R. We build on the Dyna formalism ( Sutton , 1991 ) for MBRL , and more specifically , the recently proposed HC-Dyna ( Pan et al. , 2019 ) as shown in Algorithm 1 . HCDyna provides a special approach to Search-Control ( SC ) —the mechanism of generating states or state-action pairs from which to query the model to get the next states and rewards . HC-Dyna ’ s search-control mechanism generates states by Hill Climbing ( HC ) on some criterion function h ( · ) . The term HC is used for generality as the vanilla gradient ascent is modified to resolve certain challenges ( Pan et al. , 2019 ) . The algorithmic framework maintains two buffers : the conventional ER buffer storing experiences ( i.e. , an experience/transition has the form of ( st , at , st+1 , rt+1 ) ) and a search-control queue storing the states acquired by search-control mechanisms . At each time step t , a real experience ( st , at , st+1 , rt+1 ) is collected and stored into ER buffer . Then the HC search-control process starts to collect states and store them into the search-control queue . A hypothetical experience is obtained by first selecting a state s from the search-control queue , then selecting an action a according to the current policy , and then querying the model to get the next state s′ and reward r to form an experience ( s , a , s′ , r ) . These hypothetical transitions are combined with real experiences into a single mini-batch to update the training parameters . The n updates , performed before taking the next action , are called planning updates ( Sutton & Barto , 2018 ) , as they improve the value/policy by using Algorithm 1 HC-Dyna : Generic framework Input : Hill Climbing ( HC ) criterion function h : S 7→ R , batch-size b ; Initialize empty searchcontrol queue Bsc ; empty ER buffer Ber ; initialize policy and model P ; HC stepsize αh ; minibatch size b ; environment P ; mixing rate ρ decides the proportion of hypothetical experiences in a mini-batch . for t = 1 , 2 , . . . do Add ( st , at , st+1 , rt+1 ) to Ber while within some budget time steps do s← s+ αh∇sh ( s ) // HC for search-control Add s into Bsc // n planning updates/steps for n times do B ← ∅ // initialize an empty mini-batch B for bρ times do Sample s ∼ Bsc , on-policy action a Sample s′ , r ∼ P ( s , a ) B ← ( s , a , s′ , r ) Sample b ( 1− ρ ) experiences from Ber , add to B Update policy/value on mixed mini-batch B a model . The choice of pairing states with on-policy actions to form hypothetical experiences has been reported to be beneficial ( Gu et al. , 2016 ; Pan et al. , 2018 ; Janner et al. , 2019 ) . Two instances have been proposed for h ( · ) : the value function v ( s ) ( Pan et al. , 2019 ) and the gradient magnitude ||∇sv ( s ) || ( Pan et al. , 2020 ) . The former is used as a measure of the utility of a state : doing HC on the learned value function should find high-value states without being constrained by the physical environment dynamics . The latter is considered as a measure of the value approximation difficulty , then doing HC provides additional states whose values are difficult to learn . The two suffer from several issues as we discuss in the Appendix A.1 . This paper will introduce a HC search-control method motivated by overcoming the limitations of the prioritized ER . 3 A DEEPER LOOK AT ERROR-BASED PRIORITIZED SAMPLING . In this section , we provide theoretical motivation for error-based prioritized sampling by showing its equivalence to optimizing a cubic power objective with uniform sampling in a supervised learning setting . We prove that optimizing the cubic objective provides a faster convergence rate during early learning . Based on the insight , we discuss two limitations of the prioritized ER : 1 ) outdated priorities and 2 ) insufficient coverage of the sample space . We then empirically study the such limitations .

Summary: This paper proposes an alternative method for performing Prioritized Experience Replay (PER), which avoids issues of inadequate state coverage and staleness in priority scores. Their method is based on a result from stochastic langevin dynamics, which shows their specific stochastic gradient Langevin dynamics (SGLD) update to online states leads to the ground-truth PER distribution in the limit. Updating collected states this way and adding them to an additional "stochastic-control queue" from which sampled minibatch transitions are mixed with those sampled from a standard experience replay buffer. Empirically, this paper shows this SGLD-based method leads to a closer approximation to the ideal prioritized experience replay distribution over recent transitions than standard PER.