Learning World Graph Decompositions To Accelerate Reinforcement Learning

1 INTRODUCTION . Many real-world applications , e.g. , self-driving cars and in-home robotics , require an autonomous agent to execute different tasks within a single environment that features , e.g . high-dimensional state space , complex world dynamics or structured layouts . In these settings , model-free reinforcement learning ( RL ) agents often struggle to learn efficiently , requiring a large amount of experience collections to converge to optimal behaviors . Intuitively , an agent could learn more efficiently by focusing its exploration in task-relevant regions , if it has knowledge of the high-level structure of the environment . We propose a method to 1 ) learn and 2 ) use an environment decomposition in the form of a world graph , a task-agnostic abstraction . World graph nodes are waypoint states , a set of salient states that can summarize agent trajectories and provide meaningful starting points for efficient exploration ( Chatzigiorgaki & Skodras , 2009 ; Jayaraman et al. , 2018 ; Ghosh et al. , 2018 ) . The directed and weighted world graph edges characterize feasible traversals among the waypoints . To leverage the world graph , we model hierarchical RL ( HRL ) agents where a high-level policy chooses a waypoint state as a goal to guide exploration towards task-relevant regions , and a low-level policy strives to reach the chosen goals . Our framework consists of two phases . In the task-agnostic phase , we obtain world graphs by training a recurrent variational auto-encoder ( VAE ) ( Chung et al. , 2015 ; Gregor et al. , 2015 ; Kingma & Welling , 2013 ) with binary latent variables ( Nalisnick & Smyth , 2016 ) over trajectories collected using a random walk policy ( Ha & Schmidhuber , 2018 ) and a curiosity-driven goal-conditioned policy ( Ghosh et al. , 2018 ; Nair et al. , 2018 ) . World graph nodes are states that are most frequently selected by the binary latent variables , while edges are inferred from empirical transition statistics between neighboring waypoints . In the task-specific phase , taking advantage of the learned world graph for structured exploration , we efficiently train an HRL model ( Taylor & Stone , 2009 ) . In summary , our main contributions are : • A task-agnostic unsupervised approach to learn world graphs , using a recurrent VAE with binary latent variables and a curiosity-driven goal-conditioned policy . • An HRL scheme for the task-specific phase that features multi-goal selection ( Wide-thenNarrow ) and navigation via world graph traversal . • Empirical evaluations on multiple tasks in complex 2D grid worlds to validate that our framework produces descriptive world graphs and significantly improves both sample efficiency and final performance on these tasks over baselines , especially thanks to transfer learning from the unsupervised phase and world graph traversal . 2 RELATED WORK . An understanding of the environment and its dynamics is essential for effective planning and control in model-based RL . For example , a robotics agent often locates or navigates by interpreting a map ( Lowry et al. , 2015 ; Thrun , 1998 ; Angeli et al. , 2008 ) . Our exploration strategy draws inspiration from active localization , where robots are actively guided to investigate unfamiliar regions ( Fox et al. , 1998 ; Li et al. , 2016 ) . Besides mapping , recent works ( Azar et al. , 2019 ; Ha & Schmidhuber , 2018 ; Guo et al. , 2018 ) learn to represent the world with generative latent states ( Tian & Gong , 2017 ; Haarnoja et al. , 2018 ; Racanière et al. , 2017 ) . If the latent dynamics are also extrapolated , the latent states can assist planning ( Mnih et al. , 2016a ; Hafner et al. , 2018 ) or model-based RL ( Gregor & Besse , 2018 ; Kaiser et al. , 2019 ) . While also aiming to model the world , we approach this as abstracting both the structure and dynamics of the environment in a graph representation , where nodes are states from the environment and edges encode actionable efficient transitions between nodes . Existing works ( Metzen , 2013 ; Mannor et al. , 2004 ; Eysenbach et al. , 2019 ; Entezari et al. , 2010 ) have shown benefits of such graph abstractions but typically select nodes only subject to a good coverage the observed state space . Instead , we identify a parsimonious subset of states that can summarize trajectories and provide more useful intermediate landmarks , i.e . waypoints , for navigating complex environments . Our method for estimating waypoint states can be viewed as performing automatic ( sub ) goal discovery . Subgoal and subpolicy learning are two major approaches to identify a set of temporally-extended actions , “ skills ” , that allow agents to efficiently learn to solve complex tasks . Subpolicy learning identifies policies useful to solve RL tasks , such as option-based methods ( Daniel et al. , 2016 ; Bacon et al. , 2017 ) and subtask segmentations ( Pertsch et al. , 2019 ; Kipf et al. , 2018 ) . Subgoal learning , on the other hand , identifies “ important states ” to reach ( Şimşek et al. , 2005 ) . Previous works consider various definitions of “ important ” states : frequently visited states during successful task completions ( Digney , 1998 ; McGovern & Barto , 2001 ) , states introducing the most novel information ( Goyal et al. , 2019 ) , bottleneck states connecting densely-populated regions ( Chen et al. , 2007 ; Şimşek et al. , 2005 ) , or environment-specific heuristics ( Ecoffet et al. , 2019 ) . Our work draws intuition from unsupervised temporal segmentation ( Chatzigiorgaki & Skodras , 2009 ; Jayaraman et al. , 2018 ) and imitation learning ( Abbeel & Ng , 2004 ; Hussein et al. , 2017 ) . We define “ important ” states ( waypoints ) as the most critical states in recovering action sequences generated by some agent , which indicates that these states contain the richest information about the executed policy ( Azar et al. , 2019 ) . 3 LEARNING WORLD GRAPHS . We propose a method for learning a world graph Gw , a task-agnostic abstraction of an environment that captures its high-level structure and dynamics . In this work , the primary use of world graphs is to accelerate reinforcement learning of downstream tasks . The nodes of Gw , denoted by a set of waypoints states sp ∈ Vp , are generically “ important ” for accomplishing tasks within the environment , and therefore useful as starting points for exploration . Our method identifies such waypoint states from interactions with the environment . In addition , we embed feasible transitions between nearby waypoint states as the edges of Gw . In this work , we define important states in the context of learning Gw ( see Section 2 for alternative definitions ) . That is , we wish to discover a small set of states that , when used as world graph nodes , concisely summarize the structure and dynamics of the environment . Below , we describe 1 ) how to collect state-action trajectories and an unsupervised learning objective to identify world graph nodes , and 2 ) how the graph ’ s edges ( i.e. , how to transition between nodes ) are formed from trajectories . 3.1 WAYPOINT STATE IDENTIFICATION . The structure and dynamics of an environment are implicit in the state-action trajectories observed during exploration . To identify world graph nodes from such data , we train a recurrent variational autoencoder ( VAE ) that , given a sequence of state-action pairs , identifies a subset of the states in the sequence from which the full action sequence can be reconstructed ( Figure 2 ) . In particular , the VAE infers binary latent variables that controls whether each state in the sequence is used by the generative decoder , i.e. , whether a state is “ important ” or not . Binary Latent VAE The VAE consists of an inference , a generative and a prior network . These are structured as follows : the input to the inference network qφ is a trajectory of state-action pairs observed from the environment τ= { ( st , at ) } Tt=0 , with s= { st } Tt=0 and a= { at } Tt=0 denoting the state and action sequences respectively . The output of the inference network is the approximated posterior over a sequence z= { zt } Tt=0 of binary latent variables , denoted as qφ ( z|a , s ) . The generative network pθ computes a distribution over the full action sequence a using the masked state sequence , where st is masked if zt=0 ( we fix z0=zT=1 during training ) , denoted as pθ ( a|s , z ) . Finally , a state-conditioned pψ ( zt|st ) given by the prior network pψ for each st encodes the empirical average probability that state st is activated for reconstruction . This choice encourages inference to select within a consistent subset of states for use in action reconstruction . In particular , the waypoint Algorithm 1 : Identifying waypoint states Vp and learning a goal-conditioned policy πg Result : Waypoint states Vp and a goal-conditioned policy πg Initialize network parameters for the recurrent variational inference model V Initialize network parameters for the goal-conditioned policy πg Initialize Vp with the initial position of the agent , i.e . Vp = { s0 = ( 1 , 1 ) } while VAE reconstruction error has not converged do for n← 1 to N do Sample random waypoint sp ∈ Vp Navigate agent to sp and perform T -step rollout using a randow walk policy : τ rn ← { ( s0 = sp , a0 ) , ... , ( sT , aT ) } gn ← sT Navigate agent to sp and perform T -step rollout using πg with goal gn : τπn ← { ( s0 = sp , a0 ) , ... , ( sT , aT ) } at∼πg ( ·|st , gn ) Re-label πg rewards with action reconstruction error as curiosity bonus : rπn ← { 1st+1=gn − λ · pθ ( at|s , z ) } Tt=0 end Perform policy gradient update of πg using τπ and rπ Update V using τ r and τπ Update Vp as set of states with largest prior mean αsαs+βs . end states Vp are chosen as the states with the largest prior means and during training , once every few iterations , Vp is updated based on the current prior network . Objective Formally , we optimize the VAE using the following evidence lower bound ( ELBO ) : ELBO = Eqφ ( z|a , s ) [ log pθ ( a|s , z ) ] −DKL ( qφ ( z|a , s ) |pψ ( z|s ) ) . ( 1 ) To ensure differentiablity , we apply a continuous relaxation over the discrete zt . We use the Beta distribution pψ ( zt ) = Beta ( αt , βt ) for the prior and the Hard Kumaraswamy distribution qψ ( zt|a , z ) = HardKuma ( α̃t , β̃t ) for the approximate posterior , which resembles the Beta distribution but is outside the exponential family ( Bastings et al. , 2019 ) . This choice allows us to sample 0s and 1s without sacrificing differentiability , accomplished via the stretch-and-rectify procedure ( Bastings et al. , 2019 ; Louizos et al. , 2017 ) and the reparametrization trick ( Kingma & Welling , 2013 ) . Lastly , to prevent the trivial solution of using all states for reconstruction , we use a secondary objective L0 to regularize the L0 norm of z at a targeted value µ0 ( Louizos et al. , 2017 ; Bastings et al. , 2019 ) , the desired number of selected states out of T steps , e.g . for when T = 25 , we set µ0 = 5 , meaning ideally 5 out of 25 states are activated for action reconstruction . Another term LT to encourage temporal separation between selected states by targeting the number of 0/1 switches among z at 2µ0 : L0 = ∥∥Eqφ ( z|s , a ) [ ‖z‖0 ] − µ0∥∥2 , LT = ∥∥∥∥∥Eqφ ( z|s , a ) [ T∑ t=0 1 [ zt 6= zt+1 ] ] − 2µ0 ∥∥∥∥∥ 2 . ( 2 ) See Appendix A for details on training the VAE with binary zt , including integration of the Hard Kumaraswamy distribution and how to regularize the statistics of z .

Summary: This paper proposes a novel approach to hierarchical reinforcement learning approach by first learning a graph decomposition of the state space through a recurrent VAE and then use the learned graph to efficiently explore the environment. The algorithm is separated into 2 stages where in the first stage random walk and goal conditioned policy is used to explore the environment and simultaneous use a recurrent binary VAE to compress the trajectory. The inference network is given the observation and action and the reconstruction is to, given the hidden state or hidden state+observation, reconstruct the action taken. The approximate posterior takes on the form of a hard Kumaraswamy distribution which can differentiably approximate a binary variable; when the approximate posterior is 0, the decoder must reconstruct the action using the hidden state alone. The nodes of the world graph are roughly states that are used to reconstruct the trajectories in the environment. After the graph is constructed, the agent can use a combination of high-level policy and classical planning to solve tasks with sparse reward.