Inverse Dynamics Pretraining Learns Good Representations for Multitask Imitation

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.

Strengths

1. The authors support their claims with both compelling empirical evidence and novel theoretical arguments. 2. 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.

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. [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. [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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

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.

Summary

This paper analysis the downstream fine-tuned performance on various multi-task settings after pretaining with different objectives including - 1) Inverse dynamics 2) Contrastive 3) Forward Dynamics - Explicit 4) Forward Dynamics - Implicit 5) Behaviour Cloning They 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.

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.

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. From 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.

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? - 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? SGI - https://arxiv.org/abs/2106.04799 SPR - https://arxiv.org/abs/2007.05929

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes, the authors have addressed limitations. If 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.

Summary

The paper studies how to effectively pretraining representations in imitation learning where we have acess to a pretraining dataset with multi-task demonstrations with an unobserved latent context variable for each task and limited amount of finetuning demonstrations for transferring to novel context. The goal in this setup is to learn good low dimensional representations of high dimensional input (e.g. visual) to enable transfer to nove context for finetuning. The paper claims, inverse dynamics modeling is a well-suited objective for imitation learning pretraining setting and provide empirical evidence supportng the claim. In addition, the paper also derives a theoretical analysis using simple general environmental model.

Strengths

1. The paper presents an extensive set of experiments for comparing different methods of pretraining policy representations using multiple training methods like Inverse Dynamics, Behavior Cloning, Forward dynamics (implicit and explicit), training policy from scratch , using pretrained visual representations and contrastive pretraining. 2. The results are presented on a diverse suite of tasks which include simpler environments like PointMass to slightly more difficult environments like MetaWorld or Kitchen environments. 3. The paper evaluates, how well these representations perform with respect to dataset size used for pretraining and finetuning, how well policies learnt from these mthods generalize to in and out of distribution evaluation tasks. Authors find that ID pretraining leads to better performing representations on average across benchmark. The paper also presents a nice analysis on how well the representations learned can predict state and actions. Interestingly ID performs quite well at predict all of this info vs any other method. 4. Finally, the paper also presents a theoretical analysis to support the empirical results using a simplified model 5. 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 also has results on MetaWorld, Franka Kitchen tasks and it seems like the results for kitchen environment with R3M are quite poor based on figures 11 & 12 from appendix. Can authors please clarify if the experiments use the same setup as R3M paper? Also, it might good to provide further insights into why this is the case. 2. The evaluation of these methods is shown on different simple simulated environments but it might also be interesting to compare these results on more photorealistic simulation environments like Habitat or AI2Thor for tasks like manipulation. 3. The performance of training from scratch baseline seems to almost always be equal to performance of policies trained using ID pretrained representations as we scale the training dataset size. This seems a bit concerning as pretraining also requires large amounts of data for any of these methods. Do authors have a specific set of tasks where they observe sample efficiency? i.e. ID pretrained representations converge to similar performance early on during finetuning? It'd be good to outline the gains achieved apart from task performance to quantify importance of pretrained representations.

Questions

The paper doesn't present any novel method but presents interesting and valuable analysis that is worth sharing with the community. As outlined in weakness section if authors can discuss about the mentioned questions.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Scaling these methods to much more complex high dimensional and long-horizon tasks might not be trivial given the compute requirement. No novelty in contributing a new method but the paper presents interesting analysis

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.

Strengths

- Crystal-clear writing that clearly explains the setup and the results in a well-structured manner. - Extensive experimental results that show inverse dynamics modelling can be better than other objectives in a controlled way.

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. - 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. - 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? - 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). [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. [2] James, Stephen, et al. "Rlbench: The robot learning benchmark & learning environment." IEEE Robotics and Automation Letters 5.2 (2020): 3019-3026. [3] Zhu, Yuke, et al. "robosuite: A modular simulation framework and benchmark for robot learning." arXiv preprint arXiv:2009.12293 (2020).

Questions

- Please address my concerns & questions in Weaknesses. - 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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

I couldn't find the discussion on the limitation of this paper in the submitted main draft.