Original: M ODEL -B ASED
Revised: M ODEL -B ASED

Original: V ISUAL
Revised: V ISUAL

Original: P LANNING
Revised: P LANNING

Original: WITH
Revised: WITH

Original: S ELF -S UPERVISED
Revised: S ELF -S UPERVISED

Original: F UNCTIONAL D ISTANCES
Revised: F UNCTIONAL D ISTANCES

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: A generalist robot must be able to complete a variety of tasks in its environment.
Revised: A generalist robot must be able to complete a variety of tasks in its environment.

Original: One appealing way to specify each task is in terms of a goal observation.
Revised: One appealing way to specify each task is in terms of a goal observation.

Original: However, learning goal-reaching policies with reinforcement learning remains a challenging problem, particularly when rewards are not provided and  distances in the observation space are not meaningful.
Revised: However, learning goal-reaching policies with reinforcement learning remains a challenging problem, particularly when rewards are not provided and (cid:96) 2 distances in pixel space are not meaningful.

Original: Learned dynamics models are a promising approach for learning about the environment without rewards or task-directed data, but planning to reach goals with such a model requires a notion of functional similarity between observations and goal states.
Revised: Learned dynamics models are a promising approach for learning about the environment without rewards or task-directed data, but planning to reach goals with such a model requires a notion of functional similarity between observations and goal states.

Original: We present a self-supervised method for model-based visual goal reaching, which uses both a visual dynamics model as well as a dynamical distance function learned using model-free reinforcement learning.
Revised: We present a self-supervised method for model-based visual goal reaching, which uses both a visual dynamics model as well as a dynamical distance function learned using model-free reinforcement learning.

Original: This approach trains entirely using ofﬂine, unlabeled data, making it practical to scale to large and diverse datasets.
Revised: This approach trains entirely using ofﬂine, unlabeled data, making it practical to scale to large and diverse datasets.

Original: On several challenging robotic manipulation tasks with only ofﬂine, unlabeled data, we ﬁnd that our algorithm compares favorably to prior model-based and model-free reinforcement learning methods.
Revised: On several challenging robotic manipulation tasks with only ofﬂine, unlabeled data, we ﬁnd that our algorithm compares favorably to prior model-based and model-free reinforcement learning methods.

Original: In ablation experiments, we additionally identify important factors for learning effective distances.
Revised: In ablation experiments, we additionally identify important factors for learning effective distances.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Designing general-purpose robots that can perform a wide range of tasks remains an open problem in AI and robotics.
Revised: Designing general-purpose robots that can perform a wide range of tasks remains an open problem in AI and robotics.

Original: Reinforcement learning (RL) represents a particularly promising tool for learning robotic behaviors when skills can be learned one at a time from user-deﬁned reward functions.
Revised: Reinforcement learning (RL) represents a particularly promising tool for learning robotic behaviors when skills can be learned one at a time from user-deﬁned reward functions.

Original: However, general-purpose robots will likely require large and diverse repertoires of skills, and learning individual tasks one at a time from manually-speciﬁed rewards is onerous and time-consuming.
Revised: However, general-purpose robots will likely require large and diverse repertoires of skills, and learning individual tasks one at a time from manually-speciﬁed rewards is onerous and time-consuming.

Original: How can we design learning systems that can autonomously acquire general-purpose knowledge that allows them to solve many different downstream tasks?
Revised: How can we design learning systems that can autonomously acquire general-purpose knowledge that allows them to solve many different downstream tasks?

Original: To address this problem, we must resolve three questions.
Revised: To address this problem, we must resolve three questions.

Original: (1) How can this robot be commanded to perform speciﬁc downstream tasks?
Revised: (1) How can this robot be commanded to perform speciﬁc downstream tasks?

Original: A simple and versatile choice is to deﬁne tasks in terms of desired outcomes, such as an example observation of the completed task.
Revised: A simple and versatile choice is to deﬁne tasks in terms of desired outcomes, such as an example observation of the completed task.

Original: (2) What types of data should this robot learn from?
Revised: (2) What types of data should this robot learn from?

Original: 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 can be achieved by learning from large and diverse task-agnostic datasets.
Revised: 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 can be achieved by learning from large and diverse task-agnostic datasets.

Original: In the context of RL, this means we need ofﬂine methods that can use all sources of prior data, even in the absence of reward labels .
Revised: In the context of RL, this means we need ofﬂine methods that can use all sources of prior data, even in the absence of reward labels .

Revised: While ofﬂine datasets can be comparatively limited, for example in their coverage of the state space and complexity of behaviors, they are often practical to use in real-world settings (Levine et al., 2020).

Original: (3) What exactly should the robot learn from this data to enable goal-reaching?
Revised: (3) What exactly should the robot learn from this data to enable goal-reaching?

Original: Similar to prior work (Botvinick & Weinstein, 2014;
Revised: Similar to prior work (Botvinick & Weinstein, 2014;

Original: Watter et al., 2015; Finn & Levine, 2017; Ebert et al., 2018b), we note that policies and value functions are speciﬁc to a particular task, while a predictive model captures the physics of an agent’s environment and thus can be used for solving almost any task.
Revised: Watter et al., 2015; Finn & Levine, 2017; Ebert et al., 2018b), we note that policies and value functions are speciﬁc to a particular task, while a predictive model captures the physics of an agent’s environment and thus can be used for solving almost any task.

Original: This makes model learning particularly effective for learning from large and diverse datasets, which do not necessarily contain successful behaviors.
Revised: This makes model learning particularly effective for learning from large and diverse datasets, which do not necessarily contain successful behaviors.

Original: While model-based approaches have demonstrated promising results, including for vision-based tasks in real-world robotic systems (Ebert et al., 2018a;
Revised: While model-based approaches have demonstrated promising results, including for vision-based tasks in real-world robotic systems (Ebert et al., 2018a;

Original: Finn & Levine, 2017), such methods face two major challenges.
Revised: Finn & Levine, 2017), such methods face two major challenges.

Original: First, predictive models on raw images are only effective for short horizons, as uncertainty accumulates far into the future (Denton & Fergus, 2018; Finn et al., 2016; Hafner et al., 2019b; Babaeizadeh et al., 2017).
Revised: First, predictive models on raw images are only effective for short horizons, as uncertainty accumulates far into the future (Denton & Fergus, 2018; Finn et al., 2016; Hafner et al., 2019b; Babaeizadeh et al., 2017).

Original: Second, using such models for planning toward goals requires a notion of similarity between images.
Revised: Second, using such models for planning toward goals requires a notion of similarity between images.

Original: While prior methods have utilized latent variable models (Watter et al., 2015; Nair et al., 2018), (cid:96) 2 pixel-space distance (Nair & Finn, 2019), and other heuristic measures of similarity (Ebert et al., 2018b), these metrics only capture visual similarity.
Revised: While prior methods have utilized latent variable models (Watter et al., 2015; Nair et al., 2018), (cid:96) 2 pixel-space distance (Nair & Finn, 2019), and other heuristic measures of similarity (Ebert et al., 2018b), these metrics only capture visual similarity.

Original: To enable reliable control with predictive models, we instead need distances that are aware of dynamics.
Revised: To enable reliable control with predictive models, we instead need distances that are aware of dynamics.

Original: In this paper, we propose Model-Based RL with Ofﬂine Learned Distances (MBOLD), which aims to address both of these challenges by learning predictive models together with image-based distance functions that reﬂect functionality, from ofﬂine, unlabeled data.
Revised: In this paper, we propose Model-Based Control with Ofﬂine Learned Distances (MBOLD), which aims to address both of these challenges by learning predictive models together with image-based distance functions that reﬂect functionality, from ofﬂine, unlabeled data.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Instead, we utilize approximate dynamic programming for distance estimation.
Revised: Instead, we utilize approximate dynamic programming for distance estimation.

Original: 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 ofﬂine setting and show that approximate dynamic programming techniques derived from Qlearning style Bellman updates can learn effective shortest path dynamical distances.
Revised: 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 ofﬂine setting and show that approximate dynamic programming techniques derived from Qlearning style Bellman updates can learn effective shortest path dynamical distances.

Original: Although this procedure resembles model-free reinforcement learning, we ﬁnd empirically that it does not by itself produce useful policies.
Revised: Although this procedure resembles model-free reinforcement learning, we ﬁnd empirically that it does not by itself produce useful policies.

Original: 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.
Revised: 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.

Original: The primary contribution of this work is an ofﬂine, self-supervised approach for solving arbitrary goal-reaching tasks by combining planning with predictive models and learned dynamical distances.
Revised: The primary contribution of this work is an ofﬂine, self-supervised approach for solving arbitrary goal-reaching tasks by combining planning with predictive models and learned dynamical distances.

Original: To our knowledge, our method is the ﬁrst to directly combine predictive models on images with dynamical distance estimators on images, entirely from random, ofﬂine data without reward labels.
Revised: To our knowledge, our method is the ﬁrst to directly combine predictive models on images with dynamical distance estimators on images, entirely from random, ofﬂine data without reward labels.

Original: Through our experimental evaluation on challenging simulated robotic object manipulation tasks, we ﬁnd that our method can outperform previously introduced reward speciﬁcation 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.
Revised: Through our experimental evaluation on challenging simulated robotic object manipulation tasks, we ﬁnd that our method can outperform previously introduced reward speciﬁcation 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.

Original: We also ﬁnd that combining
Revised: We also ﬁnd that combining

Original: Q-functions with planning improves dramatically over policies directly learned with model-free RL.
Revised: Q-functions with planning improves dramatically over policies directly learned with model-free RL.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Ofﬂine and Model-based RL: A number of prior works have studied the problem of learning behaviors from existing ofﬂine datasets.
Revised: Ofﬂine and Model-based RL: A number of prior works have studied the problem of learning behaviors from existing ofﬂine datasets.

Original: While recent progress has been made in applying model-free RL techniques to this problem of “ofﬂine” 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 ofﬂine model-based
Revised: While recent progress has been made in applying model-free RL techniques to this problem of “ofﬂine” 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 ofﬂine model-based

Original: RL (Lowrey et al., 2018; Kidambi et al., 2020; Yu et al., 2020; Argenson & Dulac-Arnold, 2020), where the agent learns a predictive model of the world from data.
Revised: RL (Lowrey et al., 2018; Kidambi et al., 2020; Yu et al., 2020; Argenson & Dulac-Arnold, 2020), where the agent learns a predictive model of the world from data.

