Model-Based Visual Planning with Self-Supervised Functional Distances

1 INTRODUCTION . Designing general-purpose robots that can perform a wide range of tasks remains an open problem in AI and robotics . Reinforcement learning ( RL ) represents a particularly promising tool for learning robotic behaviors when skills can be learned one at a time from user-defined reward functions . However , general-purpose robots will likely require large and diverse repertoires of skills , and learning individual tasks one at a time from manually-specified rewards is onerous and time-consuming . How can we design learning systems that can autonomously acquire general-purpose knowledge that allows them to solve many different downstream tasks ? To address this problem , we must resolve three questions . ( 1 ) How can the robot be commanded to perform specific downstream tasks ? A simple and versatile choice is to define tasks in terms of desired outcomes , such as an example observation of the completed task . ( 2 ) What types of data should this robot learn from ? In settings where modern machine learning attains the best generalization results ( Deng et al. , 2009 ; Rajpurkar et al. , 2016 ; Devlin et al. , 2018 ) , a common theme is that excellent generalization is achieved by learning from large and diverse task-agnostic datasets . In the context of RL , this means we need offline methods that can use all sources of prior data , even in the absence of reward labels . As collecting new experience on a physical robot is often expensive , offline data is often more practical to use in real-world settings ( Levine et al. , 2020 ) . ( 3 ) What should the robot learn from this data to enable goal-reaching ? Similar to prior work ( Botvinick & Weinstein , 2014 ; Watter et al. , 2015 ; Finn & Levine , 2017 ; Ebert et al. , 2018b ) , we note that policies and value functions are specific to a particular task , while a predictive model captures the physics of the environment independently of the task , and thus can be used for solving almost any task . This makes model learning particularly effective for learning from large and diverse datasets , which do not necessarily contain successful behaviors . While model-based approaches have demonstrated promising results , including for vision-based tasks in real-world robotic systems ( Ebert et al. , 2018a ; Finn & Levine , 2017 ) , such methods face two major challenges . First , predictive models on raw images are only effective over short horizons , as uncertainty accumulates far into the future ( Denton & Fergus , 2018 ; Finn et al. , 2016 ; Hafner et al. , 2019b ; Babaeizadeh et al. , 2017 ) . Second , using such models for planning toward goals requires a notion of similarity between images . While prior methods have utilized latent variable models ( Watter et al. , 2015 ; Nair et al. , 2018 ) , ` 2 pixel-space distance ( Nair & Finn , 2020 ) , and other heuristic measures of similarity ( Ebert et al. , 2018b ) , these metrics only capture visual similarity . To enable reliable control with predictive models , we instead need distances that are aware of dynamics . In this paper , we propose Model-Based RL with Offline Learned Distances ( MBOLD ) , which aims to ad- dress both of these challenges by learning predictive models together with image-based distance functions that reflect functionality , from offline , unlabeled data . The learned distance function estimates of the number of steps that the optimal policy would take to transition from one state to another , incorporating not just visual appearance , but also an understanding of dynamics . However , to learn dynamical distances from task-agnostic data , supervised regression will lead to overestimation , since the paths in the data are not all optimal for any task . Instead , we utilize approximate dynamic programming for distance estimation . While prior work has studied such methods to learn goal-conditioned policies in online model-free RL settings ( Eysenbach et al. , 2019 ; Florensa et al. , 2019 ) , we extend it to the offline setting and show that approximate dynamic programming techniques derived from Q-learning style Bellman updates can learn effective shortest path dynamical distances . Although this procedure resembles model-free reinforcement learning , we find empirically that it does not by itself produce useful policies . Instead , our method ( Fig . 1 ) combines the strengths of dynamics models and distance functions , using the predictive model to plan over short horizons , and using the learned distances to provide a global cost that captures progress toward distant goals . The primary contribution of this work is an offline , self-supervised approach for solving arbitrary goal-reaching tasks by combining planning with predictive models and learned dynamical distances . To our knowledge , our method is the first to directly combine predictive models on images with dynamical distance estimators on images , entirely from random , offline data without reward labels . Through our experimental evaluation on challenging robotic object manipulation tasks , including simulated object relocation and real-world drawer manipulation , we find that our method can outperform previously introduced reward specification methods for visual model-based control with a relative performance improvement of at least 50 % across all tasks , and compares favorably to prior work in model-based and model-free RL . We also find that combining Q-functions with planning improves dramatically over policies directly learned with model-free RL . 2 RELATED WORK . Offline and Model-based RL : A number of prior works have studied the problem of learning behaviors from existing offline datasets . While recent progress has been made in applying model-free RL techniques to this problem of offline or batch RL ( Fujimoto et al. , 2019 ; Wu et al. , 2019 ; Kumar et al. , 2019 ; 2020 ; Nair et al. , 2020b ) , one approach that has shown promise is offline model-based RL ( Lowrey et al. , 2019 ; Kidambi et al. , 2020 ; Yu et al. , 2020 ; Argenson & Dulac-Arnold , 2020 ) , where the agent learns a predictive model of the world from data . Such model-based methods have seen success both in the offline and online RL settings , and have a rich history of being effective for planning ( Deisenroth & Rasmussen , 2011 ; Watter et al. , 2015 ; McAllister & Rasmussen , 2016 ; Chua et al. , 2018 ; Amos et al. , 2018 ; Hafner et al. , 2019b ; Nagabandi et al. , 2018 ; Kahn et al. , 2020 ; Dong et al. , 2020 ) or policy optimization ( Sutton , 1991 ; Weber et al. , 2017 ; Ha & Schmidhuber , 2018 ; Janner et al. , 2019 ; Wang & Ba , 2019 ; Hafner et al. , 2019a ) . However , the vast majority of these prior works consider the single task setting where the agent aims to maximize a single task reward . In contrast , in this work we circumvent the need for task rewards by adopting a selfsupervised multi-task approach , where a single learned model is used to perform a variety of tasks , specified in a flexible and general way by desired outcomes – i.e. , goal images . Self-supervised goal reaching : While the standard RL problem involves optimizing for a taskspecific reward , an alternative and potentially more general formulation involves learning a generic goal reaching policy , without task-specific reward labels . In fact , a number of prior works learn goal-conditioned policies using model-free RL ( Kaelbling , 1993 ; Nair et al. , 2018 ; Mandlekar et al. , 2019 ; Nair et al. , 2020a ) , or variants of goal-conditioned behavioral cloning ( GCBC ) ( Ghosh et al. , 2019 ; Ding et al. , 2019 ; Lynch et al. , 2020 ) . In our experiments , we show that our method outperforms both model-free approaches and goal-conditioned behavioral cloning . A number of methods combine model-free and model-based elements by planning over a graph representation ( Eysenbach et al. , 2019 ; Nasiriany et al. , 2019 ; Savinov et al. , 2018 ; Liu et al. , 2020 ) . Such methods can struggle in higher dimensions , where constructing graphs that adequately cover the space may require an excessive number of samples . We compare to these methods in our experiments . Similarly to Finn & Levine ( 2017 ) ; Ebert et al . ( 2018b ) ; Nair & Finn ( 2020 ) ; Yen-Chen et al . ( 2019 ) ; Suh & Tedrake ( 2020 ) , our method uses an action-conditioned video prediction model to generate plans . However , these prior methods generally utilize hand-crafted image similarity reward measures such as ` 2 pixel-error ( Ebert et al. , 2018a ; Nair & Finn , 2020 ) and pixel-flow prediction ( Finn & Levine , 2017 ) . In complex scenes , this can become a major bottleneck : predictions degrade rapidly further in the future , making an informative image similarity metric critical for effective planning . We propose to learn functional similarity metrics in terms of dynamical distances , which we find can be combined with predictive models to attain significantly improved results . Dynamical distance learning : Our method learns dynamical distances – distances that represent shortest paths – from offline data . In the literature , dynamical distances have been learned via direct regression using online data ( Hartikainen et al. , 2019 ) , representation learning ( Warde-Farley et al. , 2018 ; Yu et al. , 2019b ) , or via Q-learning by relabeling goals ( Eysenbach et al. , 2019 ; Florensa et al. , 2019 ) . While these last two works are most similar to ours , in that they also employ approximate dynamic programming to learn distances , our method directly combines these dynamical distances with visual predictive models and planning . Lastly , while prior work has also explored combining model-based planning with value functions ( Zhong et al. , 2013 ; Lowrey et al. , 2019 ; Hafner et al. , 2019a ; Schrittwieser et al. , 2019 ; Argenson & Dulac-Arnold , 2020 ) , these works consider the single task domain with a reward function , while our learned value function considers the multi-task goal reaching domain from entirely random , offline data without reward labels . 3 THE SELF-SUPERVISED OFFLINE RL PROBLEM STATEMENT . In this section , we introduce notation and define the problem setting . We will employ a Markov decision process ( MDP ) with state observations st ∈ S and actions at ∈ A , both indexed by time t ∈ 0 , 1 , · · · , H , where H denotes the maximum episode length . The initial state is sampled from an initial state distribution s0 ∼ p0 ( s0 ) , and subsequent states are sampled according to Markovian dynamics : st+1 ∼ p ( st+1 | st , at ) . Actions are sampled at ∼ π ( at | st , sg ) from a policy that is conditioned on both the current state and a goal state sg ∈ S. In our experiments , both the state and goal are images ( i.e. , S = RH×W×3 ) . We tackle offline learning in this setting , assuming access to a fixed dataset D consisting of trajectories { s0 , a0 , s1 , ... sT } of the agent interacting with the environment . This data can include any environment interactions , from expert demonstrations to trajectories which are not particularly successful at any task . In our experiments , we use data collected using a random policy , which is inexpensive to obtain . The agent does not have access to the environment to collect additional training data . Given this dataset , the objective is to determine the optimal goal-conditioned policy π ? ( at | st , sg ) , under which the agent is able to transition to any goal state sg from any starting state st in the minimum number of time steps possible . Note that unlike in the standard formulation of the RL problem , the agent does not receive any reward signal from its environment .

Summary: This paper proposes to learn functional distances to varying goals concurrently with a latent dynamics model from images. A network is trained to predict the Q-value of state action pairs for a sparse reward at the goal. This way the Q-function represents the shortest path distance to the goal. The paper also proposes a specific learning scheme for training the model from random rollouts in the environment. The learned distance function is used with the dynamics model for planning from start to goal images using model-predictive control (CEM). The approach is evaluated on simulated reaching and object pushing tasks. It is compared with state-of-the-art methods and demonstrates improved performance in some of the tasks.