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: The paper proposes a new sampling strategy to improve the sample efficiency of Q-learning based methods. Bootstrapping being a fundamental characteristic of TD methods, the accuracy of the next state's value is critical in updating the current state. Previous work preform a similar observation, leading to methods such as backward value iteration. However, performing reverse sweeps is challenging in high dimensional tasks. As such, the authors propose to iteratively construct a graph representing the underlying MDP on which reverse sampling is executed. The authors report considerable improvement on grid-like tasks when compared to baselines that propose sampling schemes. The authors also provide numerous additional experiments attempting to elucidate the merits and limitations of the approach.