Model-augmented Prioritized Experience Replay

1 INTRODUCTION . Experience replay ( Lin , 1992 ; Mnih et al. , 2015 ) , which provides experiences that different policies may collect , is an essential component of policy training in reinforcement learning ( RL ) . By utilizing many past experiences in a large buffer , experience replay can stabilize and improve policy training . As a result , experience replay with the buffer has significantly contributed to the success of RL on various tasks ( Fujimoto et al. , 2018 ; Haarnoja et al. , 2018a ; b ; Mnih et al. , 2015 ) . Due to its importance , various methods for calculating priority scores of experiences in the buffer have been proposed to sample experiences efficiently ( Schaul et al. , 2016 ; Zha et al. , 2019 ; Sinha et al. , 2020 ) . Since well-trained critic networks lead to effective policy learning , one of the most popular methods is to utilize predefined metrics for prioritizing experiences based on the temporal difference error ( TD-error ) , directly related to the loss of critic networks ( Schaul et al. , 2016 ) . The prioritized experience replay based on TD-error ( PER ) has indeed proved its effectiveness in Q-learning ( Hessel et al. , 2018 ; Schaul et al. , 2016 ) . However , measuring Q-values requires to predict the expectation of returns , which can be obtained after multi-steps , so learning to predict Q-values generally needs a lot of interactions with an environment . Due to the difficulty of the multi-step estimation , sampling based only on high TD-errors may be often ineffective or even degrade the sample-efficiency of the RL framework on some tasks , compared to the uniform sampling , e.g. , see Figure 6 in ( Hessel et al. , 2018 ) . Moreover , Zha et al . ( 2019 ) ; Wang & Ross ( 2019 ) ; Sinha et al . ( 2020 ) show that the effectiveness of PER is questionable under some settings other than Q-learning , e.g. , policy-based methods . In particular , Zha et al . ( 2019 ) reported that transitions with low TD-errors could be suitable for training a policy for some tasks . To overcome this issue , several works attempt to design learning-based prioritizing strategies ( Zha et al. , 2019 ; Sinha et al. , 2020 ; Oh et al. , 2021 ) utilizing various features , including transition ’ s elements and quantities driven by critic networks . Here , our motivating observation is that all learnable features used in the existing schemes are from critic networks . For instance , although Zha et al . ( 2019 ) utilizes multiple features , e.g. , reward , TD-error , and timestep , TD-error is only learnable and is obtained from critic networks . In this paper , we aim for developing a new prioritizing strategy by employing learnable auxiliary features outside critic networks , but can support TD-errors for efficient critics training . To this end , we consider components in Model-based RL ( MbRL ) . To be specific , we first modify the critic network by additionally predicting the reward and the transition with shared weights , which we call Model-augmented Critic Network ( MaCN ) . Then , we propose a modified experience replay scheme , coined Model-augmented Prioritized Experience Replay ( MaPER ) , which prioritizes the past experiences having both high model estimation errors and the TD-errors ( see Figure 1. ( a ) ) . Namely , MaPER encourages MaCN to predict both the Q-value and the model ( i.e. , reward and transition ) well . Here , our intuition is two-fold : • We expect that prioritizing samples additionally with high model estimation errors helps decrease the model errors at the early stage of training , even if MaCN ’ s estimation of the Q-value is poor . • Once MaCN has small model errors , MaPER prefers samples of high TD-error . Here , the learned representation to predict the model well is also useful to predict Q-value , as Q-value is decided by the model ( Sutton & Barto , 2018 ) . One may suggest to train MaCN without using MaPER , but this is not effective ( see Figure 4 ) . This is because MaPER brings a type of curriculum learning effect ( Bengio et al. , 2009 ) to learn Q-value well by MaCN ; it gradually increases the complexity of samples for training , from high model estimation errors to high TD-errors . Indeed , Figure 1 . ( b ) shows that TD-errors start to abruptly reduce ( around 0.05 × 106 steps ) after model errors sufficiently reduce . Figure 1 . ( b ) - ( c ) also shows that MaPER improves both the decay rate of TD-errors and the convergence rate of Q-values to returns . To summarize , our contributions are as follows : • We propose Model-augmented Prioritized Experience Replay ( MaPER ) , which computes the priority scores of the past experiences additionally considering the model estimation error . MaPER helps reduce TD-errors much faster than learning the Q-values alone , as shown in Figure 1. ( b ) - ( c ) . This results in large improvement in the return with the comparable number of parameters of the original critic networks . • MaPER can be seamlessly integrated into any modern off-policy RL frameworks with critic networks , including both MfRL ( e.g . SAC ( Haarnoja et al. , 2018a ) , TD3 ( Fujimoto et al. , 2018 ) and Rainbow ( van Hasselt et al. , 2019 ) ) and MbRL ( e.g . MBPO ( Janner et al. , 2019 ) ) methods with negligible memory or computational overhead . We remark that to the best of our knowledge , none of the existing methods improve RL ’ s experience replay by employing model learning although there exist some relevant works for different purposes , e.g. , to augment rewards for improving the exploration ( Pathak et al. , 2017 ; Shelhamer et al. , 2016 ) or to enhance representation learning for pixel-based observations ( Lin et al. , 2019 ; Agarwal et al. , 2021 ; Stooke et al. , 2021 ; Lee et al. , 2020b ; Schrittwieser et al. , 2020 ; Hessel et al. , 2021 ; Hafner et al. , 2020b ; Srinivas et al. , 2020 ; Guo et al. , 2020 ; Lee et al. , 2019 ) . 2 METHOD . 2.1 PRELIMINARY . In an RL framework , an agent interacts with an environment which consists of reward and transition distributions , R : S × A → ∆ ( R ) and T : S × A → ∆ ( S ) , for state and action spaces S and A , where ∆ ( · ) denotes a set of distributions . The agent selects an action at ∼ π ( ·|st ) by its policy π on the current state st ∈ S to receive the next state st+1 ∼ T ( st , at ) and a reward rt ∼ R ( st , at ) over discrete timesteps t. The agent ’ s objective is to learn a policy π that maximizes Es0∼ρ , π [ ∑∞ k=0 γ krk ] , where ρ is a distribution of initial state s0 at each episode and ∑∞ k=0 γ krt+k is the discounted cumulative rewards with a discount factor γ ∈ [ 0 , 1 ) . Although our method is generally applicable to off-policy RL methods which have Q-networks , i.e. , critics , throughout this section , we focus on off-policy actor-critic RL algorithms which consist of the policy ( i.e. , actor ) πΘ ( a|s ) and critic networks Qθ ( s , a ) with an experience replay buffer B , where Θ and θ are their parameters , respectively . The most commonly used loss for Qθ ( s , a ) ( Mnih et al. , 2015 ; Haarnoja et al. , 2018a ; Fujimoto et al. , 2018 ) is : E ( st , at , rt , st+1 ) ∼B [ ∥δtQ πΘ θ ∥MSE ] , ( 1 ) where ∥·∥MSE is the mean-square error and δtQπΘθ is the Temporal Difference error ( TD-error ) defined as follows : for a given transition ( st , at , rt , st+1 ) , δtQ πΘ θ = δQ πΘ θ ( st , at , rt , st+1 ) = rt + γEa′∼πΘ ( ·|st+1 ) [ Qθ ( st+1 , a ′ ) ] −Qθ ( st , at ) . ( 2 ) We can interpret this value as a measure of how surprising or ‘ unexpected ’ the transition is . 2.2 MODEL-AUGMENTED PRIORITIZED EXPERIENCE REPLAY ( MAPER ) . Model-augmented Critic Network ( MaCN ) . An environment is depicted as reward and transition distributions , R ( s , a ) and T ( s , a ) . We observe that these two distributions and Qθ ( s , a ) have the same input domain , i.e. , S ×A , the product of state and action spaces . Motivated by it , we slightly modify Qθ ( s , a ) into Cθ ( s , a ) to additionally predict the environment by a reward model Rθ ( s , a ) and a transition model Tθ ( s , a ) in parallel via parameter sharing . We refer to Cθ ( s , a ) as a Modelaugmented Critic Network ( MaCN ) such that Cθ = ( Qθ , Rθ , Tθ ) . Here , Qθ measures Q-values with the reward estimationRθ as ( Hessel et al. , 2021 ) to stimulate the exploration ( see Table C.1 in the supplementary material ) . In other words , we formulate the loss for Qθ using the estimated rewards , that is similar to Eq . ( 1 ) : LQθ = E ( st , at , st+1 ) ∼B [ ∥δtQ πΘ θ ∥MSE ] , ( 3 ) where δtQπΘθ = δQ πΘ θ ( st , at , rt , st+1 ) = Qθ ( st , at ) − ( Rθ ( st , at ) + γEa′∼πΘ ( · , st+1 ) [ Qθ ( st+1 , a ′ ) ] ) . ( 4 ) Here , Rθ is detached in the backward propagation of the loss LQθ . Then the loss for Cθ is LCθ = ξ1LQθ + ξ2LRθ + ξ3LTθ , ( 5 ) where LRθ and LT θ are the losses for Rθ and Tθ , respectively , with adaptively changing positive coefficients ξ1 , ξ2 , and ξ3 by employing a dynamic method in ( Liang & Zhang , 2020 ; Liu et al. , 2019b ) . In this paper , we use the following losses forRθ and Tθ : LRθ = E ( st , at , rt , st+1 ) ∼B [ ∥δtRθ∥MSE ] , LT θ = E ( st , at , rt , st+1 ) ∼B [ ∥δtTθ∥MSE ] , Algorithm 1 Model-augmented Prioritized Experience Replay based on Actor-Critic Methods 1 : Initialize the model-augmented critic network ’ s parameters θ , the actor ’ s parameters Θ , a replay buffer B ← ∅ , priority set PB ← ∅ , and the batch size m. 2 : for each timestep t do 3 : Choose at from the actor and collect a transition ( st , at , rt , st+1 ) from the environment . 4 : Update replay buffer B ← B ∪ { ( st , at , rt , st+1 ) } and priority set PB ← PB ∪ { maxB PB } . 5 : for each gradient step do 6 : Sample an index I by the given set PB and Eq . ( 8 ) with |I| = m. 7 : Calculate weights { wi } i∈I by Eq . ( 9 ) 8 : Learn θ by Eq . ( 5 ) and Θ by { Bi } i∈I ⊂ B with corresponding weights { wi } i∈I ( Eq . ( 9 ) ) . 9 : Update a priority set { σi } i∈I by Eq . ( 10 ) 10 : end for 11 : end for where δtRθ = Rθ ( st , at ) − Er∼R ( st , at ) [ r ] , δtTθ = Tθ ( st , at ) − Es∼T ( st , at ) [ s ] . Formulation of MaPER . We first explain the concept of PER ( Schaul et al. , 2016 ) since we will augment it with the model-based components . Let [ n ] be defined as the set { 1 , · · · , n } for a positive integer n. Without loss of generality , we can 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 ) , ( 6 ) with a function κ ( i ) from the index of B to a corresponding timestep . PER calculates each Bi ’ s priority σi as the recently computed TD-error for itself and collects a set of priority scores : PB = { σ1 , · · · , σ|B| } , ( 7 ) where each priority σi is updated whenever the corresponding transition is sampled for training the actor and critic networks . The TD-error ( Eq . ( 2 ) ) is the most frequently used quantity to make the priority set ( Eq . ( 7 ) ) ( Schaul et al. , 2016 ; Brittain et al. , 2019 ; Hessel et al. , 2018 ; van Hasselt et al. , 2019 ) . Then the sampling strategy of PER is to determine an index set I in [ |B| ] from the probability pi of i-th transition defined by the priority set : pi = σαi∑ k∈ [ |B| ] σ α k , ( 8 ) with a hyper-parameter α > 0 . Finally , we calculate the importance-sampling weights as follows : wi = ( 1 |B|pi ) β , ( 9 ) where β > 0 is also a hyper-parameter to compensate the bias of probabilities . We design a new priority score equation with MaCN . Since our method uses Cθ ( s , a ) instead of the original critic network , we modify the rule in obtaining the priority set in Eq . ( 7 ) for experiences in the buffer B accordingly . To this end , we compute the priority of each transition as the sum of the TD-errors and model errors , which is used to compute probabilities and weights for corresponding transitions by Eqs . ( 8 ) - ( 9 ) , as follows : σi = ξ1∥δκ ( i ) QπΘθ ∥MSE + ξ2∥δκ ( i ) Rθ∥MSE + ξ3∥δκ ( i ) Tθ∥MSE , ( 10 ) where κ is the map introduced in Eq . ( 6 ) . By using the priority scoring in Eq . ( 10 ) , the buffer can sample various experiences that are useful to MaCN for both long-term and short-term viewpoints . The detailed descriptions of our framework MaPER are provided in Algorithm 1 . Finally , we remark that the computational overhead by MaPER is negligible , compared to the vanilla PER . Suppose that a critic network has a linear layer with N hidden units as the final layer . Then one can compute the difference in the number of parameters between the MaCN and the original critic network : N ( 1 + |S| ) , where S is a space of ( embedded ) states in critic networks . Notice that the number of additional parameters depends only on the final hidden units in the critic network . Therefore , computing costs for the additional parameters is minor compared to computing costs for both the total parameters in all networks and prioritized experience sampling .

Summary: Authors introduced Model-augmented Prioritized Experience Replay (MaPER) as a variation to the canonical prioritized sweeping. The main difference is that they extended the critic network on its last layer to also predict the reward and transition model. Then they changed the prioritization of sample for replay to also include the error in the reward and transition model, in addition to the TD-error. Experimental results in various MuJoC domains showed the performance boost introduced by MaPER on top of existing state-of-the-art techniques. Furthermore Authors carried an ablation study, investigation the impact of 1) incorporating the MaPER while all methods used the same network 2) incorporating various error values for prioritization, and 3) increasing the network size.