SlowGuess/NeurIPS_2023_OCR huggingface dataset:
David Brandfonbrener* New York University
Ofir Nachum Google
In recent years, domains such as natural language processing and image recognition have popularized the paradigm of using large datasets to pretrain representations that can be effectively transferred to downstream tasks. In this work we evaluate how such a paradigm should be done in imitation learning, where both pretraining and finetuning data are trajectories collected by experts interacting with an unknown environment. Namely, we consider a setting where the pretraining corpus consists of multitask demonstrations and the task for each demonstration is set by an unobserved latent context variable. The goal is to use the pretraining corpus to learn a low dimensional representation of the high dimensional (e.g., visual) observation space which can be transferred to a novel context for finetuning on a limited dataset of demonstrations. Among a variety of possible pretraining objectives, we argue that inverse dynamics modeling – i.e., predicting an action given the observations appearing before and after it in the demonstration – is well-suited to this setting. We provide empirical evidence of this claim through evaluations on a variety of simulated visuomotor manipulation problems. While previous work has attempted various theoretical explanations regarding the benefit of inverse dynamics modeling, we find that these arguments are insufficient to explain the empirical advantages often observed in our settings, and so we derive a novel analysis using a simple but general environment model.
Pipelines in image recognition and natural language processing commonly use large datasets to pretrain representations that are then transferred to downstream tasks where data is limited [Devlin et al., 2018, Chen et al., 2020, Radford et al., 2021]. In this paper, we consider how this paradigm can be applied to imitation learning [Pomerleau, 1991, Ho and Ermon, 2016, Kostrikov et al., 2019]. In contrast to supervised learning domains where datasets consist of input-output pairs, imitation learning datasets consist of trajectories with both the input-output mapping to be learned (namely, observation-action pairs) as well as information about the dynamics of the environment. Given this additional structure, it is worthwhile to study pretraining approaches that can incorporate this structure to improve beyond methods from traditional supervised learning domains.
To formalize the precise notion of transfer between pretraining and finetuning phases, we consider a multitask imitation setting where the environment (i.e., the transition dynamics) is fixed and data is comprised of trajectories of task experts acting in this environment. A task is defined by a latent context variable that is observed by an expert demonstrator, but is not contained in the dataset, as shown in Figure 1. During pretraining, we have access to a large number of trajectories from various tasks, while during finetuning we have access to a small number of trajectories from a single task. The goal is thus to use the pretraining dataset to learn representations that contain information about the environment that facilitates efficient learning of the finetuning task.
* david.brandfonbrener@nyu.edu
A number of existing works have proposed objectives for representation learning that are applicable in this setting [Schwarzer et al., 2021, Stooke et al., 2021, Yang and Nachum, 2021, Yang et al., 2023], and we consider a variety of algorithms and modes of analysis to determine which approach is the most promising. Algorithmically, we consider four generic classes of objectives for pretraining: inverse dynamics, behavior cloning, forward dynamics, and static observation modeling (Figure 1). We conduct two types of analysis. First, we conduct an extensive empirical evaluation and introspection of the candidate algorithms along with several strong baselines. Second, we present a simple but general theoretical model of the multitask representation learning problem and analyze the relative merits of the candidate algorithms under this model.
Our main results from these analyses are summarized as follows:
1. Across a broad array of experiments from visual observations in six environments, out of all approaches considered, inverse dynamics is the only one that consistently outperforms the baseline of training a model from scratch. The performance of inverse dynamics even matches that of finetuning from ground truth low-dimensional states on in-distribution contexts. Moreover, we find that inverse dynamics scales the best with pretraining dataset size and most effectively maintains relevant information about the observation in its learned representation.
2. In our simplified model of representation learning, we show that inverse dynamics pretraining efficiently recovers the ideal representation while behavior cloning can suffer from confounding and forward dynamics can suffer from poor sample efficiency. These results provide intuition for the empirical results and motivate why inverse dynamics pretraining is so performant and robust.
(a) Multitask imitation.
(b) ID
(d) FD
Figure 1: (a) A graphical models of the setting. Shaded nodes indicate observed variables. The expert behavior (i.e.,  $o_i \rightarrow a_i$ ) is determined by an unobserved context variable  $c$  while the transition dynamics (i.e.,  $(o_i, a_i) \rightarrow o_{i+1}$ ) are determined by the environment dynamics. (b)-(e) illustrate the candidate algorithms. We use blue to indicate inputs to the algorithm and green to indicate prediction targets. ID = inverse dynamics, BC = behavior cloning, FD = forward dynamics, Cont = contrastive learning. (f) Shows success of policies finetuned on top of various representations averaged across all datasets in our suite for default dataset sizes. Inverse dynamics (shown in green) is the only method to substantially outperform the baseline of training from scratch (shown in black). Further details about the experimental protocol and results are in Sections 4 and 5.
(c) BC
(e) Cont
(f) Experimental results
As explained above, pretraining a representation has become a dominant paradigm in computer vision and natural language processing [Devlin et al., 2018, Chen et al., 2020, Radford et al., 2021]. Determining how to best leverage similar pretraining techniques in decision making problems is an important step towards extending the success of supervised learning into more temporally extended problems like those in robotics [Yang et al., 2023].
Prior work proposes several possible pretraining objectives for learning features for decision-making (and illustrated in Figure 1). First, inverse dynamics modeling has been proposed in several settings, although never as a representation learning algorithm for multitask imitation. Most directly related to our work is Efroni et al. [2021], Lamb et al. [2022] which use multi-step inverse dynamics for feature extraction for exploration in reinforcement learning (RL) in the presence of exogenous noise. Later work from Islam et al. [2022] extended this approach to offline RL. Less closely related are Pathak et al. [2017] which uses inverse dynamics in the context of exploration and Baker et al. [2022], Venuto et al. [2022] which use an inverse dynamics model to label video data with actions for imitation.
Another, perhaps simpler approach is to use behavior cloning as a pretraining algorithm. Arora et al. [2020] shows that this can be a well-motivated approach to pretraining a representation when the task variable is observed. Other work uses behavior cloning objectives to pretrain representations of temporally extended actions [Ajay et al., 2020] or priors for offline RL [Zang et al., 2022].
A third approach is to model the forward dynamics of the system as a pretraining objective. Most directly related to our work, Nachum and Yang [2021] show that this is a well-motivated technique for imitation learning and provide empirical evidence on single task atari games, but do not compare to inverse dynamics. This technique has also been explored in empirical work for online and offline RL [Schwarzer et al., 2021, Laskin et al., 2020, Aytar et al., 2018, Lee et al., 2022b, Wu et al., 2023].
Finally, a method which we will refer to as static observation modeling does not leverage information about dynamics and rather directly uses self-supervised methods from computer vision [Pari et al., 2021, Chen et al., 2020, Grill et al., 2020]. This approach does not take advantage of any additional structure in an imitation learning setting, but has nevertheless worked well in some settings.
Several empirical studies of representation learning for decision-making already exist. Most closely related to this work, [Chen et al., 2022] conducts an empirical evaluation of representations for imitation and finds that none of them consistently outperform training directly from pixels. However, this prior work (a) considers much larger finetuning datasets which can dramatically reduce the benefits of pretraining, and (b) considers different environments than we do, where the gap between pretraining and finetuning tasks is less controlled. Another line of work like Nair et al. [2022] attempts to pretrain general representations using large human-collected video datasets like Ego4d [Grauman et al., 2022]. In contrast, we focus on a more carefully controlled (albeit smaller scale) experimental settings where we can derive a more clear understanding of the relative merits of different pretraining objectives. Another empirical study from Stooke et al. [2021] considers representations in online reinforcement learning. Meanwhile, Yang and Nachum [2021] considers representations for imitation but does not consider image-based or multitask problems. Moreover, none of these works includes a theoretical understanding for the findings presented therein.
A further discussion of pretraining in the context of imitation can be found in Appendix A.
3 Problem setup
Here we present the formal setup for our problem setting of reward free pretraining from multitask expert data. We formalize this as a contextual MDP with rich (i.e., visual) observations where the latent context determines the initial state and reward functions.
Environment. We model the environment as a contextual MDP with context-independent dynamics:
$$
c \sim P _ {c}, \quad o _ {0} \sim \rho_ {c}, \quad r _ {i} = r _ {c} \left(s _ {i}, a _ {i}\right), \quad o _ {i + 1} \sim T \left(o _ {i}, a _ {i}\right). \tag {1}
$$
Importantly, we consider the context variable  $c$  and rewards  $r_c$  to be latent, i.e., they are not available during training, and only used to evaluate a learned policy. At a high level, this captures the setting where the task (defined by the context variable) may change, but the dynamics of the world do not. For example, the context variable could be a continuous variable like a goal position that the expert is navigating towards or a discrete variable representing a behavior like locking a door.
Data generation. Data is generated by executing policies  $\pi$  that map observations to actions in the environment. We consider two different datasets for any given problem. First there is a large multi-context pretraining dataset that will be used for representation learning, specifically to learn an observation encoder. Second, there is a small single-context finetuning dataset for policy learning on top of the pretrained representation. The multi-context pretraining data is generated as follows:
$$
D _ {p r e} = \left\{\tau_ {i} \right\} _ {i = 1} ^ {N _ {p r e}}: c \sim P _ {c}, \tau = (o _ {0}, a _ {0}, o _ {1} \dots) \sim P ^ {\pi_ {c}}, \pi_ {c} \approx \pi_ {c} ^ {*} = \arg \max _ {\pi} J _ {r _ {c}} (\pi), \tag {2}
$$
where  $J_{r_c}(\pi)$  denotes the expected return of  $\pi$  when the reward is  $r_c$ . Note that the demonstration policy has access to the latent context  $c$ , but this latent context is not observed in the data.
Then the single-context finetuning data is generated for context  $c_{fine}$  as follows:
$$
D _ {f i n e} = \left\{\tau_ {i} \right\} _ {i = 1} ^ {N _ {f i n e}}: \tau = \left(o _ {0}, a _ {0}, o _ {1} \dots\right) \sim P ^ {\pi_ {c f i n e}}. \tag {3}
$$
Pretraining. The goal of the paper is to analyze different methods for pretraining feature extractors. Training of the encoders  $\phi$  to minimize a loss  $\ell$  proceeds as follows:
$$
\hat {\phi}: \mathcal {O} \rightarrow \mathbb {R} ^ {d} = \arg \min  _ {\phi} \underset {D _ {p r e}} {\mathbb {E}} [ \ell (\phi , \tau_ {i}) ]. \tag {4}
$$
A full description of the losses  $\ell$  used by different algorithms will come in Section 4.2. For simplicity (and in keeping with prior work [Nachum and Yang, 2021, Chen et al., 2022]) we will consider  $\ell$  to only be a function of transitions  $(\sigma_i^j,a_i^j,\sigma_i^{j'})$  rather than full trajectories to leverage the Markovian structure. We also run some ablations of including multistep information in Appendix B and find little difference.
Finetuning. Features are evaluated by finetuning a small policy head on top of the frozen features:
$$
\hat {\pi} _ {\hat {\phi}}: \mathbb {R} ^ {d} \rightarrow \mathcal {A} = \arg \min  _ {\pi} \underset {D _ {f i n e}} {\mathbb {E}} [ \ell (\pi , a _ {i} ^ {j}, \hat {\phi} (o _ {i} ^ {j})) ]. \tag {5}
$$
In all of our experiments,  $\ell$  is the mean squared error loss for behavior cloning. We elect to use frozen features to allow for simple and clear evaluation of the representations. This is in keeping with prior work on representations for imitation [Nachum and Yang, 2021, Chen et al., 2020, Nair et al., 2022] as well as computer vision [Chen et al., 2020].
Evaluation. Finally, we evaluate the finetuned policy by performing rollouts in the finetuning environment with context  $c_{fine}$  to estimate  $J_{r_{c_{fine}}}(\hat{\pi}_{\hat{\phi}})$ . In our tasks we usually consider  $r_{c_{fine}}$  to be a binary indicator of successful completion of the finetuning task.
4.1 Environments and Datasets
We design a suite of tasks and datasets to probe the capabilities of various representation learners for downstream imitation. We focus on robotic manipulation from vision as this is an important sequential decision making task that depends on learning task-relevant visual representations where pretraining deep visual feature extractors is a popular approach. Our suite consists of six different pretraining datasets on varied tasks and of varied size. Each pretraining dataset has several associated finetuning datasets and simulation environments that allow for online evaluation of learned policies.
(a) Pointmass
(b)Pick  $^+$  place
(c) Door
Figure 2: Our six datasets: (a) Pointmass navigation with latent goals. (b) Pick and place with latent goals. (c) Multitask manipulation of a door. (d) Sequential kitchen manipulation. (e) Multitask manipulation of diverse objects, where we consider two different train-eval splits ML45 and R3M.
(e) Metaworld (ML45/R3M)
All tasks are performed from visual inputs, as shown in Figure 2. Each pair of pretraining-finetuning datasets requires a slightly different type of generalization as dictated by the different types of context variable. Specifically, in the pointmass and pick+place datasets the context variable is a latent goal position, while in the door and metaworld datasets the context variable is a discrete identity of a desired behavior, and in the kitchen datasets the context variable is a discrete ordered sequence of subtasks. The datasets are described in full detail in Appendix C.
4.2 Algorithms
We consider nine different representations across our suite of experiments. These representations include baseline and skyline/oracle performance as well as five representations that are pretrained on
Table 1: Description of the different datasets used in the experiments. Dataset sizes are measured in number of trajectories ( $N_{traj}^{pre}$  for pretraining and  $N_{traj}^{fine}$  for finetuning) and given as ranges with default values in bold. Trajectory lengths vary from 50 to 400 steps. These default sizes may vary in each experiment when indicated. Each datasets contains a certain number of latent contexts ( $N_{context}^{pre}$  and  $N_{context}^{fine}$ ). For each finetuning context, we sample datasets with  $N_{seed}^{fine}$  different seeds.
Environment	Npretraj	Nfinetraj	Nprecontext	Nfinecontext	Nfineseed
Pointmass	(1e1, 1e2, 1e3)	(1, 2, 5, 10)	Npretraj	5	1
Pick + place	(1e1, 1e2, 1e3)	(2, 5, 10, 20)	Npretraj	5	1
Door	(1e1, 1e2, 1e3)	(2, 5, 10, 20)	3	1	5
Kitchen	(50, 150, 450)	(2, 5, 10, 15)	21	3	5
MW-ML45	(1e2, 1e3, 1e4)	(2, 5, 10, 20)	45	5	5
MW-R3M	(1e2, 1e3, 1e4)	(2, 5, 10, 20)	45	5	5
our own pretraining datasets described above. Each of the representations will be referred to by its bolded name after it is described.
All algorithms (except for the Imagenet and R3M baselines) share the exact same encoder architecture to control as best we can for variation in architecture between methods. Each method is pretrained for the same number of gradient steps. Additional training details can be found in Appendix C.
Skyline/oracle. As a skyline or oracle representation we directly use the low dimensional states (States) from the simulator. Depending on the task, this representation includes the position of the robot, position of the object to be manipulated, and/or position of the goal. A full description of the per environment state variables can be found in Appendix C.
Baselines. We consider three baseline representations that are not trained on our pretraining datasets. The first is to directly use the pixels with image augmentations (Pixels + Aug) to train an encoder and a policy from scratch on the finetuning data. It is essential to use the augmentations to ensure that this a strong baseline. The second is features of a ResNet18 pretrained onImagenet (Imagenet). The last consists of the features of a ResNet18 that is specifically pretrained for robotic manipulation by Nair et al. [2022] on the Ego4d dataset (R3M).
Inverse dynamics. The primary representation learning objective that we consider is inverse dynamics (ID) which models the distribution  $P(a|o,o')$  using an architecture that first encodes  $o,o'$  with an encoder  $\phi$  and then predicts  $a$  with a small MLP  $f$ :
$$
\phi_ {I D} ^ {*} = \arg \min  _ {\phi} \min  _ {f} \mathbb {E} _ {o, a, o ^ {\prime}} [ (a - f (\phi (o), \phi (o ^ {\prime}))) ^ {2} ]. \tag {6}
$$
Behavior cloning. A simpler alternative objective is to directly apply behavior cloning (BC) to the multitask actions in the pretraining dataset conditioned on the observations using MSE loss. The learner is parameterized as an encoder  $\phi$  followed by a small MLP  $\pi$ :
$$
\phi_ {B C} ^ {*} = \arg \min  _ {\phi} \min  _ {\pi} \mathbb {E} _ {o, a} [ (a - \pi (\phi (o))) ^ {2} ]. \tag {7}
$$
Forward dynamics. We consider two representation learners that predict the forward dynamics of the system. The first is explicit forward dynamics (FD-e) which explicitly constructs a model of the forward dynamics in the space of observations by encoding the current observation and then attempting to reconstruct the next observation  $o^{\prime}$  using a decoder  $d$ :
$$
\phi_ {E F D} ^ {*} = \arg \min  _ {\phi} \min  _ {d} \mathbb {E} _ {o, a, o ^ {\prime}} [ (o ^ {\prime} - d (\phi (o), a)) ^ {2} ]. \tag {8}
$$
The second objective is implicit forward dynamics (FD-i) which implicitly constructs a model of the forward dynamics using contrastive learning. Explicitly, we consider a form of contrastive learning where an energy function is defined as the inner product of L2-normalized projected embeddings (given by projection MLPs  $f_{1}$ ,  $f_{2}$ ) which is then passed into an InfoNCE loss:
$$
E \left(o, a, o ^ {\prime}\right) = \exp \left(f _ {1} \left(\phi (o), a\right) ^ {\top} f _ {2} \left(\phi \left(o ^ {\prime}\right)\right)\right), \tag {9}
$$
$$
\phi_ {I F D} ^ {*} = \arg \min  _ {\phi} \min  _ {f _ {1}, f _ {2}} \mathbb {E} _ {o, a, o ^ {\prime}} [ - \log (E (o, a, o ^ {\prime})) + \log \mathbb {E} _ {\bar {o} ^ {\prime}} [ E (o, a, \bar {o} ^ {\prime}) ] ]. \tag {10}
$$
Static observation modeling Finally, we consider a baseline that simply models  $P(o)$ . Rather than modeling this explicitly with reconstruction, we use a contrastive loss (Cont) where we use image augmentations to construct pairs of  $o$  and  $\bar{o}$  that do not rely on the dynamics of the environment at all. Again we use the InfoNCE loss, in what can be seen as a variant of SimCLR:
$$
E \left(o, o _ {a u g}\right) = \exp \left(f \left(\phi (o)\right) ^ {\top} \pi \left(\phi \left(o _ {a u g}\right)\right)\right), \tag {11}
$$
$$
\phi_ {C o n t} ^ {*} = \arg \min  _ {\phi} \min  _ {f} \mathbb {E} _ {o, o _ {a u g}} [ - \log (E (o, o _ {a u g})) + \log \mathbb {E} _ {\bar {o} _ {a u g}} [ E (o, \bar {o} _ {a u g}) ] ]. \tag {12}
$$
We want to determine which representation learning objective is best, but the precise answer will depend on the situation. To get a clearer understanding of this sometimes ambiguous performance we conduct a variety of controlled experiments on our diverse suite of datasets. We focus on the following questions to guide our empirical analysis:
1. How do factors of the datasets impact performance of algorithms?
2. How are the learned representations similar to and different from each other?
Note: we will focus on presenting aggregate statistics across all datasets in the main text, but full results can be found in Appendix B. Full details about the methodology are in Appendix C and code is at https://github.com/davidbrandfonbrener/imitation_pretraining.
5.1 Impact of dataset on representation learning performance
Scaling with data size. The performance of each algorithm can be highly sensitive to both pretraining and finetuning sizes. Thus, instead of producing one simple summary statistic, we sweep over both the size of the finetuning data (for default pretraining size) and size of the pretraining data (for default finetuning size). The results of these sweeps are presented in Figure 3.
Figure 3: Average success rate after finetuning averaged across datasets, contexts, and seeds. Error bars show the standard error across contexts and seeds, averaged across datasets. The plots show sweeps across finetuning size with default pretraining size (left) or pretraining size with default finetuning size (right) measured in units according to Table 1. Methods that do not depend on pretraining size are shown as horizontal lines.
The sweeps both suggest that inverse dynamics outperforms the alternatives. First, on the finetuning size sweep, we see that the ID line is the only one that consistently outperforms training from scratch on Pixels + Aug. This gap is largest at small finetuning sizes, which are perhaps the most interesting case since that is when we expect pretraining to be useful. Second, the pretraining size sweep indicated that ID is scaling the most efficiently with pretraining size. Further results, including breakdowns across each dataset can be found in Appendix B.
In distribution vs. out of distribution eval tasks. The way that our datasets are constructed, the door, kitchen, metaworld-ml45, and metaworld-r3m datasets only have a finite number of possible contexts that is much smaller than the number of pretraining trajectories. For our default datasets, we elected to construct a train-test split of contexts to ensure that the contexts used for finetuning are
not seen during pretraining. As a result, the default finetuning tasks can be in some sense "out of distribution", measuring extrapolation as opposed to in-distribution generalization. For example, in the door dataset, we pretrain on door opening, closing, and unlocking (with varied door position) and then finetune on door locking (again with varied position).