Original: Such model-based methods have seen success both in the ofﬂine and online RL settings, and have a rich history of being effective for planning (Deisenroth & Rasmussen, 2011; Watter et al., 2015; McAllister & Rasmussen, 2016;
Revised: Such model-based methods have seen success both in the ofﬂine and online RL settings, and have a rich history of being effective for planning (Deisenroth & Rasmussen, 2011; Watter et al., 2015; McAllister & Rasmussen, 2016;

Original: Chua et al., 2018; Amos et al., 2018; Hafner et al., 2019b; Nagabandi et al., 2018; Kahn et al., 2020) or policy optimization (Sutton, 1991;
Revised: Chua et al., 2018; Amos et al., 2018; Hafner et al., 2019b; Nagabandi et al., 2018; Kahn et al., 2020) or policy optimization (Sutton, 1991;

Original: Weber et al., 2017; Ha & Schmidhuber, 2018; Janner et al., 2019; Wang & Ba, 2019; Hafner et al., 2019a).
Revised: Weber et al., 2017; Ha & Schmidhuber, 2018; Janner et al., 2019; Wang & Ba, 2019; Hafner et al., 2019a).

Original: Critically, however, the vast majority of these prior works consider the single task setting where the agent aims to maximize a single task reward.
Revised: Critically, however, the vast majority of these prior works consider the single task setting where the agent aims to maximize a single task reward.

Original: 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, speciﬁed in a ﬂexible and general way by desired outcomes – i.e., goal images.
Revised: 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, speciﬁed in a ﬂexible and general way by desired outcomes – i.e., goal images.

Original: Self-supervised goal reaching : While the standard RL problem involves optimizing for a taskspeciﬁc reward, an alternative and potentially more general formulation involves learning a generic goal reaching policy, without task-speciﬁc reward labels.
Revised: Self-supervised goal reaching : While the standard RL problem involves optimizing for a taskspeciﬁc reward, an alternative and potentially more general formulation involves learning a generic goal reaching policy, without task-speciﬁc reward labels.

Original: 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)
Revised: 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)

Original: (Ghosh et al., 2019; Ding et al., 2019; Lynch et al., 2020).
Revised: (Ghosh et al., 2019; Ding et al., 2019; Lynch et al., 2020).

Original: In our experiments, we show that our method outperforms both model-free approaches and goal-conditioned behavioral cloning.
Revised: In our experiments, we show that our method outperforms both model-free approaches and goal-conditioned behavioral cloning.

Original: 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;
Revised: 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;

Original: Savinov et al., 2018; Liu et al., 2020).
Revised: Savinov et al., 2018; Liu et al., 2020).

Original: Such methods can struggle in higher dimensions, where constructing graphs that adequately cover the space may require an excessive number of samples.
Revised: Such methods can struggle in higher dimensions, where constructing graphs that adequately cover the space may require an excessive number of samples.

Original: We compare to these methods in our experiments.
Revised: We compare to these methods in our experiments.

Original: Similarly to
Revised: Similarly to

Original: Finn & Levine (2017); Ebert et al.
Revised: Finn & Levine (2017); Ebert et al.

Original: (2018b); Nair & Finn (2019); Yen-Chen et al.
Revised: (2018b); Nair & Finn (2019); Yen-Chen et al.

Original: ; Suh &
Revised: ; Suh &

Original: Tedrake (2020), our method uses an action-conditioned video prediction model to generate plans.
Revised: Tedrake (2020), our method uses an action-conditioned video prediction model to generate plans.

Original: However, these prior methods generally utilize hand-crafted image similarity reward measures such as (cid:96) 2 pixel-error (Ebert et al., 2018a; Nair & Finn, 2019) and pixel-ﬂow prediction (Finn & Levine, 2017).
Revised: However, these prior methods generally utilize hand-crafted image similarity reward measures such as (cid:96) 2 pixel-error (Ebert et al., 2018a; Nair & Finn, 2019) and pixel-ﬂow prediction (Finn & Levine, 2017).

Original: 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.
Revised: 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.

Original: We propose to learn functional similarity metrics in terms of dynamical distances, which we ﬁnd can be combined with predictive models to attain signiﬁcantly improved results.
Revised: We propose to learn functional similarity metrics in terms of dynamical distances, which we ﬁnd can be combined with predictive models to attain signiﬁcantly improved results.

