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 takes a deep look at prioritized experience replay, a popular technique in deep reinforcement learning. The paper gives insights on why error-based prioritized experience replay can help when the importance ratio is unused. This paper also pointed out two limitations of prioritized experience replay, which are outdated priorities and insufficient coverage of state space, and the author proposed to use SGLD to solve the limitations. Experiments show that the proposed method leads to a good coverage of state space and improve the return of the training algorithm.