Learning to Sample with Local and Global Contexts  in Experience Replay Buffer

1 INTRODUCTION . Experience replay ( Mnih et al. , 2015 ) , which is a memory that stores the past experiences to reuse them , has become a popular mechanism for reinforcement learning ( RL ) , since it stabilizes training and improves the sample efficiency . The success of various off-policy RL algorithms largely attributes to the use of experience replay ( Fujimoto et al. , 2018 ; Haarnoja et al. , 2018a ; b ; Lillicrap et al. , 2016 ; Mnih et al. , 2015 ) . However , most off-policy RL algorithms usually adopt a unique random sampling ( Fujimoto et al. , 2018 ; Haarnoja et al. , 2018a ; Mnih et al. , 2015 ) , which treats all past experiences equally , so it is questionable whether this simple strategy would always sample the most effective experiences for the agents to learn . Several sampling policies have been proposed to address this issue . One of the popular directions is to develop rule-based methods , which prioritize the experiences with pre-defined metrics ( Isele & Cosgun , 2018 ; Jaderberg et al. , 2016 ; Novati & Koumoutsakos , 2019 ; Schaul et al. , 2016 ) . Notably , since TD-error based sampling has improved the performance of various off-policy RL algorithms ( Hessel et al. , 2018 ; Schaul et al. , 2016 ) by prioritizing more meaningful samples , i.e. , high TD-error , it is one of the most frequently used rule-based methods . Here , TD-error measures how unexpected the returns are from the current value estimates ( Schaul et al. , 2016 ) . However , such rule-based sampling strategies can lead to sampling highly biased experiences . For instance , Figure 1 shows randomly selected 10 transitions among 64 transitions sampled using certain 1Code is available at https : //github.com/youngmin0oh/NERS metrics/rules under a policy-based learning , soft actor critic ( SAC ) ( Haarnoja et al. , 2018a ) , on Pendulum-v0 after 30,000 timesteps , which goal is to balance the pendulum to make it stay in the upright position . We observe that sampling by the TD-error alone mostly selects initial transitions ( see Figure 1 ( a ) ) , where the rods are in the downward position , since it is difficult to estimate Q-value on them . Conversely , the sampled transitions by Q-value describe rods in the upright position ( see Figure 1 ( b ) ) , which will provide high returns to agents . Both can largely contribute to the update of the actor and critic since the advantage term and mean-square of TD-errors are large . Yet , due to the bias , the agent trained in such a manner will mostly learn what to do in a specific state , but will not learn about others that should be experienced for proper learning of the agent . Therefore , such biased ( and redundant ) transitions may not lead to increased sample efficiency , even though each sampled transition may be individually meaningful . On the other hand , focusing only on the diversity of samples also has an issue . For instance , sampling uniformly at random is able to select out diverse transitions including intermediate states such as those in the red boxes of Figure 1 ( c ) , where the rods are in the horizontal positions which are necessary for training the agents as they provide the trajectory between the two types of states . However , the transitions are occasionally irrelevant for training both the policy and the Q networks . Indeed , states in the red boxes of Figure 1 ( c ) possess both lowQ-values and TD-errors . Their low TD-errors suggest that they are not meaningful for the update of Q networks . Similarly , low Q-values can not be used to train the policy what good actions are . Motivated by the aforementioned observations , we aim to develop a method to sample both diverse and meaningful transitions . To cache both of them , it is crucial to measure the relative importance among sampled transitions since the diversity should be considered in them , not all in the buffer . To this end , we propose a novel neural sampling policy , which we refer to Neural Experience Replay Sampler ( NERS ) . Our method learns to measure the relative importance among sampled transitions by extracting local and global contexts from each of them and all sampled ones , respectively . In particular , NERS is designed to take a set of each experience ’ s features as input and compute its outputs in an equivariant manner with respect to the permutation of the set . Here , we consider various features of transition such as TD-error , Q-value and the raw transition , e.g. , expecting to sample intermediate transitions as those in blue boxes of Figure 1 ( c ) ) efficiently . To verify the effectiveness of NERS , we validate the experience replay with various off-policy RL algorithms such as soft actor-critic ( SAC ) ( Haarnoja et al. , 2018a ) and twin delayed deep deterministic ( TD3 ) ( Fujimoto et al. , 2018 ) for continuous control tasks ( Brockman et al. , 2016 ; Todorov et al. , 2012 ) , and Rainbow ( Hessel et al. , 2018 ) for discontinuous control tasks ( Bellemare et al. , 2013 ) . Our experimental results show that NERS consistently ( and often significantly for complex tasks having high-dimensional state and action spaces ) outperforms both the existing the rule-based ( Schaul et al. , 2016 ) and learning-based ( Zha et al. , 2019 ) sampling methods for experience replay . In summary , our contribution is threefold : • To the best of our knowledge , we first investigate the relative importance of sampled transitions for the efficient design of experience replays . • For the purpose , we design a novel permutation-equivariant neural sampling architecture that utilizes contexts from the individual ( local ) and the collective ( global ) transitions with various features to sample not only meaningful but also diverse experiences . • We validate the effectiveness of our neural experience replay on diverse continuous and discrete control tasks with various off-policy RL algorithms , on which it consistently outperforms both existing rule-based and learning-based sampling methods . 2 NEURAL EXPERIENCE REPLAY SAMPLER . We consider a standard reinforcement learning ( RL ) framework , where an agent interacts with an environment over discrete timesteps . Formally , at each timestep t , the agent receives a state st from the environment and selects an action at based on its policy π . Then , the environment returns a reward rt , and the agent transitions to the next state st+1 . The goal of the agent is to learn the policy π that maximizes the return Rt = ∑∞ k=0 γ krt+k , which is the discounted cumulative reward from the timestep t with a discount factor γ ∈ [ 0 , 1 ) , at each state st . Throughout this section , we focus on off-policy actor-critic RL algorithms with a buffer B , which consist of the policy πψ ( a|s ) ( i.e. , actor ) and Q-function Qθ ( s , a ) ( i.e. , critic ) with parameters ψ and θ , respectively . 2.1 OVERVIEW OF NERS . We propose a novel neural sampling policy f with parameter φ , called Neural Experience Replay Sampler ( NERS ) . It is trained for learning to select important transitions from the experience replay buffer for maximizing the actual cumulative rewards . Specifically , at each timestep , NERS receives a set of off-policy transitions ’ features , which are proportionally sampled in the buffer B based on priorities evaluated in previous timesteps . Then it outputs a set of new scores from the set , in order for the priorities to be updated . Further , both the sampled transitions and scores are used to optimize the off-policy policy πψ ( a|s ) and action-value function Qθ ( s , a ) . Note that the output of NERS should be equivariant of the permutation of the set , so we design its neural architecture to satisfy the property . Next , we define the reward rre as the actual performance gain , which is defined as the difference of the expectation of the sum of rewards between the current and previous evaluation policies , respectively . Figure 2 shows an overview of the proposed framework , which learns to sample from the experience replay . In the following section , we describe our method of learning the sampling policy for experience replay and the proposed network architecture in detail . 2.2 DETAILED COMPONENTS OF NERS . Input observations . Throughout this paper , we denote the set { 1 , · · · , n } by [ n ] for positive integer n. Without loss of generality , suppose that the replay buffer B stores the following information as its i-th transition Bi = ( sκ ( i ) , aκ ( i ) , rκ ( i ) , sκ ( i ) +1 ) where κ ( i ) is a function from the index of B to a corresponding timestep . We use a set of priorities PB = { σ1 , · · · , σ|B| } that is updated whenever sampling transitions for training the actor and critic . One can sample an index set I in [ |B| ] with the probability pi of i-th transition as follows : pi = σαi∑ k∈ [ |B| ] σ α k , ( 1 ) Algorithm 1 Training NERS : batch size m and sample size n Initialize NERS parameters φ , a replay buffer B ← ∅ , priority set PB ← ∅ , and index set I ← ∅ for each timestep t do Choose at from the actor and collect a sample ( st , at , rt , st+1 ) from the environment Update replay buffer B ← B ∪ { ( st , at , rt , st+1 ) } and priority set PB ← PB ∪ { 1.0 } for each gradient step do Sample an index I by the given set PB and Eq . ( 1 ) with |I| = m Calculate a score set { σk } k∈I and weights { wi } i∈I by Eq . ( 4 ) and Eq . ( 5 ) , respectively Train the actor and critic using batch { Bi } i∈I ⊂ B and corresponding weights { wi } i∈I Collect I ← I ⋃ I and update PB ( I ) by the score set { σk } k∈I end for for the end of an episode do Choose a subset Itrain from I uniformly at random such that |Itrain| = n Calculate rre as in Eq . ( 6 ) Update sampling policy φ using the gradient ( 7 ) with respect to Itrain Empty I , i.e. , I ← ∅ end for end for with a hyper-parameter α > 0 . Then , we define the following sequence of features for { Bi } i∈I : D ( B , I ) = { sκ ( i ) , aκ ( i ) , rκ ( i ) , sκ ( i ) +1 , κ ( i ) , δκ ( i ) , rκ ( i ) + γmax a Qθ̂ ( sκ ( i ) + a ) } i∈I , ( 2 ) where γ is a discount factor , θ̂ is the target network parameter , and δκ ( i ) is the TD-error defined as follows : δκ ( i ) = rκ ( i ) + γmax a Qθ̂ ( sκ ( i ) +1 , a ) −Qθ ( sκ ( i ) , aκ ( i ) ) . The TD-error indicates how ‘ surprising ’ or ‘ unexpected ’ the transition is ( Schaul et al. , 2016 ) . Note that the input D ( B , I ) contains various features including both exact values ( i.e. , states , actions , rewards , next states , and timesteps ) and predicted values in the long-term perspective ( i.e. , TD-errors and Q-values ) . We abbreviate the notation D ( B , I ) = D ( I ) for simplicity . Utilizing various information is crucial in selecting diverse and important transitions ( see Section 3 ) . Architecture and action spaces . Now we explain the neural network structure of NERS f . Basically , f takes D ( I ) as an input and generate their scores , where these values are used to sample transitions proportionally . Specifically , f consists of fl , fg , and fs called learnable local , global and score networks with output dimensions dl , dg , and 1 . The local network is used to capture attributes in each transition by fl ( D ( I ) ) = { fl,1 ( D ( I ) ) , · · · fl , |I| ( D ( I ) ) } ∈ R|I|×dl , where fl , k ( D ( I ) ) ∈ Rdl ( k ∈ [ |I| ] ) . The global network is used to aggregate collective information of transitions by taking fg avg ( D ( I ) ) = ∑ fg ( D ( I ) ) |I| ∈ R 1×dg , where fg ( D ( I ) ) ∈ R|I|×dg . Then by concatenating them , one can make an input for the score network fs as follows : Dcat ( I ) : = { fl,1 ( D ( I ) ) ⊕ fgavg ( D ( I ) ) , · · · , fl , |I| ( D ( I ) ) ⊕ fgavg ( D ( I ) ) } ∈ R|I|× ( dl+dg ) , ( 3 ) where ⊕ denotes concatenation . Finally , the score network generates a score set : fs ( D cat ( I ) ) = { σi } i∈I ∈ R |I| . ( 4 ) One can easily observe that fs is permutation-equivariant with respect to input D ( I ) . The set { σi } i∈I is used to update the priorities set P for transitions corresponding to I by Eq . ( 1 ) and to compute importance-sampling weights for updating the critic , compensating the bias of probabilities ( Schaul et al. , 2016 ) ) : wi = ( 1 |B|p ( i ) ) β , ( 5 ) where β > 0 is a hyper-parameter . Then the agent and critic receive training batch D ( I ) and corresponding weights { wi } i∈I for training , i.e. , the learning rate for training sample Bi is set to be proportional to wi . Due to this structure satisfying the permutation-equivariant property , one can evaluate the relative importance of each transition by observing not only itself but also other transitions . Reward function and optimizing sampling policy . We update NERS at each evaluation step . To optimize our sampling policy , we define the replay reward rre of the current evaluation as follows : for policies π and π′ used in the current and previous evaluations as in ( Zha et al. , 2019 ) , rre : = Eπ ∑ t∈ { timesteps in an episode } rt − Eπ′ ∑ t∈ { timesteps in an episode } rt . ( 6 ) The replay reward is interpreted as measuring how much actions of the sampling policy help the learning of the agent for each episode . Notice that rre only observes the difference of the mean of cumulative rewards between the current and previous evaluation policies since NERS needs to choose transitions without knowing which samples will be added and how well agents will be trained in the future . To maximize the sample efficiency for learning the agent ’ s policy , we propose to train the sampling policy to selects past transitions in order to maximize rre . To train NERS , one can choose Itrain that is a subset of a index set I for totally sampled transitions in the current episode . Then we use the following formula by REINFORCE ( Williams , 1992 ) : ∇φEItrain [ rre ] = EItrain [ rre ∑ i∈Itrain ∇φ log pi ( D ( Itrain ) ) ] , ( 7 ) where pi is defined in Eq . ( 1 ) . The detailed description is provided in Algorithm 1 . While ERO Zha et al . ( 2019 ) uses a similar replay-reward ( Eq . 6 ) , there are a number of fundamental differences between it and our method . First of all , ERO does not consider the relative importance between the transitions as NERS does , but rather learns an individual sampling rate for each transition . Moreover , they consider only three types of features , namely TD-error , reward , and the timestep , while NERS considers a larger set of features by considering more informative features that are not used by ERO , such as raw features , Q-values , and actions . However , the most important difference between the two is that ERO performs two-stage sampling , where they first sample with the individually learned Bernoulli sampling probability for each transition , and further perform random sampling from the subset of sampled transitions . However , with such a strategy , the first-stage sampling is highly inefficient even with moderate size experience replays , since it should compute the sampling rate for each individual instance . Accordingly , its time complexity of the first-stage sampling depends finally on the capacity of the buffer B , i.e. , O ( |B| ) . On the contrary , NERS uses a sum-tree structure as in ( Schaul et al. , 2016 ) to sample transitions with priorities , so that its time complexity for sampling depends highly on O ( log |B| ) . Secondly , since the number of experiences selected from the first stage sampling is large , it may have little or no effect , making it to behave similarly to random sampling . Moreover , ERO updates its network with the replay reward and experiences that are not sampled from two-stage samplings but sampled by the uniform sampling at random ( see Algorithm 2 in Zha et al . ( 2019 ) ) . In other words , samples that are never selected affect the training of ERO , while NERS updates its network solely based on the transitions that are actually selected by itself .

Summary: Observing that the existed ER-based sampling method may introduce bias or redundancy in sampled transitions, the paper proposes a new sampling method in the ER learning setting. The idea is to take into consideration the context, i.e. many visited transitions, rather than a single one, based on which one can measure the relative importance of each transition. Specifically, the weights of transitions are also learned through a Reinforce agent and hence the sampling distribution is learned to directly improve sample efficiency.