Original: Dynamical distance learning : Our method learns dynamical distances – distances that represent shortest paths – from ofﬂine data.
Revised: Dynamical distance learning : Our method learns dynamical distances – distances that represent shortest paths – from ofﬂine data.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: Lastly, while prior work has also explored combining model-based planning with value functions (Zhong et al., 2013; Lowrey et al., 2018; 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, ofﬂine data without reward labels.
Revised: Lastly, while prior work has also explored combining model-based planning with value functions (Zhong et al., 2013; Lowrey et al., 2018; 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, ofﬂine data without reward labels.

Original:  T HE S ELF -S UPERVISED O FFLINE RL P ROBLEM S TATEMENT
Revised:  T HE S ELF -S UPERVISED O FFLINE RL P ROBLEM S TATEMENT

Original: In this section, we introduce notation and deﬁne the problem setting.
Revised: In this section, we introduce notation and deﬁne the problem setting.

Original: We will employ a Markov decision process (MDP) with state observations s t ∈ S and actions a t ∈ A , both indexed by time t ∈ 0 , 1 , · · · , H , where H denotes the maximum episode length.
Revised: We will employ a Markov decision process (MDP) with state observations s t ∈ S and actions a t ∈ A , both indexed by time t ∈ 0 , 1 , · · · , H , where H denotes the maximum episode length.

Original: The initial state is sampled from an initial state distribution s 0 ∼ p 0 ( s 0 ) , and subsequent states are sampled according to Markovian dynamics: s t +1 ∼ p ( s t +1 | s t , a t ) .
Revised: The initial state is sampled from an initial state distribution s 0 ∼ p 0 ( s 0 ) , and subsequent states are sampled according to Markovian dynamics: s t +1 ∼ p ( s t +1 | s t , a t ) .

Original: Actions are sampled a t ∼ π ( a t | s t , s g ) from a policy that is conditioned on both the current state and a goal state s g ∈ S .
Revised: Actions are sampled a t ∼ π ( a t | s t , s g ) from a policy that is conditioned on both the current state and a goal state s g ∈ S .

Original: In our experiments, both the state and goal are images (i.e., S =
Revised: In our experiments, both the state and goal are images (i.e., S =

Original: R H
Revised: R H

Original: × W × 3 ).
Revised: × W × 3 ).

Original: We tackle ofﬂine learning in this setting, assuming access to a ﬁxed dataset D consisting of trajectories { s 0 , a 0 , s 1 , ...
Revised: We tackle ofﬂine learning in this setting, assuming access to a ﬁxed dataset D consisting of trajectories { s 0 , a 0 , s 1 , ...

Original: s T } of the agent interacting with the environment.
Revised: s T } of the agent interacting with the environment.

Original: This data can include any environment interactions, from expert demonstrations to trajectories which are not particularly successful at any task.
Revised: This data can include any environment interactions, from expert demonstrations to trajectories which are not particularly successful at any task.

Original: In our experiments, we use data collected using a random policy, which is inexpensive to obtain.
Revised: In our experiments, we use data collected using a random policy, which is inexpensive to obtain.

Original: The agent does not have access to the environment to collect additional training data.
Revised: The agent does not have access to the environment to collect additional training data.

Original: Given this dataset, the objective is to determine the optimal goal-conditioned policy [Table] ( a t | s t , s g ) , under which the agent is able to transition to any goal state s g from any starting state s t in the minimum number of time steps possible.
Revised: Given this dataset, the objective is to determine the optimal goal-conditioned policy π (cid:63) ( a t | s t , s g ) , under which the agent is able to transition to any goal state s g from any starting state s t in the minimum number of time steps possible.

Original: Note that unlike in the standard formulation of the RL problem, the agent does not receive any reward signal from its environment.
Revised: Note that unlike in the standard formulation of the RL problem, the agent does not receive any reward signal from its environment.

Original:  M ODEL -B ASED V ISUAL G OAL -R EACHING
Revised:  M ODEL -B ASED V ISUAL G OAL -R EACHING

Original: In this section we will introduce our method, MBOLD, for ofﬂine, goal-conditioned reinforcement learning.
Revised: In this section we will introduce our method, MBOLD, for ofﬂine, goal-conditioned reinforcement learning.

Original: MBOLD, illustrated in Fig.
Revised: MBOLD, illustrated in Fig.

Original: 2, is composed of two neural networks: a predictive model and a learned distance function.
Revised: 2, is composed of two neural networks: a predictive model and a learned distance function.

Original: The video-predictive dynamics model allows the agent to predict the result of hypothetical sequences of actions.
Revised: The video-predictive dynamics model allows the agent to predict the result of hypothetical sequences of actions.

Original: However, this model cannot accurately predict far into the future, and has no notion of whether the predicted outcomes are desirable.
Revised: However, this model cannot accurately predict far into the future, and has no notion of whether the predicted outcomes are desirable.

Original: Thus, we also learn a distance function, which will estimate the timestep length of the shortest path between a 
Revised: Thus, we also learn a distance function, which will estimate the timestep length of the shortest path between a predicted state and a given goal.

Revised: Both networks are trained on the same ofﬂine dataset pre-collected by a (possibly random) policy.

Original: At test-time, we use the learned dynamics model and distance function for model-predictive control (MPC).
Revised: At test-time, we use the learned dynamics model and distance function for model-predictive control (MPC).

Original: MBOLD predicts future states for candidate action sequences using the learneddynamics model, and usesthe learned distance function todetermine which action sequence will lead the agent closestto the goal. The ﬁrst of the actions is then executed, and planning repeats upon receiving the subsequent observationfrom the environment.
Revised: MBOLD predicts future states for candidate action sequences using the learned model, and uses distance function to closest ﬁrst of the actions is then executed, and planning repeats upon receiving the subsequent observation the environment.

Original: The remainder of this section describes how we learn the dynamics model and distance function, and use them to perform control.
Revised: The remainder of this section describes how we learn the dynamics model and distance function, and use them to perform control.

Original: Dynamics learning.
Revised: Dynamics learning.

Original: Our method learns environment dynamics in order to solve for actions during test time, without explicit task reward signal during training time.
Revised: Our method learns environment dynamics in order to solve for actions during test time, without explicit task reward signal during training time.

Original: In theory, MBOLD is able to use any image-based forward model, including latent variable models like PlaNet (Hafner et al., 2019b).
Revised: MBOLD can use arbitrary imagebased forward models, including latent variable models like PlaNet (Hafner et al., 2019b) – the particular choice of model is a design decision when implementing our method.

Original: We implement a deterministic, convolutional video prediction model adapted from SAVP (Lee et al., 2018).
Revised: We implement a deterministic, convolutional video prediction model adapted from SAVP (Lee et al., 2018).

Original: The network takes as input the current observation s t and a sequence of h actions a t : t + h −and returns a prediction for the next h image observations, ˆ f θ ( s t , a t : t + h − 1 )=
Revised: The network takes as input the current observation s t and a sequence of h actions a t : t + h − 1 and returns a prediction for the next h image observations, ˆ f θ ( s t , a t : t + h − 1 )

Original:  { ˆ s t +1 , . . .
Revised: = { ˆ s t +1 , . . .

Original: , ˆ s t
Revised: , ˆ s t

Original: + h } .
Revised: + h } .

Original: We train this model to minimize the (cid:96) 2 image reconstruction loss: 
Revised: We train this model to minimize the (cid:96) 2 image reconstruction loss: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Distance learning.
Revised: Distance learning.

Original: Our method also learns a dynamical distance function, so that it can evaluate a functional sense of distance from predicted states to the goal state, to be used as a planning cost.
Revised: Our method also learns a dynamical distance function, so that it can evaluate a functional sense of distance from predicted states to the goal state, to be used as a planning cost.

Original: However, the environment does not provide a reward signal that might be used to deduce these distances.
Revised: However, the environment does not provide a reward signal that might be used to deduce these distances.

Original: Even worse, the ofﬂine dataset likely contains sub-optimal trajectories, so our method may not even have access to examples of shortest path trajectories between states.
Revised: Even worse, the ofﬂine dataset likely contains sub-optimal trajectories, so our method may not even have access to examples of shortest path trajectories between states.

Original: Our key observation is that a  Q-function trained on a modiﬁed MDP with an indicator cost function yields values that correspond to shortest path distances in the original environment.
Revised: Our key observation is that a goal-conditioned Q-function trained on a modiﬁed MDP with an indicator cost function yields values that correspond to shortest path distances in the original environment.

Original: Thus, Q-learning-like dynamic programming methods can recover optimal distance functions even from sub-optimal data.
Revised: Thus, Q-learning-like dynamic programming methods can recover optimal distance functions even from sub-optimal data.

Original: More precisely, we create an MDP by augmenting environment trajectories with the reward function r ( s, a, s t +1 , g )
Revised: More precisely, we create an MDP by augmenting environment trajectories with the reward function r ( s t , a, s t +1 , g )

Original: = 1 s t+1 =
Revised: = 1 s t+1 =

Original: s g , adding a discount factor of γ , and considering episodes terminated once they reach the goal state.
Revised: s g , adding a discount factor of γ , and considering episodes terminated once they reach the goal state.

Revised: Note that s t , s t +1 , and g all represent images, so the reward is only given when the next state and goal images exactly match.

Original: Goals are sampled according to a distribution on S , which we will discuss later.
Revised: During training, goals are sampled according to a distribution on S , which we will discuss later.

Original: Note that if γ < 1 , the Q-values for a policy that maximizes expected discounted returns in this MDP can be directly mapped to shortest path distances.
Revised: If γ < 1 , the Q-values for a policy that maximizes expected discounted returns in this MDP can be directly mapped to shortest path distances.

Original: Speciﬁcally, in discrete state environments, the optimal Q-function can be written as Q ( s, a, g ) =
Revised: Speciﬁcally, in discrete state environments, the optimal Q-function can be written as Q ( s, a, g ) =

Original: γ d ( s,a,g ) , where d ( s, a, g ) is a shortest path distance between s and g after taking action a .
Revised: γ d ( s,a,g ) , where d ( s, a, g ) is a shortest path distance between s and g after taking action a .

Original: Similarly, we can recover d ( s, a, g ) = log γ Q ( s, a, g ) .
Revised: Similarly, we can recover d ( s, a, g ) = log γ Q ( s, a, g ) .

Original: Ultimately, our Q-learning approach corresponds to the following Bellman error optimization objective:
Revised: Ultimately, our Q-learning approach corresponds to the following Bellman error optimization objective:

Original: [Equation]
Revised: [Equation]

Original: Because the policy used for collecting the ofﬂine dataset may not be directed towards any particular goals, we must somehow select goals  for each transition to optimize this objective.
Revised: In practice, we use a deep neural network to represent the Q-function. During training, we sample transitions ( s t , a t , s t +1 , g ) to optimize the objective in Equation 2. The ﬁrst three components ( s t , a t , s t +1 ) of the tuple can be sampled randomly from the dataset, but because trajectories from the ofﬂine dataset may not be directed towards any particular goals, we must somehow select goals g for each transition. Selecting goals for relabeling transitions.

Original: Na¨ıvely choosing g , say by sampling random states uniformly from the dataset, will provide an extremely sparse reward signal, as two random states will almost never be exactly identical.
Revised: Na¨ıvely choosing g , say by sampling random states uniformly from the dataset, will provide an extremely sparse reward signal, as two random state images will almost never be exactly identical.

Original: While the sparse reward problem could be mitigated by  selectively sampling as goals the states that were actually reached in future time steps along the same trajectory  (Kaelbling, 1993; Andrychowicz et al., 2017),  relabeling all goals in this way creates a major issue: since the distance function would only be trained on goals that were actually reached, it would systematically underestimate the distance to unreachable goals.
Revised: The sparse reward problem can be mitigated by hindsight relabeling, that is, selectively sampling as goals the states that were actually reached in future time steps along the same trajectory as s t (Kaelbling, 1993; Andrychowicz et al., 2017), so we adopt this procedure to make training feasible. To sample goals for a transition at time step t , we sample a discrete time offset ∆ ∼ Geom ( p ) , where p ∈ [0 , 1] is a hyper-parameter, and use the state at time t + ∆ as the goal. Note that if ∆ = 1 , the reward for this transition is 1 , avoiding the sparsity issue. However, relabeling all transitions in this way creates a major issue: since the distance function would only be trained on goals that were actually reached, it would systematically underestimate the distance to unreachable goals.

Original: Put another way, goals that were not reached from s t would be out-of-distribution goals for the resulting Q-function.
Revised: Put another way, goals that were not reached from s t would be out-of-distribution goals for the resulting Q-function.

Original: We found this to result in poor performance.
Revised: We found this to result in poor performance.

Original: In practice, prior work (Kaelbling, 1993; Andrychowicz et al., 2017) actually relabels with a mixture of reached  and commanded goals, resulting in both “positive” and “negative” goals.
Revised: In practice, prior work (Kaelbling, 1993; Andrychowicz et al., 2017) actually relabels with a mixture of reached goals and commanded but not necessarily reached goals.

Original: To sample positive goals for a transition at time step t , we sample a time offset ∆ ∼ Geom ( p ) , where p ∈ [0 , 1] is a hyper-parameter, and use the state at time t + ∆ as the goal.

Original: Note that if ∆ = 1 , the reward for this transition is 1 , avoiding the sparsity issue.

Original: While prior methods can obtain negative goals based on those commanded in online data, in the ofﬂine setting we require a procedure to produce state-goal pairs that are distant, yet informative.
Revised: While prior methods can obtain negative goals based on those commanded in online data, in the ofﬂine setting we require a procedure to produce state-goal pairs that are distant, yet informative.

Original: Randomly selecting dataset states will lead to pairs of images that are clearly distant with high probability (e.g. pairs in which all objects and the robot have been moved).
Revised: Randomly selecting dataset states will lead to pairs of images that are clearly distant with high probability (e.g. pairs in which all objects and the robot have been moved).

Original: Therefore, we would like a goal sampling procedure that produces less obvious examples of distant states, which are more informative for training.
Revised: Therefore, we would like a goal sampling procedure that produces less obvious examples of distant states, which are more informative for training.

Original: Hard negative mining is one example of such a procedure, where pairs are selected based on the model’s predictions, but is computationally expensive with large datasets.
Revised: Hard negative mining is one example of such a procedure, where pairs are selected based on the model’s predictions, but is computationally expensive with large datasets.

Original: Instead, we build upon the intuition that distance functions are likely to pay excessive attention to fully actuated factors in the state, such as the position of the robot’s arm, because they are strongly predictive of distances.
Revised: Instead, we build upon the intuition that distance functions are likely to pay excessive attention to fully actuated factors in the state, such as the position of the robot’s arm, because they are strongly predictive of distances.

Original: We propose sampling pairs of states where the actuated components are similar.
Revised: We propose sampling “negative” goal states g which have similar actuated components to the state s t .

Original: When randomly sampling pairs of states under this constraint, the underactuated dimensions (e.g. the objects), which are generally not known, are likely to have distinct positions.
Revised: When randomly sampling pairs of states under this constraint, the underactuated dimensions (e.g. the objects), which are generally not known, are likely to have distinct positions.

Original: Hence, these data points can serve as informative hard negatives that encourage the model to pay more attention to the difﬁcult, underactuated parts of the state.
Revised: Hence, these data points can serve as informative hard negatives that encourage the model to pay more attention to the difﬁcult, underactuated parts of the state.

Original: Unlike hard negative mining, this sampling approach is computationally inexpensive, as it does not rely on the current distance function, and practical, as actuated components of the state can typically be measured through encoders on the actuator.
Revised: Unlike hard negative mining, this sampling approach is computationally inexpensive, as it does not rely on the current distance function, and practical, as actuated components of the state can typically be measured through encoders on the actuator.

Original: In practice, we sample these  goals from observations across all dataset trajectories via nearestneighbors search, using arm joint (cid:96) 2 distance as the similarity key.
Revised: In practice, we sample these “negative” goals from observations across all dataset trajectories via nearest-neighbors search, using arm joint (cid:96) 2 distance as the similarity key.

Original: Note that this does assume proprioceptive state information from the agent (e.g. robot joint angles), which is almost always available in real-world robotics settings, but does not require knowledge about object positions or other ground-truth environment information.
Revised: Note that this does assume proprioceptive state information from the agent (e.g. robot joint angles), which is almost always available in real-world robotics settings, but does not require knowledge about object positions or other ground-truth environment information.

Revised: Therefore, we additionally measure actuated components of the state during data collection for this goal sampling strategy, but only the sampled images are used in the distance learning objective.

Original: See Appendix A.1 for details.
Revised: See Appendix A.1 for details.

Revised: We relabel half of the transitions in each training batch with reached goals and the other half with “negative” goals with similar actuated components, ﬁnding through ablation experiments that this combination achieves stronger performance compared to using just reached goals in our evaluation environments.

Revised: In other domains, more careful consideration is required to determine if the assumptions which motivate this “negative” goal sampling strategy are satisﬁed.

Original: Control via MBOLD.
Revised: Control via MBOLD.

Original: At test-time, the learned distance function and dynamics model are used together to solve control tasks via MPC.
Revised: At test-time, the learned distance function and dynamics model are used together to solve control tasks via MPC.

Original: In other words, the dynamics model predicts how candidate actions will affect the environment, and the distance model rates predicted sequences based on which bring the agent closest to the user-deﬁned goal state.
Revised: In other words, the dynamics model predicts how candidate actions will affect the environment, and the distance model rates predicted sequences based on which bring the agent closest to the user-deﬁned goal state.

Original: This “rating” mechanism works as follows: given the current state s t , goal state s g , candidate actions a t : t + h − 1 , and predicted future states ˆ f θ ( s t , a t : t + h − 1 ) from the learned dynamics model, the learned distance function calculates
Revised: This “rating” mechanism works as follows: given the current state s t , goal state s g , candidate actions a t : t + h − 1 , and predicted future states ˆ f θ ( s t [Equation] t [Equation]

Revised: t

Original: [Equation]
Revised: [Equation]

Original: This equation acts as an objective function for MPC.
Revised: 1 ) from the learned dynamics model, the learned distance function calculates [Equation] In practice, the maximization over α is estimated by an actor network learned simultaneously with the Q-function. V ( a t : t + h − 1 ) acts as an objective function for MPC.

Original: Plainly, the controller’s goal is to ﬁnd candidate actions a t : t + h − 1 which minimize the dynamical distance to the goal h steps into the future.
Revised: Plainly, the controller’s goal is to ﬁnd candidate actions a t : t + h − 1 which minimize the dynamical distance to the goal h steps into the future.

Original: After this process completes, the best action is executed by the agent.
Revised: After this process completes, the best action is executed by the agent.

Original: Note that this controller re-plans after every action taken in the environment (i.e every timestep), in order to prevent errors in dynamics prediction from compounding.
Revised: Note that this controller re-plans after every action taken in the environment (i.e every timestep), in order to prevent errors in dynamics prediction from compounding.

Original: MPC Controller Algorithm.
Revised: MPC Controller Algorithm.

Original: MBOLD uses the CEM algorithm (De Boer et al., 2005) to optimize the objective in Equation 3.
Revised: MBOLD uses the CEM algorithm (De Boer et al., 2005) to optimize the objective in Equation 3.

Original: It begins by sampling N random trajectories from a prior multi-variate
Revised: It begins by sampling N random trajectories from a prior multi-variate

Original: Gaussian distribution.
Revised: Gaussian distribution.

Original: Then, the top K actions which score highest according to V ( a t : t + h − 1 ) are selected as candidates.
Revised: Then, the top K actions which score highest according to V ( a t : t + h − 1 ) are selected as candidates.

Original: A new Gaussian distribution is ﬁt on these candidates, and the loop starts over again by sampling fresh actions from this distribution.
Revised: A new Gaussian distribution is ﬁt on these candidates, and the loop starts over again by sampling fresh actions from this distribution.

Original: After I iterations, the loop ﬁnishes and returns the best action found so far.
Revised: After I iterations, the loop ﬁnishes and returns the best action found so far.

Original: See Appendix A.2 for full CEM implementation details.
Revised: See Appendix A.2 for full CEM implementation details.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Our experiments aim to answer two questions: (1) How does MBOLD compare to prior modelbased and model-free methods when learning to reach goals from task-agnostic ofﬂine data?
Revised: Our experiments aim to answer two questions: (1) How does MBOLD compare to prior modelbased and model-free methods when learning to reach goals from task-agnostic ofﬂine data?

Original: How do different dynamical distance learning methods compare to MBOLD in terms of providing effective distance functions for planning?
Revised: How do different dynamical distance learning methods compare to MBOLD in terms of providing effective distance functions for planning?

Original: We evaluate our method, prior methods, and baselines on three simulated tasks with visual observations: (1) a simple reaching task that requires moving a Sawyer 7-DoF arm to a goal location, which provides a way to validate implementations of all methods, (2) object pushing , in which a Sawyer arm must relocate an object to a particular goal location, in environments with 1 or 3 objects, and (3) door sliding , which requires repositioning a sliding door with a Franka 7-DoF arm.
Revised: We evaluate our method, prior methods, and baselines on three simulated tasks with visual observations: (1) a simple reaching task that requires moving a Sawyer 7-DoF arm to a goal location, which provides a way to validate implementations of all methods, (2) object pushing , in which a Sawyer arm must relocate an object to a particular goal location, in environments with 1 or 3 objects, and (3) door sliding , which requires repositioning a sliding door with a Franka 7-DoF arm.

Original: These tasks are challenging because they require long-horizon planning without access to intermediate rewards.
Revised: These tasks are challenging because they require long-horizon planning without access to intermediate rewards.

Original: For each task, we deﬁne the action space
Revised: For each task, we deﬁne the action space

Original: A such that actions control the Cartesian position of the robot’s end-effector, as well as the robot’s gripper.
Revised: A such that actions control the Cartesian position of the robot’s end-effector, as well as the robot’s gripper.

Original: We randomly generate a set of 100 test goals, consisting of a goal image and starting state, for each task, on which all methods are tested.
Revised: We randomly generate a set of 100 test goals, consisting of a goal image and starting state, for each task, on which all methods are tested.

Original: A trial is considered successful if the ﬁnal distance to the goal of each relevant object, e.g. slide position for the door sliding task, ends below a given threshold.
Revised: A trial is considered successful if the ﬁnal distance to the goal of each relevant object, e.g. slide position for the door sliding task, ends below a given threshold.

Original: For the object relocation task, we evaluate each method on two scenes, containing one and three objects.
Revised: For the object relocation task, we evaluate each method on two scenes, containing one and three objects.

Original: All evaluation goals require the robot to move one of the objects, with the others serving as distractors.
Revised: All evaluation goals require the robot to move one of the objects, with the others serving as distractors.

Original: We also study two levels of difﬁculty: “regular,” where goals are generated from random trajectories in which the object moves a certain minimum distance, and “hard,” where the arm is additionally enforced to be distant from the object in the goal observation, requiring the robot to push the object and then withdraw the arm.
Revised: We also study two levels of difﬁculty: “regular,” where goals are generated from random trajectories in which the object moves a certain minimum distance, and “hard,” where the arm is additionally enforced to be distant from the object in the goal observation, requiring the robot to push the object and then withdraw the arm.

Original: We depict the tasks in Fig.
Revised: We depict the tasks in Fig.

Original: 3 (left) and provide full experimental details in Appendix A.3.
Revised: 3 (left) and provide full experimental details in Appendix A.3.

Original: For all tasks, we generate an ofﬂine dataset by running random policies for 1e4 episodes oftimesteps each.
Revised: For all tasks, we generate an ofﬂine dataset by running random policies for 1e4 episodes oftimesteps each.

Original: We provide only this ofﬂine dataset to all methods, with no online training.
Revised: We provide only this ofﬂine dataset to all methods, with no online training.

Original: At test time, the agent only receives the goal image and current observation at each step, and no intermediate rewards besides those that it computes itself.
Revised: At test time, the agent only receives the goal image and current observation at each step, and no intermediate rewards besides those that it computes itself.

Original: Comparative evaluation.
Revised: Comparative evaluation.

Original: We compare MBOLD to prior work in model-based and model-free RL.
Revised: We compare MBOLD to prior work in model-based and model-free RL.

Original: As MBOLD uses purely ofﬂine data and does not require rewards from the environment, we make modiﬁcations to these methods where necessary to provide a fair comparison.
Revised: As MBOLD uses purely ofﬂine data and does not require rewards from the environment, we make modiﬁcations to these methods where necessary to provide a fair comparison.

Original: Many of these prior methods (though not all) require the environment to provide a ground truth reward signal.
Revised: Many of these prior methods (though not all) require the environment to provide a ground truth reward signal.

Original: In this case, we provide these methods with simple “uninformative” rewards, following prior work (Nair et al., 2018), which consist of the MSE between the current and goal image.
Revised: In this case, we provide these methods with simple “uninformative” rewards, following prior work (Nair et al., 2018), which consist of the MSE between the current and goal image.

Original: Many of these methods were initially presented in the online setting.
Revised: Many of these methods were initially presented in the online setting.

Original: The ofﬂine setting is harder for RL methods (Fujimoto et al., 2019; Wu et al., 2019; Kumar 
Revised: The ofﬂine setting is harder for RL methods (Fujimoto et al., 2019; Wu et al., 2019; Kumar et al., 2019), partially explaining their poor performance.

Original: et
Revised: See Appendix B for details on all baselines. We compare MBOLD to the following methods:

Original: • Reinforcement Learning with Imagined Goals (RIG) (Nair et al., 2018) : RIG is a model-free RL method for visual goal-reaching. 
Revised: • Reinforcement Learning with Imagined Goals (RIG) ( Nair et al., ) : RIG is a model-free RL method for visual goal-reaching. Unlike the other methods, we

Original: Unlike the other methods, we still allow RIG to collect additional online data to train its policy.
Revised: still allow RIG to collect additional online data to train its policy.

Original: • Dreamer (Hafner et al., 2019a) : Dreamer, a model-based method for image-based tasks, also uses a combination of value functions and planning, but uses online data collection and, crucially, ground truth reward signals.
Revised: • Dreamer (Hafner et al., 2019a) : Dreamer, a model-based method for image-based tasks, also uses a combination of value functions and planning, but uses online data collection and, crucially, ground truth reward signals.

Original: We adapt Dreamer for the ofﬂine, reward-free setting.
Revised: We adapt Dreamer for the ofﬂine, reward-free setting.

Original: • Dreamer (cid:96) 2 arm distance : We additionally compare with an “oracle” version of Dreamer that uses privileged information about the ground-truth position of the arm.
Revised: • Dreamer (cid:96) 2 arm distance : We additionally compare with an “oracle” version of Dreamer that uses privileged information about the ground-truth position of the arm.

Original: • Search on the Replay Buffer (SoRB)
Revised: • Search on the Replay Buffer (SoRB)

Original: (Eysenbach et al., 2019) : SoRB performs planning on a graph constructed using learned distances, learned without a reward function.
Revised: (Eysenbach et al., 2019) : SoRB performs planning on a graph constructed using learned distances, learned without a reward function.

Original: • Goal-Conditioned Behavior Cloning : We train a behavior cloning model using goals sampled from observations achieved further in a given trajectory.
Revised: • Goal-Conditioned Behavior Cloning : We train a behavior cloning model using goals sampled from observations achieved further in a given trajectory.

Original: This can be viewed as an ofﬂine variant of GCSL (Ghosh et al., 2019) or a non-recurrent version of Lynch et al.
Revised: This can be viewed as an ofﬂine variant of GCSL (Ghosh et al., 2019) or a non-recurrent version of Lynch et al.

Original: • Visual Foresight
Revised: • Visual Foresight

Original: (Ebert et al., 2018b) :
Revised: (Ebert et al., 2018b) :

Original: Visual Foresight also plans with an action-conditioned video prediction model, but uses (among other choices) (cid:96) 2 pixel error as a cost function.
Revised: Visual Foresight also plans with an action-conditioned video prediction model, but uses (among other choices) (cid:96) 2 pixel error as a cost function.

Original: Since all methods are trained from ofﬂine data with no additional environment interaction, we present ﬁnal performance on the test goals as a bar graph, rather than learning curves.
Revised: Since all methods are trained from ofﬂine data with no additional environment interaction, we present ﬁnal performance on the test goals as a bar graph, rather than learning curves.

Original: The comparison on the simple reaching task is shown in Figure 4, and suggests that on this task, many of the methods perform quite well.
Revised: The comparison on the simple reaching task is shown in Figure 4, and suggests that on this task, many of the methods perform quite well.

Original: However, on the substantially more complex tasks, shown in Figure 3, we see clearer differentiation between the different algorithms.
Revised: However, on the substantially more complex tasks, shown in Figure 3, we see clearer differentiation between the different algorithms.

Original: On harder object pushing tasks, MBOLD attains the best performance, by a considerable margin.
Revised: On harder object pushing tasks, MBOLD attains the best performance, by a considerable margin.

Original: Interestingly, simple goalconditioned behavioral cloning actually represents one of the strongest baselines on this task.
Revised: Interestingly, simple goalconditioned behavioral cloning actually represents one of the strongest baselines on this task.

Original: On the hardest door sliding task, our method attains the best performance by a large margin.
Revised: On the hardest door sliding task, our method attains the best performance by a large margin.

Original: Qualitative analysis.
Revised: Qualitative analysis.

Original: In this section, we examine the distance functions learned by MBOLD, and show qualitatively that our learned distances better model the dependence of functional separation between two states on the relative positions of objects in their scenes.
Revised: In this section, we examine the distance functions learned by MBOLD, and show qualitatively that our learned distances better model the dependence of functional separation between two states on the relative positions of objects in their scenes.

Original: Figure 6 presents heatmaps of predicted distances for a ﬁxed goal image on the object pushing task, as the initial observation is varied based on object position.
Revised: Figure 6 presents heatmaps of predicted distances for a ﬁxed goal image on the object pushing task, as the initial observation is varied based on object position.

Original: The robot arm is set to the same position in each initial image.
Revised: The robot arm is set to the same position in each initial image.

Original: We see that the Q-function is able to learn a relatively well-shaped distance which accounts for the object position.
Revised: We see that the Q-function is able to learn a relatively well-shaped distance which accounts for the object position.

Original: We additionally visualize baseline distance models for comparison.
Revised: We additionally visualize baseline distance models for comparison.

Original: First, we look at an ablation of our distance model, which is trained via regression to map pairs of states randomly sampled from a given dataset trajectory to the number of timesteps separating them in that trajectory, and can be viewed as an ofﬂine variant of DDL (Hartikainen et al., 2019).
Revised: First, we look at an ablation of our distance model, which is trained via regression to map pairs of states randomly sampled from a given dataset trajectory to the number of timesteps separating them in that trajectory, and can be viewed as an ofﬂine variant of DDL (Hartikainen et al., 2019).

Original: We call this scheme that effectively predicts random walk distances “temporal distance regression.”
Revised: We call this scheme that effectively predicts random walk distances “temporal distance regression.”

Revised: The second baseline we compare to is pixel-wise mean-squared error, a simple and popular metric.

Original: This produces more sharply peaked distances than the Q-function, and is less desirable as a reward signal during planning, as we ﬁnd through our ablation experiments.
Revised: We ﬁnd that the temporal distance regression model produces more sharply peaked distances than the Q-function and is less desirable as a reward signal during planning, as we ﬁnd through our ablation experiments.

Original: The pixel-wise MSE metricadditionally produces low distances near the goal object position, but is impacted by occlusions of the objects as well as the position of the visually pronounced arm.
Revised: The pixelwise MSE metric produces low distances near the goal object position, but is impacted by occlusions of the objects as well as the position of the visually pronounced arm.

Original: While this analysis does not necessarily directly correspond to control performance, as it ignores the movement of the robot, it demonstrates that our learned distances are aware of the functional similarity of nearby object positions, despite the fact that they are learned entirely from images with actions corresponding to the movement of the arm , not the object .
Revised: While this analysis does not necessarily directly correspond to control performance, as it ignores the movement of the robot, it demonstrates that our learned distances are aware of the functional similarity of nearby object positions, despite the fact that they are learned entirely from images with actions corresponding to the movement of the arm , not the object .

Original: Ablations.
Revised: Ablations.

Original: Our ablation studies aim to answer the following questions: (1) How does Q-learning for learning dynamical distances compare to alternative distance metrics, such as distance in the latent space of a VAE, or dynamical distances learned using direct regression on temporal distances found in random data?
Revised: Our ablation studies aim to answer the following questions: (1) How does Q-learning for learning dynamical distances compare to alternative distance metrics, such as distance in the latent space of a VAE, or dynamical distances learned using direct regression on temporal distances found in random data?

Original: (2) How important is mining negative transitions to the performance of our method?
Revised: (2) How important is mining negative transitions to the performance of our method?

Original: (3) How beneﬁcial is it to combine the learned distance function with planning through a predictive model, as compared to directly acting using the learned policy, as in standard model-free ofﬂine RL?
Revised: (3) How beneﬁcial is it to combine the learned distance function with planning through a predictive model, as compared to directly acting using the learned policy , as in standard model-free ofﬂine RL

Original: To answer the ﬁrst two questions, we perform experiments in the object pushing domain.
Revised: To answer the ﬁrst two questions, we perform experiments in the object pushing domain.

Original: We evaluate alternative distance metrics for visual planning, by duplicating the planning setup, using the same dynamics model, and only modifying the metric used for scoring candidate trajectories.
Revised: We evaluate alternative distance metrics for visual planning, by duplicating the planning setup, using the same dynamics model, and only modifying the metric used for scoring candidate trajectories.

Original: The ﬁrst distance we consider is Euclidean distance in the latent space of a VAE, that is, d ( s, g )
Revised: The ﬁrst distance we consider is Euclidean distance in the latent space of a VAE, that is, d ( s, g )

Original: = (cid:107) e ( s ) − e
Revised: = (cid:107) e ( s ) − e

Original: ( g ) (cid:107) 2 , where e is a learned encoder, which resembles the reward function used in prior work on image-based goal reaching (Nair et al., 2018).
Revised: ( g ) (cid:107) 2 , where e is a learned encoder, which resembles the reward function used in prior work on image-based goal reaching (Nair et al., 2018).

Original: The second is the direct temporal distance regression model described previously.
Revised: The second is the direct temporal distance regression model described previously.

Original: As shown in Figure 5, Q-function distances outperform alternative distances on all of the object pushing tasks.
Revised: As shown in Figure 7, Q-function distances outperform alternative distances on all of the object pushing tasks.

Original: While the temporal distance regression scheme provides competitive performance in some settings, it often provides overestimates of distances between states rather than shortest paths, as shown qualitatively in Figure 6.[Table]
Revised: While the temporal distance regression scheme provides competitive performance in some settings, it often provides overestimates of distances between states rather than shortest paths, as shown qualitatively in Figure 6.

Original: We also ﬁnd that the negative transition mining scheme also consistently improves performance, and is particularly important for the “hard” tasks.
Revised: We also ﬁnd that the negative transition mining scheme also consistently improves performance, and is particularly important for the “hard” tasks.

Original: We hypothesize this is because augmenting the training data in this way causes learned distance functions to better take into account the positions of objects in the scene, rather than just visually prominent components such as arm position.
Revised: We hypothesize this is because augmenting the training data in this way causes learned distance functions to better take into account the positions of objects in the scene, rather than just visually prominent components such as arm position.

Original: To address the third question, we compare our method, which uses learned distances for planning, to the policy discovered when performing Q-learning to learn dynamical distances.
Revised: To address the third question, we compare our method, which uses learned distances for planning, to the policy discovered when performing Q-learning to learn dynamical distances.

Original: As shown in Table 1, the policy learned directly from ofﬂine RL alone is greatly outperformed by MBOLD.
Revised: As shown in Table 1, the policy learned directly from ofﬂine RL alone is greatly outperformed by MBOLD.

Original: We hypothesize that this is due to challenges in advantage learning from ofﬂine data with extremely sparse rewards.
Revised: We hypothesize that this is due to challenges in advantage learning from ofﬂine data with extremely sparse rewards.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original:  We presented a self-supervised approach to tackling goal-reaching tasks, which learns to reach unseen visual goals given only an ofﬂine , random dataset without reward labels .
Revised: [Table] We presented a self-supervised approach to tackling goal-reaching tasks, which learns to reach unseen visual goals given only an ofﬂine , random dataset without reward labels .

Original: Our method combines the strengths of predictive models and learned dynamical distances, where a predictive model can provide effective predictions for planning actions over short horizons, while dynamical distances can provide a useful planning cost that captures distance to goals over longer horizons.
Revised: Our method combines the strengths of predictive models and learned dynamical distances, where a predictive model can provide effective predictions for planning actions over short horizons, while dynamical distances can provide a useful planning cost that captures distance to goals over longer horizons.

Original: By performing visual model predictive control with a learned visual dynamics model and a goal conditioned Q-function as the planning cost, we ﬁnd that our method is able to perform goal reaching tasks more effectively than model-based planning approaches that utilize other reward speciﬁcation techniques, as well as purely model-free methods.
Revised: By performing visual model predictive control with a learned visual dynamics model and a goal conditioned Q-function as the planning cost, we ﬁnd that our method is able to perform goal reaching tasks more effectively than model-based planning approaches that utilize other reward speciﬁcation techniques, as well as purely model-free methods.

Revised: Our method can learn from arbitrary ofﬂine datasets, from those collected using expert demonstrators or skilled policies which may enable more complex behaviors, to randomly collected data as in our experiments.

Original: By leveraging ofﬂine data collected without a speciﬁc goal in mind, our method may make it possible to utilize large, unstructured, open-world robotic manipulation datasets.
Revised: By leveraging ofﬂine data collected without a speciﬁc goal in mind, our method may make it possible to utilize large, unstructured, open-world robotic manipulation datasets.

Original: Scaling up this method to real-world systems and large data sources therefore represents a particularly exciting direction for future work, which may broaden the capabilities and generality of robotic systems.
Revised: Scaling up this method to real-world systems and large data sources therefore represents a particularly exciting direction for future work, which may broaden the capabilities and generality of robotic systems.

Original: A MBOLD I MPLEMENTATION D ETAILS A. D ISTANCE F UNCTION
Revised: A MBOLD I MPLEMENTATION D ETAILS A. D ISTANCE F UNCTION

Original: This section explains the implementation details for our distance function.
Revised: This section explains the implementation details for our distance function.

Original: Following prior work (Fujimoto et al., 2018), we learn two independent Q-functions and use the minimum for performing Bellman backups.
Revised: Following prior work (Fujimoto et al., 2018), we learn two independent Q-functions and use the minimum for performing Bellman backups.

Original: Recall that we sampled goals from two distributions: future states in the same trajectories, and states from different trajectories where the robot arm was in a similar position.
Revised: Recall that we sampled goals from two distributions: future states in the same trajectories, and states from different trajectories where the robot arm was in a similar position.

Original: To implement the second strategy, we ﬁt a k -nearest neighbors graph on 200000 (about 60% of total) dataset observations, and use the (cid:96) 2 arm joint distance as the similarity key.
Revised: To implement the second strategy, we ﬁt a k -nearest neighbors graph on 200000 (about 60% of total) dataset observations, and use the (cid:96) 2 arm joint distance as the similarity key.

Original: Each batch contains equal numbers of transitions generated from each goal sampling method.
Revised: Each batch contains equal numbers of transitions generated from each goal sampling method.

Original: For computational efﬁciency, we implement the k -NN search using the GPU-enabled FAISS library (Johnson et al., 2017).
Revised: For computational efﬁciency, we implement the k -NN search using the GPU-enabled FAISS library (Johnson et al., 2017).

Original: We also modify the reward speciﬁcation scheme by providing a small positive reward at each step where the goal is not reached, and then a large positive reward upon reaching the goal.
Revised: We also modify the reward speciﬁcation scheme by providing a small positive reward at each step where the goal is not reached, and then a large positive reward upon reaching the goal.

Original: Speciﬁcally, we choose to give a reward of 1 by default and 10 when the goal is reached (compared to 0 andrespectively as presented in the discussion in Section 4), although we do not extensively tune this parameter.
Revised: Speciﬁcally, we choose to give a reward of 1 by default and 10 when the goal is reached (compared to 0 andrespectively as presented in the discussion in Section 4), although we do not extensively tune this parameter.

Original: We ﬁnd that it does not affect performance in a statistically signiﬁcant way (results for each reward choice are within 1 standard deviation of one another) to choose this reward over the (0 , 1) rewards.
Revised: We ﬁnd that it does not affect performance in a statistically signiﬁcant way (results for each reward choice are within 1 standard deviation of one another) to choose this reward over the (0 , 1) rewards.

Original: Note that this does not change the interpretation of the Q-function as a shortest path distance, merely slightly complicating the conversion calculations from Q-values to distances in timesteps.
Revised: Note that this does not change the interpretation of the Q-function as a shortest path distance, merely slightly complicating the conversion calculations from Q-values to distances in timesteps.

Original: Finally, we add an additional loss term to perform conservative Q-learning (CQL) (Kumar et al., 2020), a method for ofﬂine model-free RL, which penalizes Q-values of randomly selected actions and increases Q-values of in-dataset actions.
Revised: Finally, we add an additional loss term to perform conservative Q-learning (CQL) (Kumar et al., 2020), a method for ofﬂine model-free RL, which penalizes Q-values of randomly selected actions and increases Q-values of in-dataset actions.

Original: We use the Lagrangian version of CQL to automatically tune the weighting term, and detail the parameters below.
Revised: We use the Lagrangian version of CQL to automatically tune the weighting term, and detail the parameters below.

Original: We ﬁnd using CQL improves performance on the door sliding task from a mean success rate of 41% to 58% , but does not signiﬁcantly impact performance on the others.
Revised: We ﬁnd using CQL improves performance on the door sliding task from a mean success rate of 41% to 58% , but does not signiﬁcantly impact performance on the others.

Original: The Q-function network architecture consists of convolutional and fully connected layers.
Revised: The Q-function network architecture consists of convolutional and fully connected layers.

Original: We deﬁne a network called the convolutional encoder , which will be used throughout the appendix.
Revised: We deﬁne a network called the convolutional encoder , which will be used throughout the appendix.

Original: This takes as input an image of shape 64 × 64 × 6 , containing the starting and goal images concatenated channelwise, and consists of 4 2D convolutional layers, with [8 , 16 , 32 , 64] ﬁlters, respectively, with all with kernel size (4 , 4) and strides of (2 , 2) .
Revised: This takes as input an image of shape 64 × 64 × 6 , containing the starting and goal images concatenated channelwise, and consists of 4 2D convolutional layers, with [8 , 16 , 32 , 64] ﬁlters, respectively, with all with kernel size (4 , 4) and strides of (2 , 2) .

Original: We use Leaky ReLU activations after each intermediate convolutional layer, and batch-norm layers after the second and third Leaky ReLUs.
Revised: We use Leaky ReLU activations after each intermediate convolutional layer, and batch-norm layers after the second and third Leaky ReLUs.

Original: We ﬂatten the output of the convolutional encoder and feed the features through 6 fully-connected linear layers of 128 units each, with the ﬁnal layer outputting a single value.
Revised: We ﬂatten the output of the convolutional encoder and feed the features through 6 fully-connected linear layers of 128 units each, with the ﬁnal layer outputting a single value.

Original: Each intermediate fully-connected layer is followed by a ReLu activation and a batch-norm layer.
Revised: Each intermediate fully-connected layer is followed by a ReLu activation and a batch-norm layer.

Original: The actor network architecture ﬁrst contains the above “convolutional encoder”, whose outputs are ﬂattened and input into a 10 layer MLP with 128 fully connected units each, and ReLu activations and batch-norm layers in between.
Revised: The actor network architecture ﬁrst contains the above “convolutional encoder”, whose outputs are ﬂattened and input into a 10 layer MLP with 128 fully connected units each, and ReLu activations and batch-norm layers in between.

Original: The ﬁnal output, of dimension 4 , is passed through a tanh activation to constrain it to the normalized action space [ − 1 , 1] .
Revised: The ﬁnal output, of dimension 4 , is passed through a tanh activation to constrain it to the normalized action space [ − 1 , 1] .

Original: Additional training hyperparameters are detailed in Table 2.
Revised: Additional training hyperparameters are detailed in Table 2.

Original: A. M ODEL -P REDICTIVE C ONTROL
Revised: A. M ODEL -P REDICTIVE C ONTROL

Original: In Table 3, we describe the parameters for model-based planning in our experiments.
Revised: In Table 3, we describe the parameters for model-based planning in our experiments.

Original: These parameters are shared across all tasks and planning costs (in ablation experiments).
Revised: These parameters are shared across all tasks and planning costs (in ablation experiments).

Original: Most values are selected based on prior work (Ebert et al., 2018b).
Revised: Most values are selected based on prior work (Ebert et al., 2018b).

Original: We ﬁnd that replanning every 6 steps produces slightly better performance than replanning every 13 steps, but not by a large margin, and we do not tune this further due to computation constraints.
Revised: We ﬁnd that replanning every 6 steps produces slightly better performance than replanning every 13 steps, but not by a large margin, and we do not tune this further due to computation constraints.

Original: We sample actions using the ﬁltering scheme described in Nagabandi et al.
Revised: We sample actions using the ﬁltering scheme described in Nagabandi et al.

Original: (2020) to make sequences smoother in time.
Revised: (2020) to make sequences smoother in time.

Original: We initialize sampling distributions using each environment’s data collection parameters, as shown in Table 4.
Revised: We initialize sampling distributions using each environment’s data collection parameters, as shown in Table 4.

Original: To compute the planning cost described in Equation 3, we maximize over α by feeding in the ﬁnal predicted state to the policy network learned by TD3, and using the outputted action as the maximizer.
Revised: To compute the planning cost described in Equation 3, we maximize over α by feeding in the ﬁnal predicted state to the policy network learned by TD3, and using the outputted action as the maximizer.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: A. E NVIRONMENTS
Revised: A. E NVIRONMENTS

Original: The Sawyer environments are adapted from the Meta-World benchmark (Yu et al., 2019a)
Revised: The Sawyer environments are adapted from the Meta-World benchmark (Yu et al., 2019a)

Original: [Table]
Revised: [Table]

Original: and the door sliding environment is based off of the environment presented by Lynch et al.
Revised: and the door sliding environment is based off of the environment presented by Lynch et al.

Original: For each task, we deﬁne the 4 -dimensional action space
Revised: For each task, we deﬁne the 4 -dimensional action space

Original: A such that actions control the Cartesian position of the robot’s end-effector, as well as the robot’s gripper.
Revised: A such that actions control the Cartesian position of the robot’s end-effector, as well as the robot’s gripper.

Original: We randomly generate a set of 100 different test goals for each setting.
Revised: We randomly generate a set of 100 different test goals for each setting.

Original: Each task is deﬁned by a goal image and starting state, on which all methods are tested.
Revised: Each task is deﬁned by a goal image and starting state, on which all methods are tested.

Original: We deﬁne success for each task in terms of the ﬁnal distance to the goal of each relevant object, e
Revised: We deﬁne success for each task in terms of the ﬁnal distance to the goal of each relevant object, e

Original: [Table]
Revised: [Table]

Original: g. object position for the object repositioning task.
Revised: g. object position for the object repositioning task.

Original: A trial is considered successful if the ﬁnal distance is below a certain threshold (cid:15) manually chosen for each task, listed in the table below.
Revised: A trial is considered successful if the ﬁnal distance is below a certain threshold (cid:15) manually chosen for each task, listed in the table below.

Original: We evaluate the success rate of each method overdifferent random training seeds.
Revised: We evaluate the success rate of each method overdifferent random training seeds.

Original: We generate ofﬂine datasets for each task by running random policies for 1 e 4 episodes oftimesteps each.
Revised: We generate ofﬂine datasets for each task by running random policies for 1 e 4 episodes oftimesteps each.

Revised: In the beginning of each episode, object positions are reset uniformly randomly over the range of possible positions across each joint.

Original: The random policy actions are drawn using a ﬁltering technique, which smooths random zero-mean Gaussian samples across time.
Revised: The random policy actions are drawn using a ﬁltering technique, which smooths random zero-mean Gaussian samples across time.

Original: We apply the correlated noise scheme described by Nagabandi et al.
Revised: We apply the correlated noise scheme described by Nagabandi et al.

Original: (2020), setting the hyperparameter β = 0 .
Revised: (2020), setting the hyperparameter β = 0 .

Original: The parameters of the multi-variate Gaussian samples in each dimension are listed in Table 4. 
Revised: The parameters of the multi-variate Gaussian samples in each dimension are listed in Table 4. [Table] threshold

Original: [Table]
Revised: [Table]

Original: B C OMPARATIVE E VALUATION I MPLEMENTATION D ETAILS B. RIG
Revised: B C OMPARATIVE E VALUATION I MPLEMENTATION D ETAILS B. RIG

Original: In this section, we will discuss implementation details of our adaptation of RIG.
Revised: In this section, we will discuss implementation details of our adaptation of RIG.

Original: We begin by training a β -VAE with latent dimension 8 .
Revised: We begin by training a β -VAE with latent dimension 8 .

Original: The VAE is trained on randomly sampled states from the entire ofﬂine dataset.
Revised: The VAE is trained on randomly sampled states from the entire ofﬂine dataset.

Original: For the loss, we use a combination of a maximum likelihood term and a KL divergence term which constrains the latent space to a unit Gaussian.
Revised: For the loss, we use a combination of a maximum likelihood term and a KL divergence term which constrains the latent space to a unit Gaussian.

Original: In particular, we compute the mean pixel error, that is,
Revised: In particular, we compute the mean pixel error, that is,

Original: [Equation]
Revised: [Equation]

Original: s
Revised: s

Original: [Equation]
Revised: [Equation]

Original: 22 , where s is the original image, and ˆ s is the reconstruction, both normalized to be in [0 , 1] .
Revised: 22 , where s is the original image, and ˆ s is the reconstruction, both normalized to be in [0 , 1] .

Original: We add this to the KL divergence between the latent distribution and the unit Gaussian, with a weighting factor of 1 e − 3 on the KL penalty.
Revised: We add this to the KL divergence between the latent distribution and the unit Gaussian, with a weighting factor of 1 e − 3 on the KL penalty.

Original: The architecture of the VAE encoder consists of the “convolutional encoder” described in section
Revised: The architecture of the VAE encoder consists of the “convolutional encoder” described in section

Original: A.1, whose features are passed through two FC layers with 128 units with a ReLu activation and batch-norm layer in between.
Revised: A.1, whose features are passed through two FC layers with 128 units with a ReLu activation and batch-norm layer in between.

Original: The VAE decoder takes as input latent states into two FC layers with 128 units with a batch-norm layer and ReLu activation after each.
Revised: The VAE decoder takes as input latent states into two FC layers with 128 units with a batch-norm layer and ReLu activation after each.

Original: This is followed by the inverted architecture of the encoder, consisting of transposed 2D convolutions.
Revised: This is followed by the inverted architecture of the encoder, consisting of transposed 2D convolutions.

Original: Then, we perform model-free RL in a modiﬁed MDP, using encoded observations as a substitute for environment observations, and computing rewards as negative (cid:96) 2 distances in latent space.
Revised: Then, we perform model-free RL in a modiﬁed MDP, using encoded observations as a substitute for environment observations, and computing rewards as negative (cid:96) 2 distances in latent space.

Original: We sample random goals from the multivariate Gaussian prior ( N (0 , I ) ) at the beginning of every episode.
Revised: We sample random goals from the multivariate Gaussian prior ( N (0 , I ) ) at the beginning of every episode.

Original: We use the open-source implementation of soft actor-critic (SAC) in RLKit, and use the default SAC parameters and architecture found in the implementation, making the following modiﬁcations: We increase the number of layers of all MLP networks from 2 to 6 .
Revised: We use the open-source implementation of soft actor-critic (SAC) in RLKit, and use the default SAC parameters and architecture found in the implementation, making the following modiﬁcations: We increase the number of layers of all MLP networks from 2 to 6 .

Original: We use a maximum path length of 30 steps for consistency with our other experiments, and a discount factor of 0 .
Revised: We use a maximum path length of 30 steps for consistency with our other experiments, and a discount factor of 0 .

Original: Along with the goal sampled from the prior at the beginning of each episode, we ﬁnd that relabeling goals with the achieved observation at the end of the trajectory improves performance, and add these transitions to the replay buffer as well.
Revised: Along with the goal sampled from the prior at the beginning of each episode, we ﬁnd that relabeling goals with the achieved observation at the end of the trajectory improves performance, and add these transitions to the replay buffer as well.

Original: Note that unlike in the original RIG formulation, we do not update the weights of the learned VAE using data collected online.
Revised: Note that unlike in the original RIG formulation, we do not update the weights of the learned VAE using data collected online.

Original: We evaluate the learned policy afterepochs of training, long after environment returns plateau.
Revised: We evaluate the learned policy afterepochs of training, long after environment returns plateau.

Original: B. D REAMER
Revised: B. D REAMER

Original: Dreamer, a model-based method for image-based tasks, also uses a combination of value functions and planning.
Revised: Dreamer, a model-based method for image-based tasks, also uses a combination of value functions and planning.

Original: We adapt Dreamer from its original single-task setting to learn a goal-conditioned policy, reward predictor, and value function; however, we do not condition the dynamics model on the goal.
Revised: We adapt Dreamer from its original single-task setting to learn a goal-conditioned policy, reward predictor, and value function; however, we do not condition the dynamics model on the goal.

Original: Dreamer has been previously demonstrated only in settings where the environment provides rewards to the agent, so we modify the method to learn from unlabeled, ofﬂine data by using experience replay.
Revised: Dreamer has been previously demonstrated only in settings where the environment provides rewards to the agent, so we modify the method to learn from unlabeled, ofﬂine data by using experience replay.

Original: We ﬁnd that using an indicator reward function as in our method or a heuristically deﬁned reward function, image MSE, causes Dreamer to struggle to learn.
Revised: We ﬁnd that using an indicator reward function as in our method or a heuristically deﬁned reward function, image MSE, causes Dreamer to struggle to learn.

Original: We thus additionally demonstrate the performance of Dreamer using a manually speciﬁed arm distance reward for the
Revised: We thus additionally demonstrate the performance of Dreamer using a manually speciﬁed arm distance reward for the

Original: Sawyer reaching task.
Revised: Sawyer reaching task.

Original: We build off of the open source implementation of Dreamer by the original authors, written in
Revised: We build off of the open source implementation of Dreamer by the original authors, written in

Original: TensorFlow2 and found at https://github.com/danijar/dreamer.
Revised: TensorFlow2 and found at https://github.com/danijar/dreamer.

Original: Speciﬁcally, to modify the networks to support goal-conditioning, we add independent convolutional encoders which take the goal image as input to each network.
Revised: Speciﬁcally, to modify the networks to support goal-conditioning, we add independent convolutional encoders which take the goal image as input to each network.

Original: Each encoder consists of 2D convolution layers with [32 , 64 , 128 , 256] ﬁlters and kernel sizes of 4 to each network, and we concatenate the ﬂattened features to the inputs of each network.
Revised: Each encoder consists of 2D convolution layers with [32 , 64 , 128 , 256] ﬁlters and kernel sizes of 4 to each network, and we concatenate the ﬂattened features to the inputs of each network.

Original: We additionally increase the number of fully-connected layers for the value and actor networks from 3 and 2 respectively to 10 .
Revised: We additionally increase the number of fully-connected layers for the value and actor networks from 3 and 2 respectively to 10 .

Original: We use a discount factor of γ = 0 .
Revised: We use a discount factor of γ = 0 .

Original: All other hyperparameter values are defaults from the public implementation.
Revised: All other hyperparameter values are defaults from the public implementation.

Original: For training, we relabel trajectories sampled from the ﬁxed, ofﬂine dataset with a uniformly randomly selected observation from the trajectory as the goal.
Revised: For training, we relabel trajectories sampled from the ﬁxed, ofﬂine dataset with a uniformly randomly selected observation from the trajectory as the goal.

Original: In most of our experiments, we compute the negative pixel-wise MSE as the reward, but in one reaching experiment, we use the negative (cid:96)
Revised: In most of our experiments, we compute the negative pixel-wise MSE as the reward, but in one reaching experiment, we use the negative (cid:96)

Original: Euclidean distance between the arm end-effector position and the goal end-effector position.
Revised: Euclidean distance between the arm end-effector position and the goal end-effector position.

Original: We train for 2000 iterations for each experiment, although initial experiments in which we trained for 20 x longer did not yield improved results.
Revised: We train for 2000 iterations for each experiment, although initial experiments in which we trained for 20 x longer did not yield improved results.

Original: B. G OAL -C ONDITIONED B EHAVIOR C LONING
Revised: B. G OAL -C ONDITIONED B EHAVIOR C LONING

Original: To train a goal-conditioned behavior cloning policy, we begin by relabeling random transitions from the dataset with goals which are later achieved in those trajectories.
Revised: To train a goal-conditioned behavior cloning policy, we begin by relabeling random transitions from the dataset with goals which are later achieved in those trajectories.

Original: Speciﬁcally, we sample stategoal pairs from trajectories in the dataset by ﬁrst selecting the initial state index t i uniformly from all timesteps, and then selecting the goal state index t g uniformly from timesteps greater than t i .
Revised: Speciﬁcally, we sample stategoal pairs from trajectories in the dataset by ﬁrst selecting the initial state index t i uniformly from all timesteps, and then selecting the goal state index t g uniformly from timesteps greater than t i .

Original: We then train a neural network to predict the transition action a i given the state
Revised: We then train a neural network to predict the transition action a i given the state

Original: s i and the relabeled goal s g , using a mean-squared error loss.
Revised: s i and the relabeled goal s g , using a mean-squared error loss.

Original: The network architecture is the same as that of the actor network used in Q-learning for MBOLD, described in Appendix A.1.
Revised: The network architecture is the same as that of the actor network used in Q-learning for MBOLD, described in Appendix A.1.

Original: We train the model for 3125000 iterations ( 1000 epochs) using a batch size of 32 , and use the same optimizer and learning rate as the distance learned for MBOLD.
Revised: We train the model for 3125000 iterations ( 1000 epochs) using a batch size of 32 , and use the same optimizer and learning rate as the distance learned for MBOLD.

Original: B. S EARCH ON THE R EPLAY B UFFER
Revised: B. S EARCH ON THE R EPLAY B UFFER

Original: For SoRB, we train a distributional Q-function to represent distances as in the original paper.
Revised: For SoRB, we train a distributional Q-function to represent distances as in the original paper.

Original: Distributional RL discretizes possible value estimates into a set of bins – we use 10 for all of our experiments.
Revised: Distributional RL discretizes possible value estimates into a set of bins – we use 10 for all of our experiments.

Original: We train this distributional Q-function for 300 epochs, as in the distance function training for MBOLD.
Revised: We train this distributional Q-function for 300 epochs, as in the distance function training for MBOLD.

Original: We also use the same architecture and training scheme, altering the number of outputs to 10 bins and using the KL-divergence loss for the distributional Q-function as in Eysenbach et al.
Revised: We also use the same architecture and training scheme, altering the number of outputs to 10 bins and using the KL-divergence loss for the distributional Q-function as in Eysenbach et al.

Original: However, unlike in Eysenbach et al.
Revised: However, unlike in Eysenbach et al.

Original: (2019), we train on just the ﬁxed, ofﬂine dataset.
Revised: (2019), we train on just the ﬁxed, ofﬂine dataset.

Original: We then perform the planning portion of SoRB with the “maxdist” parameter set to 4 , after manual tuning.
Revised: We then perform the planning portion of SoRB with the “maxdist” parameter set to 4 , after manual tuning.

Original: We use a graph size of 2000 states for all experiments, due to computational constraints.
Revised: We use a graph size of 2000 states for all experiments, due to computational constraints.

Original: We ﬁnd that the policy learned through Q-learning performs very poorly at reaching subgoals, so we instead substitute the GCBC policy for this purpose.
Revised: We ﬁnd that the policy learned through Q-learning performs very poorly at reaching subgoals, so we instead substitute the GCBC policy for this purpose.

Original: We ﬁnd that this greatly improves performance across all tasks.
Revised: We ﬁnd that this greatly improves performance across all tasks.

Original: B. V ISUAL F ORESIGHT
Revised: B. V ISUAL F ORESIGHT

Original: To compare MBOLD to visual foresight, we use the same dynamics model and planning setup as in MBOLD, however, we substitute the learned dynamical distance function with the (cid:96) 2 pixel error cost used in visual foresight.
Revised: To compare MBOLD to visual foresight, we use the same dynamics model and planning setup as in MBOLD, however, we substitute the learned dynamical distance function with the (cid:96) 2 pixel error cost used in visual foresight.

Original: C A BLATION E XPERIMENTS I MPLEMENTATION D ETAILS C. VAE D ISTANCE
Revised: C A BLATION E XPERIMENTS I MPLEMENTATION D ETAILS C. VAE D ISTANCE

Original: We use the same architecture as the VAE used in the RIG comparison described in Appendix B.
Revised: We use the same architecture as the VAE used in the RIG comparison described in Appendix B.

Original: We set the latent space dimension to 256 and weight the KL divergence term using a factor of 1 e − 5 .
Revised: We set the latent space dimension to 256 and weight the KL divergence term using a factor of 1 e − 5 .

Original: We train the model for 3125000 iterations ( 1000 epochs) using a batch size of 32 , and use the same optimizer and learning rate as the distance learned for MBOLD.
Revised: We train the model for 3125000 iterations ( 1000 epochs) using a batch size of 32 , and use the same optimizer and learning rate as the distance learned for MBOLD.

Original: C. T EMPORAL D ISTANCE R EGRESSION
Revised: C. T EMPORAL D ISTANCE R EGRESSION

Original: To train the temporal distance regression model, we sample state-goal pairs from trajectories in the dataset by ﬁrst selecting the initial state index t i uniformly from all timesteps, and then selecting the goal state index t g uniformly from timesteps greater than t i .
Revised: To train the temporal distance regression model, we sample state-goal pairs from trajectories in the dataset by ﬁrst selecting the initial state index t i uniformly from all timesteps, and then selecting the goal state index t g uniformly from timesteps greater than t i .

Original: We compute the label for this pair as min( t g − t i , maxdist ) , where maxdist is a hyperparameter we set to 10 .
Revised: We compute the label for this pair as min( t g − t i , maxdist ) , where maxdist is a hyperparameter we set to 10 .

Original: The maxdist parameter helps to improve the optimality of distances on average.
Revised: The maxdist parameter helps to improve the optimality of distances on average.

Original: We train the neural network to regress this target label using an (cid:96) 2 error loss.
Revised: We train the neural network to regress this target label using an (cid:96) 2 error loss.

Original: We train the network for 3125000 iterations ( 1000 epochs) with a batch size of 32 , and use the same optimizer and learning rate as the distance learned for MBOLD.
Revised: We train the network for 3125000 iterations ( 1000 epochs) with a batch size of 32 , and use the same optimizer and learning rate as the distance learned for MBOLD.

Original: The architecture for the temporal distance regression model begins with the convolutional encoder described in Appendix B. Its ﬂattened outputs are fed into 5 fully-connected layers of 256 units each, with batch-norm and ReLu activations after each intermediate layer.
Revised: The architecture for the temporal distance regression model begins with the convolutional encoder described in Appendix B. Its ﬂattened outputs are fed into 5 fully-connected layers of 256 units each, with batch-norm and ReLu activations after each intermediate layer.

Original: C. Q-F UNCTION P OLICY
Revised: C. Q-F UNCTION P OLICY

Original: We ﬁnd that the policy directly learned by our method when learning distances performs extremely poorly.
Revised: We ﬁnd that the policy directly learned by our method when learning distances performs extremely poorly.

Original: However, performing Q-learning using random shooting over 100 uniformly random actions selected from [ − 1 , 1] 4 to optimize over actions to compute target values produces much better results when used directly as a policy, compared to using an actor network to perform this optimization as in our method.
Revised: However, performing Q-learning using random shooting over 100 uniformly random actions selected from [ − 1 , 1] 4 to optimize over actions to compute target values produces much better results when used directly as a policy, compared to using an actor network to perform this optimization as in our method.

Original: Therefore, we report results from acting according to this random shooting method.
Revised: Therefore, we report results from acting according to this random shooting method.

Original: At test time, we estimate the optimal action a (cid:63) =
Revised: At test time, we estimate the optimal action a (cid:63) =

Original: arg max a Q ( s t , a, g ) by again samplinguniformly random actions, and selecting the best one. 
Revised: arg max a Q ( s t , a, g ) by again samplinguniformly random actions, and selecting the best one. D C OMPUTATIONAL C OMPLEXITY A NALYSIS

Revised: In this section, we discuss the computation complexity of training and acting using MBOLD.

Revised: Training: Training the dynamics model takes about 30 hr while training the distance function takes about 5 hr.

Revised: These training times are dwarfed by the cost of collecting data in the real world, which could take on the order of 3-4 days in the real world (but can be reused for various tasks).

Revised: In contrast, a single RL approach only requires learning the distance function.

Revised: While this means that it takes

Revised: MBOLD signiﬁcantly longer to train than the single RL approach, note that the dynamics model can be shared across many tasks.

Revised: We train the dynamics model for 200k and distance function for 94k training steps.

Revised: A training step for the dynamics model involves one forward and backward pass through the dynamics model.

Revised: A training step for the distance function requires sampling positive and negative goals, two Q-function forward passes and a policy network forward pass to compute target values and current Q-values, and a backward pass to update model parameters.

Revised: In contrast, a single RL approach would just learn the distance function, not the dynamics model.

Revised: From the above estimates, this means that training steps for the dynamics model are around 3 times slower than training steps for the distance model.

Revised: Because the dynamics model can be used to perform many tasks, this cost is amortized over these tasks, as compared to a single RL approach.

Revised: Acting:

Revised: Selecting a sequence of actions (6 actions in our experiments) using MBOLD requires one forward pass of the dynamics model for each CEM iteration (3 total in our experiments), and one forward pass through the distance function and policy network.

Revised: Amortized over a trajectory, this amounts to about 2 sec wall clock time per action, which can be sped up by around 2x with similar performance by replanning less frequently.

Revised: For a single RL approach, each action would require just one forward pass through the policy network.

Revised: E P LANNING HORIZON ABLATION EXPERIMENTS

Revised: In this section, we investigate the effect of the planning horizon h on control performance.

Revised: After training distance functions according to Appendix A.1, we perform planning with three different settings for h on the simulated block pushing tasks.

Revised: We present the results in Figure 7.

Revised: We ﬁnd that a longer planning horizon is beneﬁcial, especially for solving more difﬁcult tasks.

Revised: We hypothesize that this is because longer planning horizons allow the planner and distance function to better distinguish promising predicted states, while the ﬁdelity of state predictions remains relatively high.

Revised: F R ANDOM OBJECT RESET ABLATION EXPERIMENTS

Revised: In this section, we perform experiments to evaluate the impact of the distribution of initial object position on task performance.

Revised: In particular, we look at the single-object Sawyer pushing task.

Revised: We collect an additional dataset with the same policy and other parameters as that used in the main comparative evaluations, but restrict the random object initialization position to be within [ − 0 .

Revised: 05]as opposed to [ − 0 . 2 , 0 . 2] 2 .

Revised: This represents a 16x reduction in the area of possible initializations.

Revised: We then train a new dynamics model and distance function from scratch and compare the control performance on the same benchmark tasks from the main comparisons.

Revised: We present the results in

Revised: Table 5.

Revised: We ﬁnd that the control performance on these tasks remain within one standard deviation despite the restriction in reset position.

Revised: [Table]