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: The paper combines optimal control and reinforcement learning (RL) for the execution of robotic manipulation tasks with variable goals. The approach learns, on the one hand, an image-based predictive model using a deep neural network and, on the other hand, a distance cost function using Q-Learning. These model and cost function are used to define the next action to execute using model predictive control (MPC). The cost function is learned from task-agnostic data by randomly generating goal states from the state space and including these goals as input variables in a Q-learning approach. This permits defining cost functions for MPC for variable goals.