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

1 INTRODUCTION . Using hypothetical experience simulated from an environment model can significantly improve sample efficiency of RL agents ( Ha & Schmidhuber , 2018 ; Holland et al. , 2018 ; Pan et al. , 2018 ; Janner et al. , 2019 ; van Hasselt et al. , 2019 ) . Dyna ( Sutton , 1991 ) is a classical MBRL architecture where the agent uses real experience to updates its policy as well as its reward and dynamics models . In-between taking actions , the agent can get hypothetical experience from the model to further improve the policy . An important question for effective Dyna-style planning is search-control : from what states should the agent simulate hypothetical transitions ? On each planning step in Dyna , the agent has to select a state and action from which to query the model for the next state and reward . This question , in fact , already arises in what is arguably the simplest variant of Dyna : Experience Replay ( ER ) ( Lin , 1992 ) . In ER , visited transitions are stored in a buffer and at each time step , a mini-batch of experiences is sampled to update the value function . ER can be seen as an instance of Dyna , using a ( limited ) non-parametric model given by the buffer ( see van Seijen & Sutton ( 2015 ) for a deeper discussion ) . Performance can be significantly improved by sampling proportionally to priorities based on errors , as in prioritized ER ( Schaul et al. , 2016 ; de Bruin et al. , 2018 ) , as well as specialized sampling for the off-policy setting ( Schlegel et al. , 2019 ) . Search-control strategies in Dyna similarly often rely on using priorities , though they can be more flexible in leveraging the model rather than being limited to only retrieving visited experiences . For example , a model enables the agent to sweep backwards by generating predecessors , as in prioritized sweeping ( Moore & Atkeson , 1993 ; Sutton et al. , 2008 ; Pan et al. , 2018 ; Corneil et al. , 2018 ) . Other methods have tried alternatives to error-based prioritization , such as searching for states with high reward ( Goyal et al. , 2019 ) , high value ( Pan et al. , 2019 ) or states that are difficult to learn ( Pan et al. , 2020 ) . Another strategy is to directly generate hypothetical experiences from trajectory optimization algorithms ( Gu et al. , 2016 ) . These methods are all supported by nice intuition , but as yet lack solid theoretical reasons for why they can improve sample efficiency . In this work , we provide new insights about how to choose the sampling distribution over states from which we generate hypothetical experience . In particular , we theoretically motivate why errorbased prioritization is effective , and provide a mechanism to generate states according to more accurate error estimates . We first prove that l2 regression with error-based prioritized sampling is equivalent to minimizing a cubic objective with uniform sampling in an ideal setting . We then show that minimizing the cubic power objective has a faster convergence rate during early learning stage , providing theoretical motivation for error-based prioritization . The theoretical understanding illuminates two issues of prioritized ER : insufficient sample space coverage and outdated priorities . To overcome the limitations , we propose a search-control strategy in Dyna that leverages a model to simulate errors and to find states with high expected error . Finally , we demonstrate the efficacy of our method on various benchmark domains and an autonomous driving application . 2 PROBLEM FORMULATION . We formalize the problem as a Markov Decision Process ( MDP ) , 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 , and takes an action at ∈ A . The environment transitions to the next state st+1 ∼ P ( ·|st , at ) , and emits 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 . Algorithm 1 HC-Dyna : Generic framework Input : hill climbing crit . h : S 7→ R , batch-size b Initialize empty search-control queue Bsc ; empty ER buffer Ber ; initialize policy and model P for t = 1 , 2 , . . . do Add ( st , at , st+1 , rt+1 ) to Ber while within some budget time steps do s← s+ αa∇sh ( s ) //hill climbing Add s into Bsc for n times do B ← ∅ for b/2 times do Sample s ∼ Bsc , on-policy action a Sample s′ , r ∼ P ( s , a ) B ← ( s , a , s′ , r ) Sample b/2 experiences from Ber , add to B Update policy on the mixed mini-batch B The agent ’ s objective is to find an optimal policy . A popular algorithm is Qlearning ( Watkins & Dayan , 1992 ) , where 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 is critical when using neural networks to estimate Qθ , as used in DQN ( Mnih et al. , 2015 ) , both to stabilize and speed up learning . MBRL has the potential to provide even further sample efficiency improvements . 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 . HC-Dyna provides a particular approach to search-control—the mechanism of generating states or state-action pairs from which to query the model to get next states and rewards ( i.e . hypothetical experiences ) . It is characterized the fact that it generates states by hill climbing on some criterion function h ( · ) . The term Hill Climbing ( HC ) is used for generality as the vanilla gradient ascent procedure is modified to resolve certain challenges ( Pan et al. , 2019 ) . Two particular choices have been proposed for h ( · ) : the value function v ( s ) from Pan et al . ( 2019 ) and the gradient magnitude ||∇sv ( s ) || from Pan et al . ( 2020 ) . The former is used as measure of the utility of visiting a state and the latter is considered as a measure of value approximation difficulty . The hypothetical experience is obtained by first selecting a state s , then typically selecting the action a according to the current policy , and then querying the model to get next state s′ and reward r. These hypothetical transitions are treated just like real transitions . For this reason , HC-Dyna combines both real experience and hypothetical experience into mini-batch updates . These n updates , performed before taking the next action , are called planning updates , as they improve the action-value estimates— and so the policy—using a model . However , it should be noted that there are several limitations to the two previous works . First , the HC method proposed by Pan et al . ( 2019 ) is mostly supported by intuitions , without any theoretical justification to use the stochastic gradient ascent trajectories for search-control . Second , the HC on gradient norm and Hessian norm of the learned value function Pan et al . ( 2020 ) is supported by some suggestive theoretical evidence , but it suffers from great computation cost and zero gradient due to the high order differentiation ( i.e. , ∇s||∇sv ( s ) || ) as suggested by the authors . This paper will introduce our novel HC search-control method motivated by overcoming the limitations of the prioritized ER method , which has stronger theoretical support than the work by Pan et al . ( 2019 ) and improved computational cost comparing with the existed work by Pan et al . ( 2020 ) . 3 A DEEPER LOOK AT ERROR-BASED PRIORITIZED SAMPLING . In this section , we provide theoretical motivation for error-based prioritized sampling . We show that prioritized sampling can be reformulated as optimizing a cubic power objective with uniform sampling . We prove that optimizing the cubic objective provides a faster convergence rate during early learning . Based on these results , we highlight that prioritized ER has two limitations 1 ) outdated priorities and 2 ) insufficient coverage of the sample space . This motivates our method in the next section to address the two limitations . 3.1 PRIORITIZED SAMPLING AS A CUBIC OBJECTIVE . In the l2 regression , we minimize the mean squared error minθ 12n ∑n i=1 ( fθ ( xi ) − yi ) 2 , for training set T = { ( xi , yi ) } ni=1 and function approximator fθ , such as a neural network . In error-based prioritized sampling , we define the priority of a sample ( x , y ) ∈ T as |fθ ( x ) − y| ; the probability of drawing a sample ( x , y ) ∈ T is typically q ( x , y ; θ ) ∝ |fθ ( x ) − y| . We employ the following form to compute the probabilities : q ( x , y ; θ ) def = |fθ ( x ) − y|∑n i=1 |fθ ( xi ) − yi| ( 1 ) We can show an equivalence between the gradients of the squared objective with this prioritization and the cubic power objective 13n ∑n i=1 |fθ ( xi ) − yi|3 . See Appendix A.3 for the proof . Theorem 1 . For a constant c determined by θ , T , we have E ( x , y ) ∼uniform ( T ) [ ∇θ ( 1/3 ) |fθ ( x ) − y|3 ] = cE ( x , y ) ∼q ( x , y ; θ ) [ ∇θ ( 1/2 ) ( fθ ( x ) − y ) 2 ] This simple theorem provides an intuitive reason for why prioritized sampling can help improve sample efficiency : the gradient direction of cubic function is sharper than that of the square function when the error is relatively large ( Figure 1 ) . Theorem 2 further characterizes the difference between the convergence rates by optimizing the mean square error and the cubic power objective , providing a solid motivation for using error-based prioritized sampling . Theorem 2 ( Fast early learning ) . Consider the following two objectives : ` 2 ( x , y ) def = 12 ( x− y ) 2 , and ` 3 ( x , y ) def = 13 |x− y| 3 . Denote δt def = |xt − y| , and δ̃t def = |x̃t − y| . Define the functional gradient flow updates on these two objectives : dxt dt = −η d { 12 ( xt − y ) 2 } dxt , dx̃t dt = −η d { 13 |x̃t − y| 3 } dx̃t . ( 2 ) Given error threshold ≥ 0 , define the hitting time t def = mint { t : δt ≤ } and t̃ def = mint { t : δ̃t ≤ } . For any initial function value x0 s.t . δ0 > 1 , ∃ 0 ∈ ( 0 , 1 ) such that ∀ > 0 , t ≥ t̃ .1 Proof . Please see Appendix A.4 . Given the same and the same initial value of x , first we derive t = 1 η · ln { δ0 } , t̃ = 1 η · ( 1 − 1 δ0 ) . Then we analyze the condition on to see when t ≥ t̃ , i.e . minimizing the square error is slower than minimizing the cubic error . The above theorem says that when the initial error is relatively large , it is faster to get to a certain low error point with the cubic objective . We can test this in simulation , with the following minimization problems : minx≥0 x2 and minx≥0 x3 . We use the hitting time formulae t = 1η · ln { δ0 } , t̃ = 1 η · ( 1 − 1 δ0 ) derived in the proof , to compute the hitting time ratio t t̃ under different initial values x0 and final error value . In Figure 1 ( c ) ( d ) , we can see that it usually takes a significantly shorter time for the cubic loss to reach a certain xt with various x0 values . 1Finding the exact value of 0 would require a definition of ordering on complex plane , which leads to 0 = − 1W ( log 1/a−1/a−πi ) and W ( · ) is a Wright Omega function , then we have t̃ ≤ t . Our theorem statement is sufficient for the purpose of characterizing convergence rate . Implications of the above theory . The equivalence from Theorem 1 inspires us to identify two limitations of the current prioritized ER method : 1 ) The equivalence requires the priorities of all samples to get updated after the training parameters get updated at each time step . 2 ) The equivalence requires the prioritized sampling distribution to be calculated on the whole training set ; in an online RL setting , at the current time step t , we only have visited samples . These visited samples provide a biased training set w.r.t . current policy which likely does not reasonably cover the state space . We will present our approach to overcome the limitations in Section 4 . In the next section , we will empirically verify our theoretical findings .

Summary: The paper proposes a new way of prioritization in experience replay and Dyna-style planning methods. In particular, it proposes to exploit a learned model to actively search for states with high expected errors. The states are then prioritized proportional to the expected errors. The authors motivate the approach by a theoretical/empirical observation that the prioritized optimization of L2 loss is equivalent to the direct optimization of cubic loss. More specifically they tackle (1) outdated priorities of training samples, (2) insufficient coverage of the sample space; which are identified as the main shortcomings of previously explored prioritization methods.