Topological Experience Replay

1 INTRODUCTION . A significant challenge in reinforcement learning ( RL ) is to overcome the need for large amounts of data . Off-policy algorithms have received great attention because of its ability to reuse data by experience replay ( Lin , 1992 ) for learning the policy . The key ingredient in off-policy methods is learning a Q-function ( Watkins & Dayan , 1992 ) that to predict the expected future sum of rewards ( i.e. , return ) at a state . In the recent advances of deep RL , the Q-function is trained by performing bootstrapping over states sampled from the experience replay buffer . In bootstrapping , the sum of the immediate reward and the succeeding state ’ s Q-value are backed up to the current state ’ s Qvalue . This learning procedure divides the value estimation into several sub-problems where each state ’ s Q-value depends only on its successor states ’ Q-value . Because of the dependency on the successor states ’ Q-value , the ordering of states used for iterative updates substantially influences the convergence speed of Q-value . An improper update order results in slow convergence . As a motivating example , consider the Markov Decision Process ( MDP ) shown in Figure 1 . Let the agent receive a positive reward when it reaches the goal state ( labeled as G ) , but zero at any other state ( note that our method does not need rewards to be 0 or 1 ) . Starting from state C , the agent can obtain the reward by visiting states in the sequence C → D → G. Now if the Q-value of taking action a at state D , Q ( D , a ) , is inaccurate , then using it to update Q ( C , a ) will lead to erroneous estimation of Q ( C , a ) . It is therefore natural to start at the terminal state G and first bootstrap Qvalues of preceding states ( D , E ) , and then for states ( A , B , C ) . Such a reverse bootstrapping order ensures that for each state , the dependent succeeding states ’ Q-values are updated before used for updating the current state ’ s value . In fact , it has been proven that for acyclic MDPs such a reverse sweep is the optimal order for performing bootstrapping ( Bertsekas et al. , 2000 ) . Other works ( Grześ & Hoey , 2013 ; Dai & Hansen , 2007 ; Dai et al. , 2011 ) have further empirically demonstrated the effectiveness of reverse sweep in cyclical MDPs . However , it is challenging to perform Q-learning with reverse sweep in high-dimensional state spaces since the predecessors of each state are often unknown . State-of-the-art Q-learning methods ( Mnih et al. , 2015 ; Lillicrap et al. , 2015 ; Haarnoja et al. , 2018 ; Fujimoto et al. , 2018 ) resort to random sampling of data from the replay buffer . Their speed of convergence can be slow as these methods disregard the importance of the update order . The speed of convergence matters because data collection is interleaved withQ-value updates . If theQ-values are incorrect , the agent may take actions that result in low rewards . Such data is not useful for improving the Q-values of states in the high-return trajectories . Therefore , the slower convergence of Q-values is directly linked to the data inefficiency . One way to reduce dependence on data is to artificially increase the speed of convergence by increasing the ratio between the Q-learning update steps and the data collection steps . However , when function approximators are used to estimate the Q-value , excessive updates on a fixed set of interaction data can lead to over-fitting and overestimation of the Q-values , resulting in worse overall performance . Our experiments in Section 5.4 empirically confirm this hypothesis . To speed up Q-function convergence in high-dimensional state spaces , we approximate the reverse sweep update scheme via stitching the trajectories stored in the replay buffer to construct a graph . The graph is built directly from high-dimensional observations ( e.g. , images ) . Each state is a vertex , and two vertices are connected with an edge if the agent transitions between them during training time . As a result , trajectories are joined when a common state appears in two different rollouts . In our framework , graph building and exploration proceed iteratively . The resulting graph provides information about the predecessor of each state , which is used to determine the update order of the Q-function based on the reverse sweep principle . The reverse sweep is initiated from a set of terminal states because bootstrapping is not required to determine the correct Q-value for these states . We call this method Topological Experience Replay ( TER ) because the update order of Qvalues is based on the topology of the state space . 2 RELATED WORKS . Experience replay Devising an experience replay strategy to improve Q-learning is an active field in deep RL . Prioritized Experience Replay ( PER ) ( Schaul et al. , 2015 ) prioritize states by the bootstrapping errors ( also called as temporal difference ( TD ) error ) ( Schaul et al. , 2015 ) . Since during training the TD error depends on a ( possibly ) erroneous estimate of the Q-values , PER can lead to slow convergence speed . Episodic Backward Update ( EBU ) ( Lee et al. , 2019 ) and Lin ( 1992 ) replays interactions in trajectories in the reverse order of state visitations in every episode . However , for instance , in Figure 1 , moving backward in a trajectory ( A , E , B , D , G ) misses the other predecessors of D ( i.e. , C ) while the Q-value of D affects that of C. DisCor ( Kumar et al. , 2020 ) re-weights the error of state transitions to approximate the optimal bootstrapping order while the conservative weighting scheme of DisCor impedes learning progress . Different from the above works , we reorder states for updates based on the states ’ Q-value dependency . Reverse sweep Similar ideas of reverse sweep were studied in different scenarios . Assuming the access to the simulators , Florensa et al . ( 2017 ) trained RL agents from initial states that are closer to the goal and are gradually chosen to be farther away . Moore & Atkeson ( 1993 ) fitted an environment dynamics model for generating the predecessors of the states of high bootstrapping errors successively backward to update the Q-values . Goyal et al . ( 2019 ) extended this algorithm by a learned generative model to perform backward rollouts from high-value states . Schroecker et al . ( 2019 ) trained an imitation learning model to match the expert ’ s predecessor state distribution generated by a backtracking model . Instead of augmenting the training dataset by manipulating simulators or generating synthetic data , our method rearranges the existing dataset to aid training . Graph for RL Prior works organizing experience into a graph focused on temporally abstracted planning . Savinov et al . ( 2018 ) learn to build a graph for navigation from images . Eysenbach et al . ( 2019 ) ; Zhang et al . ( 2021 ) ; Emmons et al . ( 2020 ) provide RL agents with sub-goals by planning over graph . Zhu et al . ( 2020 ) regularizes the Q-learning training via an additional graph Q-function defined over graph . Instead of taking advantages of temporal abstraction from graph , we use a graph to maintain states ’ dependency for improving Q-function learning . Model-based RL Graph is similar with the world model used in model-based RL , but TER uses it for a different purpose . The classic work ( Sutton , 1991 ) maintains the state transition probabilities of each state and sample states by these transition probabilities to update the Q-values . Recent model-based methods ( Hafner et al. , 2019 ; Nagabandi et al. , 2018 ; Deisenroth & Rasmussen , 2011 ) typically learn a world model ( e.g. , neural network or Gaussian process ) for generating “ imaginary ” data for training the RL agent or planning . Imaginary data augments the replay buffer , improving the agent ’ s performance by training it with data absent in the replay buffer . Instead of generating synthetic state transitions , TER employs the graph to find a better update order of each state transition in the existing replay buffer for training the value function . In other words , model-based RL augments the dataset while TER sorts the dataset . 3 PRELIMINARIES . We consider reinforcement learning ( RL ) ( Sutton & Barto , 2018 ) in an episodic discrete-time Markov decision process ( MDP ) . The objective of RL is to find the optimal policy , π∗ , that maximizes expected return E [ ∑T−1 τ=t γ τ−trτ |st = s ] ∀s ∈ S , where γ is a discount factor ( Sutton & Barto , 2018 ) , the S represents the set of all states in the MDP , and τ denotes the time step . At a time step τ , the agent takes the action aτ = π ( sτ ) , receives reward rτ = ( sτ , aτ , sτ+1 ) , and transitions to the next state sτ+1 , where R is the reward function . At the time τ of task completion ( e.g. , reaching goal states ) the episode termination indicator E ( sτ ) = 1 , otherwise E ( sτ ) = 0 . Q-learning is a popular algorithm for finding the optimal policy . It learns the Q function Q ( s , a ) = E [ ∑T−1 τ=t γ τ−trτ |st = s , at = a ] using bootstrapping operations : Q ( s , a ) ← R ( s , a , s′ ) + γmaxa′ Q ( s ′ , a′ ) where s′ is the state encountered on executing action a in the state s. The policy π ( s ) can be easily derived as : π ( s ) : = arg maxaQ ( s , a ) . When S is high-dimensional , the Q function is usually represented by a deep neural network ( Mnih et al. , 2015 ) . The interaction data collected by the agent is stored in an experience replay buffer ( Lin , 1992 ) in the form of state transitions ( st , at , rt , st+1 ) . The Q function is updated using stochastic gradient descent on batches of data randomly sampled from the replay buffer . 4 METHOD . Our method , topological experience replay ( TER ) , is central at performing reverse sweep to update the Q-function in high-dimensional state spaces . Reverse sweep is known to accelerate Q-function convergence yet requires the knowledge of predecessors of a state , which is often unknown in highdimensional state spaces . We overcome this limitation by building a graph from the replay buffer and guide the Q-function update with a reverse sweep over the graph . Section 4.1 presents the overview of our proposed algorithm . In Section 4.2 , we describe the procedure for building the graph from high-dimensional states . Next , in Section 4.3 , we describe how the graph is used to determine the update order for Q-learning . Finally , we present a batch mixing technique in Section 4.4 that guarantees TER to converge . 4.1 TOPOLOGICAL EXPERIENCE REPLAY . Assumption Similar to prior works of reverse sweep in tabular domains ( Dai & Hansen , 2007 ) , we assume that the objective of the MDPs is to reach terminal states ( i.e. , goal states ) . When an agent successfully completes the task , the agent must be at some terminal states . This assumption is required , otherwise reverse sweep will not help the agent learn optimal Q-value faster . Another assumption is that an agent must be able to visit the same state twice so that we can find the joint state between two episodes for building the graph . This is an admissible assumption since Zhu et al . ( 2020 ) shows that an agent can visit repeated states even in high-dimensional image state spaces Bellemare et al . ( 2013 ) . Overview TER iteratively interleaves between graph building and reverse sweep online during training . The agent collects experience at each step and periodically updates the Q-function . Each new transition ( s , a , r , s′ ) is mapped to an edge on the graph . When updating the Q-function , TER performs a few steps of reverse sweep for making a batch of training samples . The reverse sweep starts from the vertices associated with terminal states . Each step in the sweep gathers transitions stored at an edge for training the Q-function . The batch mixing technique mixes the transitions that are fetched during reverse sweep and randomly sampled from the replay buffer , ensuring those transitions disconnected from terminal states to be updated . The pseudo-code in Algorithm 1 details the whole procedure .

Summary: This paper proposes a novel experience replay technique motivated by the topology of the collected data. The technique changes the order in which states are sampled for a minibatch Q-function update, sampling states backwards in a breadth-first way through a trajectory graph constructed from the replay buffer. This means that terminal states, which are more likely (by assumption) to receive reward, and don't require bootstrapping for value estimation, are updated first allowing reward from these states to propagate backwards to earlier states quicker. This method shows improved performance on several MiniGrid and Sokoban tasks, outperforming other baseline experience replay methods (PER, EBU, DisCor). They perform extensive ablations on the difference between TER (their method) and EBU, showing why TER outperforms EBU.