Inverse Dynamics Pretraining Learns Good Representations for Multitask Imitation
Authors: David Brandfonbrener
Ofir Nachum
Joan Bruna
Venue: NeurIPS 2023 poster (2023)
Keywords: representation learning
imitation learning

Abstract: 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.

Inverse Dynamics Pretraining Learns Good Representations for Multitask Imitation
David Brandfonbrener' New York University
Ofir Nachum Google
Joan Bruna New York University
Abstract
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 effectively transferred to downstream tasks. In this work we evaluate how such paradigm should be done in imitation learning, where both pretraining and finetuning data are trajectories collected by experts interacting with an unknown environment. Namely, consider setting where the pretraining corpus consists multitask demonstrations and the task for each demonstration is set by an unobserved latent context variable. The goal is to the pretraining corpus to learn a low dimensional representation of the high dimensional (e.g.. visual) observation space which can be transferred to novel context for finetuning a limited dataset of demonstrations. Among a variety of possible pretraining objectives, we argue that inverse dynamics modeling i.e., predicting action given the observations appearing before and after it in the demonstration well-suited to this setting. We provide empirical evidence of this claim through evaluations on 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 simple but general environment model.
Introduction
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 al., 2018 Chen al.) 2020 Radford al.] [00T]. 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 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 multitask imitation setting where the environment (i.e., the transition dynamics) is fixed and data comprised of trajectories of task experts acting this environment. task defined by latent context variable that is observed by an expert demonstrator, but is not contained in the dataset, as shown Figure | During pretraining, we have access to large number of trajectories from various tasks, while during finetuning we have access to small number of trajectories from single task. The goal thus to use the pretraining dataset to learn representations that contain information about the environment that facilitates efficient learning of the finetuning task.
david. Srandfoobbenerrdhng
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

