Plan2Vec: Unsupervised Representation Learning by Latent Plans

1 Introduction . Unsupervised representation learning is often motivated by the goal of reducing human involvement in the learning loop , such that an algorithm can learn directly from streams of unlabeled data . Much focus has been placed on an algorithm ’ s transfer performance across supervised learning tasks . However , for control tasks , where an agent is planning on the learned representation space , performance is often poor ( Watter et al. , 2015 ) . A quick glance at the representation learned by a variational auto-encoder ( VAE ) reveals that the learned embedding often contains local patches that are quite reasonable by themselves , but the global structure of the learned embedding is often “ crumpled ” , such that Euclidean lines between points that are sufficiently far-apart either cross domain-boundaries , or otherwise detach from the support of the learned manifold ( see Fig.4 ) . This observation implies that although the VAE objective encourages embeddings that behave well locally , memorizing individual image samples by reconstruction is insufficient to attain meaningful global structure . This raises an interesting and important question : if rote memorization is insufficient , what else do we need in order to build an agent that can make good plans ? Various works in learning representations for planning attempt to address this issue by learning locally constrained generative models ( Banijamali et al. , 2017 ; Watter et al. , 2015 ; Kurutach et al. , 2018 ) . One line of work , motivated by controlling complex dynamic systems directly from high-dimensional input , attempts to learn generative models that explicitly impose a reduced local linearity constraint on the learned dynamic manifold ( Watter et al. , 2015 ; Banijamali et al. , 2017 ) . Such methods contain three major shortcomings . First , some of these formulations rely on learning a forward model , which can not be applied to datasets where action data is unavailable or ill-defined . Second , these generative models rely heavily on the inductive prior within image generation , that images nearby in pixel space are semantically similar , which limits the applicability of these methods to domains where visual similarities map well to the conceptual space . Finally , the linear constraint and the optimization objective are both local , yet making plans involves non-local concepts of distances and direction . How to learn from streams of observation data to attain a cognitive map of the problem domain without relying on the image similarity priors provided by generative models remains an open problem . In this work , we pose the problem of unsupervised learning a plannable representation as learning a map of the domain without access to the underlying sampling process and the environment . Such a map has two main properties : First , the map informs conceptual distance between any pair of observations , beyond the typical limit of a short spatiotemporal window ( Perozzi et al. , 2014 ; Caron et al. , 2018 ) . Second , this map has to be consistent with a local metric1 , which is easily attained via self-supervision . Motivated by this problem , we propose Plan2Vec , a method for unsupervised representation learning that incorporates planning as part of the learning objective . The technical challenges of this work are threefold : First , the standard formulation of reinforcement learning requires substantial human supervision in the form of meticulously shaped , dense rewards . Different tasks usually require different reward functions , making it difficult to scale across multiple tasks . The second issue is that reinforcement learning is active , as it requires access to an environment between optimization phases to receive trajectories in order to learn . Third , to plan on a continuous state and action space , one usually needs to learn a closed-form behavior policy that outputs actions , or a forward model of the environment with actions as an input . The main contributions of this paper is to overcome all three of these problems by formulating the problem of learning the global structure of a data manifold as a goal-directed reinforcement learning task . To solve the issue of offering a reward , we train a local metric function from local context without supervision , and use it as a sparse reward for reaching the goal with hind-sight relabeling . To address the necessity of active RL and extending RL to a passive setting , we remove the need for either action data , or a model of the world , by planning entirely in the latent configuration space on a graph . Goal-directed planning can be done with a 1-step greedy policy in the learned representation space . We call our method Plan2Vec , for learning an embedding space for planning . To help illustrate our method , we lead the introduction of Plan2Vec with a set of toy tasks on simulated navigation domains , and show visualization of the components and learned manifold . We then evaluate Plan2Vec under two challenging task settings : First , we show that we can learn representations on deformable objects such as a piece of rope , which is otherwise hard to model . Moreover , we show qualitative results on visual plans between pairs of rope configurations that are randomly selected from the dataset . Second , we tackle real-world navigation on StreetLearn ( Mirowski et al. , 2018 ) , where we learn to embed a map directly from videos of a car driving through the streets , with no access to the ground-truth GPS location data . We show quantitatively that under a constrained planning computation budget , the embedding that Plan2Vec learns using a globally consistent planning objective outperforms baselines that only plan with the local metric . 2 Related Works . The work most similar to ours from the manifold learning community is DeepWalk ( Perozzi et al. , 2014 ) . DeepWalk aims to embed a social graph by randomly sampling trajectories , then use skip-gram ( Mikolov et al. , 2013 ) to embed each graph node contrastively from its contrast . This is related to the contrastive learning objective we use to train our local metric function . Despite of this , the random walk DeepWalk employs to sample those trajectories is limited in terms of distance of travel . As a result it falls under the category of representation learning algorithms that only learn from a localized context . Similarly there is a strong connection between our value-iteration learning objective and the diffusion map literature . In diffusion map , the distance on the learned manifold measures the “ diffusion distance ” between two points on a graph G under a markovian transition kernel ( Socher & Hein , 2008 ) . One can consider value-iteration as a non-parametric version of diffusion 1This can be relaxed into a pseudo metric , allowing different images to have zero conceptual distances in-between . This does not affect the applicability of our approach . map using neural networks for the kernel . A critical difference is that the transition kernel in diffusion maps is not condition on a goal whereas the policy does . As a consequence , the diffusion distance fall-off exponentially as the number of steps increases , just like with DeepWalk . Locally linear embeddings ( LLE ) could be considered a “ stronger ” version of skip-gram , where linear contributions of each neighbor is preserved . However , LLE enforces global structure , and prevent volume collapse via addition of a global volume regularization term . This is similar to the variational prior in a variational auto-encoder ( VAE ) in that both lack meaningful alignment with planning semantics . Recent work in “ robust features ” point to a connection between the injected noise and the alignment between the input and output manifoldsIlyas et al . ( 2019 ) that might be an interesting direction to explore . Other works look at planning over a graph with various assumptions ( Savinov et al. , 2018 ; Zhang et al. , 2018 ; Eysenbach et al. , 2019 ) . With semi-parametric topological memory ( SPTM ) Savinov et al . ( 2018 ) the focus is on solving navigation instead of learning a vector embedding . All of these methods use Dijkstra ’ s shortest path first ( SPF ) search ( Dijkstra , 1959 ) on the graph to plan , resulting in a worst case bound of O ( |E| + |V | log |V | ) in computation time where |E| is the number of edges on the graph and |V | the number of vertices . In our experiment , we show quantitatively that the local embedding SPTM uses to make plans is insufficient if the hard-coded planner has a restricted planning budget , whereas the globally consistent representation that Plan2Vec learns via value iteration still plans well . Our result illustrates the importance for an agent to acquire such a global view of the domain and use it as heuristic for planning . There are also gradient-based planning methods that require expert trajectories , which is a strong assumption that we do not require . Universal Planning Networks and Distributional Planning Networks ( Srinivas et al. , 2018 ; Yu et al. , 2019 ) rely on supervised learning to get to the reward through a differentiable forward model , trained end-to-end by grounding through expert actions . Gupta et al . ( 2019 ) ; Tamar et al . ( 2016 ) assume that the feature vectors live on a 2D grid world with known environment dynamics and well-defined local connectivity between states . There have also been Laplacian methods to learn representations efficient for planning ( Wu et al. , 2019 ) , using a spectral graph objective for approximating the Laplacian eigenfunctions and using L2 distance in this space as a dense reward for goal-directed tasks . This learns a similar type of space that is a compression of the original graph , but assumes the original graph is fully given through acting in the environment as opposed to a disjoint set of trajectories which are connected via a learned local metric . Embed to control ( E2C ) , RCE , L-SBMP and causal InfoGAN ( Watter et al. , 2015 ; Banijamali et al. , 2017 ; Ichter & Pavone , 2018 ; Kurutach et al. , 2018 ) are a line of generative model that explicitly incorporate forward modeling in the latent space . They show that the learned representation is plannable , without directly incorporating a planner as part of their learning objective . Our goal is drastically different – Plan2Vec learns a representation by planning , as opposed to just showing one can plan with a learned representation . Plan2Vec explicitly acquires the concept of “ reachability ” conditioned on an optimal policy as part of the representation . This results in a semantically meaningful and locally consistent global structure . Another branch of work coming from the reinforcement learning community are selfsupervised or task-agnostic RL ( Florensa et al. , 2019 ; Kahn et al. , 2017 ; Pong et al. , 2019 ) . These work aim to reduce the amount of human involvement in designing reinforcement learning algorithms for individual tasks . Plan2Vec is distinguished from these proposals in that we do not aim to learn a policy distribution π ( a|o ) . Instead , we want to learn generalizable representations of the environments that makes learning such a low-level policy , or running classical control algorithms more efficient . By abstracting away the actions , Plan2Vec is able to plan over much longer horizons , as demonstrated in Sec . 5 . 3 Technical Background . We now overview methods that learn a local metric between pairs of images that are closeby , and proximal dynamic programming under a standard Markov decision process ( MDP ) formalism ( Sutton & Barto , 1998 ) . Learning a Local Metric . Intuitively , a metric is a bivariate function that gives a measure of similarity between two points . Formally , fa , b∼D : ( a , b ) → R+ is a symmetric , real-valued , and positive-definite function over its domain D×D . When distance labels are available one can learn such a function via supervised learning . In reality , however , we often need to work with sequential datasets without access to a sampling policy that is jointly optimized , in which case one can not assume long-horizon optimality in the sequences we want to learn from . As a result , the distance information between frames of observations is only good up to a limited temporal window , beyond which noise dominates . In language modeling and unsupervised representation learning domains , it is often easy to construct positive and negative examples , and pose a binary classification objective as a Noise-Contrastive Density Estimator ( NCE ) ( Gutmann & Hyvärinen , 2010 ) , LNCE = − log f ( xi , c ) ∑ x∼X f ( x , c ) , ( 1 ) where f is a convex function proportional to the density p ( x , c ) . Minimizing the NCE loss can be mapped to maximizing a lower-bound on the mutual information between the latent code c and the data distribution X ( Oord et al. , 2018 ; Hjelm et al. , 2018 ) , I ( X , c ) = E [ log P ( X|C ) P ( X ) ] ≥ log ( N ) − LNCE . ( 2 ) Rather than directly learning a representation this way ( Sermanet et al. , 2017 ) , Plan2Vec extends the standard binary NCE objective to learn a local metric function , and uses it as a reward function . Universal Value Function Approximator as a Metric . We formulate value iteration under the Markov decision process ( MDP ) formalism ( Bellman , 1957 ) . The MDP is parameterized by the tuple ⟨S , A , P , r⟩ . S and A are the sets of states and actions , P ( s′|s , a ) is the transition model of the environment , and r ( s , a , s′ ) is the reward function . An agent is represented by its policy distribution π ( a|s ) . The state value function Vπ : S → R represents the expected sum of discounted future rewards for being at state s , conditioned on the reward r and the policy π . In sample-based value iteration with neural networks , we can learn the value function by minimizing the empirical Bellman-residual δ = ∥∥V ( s ; θ ) − B∗πV ∥∥ , ( 3 ) where the Bellman optimality operator is defined as B∗πV = max at [ R ( st , at , st+1 ) + γ ∑ st+1 P ( st+1|st , at ) V ( st+1 ; θ ) ] . ( 4 ) Universal Value Function Approximators ( UVFAs ) ( Schaul et al. , 2015 ) extend this taskspecific reward to learn a “ universal ” value function by generalizing to all goals g ∈ S. The reward now conditions on the goal r ( s , a , s′ , g ) . Assuming that the goals are uniformly sampled from S and the value function is symmetric , UVFA becomes a metric on S up to a correction constant . If we further assume that the MDP is deterministic , the sample-based Bellman residual can be reduced to V ( s , g ; θ ) ← r ( s , a , s′ , g ) + γV ( s′ , g ; θ ) , ( 5 ) which we use to learn our latent space , as detailed in Sec . 4.2 . 4 Learning Representations by Latent Planning Our goal is to learn a representation that goes beyond rote memorization of the dataset , which consists of disjoint temporal sequences . Critically , we want the structure of the embedding to capture the global topology of the dataset , such that for any observation o in the domain , we can make useful inference with respect to another sample ogoal , no matter how far away ogoal is . Having access to such a global metric , ∀ o , ogoal pairs , would enable effective planning on non-trivial , high-dimensional , and/or complex topologies that are otherwise prohibitively slow . In this section we give a high-level overview of Plan2Vec , with implementation details available in Appendix ? ? . We first define local connectivity between states through a local metric function trained in a self-supervised manner with a contrastive loss . Different from Watter et al . ( 2015 ) ; Banijamali et al . ( 2017 ) and similar to Kurutach et al . ( 2018 ) , our method does not rely on dynamics of the underlying environment in the form of sampled action data , and neither do we learned a forward model , which distracts from long-range planning . Instead , this task occurs on a graph where disjoint temporal sequences are connected by the local metric function . We learn a representation space that encodes paths on the graph such that a 1-step greedy policy on the learned representation is equivalent to a path finding algorithm on the graph .

Summary: While cast slightly differently in the intro, it seems to me that this paper learns a goal-conditioned value function that is used at test time to construct visual plans by selecting an appropriate sequence of data points from the training data. Similar to prior work they learn a local distance metric without supervision using a temporal proximity objective and then construct a graph of the training data points using this metric. The main novelty that this paper introduces seems to be the idea to distill the results of planning algorithms run at training time into a global, goal-conditioned value function, which allows to reduce the required planning time at test time. The authors perform experiments on constructing visual plans for a simulated toy navigation task, a robotic rope manipulation task and the StreetLearn navigation task. The paper reports favorable results under a time-constrained test setting but does not include strong baselines that were designed for this setting.