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 proposes a new method for learning graph representations of RL problems. The proposed method, named TOMA, creates a set of *landmarks*, each of which is a representatives state associated with a cluster of states in the original problem.  The clustering is induced via a locality sensitive embedding function. With this embedding, a graph of landmarks is created by sampling state trajectories, and then mapping each state to existing vertices (if a nearby vertex already exist), or creating a new vertex otherwise. To add robustness, an ensemble of embedding functions is used rather than a single estimator. Once the graph is build, the shortest-path to goals can be computed using a search algorithm (e.g., Dijkstra's algorithm). The proposed approach is evaluated on a set of benchmarks, and compares favorably in this problem with recent state-of-the-art graph learning algorithms for RL.