A number of existing works have proposed objectives for representation learning that are applicable
in this setting [Schwarzer al. 2021 Stooke al. 2021 Yang and Nachum 2021 Yang al.) 2023). and we consider variety of algorithms and modes of analysis to determine which approach is most promising. Algorithmically, we consider four generic classes of objectives for pretraining: inverse dynamics, behavior cloning, forward dynamics, and static observation modeling (Figure D. 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 simple but general theoretical model 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 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 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 so performant and robust.
1.0
States
R3M
FD
megenet
Cont
(b) ID
(c) BC
0.4
02
(a) Multitask imitation.
(d) FD
(e) Cont
Experimental results
Figure 1: (a) A graphical models of the setting. Shaded nodes indicate observed variables. The expert behavior (i.e., 0g + a;) is determined by an unobserved context variable while the transition dynamics (i.e., (oi, as) -> 0:+11) 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 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 and
2 Related work
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. 22200 Radford et al. (021]. Determining how to best leverage similar pretraining techniques decision making problems an important step towards extending the success of supervised learning into more temporally extended
problems like those robotics Yang et al. 20023]
Prior work proposes several possible pretraining objectives for learning features for decision-makinn (and illustrated in Figure T]. 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 Efromi et al. (202T). Lamb et al. [2022) which use multi-step inverse dynamics for feature extraction for exploration
reinforcement learning (RL) in the presence of exogenous noise. Later work from Islam et al. 22022 extended this approach to offline RI. Less closely related are Pathak
al. 2011 which uses inverse dynamics the context of exploration andBaa et al. [2022] Venuto
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 a well-motivated approach to pretraining representation when the task variable is observed. Other work uses behavior cloning objectives to pretrain representations of temproally extended actions (Ajay et al., 2020] or priors for offline RI. (Zang et al. 2022].
A third approach is to model the forward dynamies of the system as a pretraining objective. Most directly related to our work, Nachum and Yang [202T] 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, method which 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., 202 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 recision-makinn already exist. Most closely related to this work, [Chen et al.) 2022] conducts 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 et al. [2022) attempts to pretrain general representations using large human-collected video datasets like Ego4d [Gramman more carefully controlled (albeit smaller scale) experimental
et al. 2022 In contrast, focus settings where we can derive a more clear understanding of the relative merits of different pretraining objectives. Another empirical study from SSookk et al. [2021] considers representations in online reinforcement learning. Meanwhile, Yang and Nachum [222]] considers representations for imitation but does not consider image-base 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 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-indepeedent dynamics:
c~ Pc.
Pe:
= r==rr az),
0++1 ~ T(o,,aa).
(I)
Importantly, consider the context variable and rewards re be latent, i.e., they 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 discrete variable representing behavior like locking a door.
Data generation. Data is generated by executing policies π that map observations to actions in the environment. We consider two different datasets for any given problem. First there large multi-context pretraining dataset that will be used for representation learning, specifically to learn an observation encoder. Second, there is small single-context finetuning dataset for policy learning on top of the pretrained representation. The multi-context pretraining data is generated as follows:
Dpre Nore, c~ Pos
T == (oo, an, 01 .... pTe,
Te RJ T arg max (r.(TT) (T), (2)
where J=.(m) (m) denotes the expected return of a when the reward re. Note that the demonstration policy has access to the latent context e, but this latent context is not observed in the data.
Then the single-context finetuning data is generated for context Cfine as follows:
Nymn IT (oo, Go, 01 ..... )) p
Dfine
(3)
3

---

Pretraining, The goal of the paper is to analyze different methods for pretraining feature extractors. Training the encoders minimize loss proceeds as follows:
RR arg min E Dpre
[e(d, mi)]
(4)
full description of the losses used by different algorithms will come in Section 4.2 For simplicity
(and in keeping with prior work (Nachum and Yang, 2021 Chen al., 2022] we will consider to
only function of transitions (o. aj of rather than full trajectories 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:
# Rdd A arg min E Dyine
[e(m, ai ((c))))
In all of our experiments, l is the mean squared error loss for behavior cloning. We elect to use frozen features 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 well computer vision [Chnn et al.) 2020
Evaluation. Finally, we evaluate the finetuned policy by performing rollouts the finetuning
environment with context Cfine to estimate Jr Jrey.ne (#t). In our tasks we usually consider Terine to be binary indicator of successful completion of the finetuning task.
4 Experimental setup
4.1
Environments and Datasets
We design suite of tasks and datasets 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 popular approach. Our suite consists of six different pretraining datasets on varied tasks and of varied size. Fach 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
(d) Kitchen
(e) Metaworld (ML45/R3M)
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 consider two different train-eval splits ML45 and R3M.
All tasks performed from visual inputs, shown in Figure Each pair of pretraining pretraining-fiietuning 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 latent goal position, while in the door and metaworld datasets the context variable a discrete identity a desired behavior, and in the kitchen datasets the context variable is 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 traj for pretraining and N Nirin 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 (NVrre
context
and Neontews). For each finetuning context, we sample datasets with Nfin different seeds.
seed
Environment
Ntraj
ttraj
Nore context
No
N Nova
Pointmass
(Ic1, 1e2, 1e3)
(I, 5, 10)
Npre
NPPr Kzoj
Pick + place
(Ie1, 1e2, 1e3)
(2, 5, 10, 20)
traj
Door
(Iel, le2, 1e3)
(2, 10, 20)
5
Kitchen
(50, 150, 450)
(2, 5,, 15)
21
5
MW-ML45
(Ie2, 1e3, 1e4)
(2, 5, 20)
45
5
MW-R3M
(Ie2, 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 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 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 our pretraining datasets. The first is to directly use the pixels with image augmentations (Pixels + Aug) to train an encoder and a policy from seratch the finetuning data. It is essential to use the augmentations ensure that this a strong baseline. The second is features of a ResNet18 pretrained on Imagenet (Imagenet). The last consists of the features of ResNet18 that 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(alo, o') using architecture that first encodes o, d' with an encoder & and then predicts a with a small MLP f:
iin = arg min min E I(a f(o(o).
(6)
+ f 0,a,0'
Behavior cloning. A simpler alternative objective is to directly apply behavior cloning (BC) to the multitask actions in the pretraining dataset conditioned the observations using MSE loss. The learner is parameterized as an encoder & followed by a small MLP n:
T = arg min min E [(a - (+(o)))
&
x
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 model of the forward dynamics in the space of observations by encoding the current observation and then attempting to reconstruct the next observation o' using a decoder d:
BEED = arg min min E [[o' - a(((o),
(8)
0.2,0'
The second objective is implicit forward dynamics (FD-i) which implicity constructs 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 MLPss f1, f2) which is then passed into an InfoNCE loss:
E(o, a, o') = (xp(f(((()) a) fz(o(o'))),
(9)
$IFD arg min E log(FF o')) + log
(10)
& faJZ a,
5

---

Static observation modeling Finally, we consider a baseline that simply models P(o). Rather than modeling this explicitly with reconstruction, we use contrastive loss (Cont) where we use image augmentations to construct pairs of and o that do not rely on the dynamies of the environment at all. Again we use the InfoNCE loss, in what can be seen as a variant of SimCLR:
Eo,, Oaug) = exp(f(((o) (d(Dang))),
(11)
@Cent arg min min E
Aolg(E((o Oaug)) + log [E(o, Gang)]l )-]-
(12)
+ / 0,0auy
Ouag
Experiments
We want to determine which representation learning objective is best, but the precise answer will depend on the situation. To get 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:
How 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 in Appendix and code at /
com/ davidbrande www.errrerrnnnneeeeeeeeeeneereee wnwreeeer/imttttiion
5.1 Impact of dataset representation learning performance
Scaling with data size. The performance of each algorithm can be highly sensitive both pretrain- ing 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
0.8
Success
0.4
0.2
0.2
0.0
0.0
M
XL
S
M
Finetune size
Pretrain Size
States
Imagenet
ID
FD-e
Cont
Pixels
Aug
R3M
BC
FD-i
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 Methods that do 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, 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, door, kitchen, metaworld-m145, and metaworld-r3m datasets only have finite number of possible contexts that much smaller than the number of pretraining trajectories. For our default datasets, we elected to construct train-test split of contexts to ensure that the contexts used for finetuning

---

not seen during pretraining. As 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).
To test the impact of this gap between pretraining and finetuning, we created alternative pretraining datasets, where we include the test contexts (but nor the test trajectories) into the pretraining data. For example, in the door domain we include door opening, closing, locking, and unlocking in the pretraining data and still finetune on only unlocking (but with heldout initial conditions). These datasets now require much more limited notion of generalization from pretraining to finetuning.
Results shown in Figure We again see that ID is the strongest performer, but now the gap is even larger. ID matches the skyline performance of training from ground truth low-dimensional simulator states. BC also shows substantially stronger performance and outperfroms training from scratch Pixels + Aug. None of the other pretraining algorithms benefit much from the substantially easier type of generalization re- quired on these datasets. This suggests that ID) and BC are uniquely able benefit in easier settings, suggesting that they are better represen- tation learners. If an algorithm is not able to out- perform training from scratch in this simplified setting, is unlikely to be a good representation learner.
0.8
0.6
Success
0.4
0.2
0.0
Finelune size (multiple defeult)
States
R3M
FD
Pixels + Aug
ID
iDi
Imagenet
Cont
Figure 4: Average performance on the four discrete context environments when the finetuning contexts are included in the pretraining data. The finetuning data contains heldout initial conditions and trajec- tories not seen during pretraining.
Fully latent vs. inferrable context variables. Looking at our dataset suite, the datasets can be divided into two groups: those where the context variable not inferrable all from the initial state (pointmass, pick+place, and kitchen), and those where the effect of the context variable on the initial state makes it possible to infer the context given the initial state (door, metaworld-min ml45, and metaworld-r3m). This split presents an interesting comparison in particular between ID and BC (the best performing algorithms from the prior experiment). FigureB]ssoon results for these algorithms and the Pixels + Aug baseline on the datasets where the context is latent.
There is a large gap between ID and BC when the con- text is fully latent. In these cases, it is impossible to tell from the current state alone what the context is and thus what the optimal action should be. As we will show in our simplified model (Section 6). in these settings BC is confounded by the latent context (in the terminology of causal inference). As a result, BC can fail to learn useful features. In contrast, ID uses the information about the future state to deconfound the learning problem and still learns good representation. Note that this gap largely disappears when the context is observable, see Appendix for further details.
0.8
0.6
Finctune size
Pixels Aug
BC
Figure 5: A comparison between ID and BC on the datasets where the context is not inferrable from the observation.
5.2 Predictive power of the representations
So far, we have focused on the success rate of the downstream finetuned policy as the main metric of comparison between algorithms. Now we will instead consider a series of experiments that assess the quality of the representations based on the ability to predict various quantities of interest from the representations. These experiments help to illustrate what information retained in the representations and how efficiently that information can be accessed.
Action prediction. First, we consider the ability to predict the expert actions in the finetuning dataset. This is directly related to the success of the finetuned policy, but avoids the variance of
7

---

performing rollouts and allows us to compare train and validation errors to evaluate the representations. Low train loss means the representations are not aliasing observations that require different actions. Meanwhile the validation loss measures the simplicity of the function that maps representations to targets, i.e. how well it generalizes.
The results in Shou the train and valida- tion loss during finetuning using the default pre- training and finetuning sizes from Table Since losses vary across datasets, we normalize by the Pixels + Aug validation loss so as be able to present averages across all datasets. We see that of the learned representations, ID has both the lowest train and validation losses, almost matching the performance Pixels Aug on train and almost matching the performance of States on validation, contrast, representations that attempt to predict forward dynamics have substantially higher train loss, indicating alias- ing of states in terms of their optimal actions. Interestingly, the Imagenet pretrained features have very low train loss, indicating lack of aliasing, but very high validation loss, indicating that the function that maps representations to actions does not generalize well.
Tain
Vel
1.25
1.00
0.75
0.50
Cont
Pixels Aug
R3M
Figure 6: Average train and validation action- prediction loss during finetuning. All losses are normalized by the Pixels + Aug validation loss to maintain consistency across environments.
State prediction. Since we perform all of experiments in simulated environments, have access to the ground truth low dimensional states. So, we can measure the ability of each representation to predict the ground truth low dimensional state and thus measure how well the representation retains information about this ground truth state. Results are in Figure 7 here we measure the train and validation loss on the pretraining distribution of the representation learning apart from the gap between pretraining and finetuning. Again normalize the losses for each dataset.
Train Val
State prediction
0
isolate the effect
Cont
Figure 7: Average state prediction the pretraining distribution. Values are normalized by the ID train loss.
Again we see that ID and BC yield the best performance. This suggests that in these datasets, pretraining objectives that attempt predict the optimal action do indeed facilitate recovery of the www-dimensional simulator state. In contrast, while the FD methods achieve approximately the same training error, they generalize much This suggests that the FD objectives are not throwing away relevant information, but are keeping around too much extraneous information about the observations, thus making the representations susceptible to overfitting. Standard contrastive learning is substantially worse, even on train error, suggesting that is throwing away important information. Extended results Appendix B
poorly.
6 Analysis
To add more theoretical understanding of the empirical results, we will consider simplified model of the data generating process based on linear dynamics in latent space. We begin by presenting the model and then show that under this model we can explain three key experimental findings: (1) inverse dynamics is able to recover the low dimensional state, (2) forward dynamics can be less efficient in some cases, and (3) BC can be confounded by the latent context. We present a high level sketch here and more details along with discussion of related theoretical work are in Appendix D
Model. Some of the key interesting properties of problems like visual manipulation that we consider empirically are that (a) the observation is very high dimensional relative to the action, (b) the actual state of the world (or simulator) can be summarized in a much lower dimensional state variable, and (c) the dynamics relatively simple if given the right representation. All of these motivating

---

properties can be captured in a simplified model that assumes linear dynamics occurring in a hidden low-dimensional state space, as presented below.
For simplicity, we will only consider one step of the dynamics represented by tuple (o, a, o', 8, 8 c) that is sampled iid from the joint distribution over those variables. Recall that we only observe (o, a, o') and that (s, s',c) e) are latent. Formally, let O Rd, S IR and A R* with d >>>>> > k. Let o O + S be the ground truth encoder, which we assumee is invertible by o Let e ~ _V(
in IR' and A, B to be any matrices in RéxÉ and Rx Then, assume that the dynamics are:
d' = の (AA(o) Ba +)).
(13)
Note that we make no assumption on the policy =5. other than that it only depends on via Φ(o). This model is similar to ones studied in the online control setting by Mhammedi al. (2020), Dean and Recht [202T], but is different from models where inverse dynamics have been studied for online
control with exogenous noise since the dynamics are entirely contained in the low dimensional state space [Efroni et al. 2021 Lamb et al. 2022
Inverse dynamics recovers the state. To get an intuition as to why inverse dynamics learning is feasible this model, note that if B+ is the pseudoinverse of B that:
6(0') B+ Ao(o) B+c.
(14)
Thus the inverse dynamics are simple linear function of the embeddings Φ(o), ó(o'). As a result, when we solve for a with least squares regression, if the encoder e representable by our function class, we will be able to recover it up to linear transformation, provided the matrix B is well- conditioned, so that the noise term B+ e does not blow up.
Forward dynamics can be less statistically efficient. Intuitively, the potential problem with learning forward dynamics is that it requires learning both an encoder and decoder while inverse dynamics only requires learning the encoder. This is not necessarily a problem priori, but we hypothesize that in practical problems of interest (like the ones in our experiments) the decoder (mapping from low dimensional state to high dimensional observation) may be more complicated than the encoder (mapping from observations to states).
W
o (x. y)
To grasp why we might expect this, note that the set of possible observations the manifold represented by the image of the decoder, i.e. Im(Ø-)). As a simple example, consider toy 2d example where the high dimensional observation is (z, fx)) E IR2 and the low dimensional state is simply x E IR as depicted in Figure & Here the encoder Φ is very simple since it just needs to recover æ, while the decoder must learn (ar). Of course this a very toy example, but we find it illustrative of the idea that it is possible that the encoder much simpler than the decoder in practice.
s
+
Figure 8: An example where the decoder is more complicated than the encoder.
BC can be confounded by the latent context. As we alluded to in the experimental section, the latent context variable can confound BC. Now we will show an example in our model where this problem arises. In this case, even with a linear encoder, infinite data, and fully expressive policy class, the Bayes optimal BC representation cannot be used to recover anything better than a random policy. This example is extreme, but shows the shortcomings of a confounded pretraining objective.
For simplicity, let f = k and e = O. Let R (x)) be the set of rotation matrices in R Let gk-1 be the unit sphere in RR. U be the uniform distribution, and 5 denote a Dirac delta. Now, assume:
c ~~ U (R(RR(x))
U (S*-)))) T. (alo) = 5/a = co(o)
(15)
Note that (0)) returns unit vector in R and that uniformly sampled rotation of a unit vector is uniformly sampled unit vector. Thus, we can marginalize over e to get:
P(alo) = P(e)=t(al))
P(c)o(a = cœ(o)] Pu(st-i)(a)) = 77:
(16)
for constant 7he equal to the reciprocal of the surface area of the unit sphere in #

---

Thus, the Bayes optimal BC policy does not depend on o at all. As result, the optimal representation learned by BC can just map every observation zero. This representation is not capable of represent- ing the optimal policy for any choice of c. However, switching to inverse dynamics pretraining where we condition on the outcome observation o' breaks the confounding and allows us to learn the true representation even without observing
Discussion
We have seen that inverse dynamics pretraining provides an effective method for learning features from multitask demonstration data. We demonstrated this across suite of datasets with visual observations and provided analysis in a simplified model to understand the strong empirical performance.
Limitations. There are still a few limitations of our work that are worth pointing out explicitly. First, in this work prioritized simulated domains with large numbers of predefined tasks and datasets with a single morphology to allow for a variety of experiments. However, it is possible that the results observed in these tasks would differ when scaled to real world tasks with additional visual diversity and physical realism. We leave this extension to future work.
Second, while our theoretical analysis provides a clear rationale for the observed empirical results in toy model, there is clearly room for better theory. Ideally, future work could present more rigorous theory that goes beyond a toy model. However, we do think that the toy model captures some of the essential characteristics of the problem and recognize that any theory must make simplifying assumptions.
Future directions. In addition to removing the limitations described above, there are many other interesting directions for future work to build
results. One direction would be extend these results to settings with suboptimal data. In this work we focus on an imitation learning setting where data is collected by expert policies across variety of tasks. In future, it would be interesting study how and if the properties of various representation learning algorithms change in the presence of suboptimal data.
It would also be interesting for future work to compare the relative merits of a broader array of pretraining techniques that go beyond representation learning. For example, methods that learn con- ditional generative models (e.g. goal-conditioning, language-conditionnn or reward-conditioning provide a different paradigm for pretraining policies instead of the feature extractors that we consider in this work.
Finally, it would be interesting to consider developing new pretraining objectives for representation learning. This could be done by combining existing objectives or developing completely new ones.
Acknowledgments
This work was completed as part of the Google Research Collabs program. We would like to thank Mahi Shafiullah, Mark Goldstein, Aahlad Puli, Siddhant Haldar, Ben Evans, Ulyana Piterbarg, and Lerrel Pinto for helpful discussions and feedback.
References
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. Opal: Offline primitive discovery for accelerating offline reinforcement learning, arXiv preprint arXiv:2010.336// 2020.
Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saun- shi. A theoretical analysis of contrastive unsupervised representation learning. arXiv preprint 1rX.v:1902.09229, 2019.
Sanjeev Arora, Simon Du, Sham Kakade, Yuping Luo, and Nikunj Saunshi. Provable representation learning for imitation learning via bi-level optimization. In International Conference on Machine Learning, pages 367-376. PMLR, 2020.
Yusuf Aytar, Tobias Pfaff, David Budden, Thomas Paine, Ziyu Wang, and Nando De Freitas. Playing hard exploration games by watching youtube. Advances in neural information processing systems, 31, 2018
10

---

Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, lurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http::t/ttttthhii com deepmede
Bowen Baker, lige Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances in Neural Information Processing Systems, 35:24639-24654, 2022.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miine and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github cwwm homm//ooooeeaaa
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprini www.x::2/220681111 2022.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pages 1597-1607. PMLR, 2020.
Xin Chen, Sam Toyer, Cody Wild, Scott Emmons, Ian Fischer, Kuang- Huung--ueille Lee, Neel Alex, Steven H Wang, Ping Luo, Stuart Russell, et al. An empirical investigation representation learning for imitation. arXiv preprint 220.::2005008886, 2022.
Zichen Jeff Cui, Yibin Wang, Nur Muhammad Mahi Shafiullah, and Lerrel Pinto. From play to policy: Conditional behavior generation from uncurated robot data. arXiv e-prints, pages arXiv-222 10, 2022.
Sarah Dean and Benjamin Recht. Certainty equivalent perception-baser control. In Learning for Dynamics and Control, pages 399-411. PMLR, 2021.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 18Xiv:/8100048055 2018.
Yiming Ding, Carlos Florensa, Pieter Abbeel, and Mariano Phielipp. Goal-conditioned imitation learning. Advances neural information processing systems, 32, 2019.
Yan Duan, Marcin Andrychowicz, Bradly Stadie, OpenAl Jonathan Ho. Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojeiech Zaremba. One-shot imitation learning. Advances in neural information processing systems, 30, 2017.
Yonathan Efroni, Dipendra Misra, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Prov- able with exogenous distractors via multistep inverse dynamics. arXiv preprint 10.00:2110.08847, 2021.
Benjamin Eysenbach, Tianjun Zhang, Ruslan Salakhutdinov, and Sergey Levine. Contrastive learning as goal-conditioned reinforcement learning. arXiv preprint arXiv:2206.07566 2022.
Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, and Pieter Abbeel. Deep spatial autoencoders for visuomotor learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA), pages 512-519. IEEE, 2016.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pages 1126-1 135. PMLR, 2017a.
11

---

Chelsea Finn, Tianhe Yu, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot visual imitation learning via meta-learnin. In Conference on robot learning, pages 357-368. PMLR, 2017b.
Xiang Fu, Ge Yang. Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions, In International Conference on Machine Learning, pages 3480-3491. PMLR, 2021.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, pages 2170-2179. PMLR, 2019.
Dibya Ghosh, Abhishek Gupta, and Sergey Levine. Learning actionable representations with goal- conditioned policies. arXiv preprint (811088/ 2018
Dibya Ghosh, Chethan Bhateja, and Sergey Levine. Reinforcement learning from passive data via latent intentions. arXiv preprint arXiv:2304.04782, 2023.
Kristen Grauman, Andrew Westbury, Eugene Byrne, Zachary Chavis, Antonino Furnari, Rohit Girdhar, Jackson Hamburger, Hao Jiang, Miao Liu, Xingyu Liu, al. Ego4d: Around the world in 3,000 hours of egocentric video. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 18995-19012, 2022.
Jcan-Basticn Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems, 33:21271 11-21271-21284 2020.
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. arXiv preprint 1rXiv:1910.11950. 2019.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2023. URL
http: httt://gittube cww/ wwt.ccoomoooc.ee goog////gl//lar
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
Riashat Islam, Manan Tomar, Alex Lamb, Yonathan Efroni, Hongyu Zang, Aniket Didolkar, Dipendra Misra, Xin Li. Harm van Seijen, Remi Tachet des Combes, al. Agent-controller representations: Principled offline rl with rich exogenous information. arXiv preprint (tti::222110006666 2022.
Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn. Bc-z: Zero-shot task generalization with robotic imitation learning. In Conference on Robot Learning, pages 991-1002. PMLR, 2022.
Ilya Kostrikov. JAXRL: Implementatiions of Reinforormmnn Learning algorithms in JAX, 10 2022.
URL https::/tttttuubbbbbbbbb com httpu//o....ooooooooooorrrrrrrra v2.
Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off-policy distribution
matching. arXiv preprint arXiv:1912.050:3 2019.
Alex Lamb, Riashat Islam, Yonathan Efroni, Aniket Didolkar, Dipendra Misra, Dylan Foster, Lekan Molu, Rajan Chari, Akshay Krishnamurthy, and John Langford. Guaranteed discovery of control- lable latent states with multi-step inverse models. arXiv preprin! arXiv: arXiv:2207.08229) 2022.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning, pages 5639-5650. PMLR, 2020.
Kuang-Huei Lee, Ofir Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadar- rama, Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et Multi-game decision transformers. Advances in Neural Information Processing Systems, 35:27921-27955 2022a.
12

---

Kuang-Huce Lee, Ofir Nachum, Tingnan Zhang. Sergio Guadarrama, Jie Tan, and Wenhao Yu. Pi-ars: Accelerating evolution-learned lisomo-looomttinn with predictive information representations. In 2022 JEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 1447-1454. IEEE, 2022b.
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In Conference on robot learning, pages 1113-1132. PMIR, 2020.
Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. Vip: Towards universal visual reward and representation via value-implicit pre-training. arXiv preprint arXiv:22 10.00033. 2022.
Zakaria Mhammedi, Dylan J Foster, Max Simchowitz, Dipendra Misra, Wen Sun, Akshay Krish- namurthy, Alexander Rakhlin, and John Langford. Learning the linear quadratic regulator from nonlinear observations. Advances in Neural Information Processing Systems, 33:14532-14543,
2020.
Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In International Conference on Machine Learning, pages 7780-7791. PMLR, 2021.
Ofir Nachum and Mengjiao Yang. Provable representation learning for imitation with contrastive fourier features. Advances in Neural Information Processing Systems, 34:30100-3011 12, 2021.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A universal visual representation for robot manipulation. arXiv preprint arXiv:2203. 12601, 2022.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint urXiv:1807.0322 2018
Jyothish Pari, Nur Muhammad Shafiullah, Sridhar Pandian Arunachalam, and Lerrel Pinto. The surprising effectiveness of representation learning for visual imitation. arXiv preprint arXiv:21 2021.
Deepak Pathak, Pulkit Agrawal, Alexci A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning, pages 2778-2787 PMLR, 2017.
Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural computation, 3(1):88-97 1991.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning. pages
PMLR, 2021.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforceemntt learning via probabilistic context variables. In International conference on machine learning, pages 5331-5340. PMLR, 2019.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint 21Vv:2205.06775 2022.
Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, R De- von Hjelm, Philip Bachman, and Aaron C Courville. Pretraining representations for data-efficient reinforcement learning. Advances in Neural Information Processing Systems, 34:1 12686-126999 2021.
Younggyo Seo, Kimin Lee, Stephen L James, and Pieter Abbeel. Reinforcement learning with action-free pre-training from videos. International Conference on Machine Learning, pages 19561-19579, PMLR, 2022.
13

---

Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel. Masked world models for visual control. In Conference on Robot Learning, pages 1332-1344. PMLR, 2023.
Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context- based representations. International Conference on Machine Learning, pages 9767-9779. PMLR, 2021.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning, pages 9870-9879. PMLR, 2021.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pages 5026-5033.
IEEE, 2012.
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm_control: Software and tasks for continuous control. Software Impacts, 6:100022, 2020.
David Venuto, Sherry Yang, Pieter Abbeel, Doina Precup. Igor Mordatch, and Ofir Nachum. Multi-environment pretraining enables transfer to action limited datasets. arXiv preprint WrXiv:2211.13337, 2022.
William Whitney, Rajat Agarwal, Kyunghyun Cho, and Abhinav Gupta. Dynamics-aware embeddings.. arXiv preprinl arXiv: 0199900093577 2019.
Philipp Wu, Arjun Majumdar, Kevin Stone, Yixin Lin, Igor Mordatch, Pieter Abbeel, and Aravind Rajeswaran. Masked trajectory models for prediction, representation, and control. arXiv preprint 2rXx::2305.02968 2023.
Mengjiao Yang and Ofrr Nachum. Representation matters: offline pretraining for sequential decision making. In International Conference on Machine Learning, pages 11784-11794. PMLR, 2021.
Mengjiao Yang, Sergey Levine, and Ofir Nachum. Trail: Near-optimal imitation learning with suboptimal data. arXiv preprint 2XXx:2110.14770, 2021.
Sherry Yang, Ofir Nachum, Yilun Du, Jason Wei, Pieter Abbeel, and Dale Schuurmans. Foun- dation models for decision making: Problems, methods, and opportunities. arXiv preprint 03t.: 2023.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 10674-10681, 2021.
Tianhe Yu, Chelsea Finn, Annie Xie, Sudeep Dasari, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot imitation from observing humans via domain-adaptiv meta-learning. arXiv preprint Xrxiv: 102vv.180..01557, 2018.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pages 1094-1100 PMLR, 2020.
Kevin Zakka, Andy Zeng, Pete Florence, Jonathan Tompson, Jeannette Bohg, and Debidatta Dwibedi. Xirl: Cross-embodiment inverse reinforcement learning. In Conference on Robot Learning, pages 537-546. PMLR, 2022.
Hongyu Zang, Xin Li, Jie Yu, Chen Liu. Riashat Islam, Remi Tachet Des Combes, and Romain Laroche. Behavior prior representation learning for offline reinforcement learning. arXiv preprint N.Niv:22/1.00866 2022.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint 2071:20006.1742, 2020.
14

---

Thomas T Zhang, Katie Kang, Bruce D Lee, Claire Tomlin, Sergey Levine, Stephen Tu, and Nikolai Matni. Multi-task imitation learning for linear dynamical systems. arXiv preprint NNNi:22/2.001/88 2022.
15

---

A Extended related work
In this paper we focus specifically on pretraining methods that learn representations of high di- mensional observations from multitask demonstration data with latent contexts for the purpose of imitation. There are many closely related problems that are studied in other papers that we did not have space to address fully in the main text that we more fully describe here. These all very interesting and complementary lines of work, but are beyond the scope of this paper.
Perhaps the largest closely related line of work focuses on learning reward -drrecteee representations in the context of reinforcement learning. This is different setting than ours and methods from there are not applicable in our setting where we do not have access to rewards. Moreover, most of these
methods do not consider multitask settings [Zhang et al., 2020 Gelada et al. 2019 Fu et al. 2021
Ghosh et al. 2018 Eysenbach et al. 2022 Sodhani et al.
2021
Another line of work seeks to learn representations of actions or sequences of actions rather than observations. This is a directly complementary line of work to the ideas presented in this paper (Ajay
ct al.
2020 Yang et al. 2021 Lynch et al. 2020 Whitney et al. 2019].
Another body of literature focuses on learning representations that can be transferred across domain and embodiment gaps and even trained directly from videos without access to actions at all [Oord
2018 Aytar et al. 2018 Seo al. 2022 Ma et al. 2022 Zakka et al. 2022 Ghosh al. 2023 In this paper, we focus on the simpler task of pretraining a representation within one MDP with
al.
consistent dynamics and access to demonstration actions, but with varied tasks. This choice allows us to make more clear comparisons between a algorithms and rigorous claims about when representations will be effective, but also limits the generality of the representations that are learned.
There are variety of new methods that rely transformer architectures construct interesting new pretraining objectives (Yang and Nachum, [021] Lee et al., 2022a Reed et al., 2022 Sco et al., 2023 Wu et al. 2023]. In this paper we focus on simple methods that can all use the same simple convolutional architecture operating on transition tuples to provide the most controlled comparison that we can. It is an interesting direction for future work to see how our insights in the Markovian case could be leveraged to inform sequence level models of partially observed problems.
Another line of work avoids pretraining representations directly and instead meta-learns a policy that adapt to new tasks (Duan et al., 2017, Finn et al. [017a|bb b) Yu et al., [018]] [Rakelly et al.) 2019, Mitchell et al. 2021]. This approach is beyond the scope of this paper which focuses on representation learning. Moreover, these meta-learning algorithms require the pretraining trajectories to be partitioned by task so that each task has multiple trajectories. Since we focus on pretraining data where we don't 't have access to the latent context, it unclear how to create these meta- meta-training datasets.
Finally, recent work has shown the promise of zero-shot generalization for multitask imitation, especially when the task identifying information is expressed in natural language to leverage advances in language models [Ding et al. 2019 Jang et al., 2022] Cui et al., 2022 Brohan et al.) 2022]. This is an exciting line of work, but beyond the scope of this project where focus data where the context information is latent. It is an interesting direction for future work to assess precisely how much performance can be improved via extra context information gauge whether is worth the cost of labeling trajectories with context information.
is an interesting direction for future work to try to better synthesize some of the findings from across this broad array of approaches to pretraining in slightly different settings.
B Extended experimental results
In this section we present the experimental results that were excluded from the main text due to space constraints. In particular, Section B.I presents representation analysis by predicting one representation from another, Section B.2 presents the per-dataset results of various sweeps over dataset size and type, Section [B.3 presents per-dataset results for representation analysis, and Section presentstts results of an ablation over multistep dynamics.
B.1 Cross-representation prediction
In the main text, we evaluated representation quality by measuring accuracy of small ML.Ps to predict either the actions on the finetuning data or the low dimensional states the pretraining
16

---

data. Here we present similar analysis, but now where we use small MLPs to predict the other representations themselves. This is interesting since it lets assess which representations contain enough information and shared structure to predict the other representations. Hypothetically, representation that is easily able to recover another representation may preferable since it retains more information.
Results presented in Figure 2 show the average across datasets of the coss-representation prediction error on validation set from the pretraining distribution (normalized by the mean prediction error on each dataset). There are several possible takeaways from this experiments. First, looking a the rows, which correspond the error when each method is used the source, we can see that in- verse dynamics generally has the lowest average error for predicting the other representations. This suggests that inverse dynamics doing a good job of recovering the information that is shared among all the representations, Second, looking at the columns, which correspond to error that
BC
Source
100
FD-e
FD+i
10-1
1D
Cont
when each representation is used as the target, BC is the most difficult to predict and inverse dynamics is second most difficult. This is somewhat surprising result, but suggests that these representations actually con- tain information that may have been thrown away (or at made least difficult to access via small MLP) within the other representations. Finally, note that the contrastive learner is both the worst source and easiest target, which consistent with the idea that those representations are losing important task-relevant information.
Target
Figure 9: Cross-representation predic- tion error of small MLP validation set from the pretraining distribution. Re- sults normalized per dataset by the mean error on that dataset and then aver- aged across datasets.
Full results on each dataset can be found in Appendix B.3 and full methodological details can be found in Appendix
B.2 Per dataset evaluation success results
In the main text and Section B. we have only presented aggregate results that average across datasets. These averages make it easier to summarize comparisons between methods, but they sacrifice the precision of how the results vary across datasets. In this section we present per dataset results for all of the relevant sweeps across dataset variations including pretraining size, finetuming size, and finetuning size when we ablate in distribution contexts or observability of the context in the observation.
Pretraining size. First, we present the full ablation over pretraining size, corresponding to the right panel of Figure The full per dataset results are shown in Figure 10)
There are several findings in the -ataset-specific results that are not visible in the aggregate reported in the main text:
First, the kitchen environment is a clear outlier mainly due to the stochasticity in the data generating process and smaller dataset size compared the others (see Appendix C.T for more detailed description of the data). As a result of the noise added to the low dimensional states, training from States actually underperforms training from Pixels + Aug. We hypothesize that this is due to some implicit regularization that arises from training from the rendered noisy observations instead of the low dimensional noisy states. Importantly, inverse dynamics is much better able to handle the stochasticity than the alternative methods given the relatively small pretraining dataset and is the only method that able to perform comparably to training from scratch.
Point mass is the only environment where the externally pretrained representations (R3M and Imagenet) substantially outperform training from Pixels + Aug and they are substantially outperformed on kitchen and the metaworld datasets. We hypothesize that this shows how it is quite difficult to transfer features across domains and see consistent benefits on challenging tasks.
17

---

Point Mass
Pick and Place
Door
1.0
0.8
0.8
0.4
0.4
0.0
0.0
0.0
Kitchen
Metaworld (ML45 split)
Metaworld (R3M split)
0.5
0.8
0.6
0.4
0.4
0.0
0.0
0.0
M
M
Pretrain size
Pretrain size
Pretrain size
Pizels Aug
Imagenet
R3M
PC
FDve
FD-4
Cont
Figure 10: The per dataset results of sweeping over pretraining size, corresponding to the right panel of Figure Error bars show standard error over seeds and contexts (as described in Table Horizontal lines indicate mean performance of algorithms that do depend on pretraining size.
Note that performance of contrastive learning is substantially better relative to the alternatives on point mass. We hypothesize that this is due to the fact that random crop augmentations are actually reasonable simulation of the dynamics in the pointmass environment specifically so that contrastive learning becomes more similar to implicit forward dynamics.
Finetuning size. Next, we present the full ablation over finetuning size, corresponding to the left panel of Figure The full per dataset results are shown in Figure II]
Again, described above, Kitchen is clear outlier due stochasticity with inverse dynamics the best performer. Inverse dynamics is also the clear winner on point mass and slight winner on pick and place. The other tasks more ambiguous with many methods performing about the same, and none substantially better than training from scratch (across all pretraining sizes). Disaggregating the results here shows how even though inverse dynamics is clearly the best in aggregate, this is not necessarily true on every dataset. As we will see Figure 12 we hypothesize that much of this weak performance be attributed to the fact that the evaluation contexts in door and the metaworld variants are truly out of distribution, making it difficult for any pretraining method to generalize.
Ablating in distribution contexts. Next, we present the full per dataset results when we ensure that all the evaluation contexts are included in the pretraining distribution, corresponding Figure 4 in the main text. The full per dataset results are shown in Figure [2
is important to compare these results to those that include out distribution evaluation contexts in Figure II First, note that the evaluation contexts point mass and pick and place were already in distribution, so they are kept the same. However, on door and the two metaworld splits there is a substantial improvement, especially for inverse dynamics and BC. This shows how these methods can benefit from being applied on tasks that are contained in the pretraining distribution. Interestingly, though the evaluation contexts now in distribution, the forward dynamics representations do not see substantial improvements and are still outperformed by training from scratch on the more challenging datasets.
18

---

Point Mass
Pick and Place
Door
1.0
1.0
1.0
0.8
0.8
0.8
0.6
0.6
0.0
0.0
5
14
XL
Kitchen
Metaworld (ML45 spliit)
Metaworid (R3M split)
1.0
0.8
0.8
0.8
0.6
0.6
0.6
0.4
0.4
0.4
0.2
Finetune size
Finetune size
Finetune Size
Imagenet
Figure 11:
The per dataset results of sweeping over finetuning size, corresponding to the left panel of Error bars show standard error over seeds and contexts (as described in Table
Figure
Place
Door
1.0
1.0
1.0
0.8
0.4
0.4
0.4
0.2
0.2
0.2
0.0
0.0
0.0
M4
XL
Kitchen
Metaworld (ML45 split)
Metaworld (R3M split)
1.0
1.0
1.0
0.8
0.9
0.2
0.2
0.2
0.0
0.0
XL 0.0
M
XL
M4
XL
----
States
Pixels Aug
Imagenet
R3M
BC
rD-e
FD4
Cont
Figure 12: The per dataset results of sweeping over finetuning size when we include the evaluation tasks in the pretraining data, corresponding to Figure 4 Error bars show standard error over seeds and contexts (as described Table TT.
19

---

Aggregating based on context observability. Finally, we present the full results for aggregations across whether the context is observable, corresponding to in the main text. Context latent in point mass, pick and place, and kitchen, but inferrable in door and both metaworld splits. The Note that these results are just grouped averages over the results results are shown Figure
13
presented in Figure 11
Latent Context
Inferrable Context
1.0
1.0
0.8
0.8
Success
0.6
0.6
0.4
0.4
0.2
0.2
0.0
XL
Finetune
Finetune
Imagenet
Pixels Aug
DC
FD-i
Figure 13: The full results of aggregating based on the observability of the context variable, corre- sponding to Figure 3 Error bars show standard error over seeds and contexts (as described in Table T then averaged across datasets.
now include the results from all algorithms and also from the environments where the context is inferrable. As reported in the main text, there is a clear gap between inverse dynamics and BC when the context is latent, likley due to confounding. Here largely disappears in the datasets where the context is inferrable and generally the disparities between methods shrink.
Compared Figure
that this gap
B.3
Per dataset representation analysis
Now we present the per dataset results of the various methods of representation analysis based on predicting different target quantities of interest: the action, the low dimensional state, and the other representations themselves.
Point Mass
Pick and Place
Door
10-1
Thain
10 1-
10
toss
10
10-1
10
10
Kitchen
Metaworid (ML45 split)
Metaworld (R3M spilt)
10-4-
10
1.0
10-1
10
1
Pioels Aug
Imaggentt
Figure 14: Full per dataset results of action prediction on the finetuning distribution.
Predicting action. First we present the per dataset results for train and validation action prediction on the finetuning datasets using the default pretraining and finetuning size. These results correspond
20

---

to Figure 6 from the main text. Unlike in the main text, here we do not do any normalization of the losses, so the losses occur at different scales on each dataset depending on how difficult the prediction task is. Results are shown in Figure 14
Point Mass
Fick and Place
Door
10
10-4
Kitchen
Metaworld (ML45 spilt)
Metaworld (R3M split)
1.39 10
1.38 10
1.36x10'1
1035x10
10-1
1-3x10-
10-3
FD-e
FD4i
Contt
Figure 15: Full per dataset results for state prediction on the pretraining distribution.
Predicting state. Next, we present the per dataset results for predicting the low dimensional state on the pretraining distribution from the various learned representations. These results correspond to Figure 7 in the main text. Again, unlike in the main text, results are normalized, they occur at different scales across environments. Results are shown in Figure [5]
Note that as mentioned before, there is stochasticity added to the low dimensional states the kitchen environment. This makes it difficult for any of the methods to substantially outperform the floor set by the noise level.
Point Mass
Pick and Place
Door
10-
10-1
10-
FD.0
FD-1
10-1
10"1
Kitchen
Metaworld (ML45 split)
Metaworld (R3M split)
RC
Source
20e
Cont
10-4
Cont
1D4
Tanget
Tanget
Target
Figure 16: Per dataset results for crosstrepresenttation prediction the pretraining distribution. Color shows the validation enror of predicting target from source.
21

---

Predicting across representations. Finally, we present the per dataset results for predicting across the different learned representations on the pretraining distribution. These results correspond to Figure 9 Again, unlike in the averaged figure, this figure is not normalized, so the scales vary across datasets. We truncate the color scale at le-4 on the low end for easier visualization.
B.4 Ablation of multistep dynamics
As mentioned in the main text, some work argues for multistep dynamics models [Efroni et al. 2021 Lamb et al. 2022]. Note that this work focuses on settings with exogenous noise which are different from the simpler settings that consider. To confirm that using multistep dynamics models does not help to learn better representations, we run an ablation of the number of steps included in the dynamics model three environments: point mass, pick and place, and door and two algorithms: inverse dynamics and implicit forward dynamics. Results are shown in Figure 17 At high level, basically find little difference when ablating the number of steps, models everywhere for simplicity.
default to using one step
Note: for inverse dynamics models, we learn a step model by predicting a, given or and Oct.k. For forward dynamics, we learn a k step model by predicting 0r+- given and Ocr+k-
Point
and Place
Swccess
5
5
N steps
N steps
N steps
Figure 17: Sweep over the number of timesteps included in the dynamics models.
C Detailed experimental methodology
In this section we present a detailed account out methodology. We also release our code that was used to perform the experiments for full transparency. We split up the description into Section C. which describes the environments and dataset generation, Section C.2 which describes the details of the pretraining pipeline, and Section C33 which describes the details of the finetuning and evaluation pipeline.
All code is at http://hhhhooooooooo////////aatt
C.1 Envionment and dataset details
Software dependencies. All of our environments are based on the MuJoCo simulator Todorov cct al. 2012). The point mass environment is derived from the DM control suite (Tunyasuvunkkool
al., 2020 The kitchen environment and dataset was introduced in Gupta al. [2019]. The rest of the environments are taken from Metaworld (Yu et al. 2020). We describe each environment detail and summarize the descriptions in Table 2
Point mass. The point mass environment consists of an actuated point mass on 2d plane. In our version, the context e RR² determines the goal location. Then, the demonstration policy गテ is a PD controller that moves the point from the current position r to the goal position c. Because the context variable is continuous, we sample an independent context for each trajectory the pretraining dataset from the uniform distribution over possible goal states. The context is fully latent and not observable the observation. The low dimensional state the 2d position and the high dimensional images are 84x84x3.
22

---

Pick and place. The pick and place task is taken from the metaworld suite. In our version, the context E R determines the goal location for the block. The demonstration policy тद is a scripted policy from the metaworld repo. We remove the goal indicator from the image in this environment so that the context fully latent and not observable from the observation. The low dimensional state is the 3d position of the gripper, 1 openness of the gripper, and 7d position and orientation of the block. The high dimensional observations are images of size 120x 120x3.
Door. The door environment also taken from the metaworld suite. In our version, the context E [4] determines the index of the environment from door-close, door-open, door-unlock, and door- lock. For our default experiments we use door-close, door-open, and door-unlock as the pretraining contexts and door-lock the eval context. For the ablation where we ensure that the eval context is in the pretraining distribution, we include door-lock in the pretraining data. The demonstration policy =5 scripted policy from the metaworld repo. Given the context, the initial position of the robot, initial position of the door, and goal position (which is visible in the observation image) are all randomized. Note, the context is inferrable since the initial position of the door and lock allow the learner to infer the context. The low dimensional state is the 3d position of the gripper. Id openness of the gripper, 7d position and orientation of two objects in the scene, and 3d goal position. The high dimensional observations are images of size 120x 120x3.
Kitchen. The kitchen environment and dataset are taken from Gupta et al. [2019]. Each trajectory contains a sequence of four tasks in a simulated kitchen collected by a human demonstrator, In our version, the context E [24] is determined by the sequence of four tasks contained within the demonstration trajectory (of which there are 24 possibilities). We evaluate on three contexts: microwave-ketee -eekke-lttllhtt light switch-slide cabinet, bottom burner-top burner-slide cabinet-hings cabinet, and kettle-botte bottom burner-top ourner-light switch. In our default setup. we pretrain on the other 21 contexts, and in the ablation of in distribution evaluation we pretrain all 24 contexts. The context is fully latent and not observable from the initial state. The low dimensional state is a 9d description of the position and 21d description of the position of objects in the kitchen. The high dimensional observations are images of size 120x 20x3..
Note: the kitchen environment is the only one that we consider that has added noise. The raw data from Gupta et al. [2019] contains gaussian noise added to the low dimensional states and actions, so this noise cannot be removed without re- re-generating the data. We render the images from the noisy states, there is also noise present in the image observations. We also evaluate in environment with the same noise added, so there is no gap between training and eval.
Metaworld (ML45 and R3M). Finally, we consider two variants of the full metaworld suite. Here the context [50] determines which metaworld task is used. We consider two different train-eval splits for our default environments. The ML45 split takes the eval tasks from the original metaworld ML45 task which are bin- bin-picking, box-close, hand-insert, door-lock, and door-unlodd The R3M split takes the eval tasks that were chosen in the R3M paper (Nair et al., [022]: assembly, bin bin-picking. button-press, opewer-open, and hammer. Given the context, the initial and goal positions are randomized. The goal position is visible in the observation. The low dimensional state is the 3d position of the gripper, 1 openness of the gripper, 7d position and orientation of (potentially) two objects in the scene, and 3d goal position. The high dimensional observations are images of size 120x 120x3.
Table 2: A summary of the description of datasets above. Inferrable refers to whether the context is observable. OOD refers to whether the evaluation context out of distribution.
Dataset
Policy
Context
Inferrable
OOD
Noise
State dim
Point mass
PD controller
R
No
No
No
Pick and place
Script
R3
No
No
No
Door
Script
[4]
Yes
Yes
No
21
Kitchen
Human
[24]
No
Yes
Yes
Metaworld-ML45
Script
[50]
Yes
Yes
No
21
Metaworld-R3M.
Script
[50]
Yes
Yes
No
23

---

C.2 Pretraining details
Software dependencies. We implement all of our training in JAX [Bradbury et al. 2018 We use flax for neural networks (Heek et al. 2023 and optax for optimization [Babuschkin et al. 2020 Our
code is loosely based on Kostrikov 2022].
Architecture. All of our pretraining algorithms share exactly the same encoder architecture to ensure that we have a fair comparison. Since our tasks are relatively simple visually, and so to allow for large scale experiments without too much compute, we use a relatively small convnet encoder. Specifically, we follow the architecture from Yarats et al. (2021) which consists of a layer convnet with 3x3 filters, number of channels of (32, 64, 128, 256), and strides of (2.2,1.1). We add modification include a spatial softmax activation (Finn et al. 2016]. which we found to be important for the manipulation tasks we consider. This is followed by linear layer to project into the embedding dimension of 64 and finally a layernorm and tanh activation to normalize the embedding. We use the gelu activation function throughout.
Now we will birefly describe the architecture used for each pretraining algorithm, following their descriptions in Section 4.2
Inverse dynamics: the inverse dynamics head is an MLP that takes in @(o), o(o') and produces an estimated action. This MLP has two hidden layers of width 256 and dropout of 0.1 during training.
BC: the BC policy head is an MLP with two hidden layers of width 256 and dropout of 0. during training.
Implicit forward dynamics: the implicit forward dynamics model uses an action encoder (a) which outputs a 64 dimensional normalized action embedding which is concatenated
to Φ(() to form ø(o, a). Then there are two projection heads fi, f2 that take in Φ(o, a) and obo') respectively and produce 64 dimensional embeddings that are normalized to have unit norm. All these networks (ø.. f1. and f2) are MLPs with two hidden layers of width 256 and the relevant input and output dimensions.
Explicit forward dynamics: the explicit forward dynamics model uses the same architecture to encode a with Φa. Then, instead of projection heads, we require a convolutional decoder to produce an image. Following Yarats al. [202]] we use an architecture that inverts the encoder, having a dense projection layer followed by channels of (256, 128, 64, 32) and strides of (1,1,2,2).
Contrastive: the contrastive network is the same as the implicit forward dynamics network except that there is no action input and o' is replaced by an augmentation of o.
Training hyperparameters. For pretraining, we split the datasets into two categories: easy (point mass, pick and place, and door) and hard (kitchen, metaworid-ml45, and meatworld-r3m), On the easy tasks we train for 100k gradient steps and on the hard tasks we train for 200k gradient steps. Batch size is 256 for all methods except explicit forward dynamics where (due to the added compute required for the decoder) we use batch size of 128 even out computational requirements across methods. All methods are trained with the adamw optimizer with learning rate 1e-3, a cosine learning rate decay schedule, and default weight decay of le-4.
Data augmentation. Following [Chen et al. 2022) and others, we note that cropping augmentations are the most important for training policies in simulated visual domains. As such, all of our pretraining algorithms (and the Pixels + Aug baseline) use random cropping augmentations, and we found this to be an important implementation detail. The one exception is explicit forward dynamics where we found it difficult reconstruct images with augmentations, so we omit them for that algorithm.
Compute resources. Pretraining was all done on an internal cluster using RTX8000 GPUs. We estimate that the final training run needed to produce the results in the paper took approximately 600 GPU hours.
C.3 Finetuning and evaluation details
Training hyperparameters. The policy always an MI.P with two hidden layers of width 256. We use gelu activation and apply dropout with probability 0.1 during finetuning. We finetune every
24

---

dataset for 10k gradient steps with batch size 256. All policies are trained with the adamw optimizer with learning rate le-3, a cosine learning rate decay schedule, and default weight decay of le-4.
As explained in Table there are several seeds and evaluation contexts for each environment. For example, for the default results in Figure we end up having a total of 80 different finetuning datasets per representation when sweeping across dataset, context, and seed that Figure is reporting aggregate results across 720 finetuning and evaluation runs.
Evaluation byperparameters. Each evaluation is run for 100 episodes in the environment to estimate the success of the policy (except for the kitchen environment where we run 50 episodes due to slow rendering of that environment).
Compute resources. Finetuning and evaluation was all done on an internal cluster on CPU (since the finetuned policy network is small and environments run on CPU). We estimate that all the finetuning and evaluation in the final runs used to produce results for the paper took approximately 2000 CPU hours.
C.4 Comparison to R3M experimental setup
There several low-level but important differences between evaluation setup and the used in the R3M paper [Narr et al. 2022] which uses some similar environments. These differences end up making the pretrained R3M representations perform worse in our evaluations than those in the original paper. For the kitchen tasks in particular, the biggest difference is that while the R3M paper considers only learning single subtasks (e.g. slide the door open, section 4.2 of the R3M paper), we consider learning sequences of subtasks (e.g. open the microwave, put the kettle on, turn on the light, and slide the door open, all in one trajectory). The R3M paper considers explicitly easier tasks. We did this because the kitchen data itself contains sequences of subtasks, not single subtasks (following the paper that introduced the kitchen dataset). For the metworld tasks, R3M chose to evaluate on a particular subset of tasks that are somewhat easier than average (this is why we consider two different splits of metaworld on with the R3M eval tasks and one with the original eval tasks from the metaworld paper). Another difference is that to focus solely on feature learning, we only pass in the image observation and not the proprioception while R3M passes in both. Again this makes the problem little bit more difficult. wW also render images at lower resolution due to computational constraints.
Finally, it important to note that R3M is attempting solve a different problem of general image representation learning that transfers across domains, while we are focusing on within domain, but cross-task generalization (which is easier analyze in a controlled way).
D Extended analysis discussion
Here provide more detailed discussion of related theoretical work.
One recent line of work focuses on learning representations for exploration [Efroni et al. 2021 Lamb al. 2022] and offline RL [Islam et al., 2022] in the presence of exogenous noise. The exogenous noise setting means that the high dimensional observations contain information that is not effected by the actions; e-g., background dynamies that appear in image observations but do not affect the task. This line of work argues that inverse dynamics modeling is the best approach to ignore exogenous noise. Our results are complementary to this line of work in showing that even settings without exogenous noise, inverse dynamics is still often preferable to alternatives for representation learning. Moreover, we consider a multitask imitation setting with latent contexts while they consider single task and reward-directed problems.
Another line of work proves that learning a forward dynamics model is a well-motivated approach for multitask imitation [Nachum and Yang/ (021). While that work does not directly compare to inverse dynamics pretraining, we find that inverse dynamics pretraining outperforms forward dynamics modeling in our settings. Moreover, while this paper shows that if our representation learns good forward dynamics model that it works well for imitation, it does not discuss how efficiently such representation can be learned. So, while both methods are mell-motivated, we find inverse dynamics modeling to be more efficient than learning the forward dynamics.
25

---

Finally, another line of work studies multitask representation learning for imitation by directly performing behavior cloning [Arora et al. 2019 Zhang al. 2022]. These methods provide positive results for the approach, but require algorithms that have access to the latent context information which must be discrete so as to learn a separate policy for every pretraining context, thus avoiding confounding. This method requires extra information and is difficult to scale to very large numbers of contexts. contrast, we find that inverse dynamics modeling is able to perform well without this extra information added complexity of learning multiple models and naturally avoids confounding by the latent context information.
26


Reviews:
[{"review_id": "YqGOkGPR2X", "reviewer": "Reviewer_JZDX", "summary": "This paper investigates the effectiveness of several popular representation learning techniques in the context of imitation learning. While there is no novel components in this paper, the main strength of this paper is that it does well in designing the experimental setups to investigate the effect of several components that might disrupt the message of experiments. And in conclusion, the paper shows that inverse dynamics modelling objective is the best objective among the ones considered in the paper.", "questions": "- Please address my concerns & questions in Weaknesses.\n- Having a recurrent architecture for BC or stacking some frames might resolve the issue of BC for inferring the latent variable. Did you consider these baselines?", "limitations": "I couldn't find the discussion on the limitation of this paper in the submitted main draft.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Crystal-clear writing that clearly explains the setup and the results in a well-structured manner.\n- Extensive experimental results that show inverse dynamics modelling can be better than other objectives in a controlled way.\n", "weaknesses": "- While the results are quite conclusive in the considered setup of imitation learning with clean expert demonstrations, it might not hold in a setup where the dataset is suboptimal so that it's difficult to learn the intention of the agent for making $a_{t}$ even with the access to consecutive two observations $o_{t}$ and $o_{t+1}$. Though I understand that this is not the main goal of this paper, having more discussion on this front, along with the discussion on the paper that discusses the sufficiency of representation learning for control [1], can be useful and strengthen the paper.\n- Current trend might not be conclusive as the considered domains are still very simple. Including some additional experiments on more challenging domains like RLBench [2] and RoboSuite [3], which are considered more real-world-ish than simple Kitchen or Meta-World domains, or even on real-world robotics domains (I understand this might not be available in most cases) can make the conclusion of this paper be much more convincing.\n- Investigation into the effect of data configuration can be more helpful for understanding when the inverse dynamics modelling can be effective. For instance, how does the trend change when the pre-training and fine-tuning domains are more simpler domains?\n- This is very minor, but including the resolution of the result figures and making them be vectorized figures (by exporting them as pdfs) in the draft can be helpful for the clarity of the paper. Also there's a typo in line 278 (the the).\n\n[1] Rakelly, Kate, et al. \"Which Mutual-Information Representation Learning Objectives are Sufficient for Control?.\" Advances in Neural Information Processing Systems 34 (2021): 26345-26357.\n\n[2] James, Stephen, et al. \"Rlbench: The robot learning benchmark & learning environment.\" IEEE Robotics and Automation Letters 5.2 (2020): 3019-3026.\n\n[3] Zhu, Yuke, et al. \"robosuite: A modular simulation framework and benchmark for robot learning.\" arXiv preprint arXiv:2009.12293 (2020).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "First, we'd like to thank the reviewer for their positive comments about the clarity of the paper and comprehensiveness of our controlled experiments.\n\nHere we will address each of the weaknesses and one additional question raised by the reviewer in turn. Hopefully these comments can provide some additional clarity. If they do, we encourage the reviewer to increase their score, and otherwise are happy to answer any follow up questions.\n\n1. (Discussion of suboptimal data) We agree with the reviewer that the results do not necessarily extend to cases of suboptimal data (as we state briefly in the discussion in lines 361-362). The main reason for this choice is to keep everything self-contained in an imitation framework. Once data is suboptimal, then it makes more sense to look towards methods like offline RL which were beyond the scope of our study. We would hope that the insights from the imitation case would transfer to RL settings, but this is left to future work. We will add a broader discussion of this issue to the paper.\n\n2. (More challenging domains) Again we agree with the reviewer that using more challenging or real-world domains would be a nice extension of the paper (as we briefly state in the discussion in lines 360-361). However, we would argue that the tasks that we consider are not toy and do provide sufficient complexity for interesting results. The main goal of this paper was to create carefully controlled and targeted experiments in diverse domains to get clear insight, and we think we have accomplished this. Moreover, we prioritized domains with large numbers of predefined tasks and datasets with a single morphology. We decided that robosuite did not have sufficiently many tasks and chose metaworld over RLBench because it had been used more frequently in related work (like R3M) and depends on mujoco rather than coppeliasim which was not supported on our compute infrastructure. That said, we definitely agree that future work to scale up these insights by building better environments and datasets and trying things on real robots is indeed a good idea. We will add a broader discussion of this issue to the paper.\n\n3. We are not sure what the reviewer meant by \"how does the trend change when the pre-training and fine-tuning domains are more simpler domains?\" Any further clarification would be much appreciated.\n\n4. (Typo and resolution) Thanks for pointing out these issues, we will fix them in the paper.\n\n5. (Question -- frame stacking) Thanks for raising this interesting question. We did try adding frame stacking at one phase of experimentation and found that it did not help, so we did not scale it up to the full experiments. If the reviewer thinks this is a crucial issue, we can attempt to go back and run these experiments across the full suite. From the theoretical side, we do not think that adding recurrence or frame stacking will resolve the dependence issue in general. Looking at the graphical model in Figure 1(a), note that even if we condition on $ o_1 $ and $ o_2$, we do not break the connection between $ c $ and $ a_2$. However, it is possible that in certain environments the history could be sufficient to uniquely determine $ c $ and thus break the dependence. That said, just because $ c $ could be inferred does not mean that we will learn the desired features which are importantly independent of $ c$. It seems that we may instead learn features that depend explicitly on $ c $ and thus would struggle to generalize to new contexts.\n\nAlso, as a note, we will explicitly add some of the above discussion about suboptimal data and more challenging domains to an explicit limitations subsection included with the discussion at the end of the paper (currently limitations are only discussed in passing, not in one central place as the reviewer noted)."}, {"review_id": "yash3kg4JT", "reviewer": "Reviewer_t1mu", "summary": "The paper studies how to effectively pretraining representations in imitation learning where we have acess to a pretraining dataset\nwith multi-task demonstrations with an unobserved latent context variable for each task and limited amount of finetuning demonstrations\nfor transferring to novel context. The goal in this setup is to learn good low dimensional representations of high dimensional input \n(e.g. visual) to enable transfer to nove context for finetuning.\n\nThe paper claims, inverse dynamics modeling is a well-suited objective for imitation learning pretraining setting and provide\n empirical evidence supportng the claim. In addition, the paper also derives a theoretical analysis using simple \ngeneral environmental model.", "questions": "The paper doesn't present any novel method but presents interesting and valuable analysis that is worth sharing with the community.\n\nAs outlined in weakness section if authors can discuss about the mentioned questions.", "limitations": "Scaling these methods to much more complex high dimensional and long-horizon tasks might not be trivial given the compute \nrequirement.\n\nNo novelty in contributing a new method but the paper presents interesting analysis", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents an extensive set of experiments for comparing different methods of pretraining policy representations using \nmultiple training methods like Inverse Dynamics, Behavior Cloning, Forward dynamics (implicit and explicit), training policy from scratch\n, using pretrained visual representations and contrastive pretraining. \n\n2. The results are presented on a diverse suite of tasks which include simpler environments like PointMass to slightly more difficult \nenvironments like MetaWorld or Kitchen environments.\n\n3.  The paper evaluates, how well these representations perform with respect to dataset size used for pretraining and finetuning,\n how well policies learnt from these mthods generalize to in and out of distribution evaluation tasks. Authors find that ID pretraining leads\nto better performing representations on average across benchmark. The paper also presents a nice analysis on how well the representations \nlearned can predict state and actions. Interestingly ID performs quite well at predict all of this info vs any other method. \n\n4. Finally, the paper also presents a theoretical analysis to support the empirical results using a simplified model\n\n\n5. The paper doesn't present any novel method but presents interesting and valuable analysis that is worth sharing with the community", "weaknesses": "1. It is not clear why pretrained robotic representations like R3M performs so poorly on almost all of the tasks. The R3M paper \nalso has results on MetaWorld, Franka Kitchen tasks and it seems like the results for kitchen environment with R3M are quite \npoor based on figures 11 & 12 from appendix. Can authors please clarify if the experiments use the same setup as R3M paper? \nAlso, it might good to provide further insights into why this is the case.\n\n2. The evaluation of these methods is shown on different simple simulated environments but it might also be interesting to compare\nthese results on more photorealistic simulation environments like Habitat or AI2Thor for tasks like manipulation.\n\n3. The performance of training from scratch baseline seems to almost always be equal to performance of policies trained using ID\npretrained representations as we scale the training dataset size. This seems a bit concerning as pretraining also requires large \namounts of data for any of these methods. Do authors have a specific set of tasks where they observe sample efficiency? i.e. \nID pretrained representations converge to similar performance early on during finetuning? It'd be good to outline the gains achieved \napart from task performance to quantify importance of pretrained representations.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "First, we would like to thank the reviewer for their thoughtful and detailed comments about the strengths of the paper on both the experimental and analytical/theoretical fronts.\n\nHere we will address each of the weaknesses raised by the reviewer in turn. Hopefully these comments can provide some additional clarity. If they do, we encourage the reviewer to increase their score, and otherwise are happy to answer any follow up questions.\n\n1. (R3M results) Thanks for raising this subtle issue. Indeed, there are several low-level but important differences between our evaluation setup and the one used in the R3M paper. For the kitchen tasks in particular, the biggest difference is that while the R3M paper considers only learning single subtasks (e.g. slide the door open, see section 4.2 of the R3M paper), we consider learning *sequences of subtasks* (e.g. open the microwave, put the kettle on, turn on the light, *and* slide the door open, all in one trajectory). The R3M paper considers explicitly easier tasks. We did this because the kitchen data itself contains sequences of subtasks, not single subtasks (following the paper that introduced the kitchen dataset). For the metworld tasks, there is a similar pattern where R3M chose to evaluate on particularly easy tasks (this is why we consider two different splits of metaworld on with the R3M eval tasks and one with the original eval tasks from the metaworld paper). Another difference is that to focus solely on feature learning, we only pass in the image observation and not the proprioception while R3M passes in both. Again this makes the problem a little bit more difficult. Finally, we also render images at a lower resolution due to computational constraints. When applying R3M we resize the images, but they will be more pixelated than the ones R3M was originally evaluated on. All of these differences likely contribute to the results that we report. We will add a discussion of these differences to the appendix in the paper. As one last point, it is important to note that R3M is attempting to solve a different problem of general image representation learning that transfers across domains, while we are focusing on within domain, but cross-task generalization (which is easier to analyze in a controlled way).\n\n2. (Photorealistic environments) We agree that evaluation in more challenging environments would be an interesting extension of our work (as we briefly state in the discussion in lines 360-361). However, we would argue that the tasks that we consider are not toy, and do provide sufficient complexity for interesting results. The main goal of this paper was to create carefully controlled and targeted experiments in simple domains to get clear insight, and we think we have accomplished this. Moreover, we prioritized domains with large numbers of predefined tasks and datasets with a single morphology that had been used in related work. We are aware of Habitat and AI2-Thor, but have not seen similar suites of tasks with demonstration data as exists for the environments that we chose (and as a result, none of the related work uses these environments either). That said, we definitely agree that future work to scale up these insights by building better environments and datasets and trying things on real robots is indeed a good idea. We will add a broader discussion of this issue to the paper.\n\n3. (Data scaling) As the reviewer notes, once we approach the limit of large finetuning data, the gap between pretraining and training from scratch disappears. We want to emphasize that this is totally expected behavior for testing transfer learning with any pretraining algorithm. As long as the training from scratch algorithm is sound, it will approach optimal performance if given enough data. The interesting case for pretraining (and the one that we focus on in our main experiments, e.g. in Figure 1), is when we only have small amounts of finetuning data. Figure 1 shows that in this regime, there is indeed a benefit to pretraining. But as the reviewer notes, we should indeed be careful to point out that these gains explicitly depend on the relative amounts of pretraining and finetuning data that are available as shown explicitly by sweeping these parameters in Figure 3. We will make this more clear in the paper.\n\nAlso, as a note, we will explicitly add some of the above discussion about more challenging domains to an explicit limitations subsection included with the discussion at the end of the paper (currently limitations are only discussed in passing, not in one centralized place).\n"}, {"review_id": "sIBQplH0Jj", "reviewer": "Reviewer_8bvo", "summary": "This paper analysis the downstream fine-tuned performance on various multi-task settings after pretaining with different objectives including - \n\n1) Inverse dynamics\n2) Contrastive\n3) Forward Dynamics - Explicit\n4) Forward Dynamics - Implicit\n5) Behaviour Cloning\n\nThey compare the performance to a model trained from scratch using image representation and augmentations (Pixels + Aug). They also compare to R3M and imagenet pretrained models as baselines. The main conclusion of this paper is that IDM works better than other pretraining objectives and also works better than models trained from scratch in small data regimes. They experimentally and theoretically show that IDM can recover the representation of the true ground truth state while other methods such as Forward dynamics and Behavior Cloning can fail to do so.\n", "questions": "- It is not clear to me what context variables mean in the context of experiments. Could you clarify with examples what would be a context variable in some datasets considered in the paper?\n- There is some evidence in literature in papers such as SGI and SPR that contrastive learning is useful for RL. Although their finetuning setup is different from this paper as they do not use imitation learning for finetuning. Do the authors have any intuition as to why contrastive learning helps there but not here?\n\nSGI - https://arxiv.org/abs/2106.04799\nSPR - https://arxiv.org/abs/2007.05929", "limitations": "Yes, the authors have addressed limitations.\n\nIf the authors use expert data for training as I pointed above, then I think they should mention that as a limitation as well as I think it might be difficult to obtain such data in many environments.\n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a comprehensive evaluation of different pretraining objectives, including a comparison to models trained from scratch using image representation and augmentations (Pixels + Aug), as well as R3M and imagenet pretrained models as baselines. The main conclusion drawn from this study is that IDM (Input-Dependent Memory) outperforms other pretraining objectives, as well as models trained from scratch, especially in scenarios with limited data availability. The strengths of this paper lie in its valuable insights and experiments that not only demonstrate the superior performance of IDM but also provide a clear understanding of why it works through various experimental analyses. These findings contribute to a better understanding of pretraining objectives and their effectiveness in recovering the representation of the true ground truth state, highlighting the limitations of alternative methods such as Forward dynamics and Behavior Cloning, which fail to achieve comparable results.\n", "weaknesses": "I think the paper lacks certain details and contexts in various places. I am not sure what is the objective used while training the policy in the finetuning phase. It seems to be that it is behavior cloning but I am not sure. Another confusion I have is the inverse dynamics objective is specified using a mean-squared error loss so it seems that the actions are continuous valued but they evaluate action predictions using binary cross-entropy loss, so are action values constrained between 0 and 1. It would be useful for readers if such details are clarified.\n\nFrom my understanding, the approach relies on expert trajectories to obtain pretraining data which seems like a limitation to me since such data may not be available in many cases. Many approaches like decision transformer (https://arxiv.org/abs/2106.01345) train using suboptimal data. It would be nice to have an analysis in the paper which uses suboptimal data and shows how it affects the different pretraining objectives. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nFirst, we would like to thank the reviewer for their comments about the strengths of the paper, namely the \"comprehensive evaluation\" and resulting \"valuable insights\".\n\nHere we will address each of the weaknesses and questions raised by the reviewer. Hopefully these comments can provide some additional clarity. If they do, we encourage the reviewer to increase their score, and otherwise are happy to answer any follow up questions.\n\n1. (Setup details) Thanks for raising these issues to our attention. Indeed, we use behavior cloning during the finetuning phase, we will clarify this more explicitly in Equation (5). For inverse dynamics, indeed we also use mean squared error. We are not sure where the reviewer sees binary cross-entropy loss in the paper since it is never used or mentioned. Actions are always continuous and always trained using mean squared error. Perhaps the confusion arises since we report task success as a percentage (valued between 0 and 1), but this is really the average reward not a measure of action prediction accuracy (where the reward is 1 for success and 0 for failure). We will clarify this in the paper.\n\n2. (Suboptimal data) We agree that extending our results to cases with suboptimal data would be an interesting direction for future work and explicitly say so briefly in the discussion in lines 361-362. However, algorithms like decision transformer that rely on access to task rewards are beyond the scope of our paper where we instead focus on the imitation setting where there is no reward information available. We made this choice so that we could create carefully controlled experiments to just test the different representation learning algorithms without the additional complexities of reward-based learning. We will add a broader discussion of this decision in the discussion section as a great direction for future work.\n\n3. (Question - context variables) Thanks for this question, the context variables are detailed in Appendix C, but we realize that we likely did not provide enough clarity in the main text. For example, in the pointmass task the context determines the 2d continuous goal location, in the kitchen task the context determines which of 24 possible sequences of subtasks is the desired behavior (e.g. open the microwave then put the kettle on then turn on the light then open the cabinet is an example of one specific context), and in the metaworld tasks the context determines which of 50 behaviors is desired (e.g. close the box is one context and lock the door is a different context). We will move some of these examples from the appendix into the main text to increase clarity.\n\n4. (Question - contrastive learning) Thanks for bringing these related works to our attention. There are substantial differences between our setup and the one in SGI or SPR beyond just imitation vs. RL (although that is indeed a major difference) that could lead to different results for contrastive learning. SPR is a single-task online method which is very different in nature from our offline multitask setup that explicitly considers the ability of representations to transfer to novel contexts (i.e. tasks). SGI is more similar since they also consider offline pretraining, however instead of considering a setting where information is transferred across tasks from multitask expert pretraining data, they consider single task learning from non-expert policies by collecting data from suboptimal (or random) policies that are attempting to solve the same task. We hypothesize that this substantial difference in dataset composition and gap between pretraining and finetuning tasks is likely the reason that they find contrastive learning to be more useful. It is interesting to note that even in this substantially different setting, the SGI paper finds inverse dynamics modeling alone to be more useful that contrastive learning alone in an ablation (Table 3), so maybe our analysis could be extended to the setting they consider in SGI in future work. We can add a discussion of these paper to the extended related work in Appendix A.\n\nThe biggest difference is that those papers do single task learning with auxiliary objectives, while we do multitask learning. Explicitly, SGI and SPR consider an online RL problem where there is only one task (i.e. one context) which determines the reward function. They do not attempt to re-use representations across tasks (i.e. for different reward functions). On the other hand, we focus on a multitask imitation setup where we are given a fixed dataset of offline experience from many different contexts (i.e. tasks) and learn one representation .\n\nAlso, as a note, we will explicitly add some of the above discussion about suboptimal data to an explicit limitations subsection included with the discussion at the end of the paper (currently limitations are only discussed in passing, not in one central place as the reviewer noted).\n"}, {"review_id": "bbGYGdJlTU", "reviewer": "Reviewer_Rymd", "summary": "This paper provides both empirical evidence and theoretical arguments for the usefulness of the inverse dynamics criteria as a pre-training objective for multitask imitation learning setups. Furthermore, the authors were able to go beyond the final conclusion and provide insightful analysis where they were able to show that the advantages of the inverse dynamics objective compared to alternative objectives such as behavioral cloning are rooted in better out-of-distribution generalization and robustness to environments with fully latent context. \n", "questions": "Can you elaborate on how your work relates to recent works on multitask imitation learning via generative pre-training on trajectories? It seem that the generative pre-training objective in these works implicitly combines the forward dynamics and behavioral cloning criteria.\n\n[1] Kuang-Huei Lee, Ofir Nachum, Mengjiao Yang, Lisa Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, Ian Fischer, Eric Jang, Henryk Michalewski, and Igor Mordatch. Multi-Game Decision Transformers, 2022.\n\n[2] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A Generalist Agent, 2022.\n", "limitations": "In its current state, the theoretical arguments lack formal claims and proofs. Thus they feel more like mathematical intuitions rather than solid theoretical results.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The authors support their claims with both compelling empirical evidence and novel theoretical arguments.  \n\n2. The paper is clearly written and accessible to readers with varying backgrounds. ", "weaknesses": "The superior performance of the inverse dynamics criteria in the examined setting is somewhat expected. Namely, unlike behavioral cloning, the latent multitask performance does not hurt it since the inverse dynamics criteria is oblivious to the latent reward. And the examined architecture seems more natural for inverse dynamics than for forward dynamics. Arguably, forward dynamics could benefit from an architecture that can fuse the observation with the action earlier.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nFirst, we would like to thank the reviewer for their positive comments about the compelling evidence, novel analysis, and clear writing in the paper.\n\nHere we will address each of the weaknesses and the questions raised by the reviewer. Hopefully these comments can provide some additional clarity. If they do, we encourage the reviewer to increase their score, and otherwise are happy to answer any follow up questions.\n\n1. (Somewhat expected results) We agree with the reviewer that the results can be explained well in hindsight, and indeed this is what our analysis section intends to do. However, we disagree that these results are obvious or expected since they are not found in the existing literature (as detailed in Section 2 and Appendix A). As a side remark, we were not initially expecting to see a significant advantage of inverse dynamics ourselves at the beginning of the project. If the reviewer could point to prior work in the literature that made these observations, we would be happy to consider it, but we do not think this is a valid weakness of the work.\n\n2. (Architectural details) We agree with the reviewer that it is difficult to make comparisons between algorithms that require different architectures. However, we disagree that our choices were in any way unreasonable. We attempted to maximize performance of all baselines with a consistent compute budget across methods. We control for architectural differences as much as possible by using the exact same image encoder architecture and feature fusion method for all networks. To compare inverse dynamics and implicit forward dynamics explicitly, both networks consist of an MLP on top of the concatenation of two encoded vectors: $ [\\phi(o), \\phi(o')] $ for inverse dynamics and $ [\\phi(o), \\phi_a(a)]$ for forward dynamics. Then the MLP head has the capacity to fuse the features. Moreover, since the image needs to be encoded with a convnet, this seems to be the standard way to incorporate action information since it cannot be fed directly into the convnet. Note for explicit forward dynamics we need to add a decoder on top of encoded features to produce an image to compute the loss. The objective requires this larger network, but we think we handled the difference in the most carefully controlled way possible. If the reviewer can point to related literature or more specific arguments about why our architectural decisions did not make sense, we would be happy to discuss further. \n\n3. (Question - generative pretraining) Thanks for this interesting question. Indeed, it does seem that generative pretraining could be thought of as a combination of BC and FD. It is not clear why creating such a combination would resolve the issues of either method in our particular setting, but this raises an interesting point that all of the objectives we consider can be combined in various ways (e.g. added together). Anecdotally, we did try adding together various objectives during development for this project but found the mixtures to perform worse than the simple objectives. In the end, in this paper we want to focus on presenting the cleanest controlled results that we can, so we opt to exclude these combinations for now, but it is definitely an important direction for future work to consider whether cleverly combining objectives could outperform the individual ones. We will add the papers that the reviewer referenced to the related work section as well as an expanded discussion of this direction of future work.\n\nAlso, as a note, we agree with the reviewer that the analysis section is not formal theory (as in theorems). It also only considers a substantially simplified model. We include it for intuition (which we believe the reviewer agrees is useful, based on the strengths section of the review). That said we will discuss how the lack of formal theorem statements is a limitation in a new limitations subsection to be included with the discussion at the end of the paper, and could present a fruitful direction for future theoretical work.\n\n"}]

Author rebuttal:

Meta-review: Inverse dynamics is the superior objective function for multitask imitation. The demonstration of this and the arguments that support it will be of wide interest and will influence other work. Reviewers were unanimous in their support. At the same time, they noted that the manuscript does not provide a proof as to why this is. This could also be interesting future work.

Multiple reviewers brought up the fact that the experiments are all in simple environments. Authors responded that they believe their experiments are sufficient. There is some question about how true this is and it would be good to highlight this explicitly in the manuscript. It may be the future work finds that the picture is not as clear in more complex domains or in some domains with specific properties.
Decision: Accept (poster)