SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To test the impact of this gap between pretraining and finetuning, we created alternative pretraining datasets, where we include the test contexts (but not 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 a much more limited notion of generalization from pretraining to finetuning.
Results are shown in Figure 4. 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 outperforms training from scratch Pixels + Aug. None of the other pretraining algorithms benefit much from the substantially easier type of generalization required on these datasets. This suggests that ID and BC are uniquely able to benefit in easier settings, suggesting that they are better representation learners. If an algorithm is not able to outperform training from scratch in this simplified setting, it is unlikely to be a good representation learner.
Fully latent vs. inferrable context variables. Looking at our dataset suite, the datasets can be divided into two groups: those where the
context variable is not inferrable at 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-ml45, and metaworld-r3m). This split presents an interesting comparison in particular between ID and BC (the best performing algorithms from the prior experiment). Figure 5 shows results for these algorithms and the Pixels + Aug baseline on the datasets where the context is latent.
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 trajectories not seen during pretraining.
Figure 5: A comparison between ID and BC on the datasets where the context is not inferrable from the observation.
There is a large gap between ID and BC when the context 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 a good representation. Note that this gap largely disappears when the context is observable, see Appendix B for further details.
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 is 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
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 Figure 6 show the train and validation loss during finetuning using the default pretraining and finetuning sizes from Table 1. Since losses vary across datasets, we normalize by the Pixels + Aug validation loss so as to be able to present averages across all datasets. We see that out of the learned representations, ID has both the lowest train and validation losses, almost matching the performance of Pixels + Aug on train and almost matching the performance of States on validation. In contrast, representations that attempt to predict forward dynamics have substantially higher train loss, indicating aliasing of states in terms of their optimal actions. Interestingly, the Imagenet pretrained features have very low train loss, indicating a lack of
aliasing, but very high validation loss, indicating that the function that maps representations to actions does not generalize well.
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 our experiments in simulated environments, we 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 so as to isolate the effect of the representation learning apart from the gap between pretraining and finetuning. Again we normalize the losses for each dataset.
Again we see that ID and BC yield the best performance. This suggests that in these datasets, pretraining objectives that attempt to predict the optimal action do indeed facilitate recovery of the low-dimensional simulator state. In contrast, while the FD methods achieve approximately the same training error, they generalize much more poorly. 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 it is throwing away important information. Extended results are in Appendix B.
Figure 7: Average state prediction error on the pretraining distribution. Values are normalized by the ID train loss.
To add a more theoretical understanding of the empirical results, we will consider a simplified model of the data generating process based on linear dynamics in a 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 are 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 a tuple  $(o,a,o^{\prime},s,s^{\prime},c)$  that is sampled iid from the joint distribution over those variables. Recall that we only observe  $(o,a,o^{\prime})$  and that  $(s,s^{\prime},c)$  are latent. Formally, let  $\mathcal{O} = \mathbb{R}^d$ ,  $S = \mathbb{R}^\ell$ , and  $\mathcal{A} = \mathbb{R}^k$  with  $d\gg \ell >k$ . Let  $\phi :\mathcal{O}\to S$  be the ground truth encoder, which we assume is invertible by  $\phi^{-1}$ . Let  $\epsilon \sim \mathcal{N}(0,\Sigma)$  in  $\mathbb{R}^\ell$  and  $A,B$  to be any matrices in  $\mathbb{R}^{\ell \times \ell}$  and  $\mathbb{R}^{\ell \times k}$ . Then, assume that the dynamics are:
$$
o ^ {\prime} = \phi^ {- 1} (A \phi (o) + B a + \epsilon). \tag {13}
$$
Note that we make no assumption on the policy  $\pi_c^*$  other than that it only depends on  $o$  via  $\phi(o)$ . This model is similar to ones studied in the online control setting by Mhammedi et al. [2020], Dean and Recht [2021], 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 in this model, note that if  $B^{+}$  is the pseudoinverse of  $B$  that:
$$
a = B ^ {+} \phi \left(o ^ {\prime}\right) - B ^ {+} A \phi (o) - B ^ {+} \epsilon . \tag {14}
$$
Thus the inverse dynamics are a simple linear function of the embeddings  $\phi(o), \phi(o')$ . As a result, when we solve for  $a$  with least squares regression, if the encoder  $\phi$  is 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^{+}\epsilon$  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 a decoder while inverse dynamics only requires learning the encoder. This is not necessarily a problem a 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).
To grasp why we might expect this, note that the set of possible observations is the manifold represented by the image of the decoder, i.e.  $\operatorname{Im}(\phi^{-1})$ . As a simple example, consider a toy 2d example where the high dimensional observation is  $(x, f(x)) \in \mathbb{R}^2$  and the low dimensional state is simply  $x \in \mathbb{R}^1$ , as depicted in Figure 8. Here the encoder  $\phi$  is very simple since it just needs to recover  $x$ , while the decoder must learn  $f(x)$ . Of course this is a very toy example, but we find it illustrative of the idea that it is possible that the encoder is much simpler than the decoder in practice.
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 a 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  $\ell = k$  and  $\epsilon = 0$ . Let  $\mathcal{R}(\mathbb{R}^{k\times k})$  be the set of rotation matrices in  $\mathbb{R}^k$ . Let  $\mathbb{S}^{k - 1}$  be the unit sphere in  $\mathbb{R}^k$ ,  $U$  be the uniform distribution, and  $\delta$  denote a Dirac delta. Now, assume:
$$
c \sim U (\mathcal {R} \left(\mathbb {R} ^ {k \times k}\right)), \quad o \sim U \left(\phi^ {- 1} \left(\mathbb {S} ^ {k - 1}\right)\right), \quad \pi_ {c} ^ {*} (a | o) = \delta [ a = c \phi (o) ] \tag {15}
$$
Note that  $\phi (o)$  returns a unit vector in  $\mathbb{R}^k$  and that a uniformly sampled rotation of a unit vector is a uniformly sampled unit vector. Thus, we can marginalize over  $c$  to get:
$$
P (a | o) = \int_ {c} P (c) \pi_ {c} ^ {*} (a | o) = \int_ {c} P (c) \delta [ a = c \phi (o) ] = P _ {U \left(\mathbb {S} ^ {k - 1}\right)} (a) = \eta_ {k}, \tag {16}
$$
for a constant  $\eta_{k}$  equal to the reciprocal of the surface area of the unit sphere in  $\mathbb{R}^k$ .
Thus, the Bayes optimal BC policy does not depend on  $o$  at all. As a result, the optimal representation learned by BC can just map every observation to zero. This representation is not capable of representing 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  $c$ .
We have seen that inverse dynamics pretraining provides an effective method for learning features from multitask demonstration data. We demonstrated this across a 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 we 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 we 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 a toy model, there is clearly room for better theory. Ideally, future work could present a 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 on our results. One direction would be to 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 a variety of tasks. In future, it would be interesting to 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 conditional generative models (e.g. goal-conditioning, language-conditioning, 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.
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.
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. Opal: Offline primitive discovery for accelerating offline reinforcement learning. arXiv preprint arXiv:2010.13611, 2020.
Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A theoretical analysis of contrastive unsupervised representation learning. arXiv preprint arXiv: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.
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 arXiv:2212.00186, 2022.
In this paper we focus specifically on pretraining methods that learn representations of high dimensional 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 are 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-directed representations in the context of reinforcement learning. This is a 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 et 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 et al., 2018, Aytar et al., 2018, Seo et al., 2022, Ma et al., 2022, Zakka et al., 2022, Ghosh et al., 2023]. In this paper, we focus on the simpler task of pretraining a representation within one MDP with consistent dynamics and access to demonstration actions, but with varied tasks. This choice allows us to make more clear comparisons between algorithms and rigorous claims about when representations will be effective, but also limits the generality of the representations that are learned.
There are a variety of new methods that rely on transformer architectures to construct interesting new pretraining objectives [Yang and Nachum, 2021, Lee et al., 2022a, Reed et al., 2022, Seo 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 can adapt to new tasks [Duan et al., 2017, Finn et al., 2017a,b, Yu et al., 2018, 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 have access to the latent context, it is unclear how to create these 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 we focus on 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 to gauge whether it is worth the cost of labeling trajectories with context information.
It 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.1 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 B.4 presents 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 MLPs to predict either the actions on the finetuning data or the low dimensional states on the pretraining
data. Here we present a similar analysis, but now where we use small MLPs to predict the other representations themselves. This is interesting since it lets us assess which representations contain enough information and shared structure to predict the other representations. Hypothetically, a representation that is easily able to recover another representation may be preferable since it retains more information.
Results presented in Figure 9 show the average across datasets of the cross-representation prediction error on a 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 at the rows, which correspond to the error when each method is used as the source, we can see that inverse dynamics generally has the lowest average error for predicting the other representations. This suggests that inverse dynamics is doing a good job of recovering the information that is shared among all the representations. Second, looking at the columns, which correspond to error when each representation is used as the target, we see that BC is the most difficult to predict and inverse dynamics is second most difficult. This is a somewhat surprising result, but suggests that these representations actually contain 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 is consistent with the idea that those representations are losing important task-relevant information.
Figure 9: Cross-representation prediction error of a small MLP on a validation set from the pretraining distribution. Results are normalized per dataset by the mean error on that dataset and then averaged across datasets.
Full results on each dataset can be found in Appendix B.3 and full methodological details can be found in Appendix C.
B.2 Per dataset evaluation success results
In the main text and Section B.1 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, finetuning 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 3. The full per dataset results are shown in Figure 10.
There are several findings in the dataset-specific results that are not visible in the aggregate reported in the main text: