Discrete Predictive Representation for Long-horizon Planning

1 INTRODUCTION . In future , we hope that robots will be able to operate in unstructured environments such as homes and hospitals , and endowed with long-horizon planning ability . Despite successes in deep reinforcement learning ( RL ) from raw observations , much progress relies on the availability of shaped reward to guide the learning ( Ng et al. , 1999 ; Mirza et al. , 2020 ) . On the other hand , over past decades , task and motion planning has been shown to solve much longer-horizon goal-directed tasks such as making a cup of coffee from torque control ( Kaelbling & Lozano-Pérez , 2011 ; Srivastava et al. , 2014 ; Toussaint , 2015 ; Wang et al. , 2018 ) . However , these methods often require pre-specified discrete abstract states , task representations and transition models , e.g. , whether the robot is holding a cup and what actions ( or perturbations ) change such an abstract state . In this paper , we aim to learn discrete representations for high-level abstract planning from video interaction data , combined with a learned short-horizon controller . Learning discrete representations from unsupervised data for planning is challenging for two reasons . First , the relationship between the optimization objective and the true task objective is not well-defined . Second , optimizing a model with a discrete layer is difficult with standard deep learning techniques . Recent methods approach the first problem using reconstruction or constrastive objectives ( Watter et al. , 2015 ; Anand et al. , 2019 ; Ha & Schmidhuber , 2018 ; Kurutach et al. , 2018 ; Hafner et al. , 2019 ; Srinivas et al. , 2020 ) ; however , the learned representations are in continuous latent space which is unstructured and difficult to combine with high-level abstract planning . While other methods show promise in learning discrete representations , they have not been applied to temporally-extended RL tasks ( Oord et al. , 2018 ; Razavi et al. , 2019 ; Risi & Stanley , 2019 ; Asai & Fukunaga , 2017 ; Stratos & Wiseman , 2020 ) . In this work , we propose Discrete Object-factorized Representations for Planning ( DORP ) – a novel framework for visual planning and control by learning discrete representations and a low-level controller . DORP learns discrete representations from images that change slowly overtime , such as whether or not the agent holds a key or which room the agent is in , along with a low-level predictive model for control . These slow features enable the agent to plan at a low frequency in longer-horizon tasks . More specifically , DORP represents an abstract state as a set of one-hot vectors , and optimizes its encoder by maximizing a mutual information lower bound between the current representations to future observations ( Oord et al. , 2018 ) . In order to train through the discrete layer , we apply the Gumbel-Softmax reparametrization trick ( Jang et al. , 2016 ; Maddison et al. , 2016 ) . Using abstract states as nodes , we build an approximate feasibility graph based on observed transition data . When provided with new start and goal images , the agent plan the shortest abstract path . Using the next abstraction as waypoint , model-predictive control maximizes the objective that is 1 if it reaches the target abstraction and 0 otherwise with a trained video prediction model . Unlike other subgoal planning work ( Savinov et al. , 2018 ; Nasiriany et al. , 2019 ; Laskin et al. , 2020 ; Liu et al. , 2020 ) , by following abstract waypoints , DORP avoids unneccessary steps to match exact waypoint states . In a set of experiments , we demonstrate that DORP learns temporally-consistent and objectfactorized representations suitable for planning . We show that these representations enable DORP to handle unseen long-horizon tasks more successfully compared to the states-of-the-arts in visual planning . Interestingly , we observe that latent representations show object-level factorization such as key-and-door . 2 PRELIMINARIES . We present background material on unsupervised representation learning and discrete optimization that our method builds on , and our problem formulation . Contrastive Predictive Coding ( CPC ) ( Oord et al. , 2018 ) learns low-dimensional representations that are most predictive of the future future high-dimensional sequential data . A non-linear encoder qθ : O → Rl , parametrized by θ , encodes the observation ot ∈ O to a latent l-dimensional vector representation zt . Let ’ s define a similarity score fk ( zt , ot+k ) = exp ( ztψqθ ( ot+k ) ) where ψ is a trainable l-by-l similarity matrix and ot+k is a future observation k steps ahead of ot . Given the query observation ot , we aim to classify the key observations – ot+k as positive and other sample õ from the dataset as negative . Formally , we optimize the loss function LCPC = −Eot , ot+k [ log fk ( zt , ot+k ) − log ∑ oj∈X fk ( zt , oj ) ] with respect to θ and ψ . This also corresponds to maximizing a lowerbound of the mutual information between the latent representation zt and the future observation ot+k . Gumbel-Softmax ( GS ) ( Jang et al. , 2016 ) is a discrete optimization technique to compute the gradient of through samples from a categorical distribution πi . GS first applies a reparametrization trick to rewrite samples zi as onehot ( argmaxj ( gj + logπi ) ) where gj are i.i.d samples from Gumbel ( 0 , 1 ) ( Gumbel , 1935 ) . GS approximates argmax with softmax making the discrete stochastic layer differentiable , i.e. , zi = softmax ( ( gi + log πi ) /τ ) . When the temperature parameter τ approaches 0 , the samples converge to the true categorical distribution . Empirically , the temperature τ starts high and is annealed closer to 0 by a certain schedule . 2.1 PROBLEM STATEMENT : GOAL-DIRECTED VISUAL PLANNING AND CONTROL . We define an unknown , fully-observable , stochastic dynamical system f which maps input observation ot ∈ O and action at ∈ A to the next observation ot+1 . Under this dynamical system , we assume a simple exploration policy πrand which can collect data that characterizes the dynamics of the system . This is known as self-supervised data or play data ( Agrawal et al. , 2016 ; Lynch et al. , 2020 ) . We consider high-dimensional observation such as images . Our goal is to learn discrete abstraction zt ∈ Z given observation ot , an abstract feasibility model , and a low-level local controller . At test time , given start os and goal og observations , we can plan a sequence of abstract states zs , z1 , ... , zg where zs and zg are the representation of os and og , and apply the low-level controller to reach these abstract waypoints and finally to og . 3 DISCRETE OBJECT-FACTORIZED REPRESENTATION FOR PLANNING . Our objective is to derive representation properties such as object-factorization and temporalconsistency from an unsupervised learning objective to facilitate long-horizon planning . When object representations are factorized , the agent can escape the need of combinatorial data configurations and can quickly generalize to unseen tasks . Also , the connectivity graph can also be memory-efficiently represented and combined with a more powerful planning algorithm . Another important property for a representation is temporal consistency , meaning any two state observations in the same abstraction should be reachable from one-another by a short sequence of actions . When this property holds in the latent space , a high-level plan can be successfully executed by a low-level controller . We propose Discrete Object-factorized Representation for Planning ( DORP ) to tackle goal-directed visual planning and control problem . DORP learns latent discrete representations , a connectivity graph for abstract planning , and a low-level model-predictive controller . We describe each subcomponent in Section 3.1 - 3.3 . 3.1 DISCRETE REPRESENTATIONS . Despite successes in applying its unsupervisedly learned representations for downstream tasks , CPC is not endowed with an inductive bias for its representations to facilitate planning ( Oord et al. , 2018 ) . Additionally , it is unclear how a discrete representations can be extracted and optimized . Toward this goal , we propose an object-factorized architecture with asymmetric Gumbel softmax for discrete optimization , and a diagonally dominating and diagonally positive similarity matrix to encourage temporal consistency . We implement our DORP architecture to encourage object-factorization and capture discrete representations . The extractor passes an anchor observation through a convolutional neural network and outputs a c-channel feature map . Each feature map is inputted to a shared encoder followed by a Gumbel softmax . For a positive and negative pair of images , however , we opt for a continuous embedding by swapping Gumbel softmax for softmax . Empirically , without this asymmetry trick , the optimization tends to converge to a poor local optima . Together these one-hots are flattened into long vectors , and their bilinear product is the similarity score between the query-key pair ( see Figure 2 ) . We propose an inductive bias in the similarity score function to encourage temporal consistency . Instead of a fully trainable matrix , we parameterize the similarity matrix ψ to be diagonally dominant and diagonally positive ( DDDP ) . This property biases the similarity score to be high when the key embedding z̃t+k is close to the query embedding zt and maximized when z̃t+k = zt . In other words , this incentivizes the representations of positive pairs to share weight matrix as eψ0 ( −σ ( ψ1 ) + αIl×l ) where Il×l is an identity matrix , α is a positive constant , σ is a sigmoid function , ψ0 is a trainable scalar , and ψ1 is a trainable l-by-l matrix . 3.2 ABSTRACT PLANNING . A critical component in planning is a connectivity graph that decides whether two discrete representations are connected . One naive solution is to build such a graph from data – all observed representations as nodes and all observed transitions as edges . However as the number of independent entities increases , the amount of data required to generate the nodes and edges in order to cover the state space increases exponentially . To solve this , we propose to reduce the planning problem by exploiting factorization ( see Figure 3 ) . 1 . First , embed the current image os and goal og as zs = qθ ( os ) and zg = qθ ( og ) . 2 . Pick a random one-hot index i s.t . the current state ’ s one-hot zis differs from the goal ’ s one-hot zig . 3 . Build a graph based on how this one-hot transits in the data ( ignoring all other one-hots ) . 4 . Plan a path zi1 , ... , z i g from z i s to z i g . If this does not exists , then we fail the task . 5 . Execute the low-level controller following the next target zik , k = 1 , · · · , g , while keeping other zj the same . 6 . If it doesn ’ t reach zik , remove the edge from the graph and redo the planning in step 4 . 7 . If it succeeds , follow the next target until it reaches zig . 8 . Repeat step 2 until all the one-hots match . 9 . Finally , execute the low-level controller to the goal image og .1 Our proposed method assumes that each latent one-hot can be manipulated independently to solve the task . This assumption can hold various settings . For example , imagine a task that requires reorganizing furniture in a spacious room , and each objects is represented as a separate one-hot . Here , we can manipulate each object separately , and since there is plenty of empty space in this scenario , we can easily plan around objects to avoid collisions . This assumption does not always hold ; along the same scenario , when the space is more cluttered , objects are more likely to collide , and the proposed fully-factorized planning may fail even though a solution exists ( incomplete ) . We extend this planning algorithm by building a graph on a random set of one-hot indices at a time in step 2 . When can choose this set to contain all the one- 1An implicit assumption of temporal consistency in the learned discrete representations is required . hots the method is equivalent to the full graph search ( complete but expensive ) . Thus , our extended planning method can trade between completeness and efficiency .

Summary: This paper tackles the problem of long horizon visual planning, with the aim of of being able to plan actions to reach distant goals. This is a well studied problem, and like prior work this method considers the setting where the agent is given an offline dataset of interaction, which it learns from to be able to reach new goals (specified by a goal image). The method first does unsupervised representation learning, where it learns a discrete representation of images (using contrastive predictive coding (CPC) with a discrete latent variable using Gumbel Softmax). Using a set of discrete latent variable, it then builds a set of graphs which connects these discrete states based on the collected experience, and derives a planning procedure to find a path in the graph which reaches the goal.