TOMA: Topological Map Abstraction for Reinforcement Learning

1 INTRODUCTION . Animals are able to discover topological map ( graph ) of surrounding environment ( O ’ Keefe and Dostrovsky , 1971 ; Moser et al. , 2008 ) , which will be used as hints for navigation . For example , previous maze experiments on rats ( O ’ Keefe and Dostrovsky , 1971 ) reveal that rats can create mental representation of the maze and use such representation to reach the food placed in the maze . In cognitive science society , researchers summarize these discoveries in cognitive map theory ( Tolman , 1948 ) , which states that animals can extract and code the structure of environment in a compact and abstract map representation . Inspired by such biological phenomenon , researchers have proposed to generate topological graph representation for Markov decision process ( MDP ) and use such graphs for planning in reinforcement learning ( RL ) . Early graph generation methods ( Mannor et al. , 2004 ) are usually prior-based , which apply some human prior to aggregate similar states to generate vertices . Recently , researchers propose some learning-based graph generation algorithms which learn such state aggregation automatically . Such methods have been proved to be better than prior-based methods ( Metzen , 2013 ) . These methods generally treat the states in a replay buffer as vertices . For the edges of the graphs , some methods like SPTM ( Savinov et al. , 2018 ) train a reachability predictor via self-supervised learning and combine it with human experience to construct the edges . Other methods like SoRB ( Eysenbach et al. , 2019 ) exploit a goal-conditioned agent to estimate the distance between vertices , based on which edges are constructed . These existing methods suffer from the following drawbacks . Firstly , these methods do not learn an abstraction for graphs and usually consider all the states in the buffer as vertices ( Savinov et al. , 2018 ) , which results in high memory and computation cost . This drawback also makes generated graph non-robust , which will degrade the planning performance . Secondly , existing methods can not be used for facilitating exploration , which is important in RL . In particular , methods like SPTM rely on human sampled trajectories to generate the graph , which is infeasible in RL exploration . Methods like SoRB require training another goal-conditioned agent . Such training procedure assumes knowledge of the environment since it requires to generate several goal-reaching tasks to train the agent . This practice is also intractable in RL exploration . In this paper , we propose a new method , called TOpological Map Abstraction ( TOMA ) , for learningbased graph generation . The main contributions of this paper are outlined as follows : • TOMA can learn to generate an abstract graph representation for MDP . Different from existing methods in which each vertex of the graph represents a state , each vertex in TOMA represents a cluster of states . As a result , compared with existing methods TOMA has much less memory and computation cost , and can generate more robust graph for planning . • TOMA can be used to facilitate exploration . In particular , we propose planning to explore , in which TOMA is used to accelerate exploration by guiding the agent towards unexplored states . A novel experience replay module called vertex memory is also proposed to improve exploration performance . • Experimental results show that TOMA can robustly generate abstract graph representation on several 2D world environments with different types of observation and can outperform previous learning-based graph generation methods to achieve the state-of-the-art performance . 2 ALGORITHM . 2.1 NOTATIONS . In this paper , we model a RL problem as a Markov decision process ( MDP ) . A MDP is a tuple M ( S , A , R , γ , P ) , where S is the state space , A is the action space , R : S × A → R is a reward function , γ is a discount factor and P ( st+1|st , at ) is the transition dynamic . ρ ( x , y ) = ‖x − y‖2 denotes Euclidean distance . G ( V , E ) denotes a graph , where V is its vertex set and E is its edge set . For any set X , we define its indicator function 1X ( x ) as follows : 1X ( x ) = 1 if x ∈ X , 1X ( x ) = 0 if x /∈ X . 2.2 TOMA . Figure 1 gives an illustration of TOMA , which tries to map states to an abstract graph . A landmark set L is a subset of S and each landmark li ∈ L is a one-to-one correspondence to a vertex vi in the graph . Each li and vi will represent a cluster of states . In order to decide which vertex a state s ∈ S corresponds to , we first use a locality sensitive embedding function φθ to calculate its latent representation z = φθ ( s ) in the embedding space Z . Then if z ’ s nearest neighbor in the embedded landmark set φθ ( L ) = { φθ ( l ) | l ∈ L } is φθ ( li ) , we will map s to vertex vi ∈ V . 2.2.1 LOCALITY SENSITIVE EMBEDDING . A locality sensitive embedding is a local distance preserving mapping φθ from state space S to an embedding space Z , which is an Euclidean space Rn in our implementation . Given a trajectory T = ( s0 , a0 , s1 , a1 , .... sn ) , we can use dij = |j − i|/r to estimate the distance between si and sj . Here r is a radius hyper-parameter to re-scale the distance and we will further explain its meaning later . In practice , however , dij is a noizy estimation for shortest distance and approximating it directly won ’ t converge in most cases . Hence , we propose to estimate which interval the real distance lies in . First , we define three indicator functions : χ1 ( x ) = 1 [ 0,1 ] ( x ) , ( 1 ) χ2 ( x ) = 1 ( 1,3 ] ( x ) , ( 2 ) χ3 ( x ) = 1 ( 3 , +∞ ) ( x ) , ( 3 ) which mark three disjoint regions [ 0 , 1 ] , ( 1 , 3 ] , ( 3 , +∞ ) , respectively . Then we define an anti-bump function ξa , b ( x ) = Relu ( −x+ a ) + Relu ( x− b ) . Here , Relu ( x ) = max ( 0 , x ) is the rectified linear unit ( ReLU ) function ( Glorot et al. , 2011 ) . With this ξ ( x ) , we can measure the deviation from the above intervals . Let L1 ( x ) = ξ−∞,1 ( x ) = Relu ( x− 1 ) , ( 4 ) L2 ( x ) = ξ1,3 ( x ) = Relu ( −x+ 1 ) + Relu ( x− 3 ) , ( 5 ) L3 ( x ) = ξ3 , +∞ ( x ) = Relu ( −x+ 3 ) , ( 6 ) and let d′ij = ρ ( φθ ( si ) , φθ ( sj ) ) denote the distance between si and sj in the embedding space . Our embedding loss is defined as L ( θ ) = E ( si , sj ) ∼Ps ( χ1 ( dij ) L1 ( d′ij ) + λ1χ2 ( dij ) L2 ( d′ij ) + λ2χ3 ( dij ) L3 ( d′ij ) ) . ( 7 ) Here Ps is a sample distribution which will be described later , λ1 and λ2 are two hyper-parameters to balance the importance of the estimation for different distances . We find that a good choice is to pick λ1 = 0.5 , λ2 = 0.2 to ensure that our model focuses on the terms with lower variance . In this equation , there are some critical components to notice : Radius r As we will see later , the hyper-parameter r will determine the granularity of each graph vertex , which we term as radius . If we define the k-ball neighborhood of s ∈ S to be Bk ( s ) = { s′ ∈ S|ρ ( φθ ( s′ ) , φθ ( s ) ) < k } , ( 8 ) then Bk ( s ) will cover more states when r is larger . During the graph generation process , we will remove redundant vertices by checking whether B1 ( li ) and B1 ( lj ) intersect too much . Re-scaling by r makes it easier to train the embedding function . Sample Distribution Ps The state pair ( si , sj ) in the loss function is sampled from a neighborhood biased distribution Ps . We will sample ( si , sj ) ( i < j ) with probability α , if j− i ≤ 4r . And we will sample ( si , sj ) ( i < j ) with probability 1− α , if j − i > 4r . We simply take α = 0.5 and the choice of α is not sensitive in our experiment . In the implementation , we use this sample distribution to draw samples from trajectory and put them into a replay pool . Then we train the embedding function by uniformly drawing samples from the pool . Anti-Bump Functions The idea of anti-bump function is inspired by the partition of unity theorem in differential topology ( Hirsch , 1997 ) , where a bunch of bump functions are used to glue the local charts of manifold together so as to derive global properties of differential manifolds . In proofs of many differential topology theorems , one crucial step is to use bump function to segregate each local chart into three disjoint regions by radius 1 , 2 and 3 , which is analogous to our method . The loss function is crucial in our method , as in experiment we find that training won ’ t converge if we replace this loss function with a commonly used L2 loss . 2.2.2 DYNAMIC GRAPH GENERATION . An abstract graph representation G ( V , E ) should satisfy the following basic requirements : • Simple : For any vi , vj ∈ V , if vi 6= vj , B1 ( li ) ∩ B1 ( lj ) should not contain too many elements . • Accurate : For any vi , vj ∈ V and vi 6= vj , 〈vi , vj〉 ∈ E if and only if the agent can travel from some s ∈ B1 ( li ) to some s′ ∈ B1 ( lj ) in a small number of steps . • Abundant : ∪i : vi∈VB1 ( li ) should cover the states as many as possible . • Dynamic : G grows dynamically , by absorbing topology information of novel states . In the following content , we show a dynamic graph generation method fulfilling such requirements . First , we introduce the basic operations in our generation procedure . The operations can be reduced to the following three categories : Initializing [ I1 : Initialize ] If V = ∅ , we will pick a landmark from currently sampled trajectories and add a vertex into V accordingly . In our implementation , this landmark is the initial state of the agent . Adding [ A1 : Add Labels ] For each state s on a trajectory , we label it with its corresponding graph vertex . Let i = argminj : vj∈V ρ ( φθ ( s ) , φθ ( lj ) ) and d = ρ ( φθ ( s ) , φθ ( li ) ) . There are three possible cases : ( 1 ) d ∈ [ 0 , 1.5 ] . We label s with vi . ( 2 ) d ∈ [ 2 , 3 ] . We consider s as an appropriate landmark candidate . Therefore , we label s with NULL but add it to a candidate queue . ( 3 ) Otherwise , s is simply labelled with NULL . [ A2 : Add Vertices ] We move some states from the candidate queue into the landmark set and update V accordingly . Once a state is added to the landmark set , we will relabel it from NULL to its vertex identifier . [ A3 Add Edges ] Let the labelled trajectory to be ( vi0 , vi1 , ... , vin ) . If we find vik and vik+1 are different vertices in the existing graph , we will add an edge 〈vik , vik+1〉 into the graph . Checking [ C1 : Check Vertices ] If ρ ( φθ ( li ) , φθ ( lj ) ) < 1.5 , then we will merge vi and vj . [ C2 : Check Edges ] For any edge 〈vi , vj〉 , if ρ ( φθ ( li ) , φθ ( lj ) ) > 3 , we will remove this edge . For efficient nearest neighbor search , we use Kd-tree ( Bentley , 1975 ) to manage the vertices . Based on the above operations , we can get our graph generation algorithm TOMA which is summarized in Algorithm 1 . Algorithm 1 Topological Map Abstraction ( TOMA ) 1 : Pool P ← ∅ . Vertex set V ← ∅ . Edge set E ← ∅ . Graph G ( V , E ) . 2 : for t = 1 , 2 , ... do 3 : Sample a trajectory T using some policy π or by random . 4 : Sample state pairs from T using distribution Ps and put them to P . 5 : Training the embedding function φθ using samples from P . 6 : Initialize G using ( I1 ) if it ’ s empty . 7 : Add vertices and edges using ( A1 ) to ( A3 ) . 8 : Check the graph using ( C1 ) to ( C2 ) . 9 : end for

Summary: This paper studies nonparametric planning in RL, building on prior work in two ways. First, the proposed method creates graph vertices corresponding to *clusters* of observations, rather than single observations. The motivation for this decision is that it leads to sparser graphs, making planning easier. The second idea is to use the graph for exploration by sampling goal states that are near the "fringe" of the graph. Experimental results suggest that the first idea increases performance (success rate and wall-clock time) on 2D navigation tasks. Additional experiments show that the second idea improves exploration.