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 presents a method that combines learning discrete representations together with planning using graph search to solve long horizon tasks from vision. The approach works by generating data via random exploration and trains a representation encoder on this data. This network extracts objects from the observations and passes them through a CNN and shared encoder to generate one-hot encodings; these encodings are concatenated to generate the discrete representation. The encoder is trained via a contrastive learning objective with a similarity matrix that encourages nearby states to share similar encodings, thereby encouraging spatial and temporal abstractions. Next, these representations are combined together with an abstract planner to generate a sequence of waypoints to the goal. This is done by creating a graph of transitions from the collected exploratory data and search within this graph — this is executed for each encoding at a time with the assumption that the task can be solved by moving each object independently of the others. Finally, a low-level controller is used for reaching the waypoints and final goal, this is done via MPC on an action-conditional predictive model that generates future observations given the current state and goal. The approach is tested on two simple planar planning tasks where the agent has to solve k-object arrangement and open a room with a key respectively.