Paper: What Makes Pre-Trained Visual Representations Successful for Robust Manipulation? (ICLR-2024)

Reviewer comment (Experiments):
The testing scenarios are not significantly different from the experimental settings, leading to conclusions that may not be solid.

Author response:
Some visual distribution shifts (such as lighting) are subtle and may only include slight changes in brightness or the direction of shadows. In spite of this, we see that lighting shifts cause the largest average drop in policy performance on FrankaKitchen when compared to more obvious changes such as changes in background texture or the presence of distractors. This suggests that the difficulty of zero-shot transfer may not correlate well with visual difference as perceived by humans.

Reviewer comment (Evaluation):
The scope of the conclusions is too narrow. While pre-trained models for visuo-motor control is a broad topic, this paper's findings are limited to the relationship between the ViT model and OOD performance under visual distribution shift, which limits in the specific setting and specific model. In this filed, there are too many different conclusions under different settings. Some slight differences may change a lot.

Author response:
One of the striking insights of our paper is the high correlation between the Jaccard index of the attention heads of a ViT model and the out-of-distribution performance of pre-trained models. In addition to this, we include an in-depth comparison of 5 ResNet models and a substantial discussion and analysis of the role of different datasets and pre-training strategies in yielding robust pre-trained models in sections 3 and 4. 

Broadly speaking, how best to pre-train visual representations is an open question in the field of learning based robotics. Much of the past work tried to train on datasets like Ego4D that include human-object interactions under the suspicion that these will be better for control. This paper shows that there are other design choices that supersede dataset choice. Moreover, we present an alternative metric to help guide the development of such models going forward. We believe this is a valuable contribution that the ICLR community could benefit from.

Reviewer comment (Novelty):
The conclusion that models pre-trained on manipulation-relevant data do not necessarily generalize better than models trained on standard pre-training datasets has been similarly addressed in previous papers.

Author response:
Could you provide a specific reference that shows this? To our knowledge, the only work that shows this is [1], which was made publicly available after our work was submitted to ICLR and is therefore considered concurrent work.

Reviewer comment (Novelty):
"We develop a benchmark for out-of-distribution generalization within FrankaKitchen and Meta-World." I think such a level of modification to the environment doesn't qualify as developing a benchmark. It's more accurate to say that some environments were constructed to verify our conclusions.

Author response:
We’ve changed the wording in that paragraph to “We construct environments to study out-of-distribution generalization within FrankaKitchen and Meta-World.” Please find the change highlighted in the updated PDF.

Reviewer comment (Experiments):
I find the spatial features in the figure intriguing. Could you conduct more analyses on how they change with varying levels of distracting factors? Additionally, I would suggest annotating Figure 1 with the value of the Jaccard Index.

Author response:
Yes, please see Figure 6 and the last paragraph of Section 5.2. From that section:  “One counter-argument to the use of Jaccard index as a metric for OOD performance is that it would be less predictive for object-level distribution shift, which would occur any time a large distractor is placed in the background of the image. In Figure 6, we plot the success rates of each ViT model as the number of objects increases and verify that the models with the higher Jaccard index actually maintain the highest performance as the number of distractors increases.”

Reviewer comment (Reproducibility):
Which pre-trained models were trained by you, and which ones are off-the-shelf?

Author response:
Consistent with past work on evaluating pre-trained visual representations [2], all of the models we study are off-the shelf.

Reviewer comment (Presentation):
The description for Figure 5 appears twice.

Author response:
Thank you for pointing these out. Please find both updated in the latest draft.

Reviewer comment (Writing):
The distinction between in-domain and out-domain generalization is blurry. As I discerned, the paper's in-domain generalization seems to pertain to training performance. This conflicts with more conventional uses of the term, such as evaluation on previously unobserved positions during training.

Author response:
Our in-domain generalization is not training performance. We run learned policies on kitchen and table positions sampled from the training distribution. The scene initializations are not seen during training. In the uploaded PDF, we’ve highlighted this distinction under the “Distribution Shifts” section.

Reviewer comment (Novelty):
The observation that self-supervised learning outperforms supervised learning in yielding generalized features is already discussed in [2], as acknowledged by the authors, making it less of a novelty.

Author response:
Could you clarify where in the paper we make this claim? Our results show that the choice to use supervision or not does not matter in comparison to metrics such as the Jaccard index. In paragraph 5 of section 4 we directly cite [2] stating how our results are a counterexample to theirs: “Parisi et al. (2022) also find that self-supervised learning outperforms supervised learning. Our results contradict this finding. Figure 4 shows that supervised training on Stylized ImageNet achieves a higher success rate in the training distribution than self-supervised training on ImageNet with a masked auto-encoding loss. These models maintain the same rank out-of-domain as well.”

Reviewer comment (Experiments):
The paper presents a hypothesis on shape bias correlating linearly with domain shift success for Resnets, but this is supported by merely three data points. This is a thin foundation for drawing substantial conclusions.

Author response:
In parts of the paper we make reference to [1] which shows that ImageNet trained CNNs are biased towards texture and that increasing shape bias improves accuracy and robustness. Our goal is not to make any new claims about shape-bias in this paper. If there are sections where we claim this please let us know which ones and we will modify them accordingly.

Reviewer comment (Experiments):
My main concern is that the paper assumes an artificially restricted setting with unrealistically limited demonstrations for a task (10, collected without varying a single thing in the background like lighting etc), on which even the training distribution performance is highly suboptimal (success rate is 40-60%). At this performance level, the policy would not be deployed even in the training environment. Typically, in such a limited demonstration regime, the demonstrations would be collected by varying placement and lighting conditions. I’m not sure (yet) that the paper actually measures what it says it measures - since the demos are so few and the training distribution policy performance is so low, the current degradation in performance under visual shift might not even be due to the representation being bad, but due to the policy network overfitting on correlations in the training data.

Author response:
Our training protocol is identical to the protocol used to evaluate pre-trained models for manipulation. In particular, we borrow training environments, experiment parameters, and code directly from R3M and recover similar performance within the train distribution as the original R3M does.  This evaluation protocol closely resembles the protocol typically used to evaluate many of the pre-trained models for manipulation.

With that said, we completely agree with the reviewer that the old way of evaluating models (that is, training with demonstrations in a single visual environment and evaluating performance in that environment) is suboptimal. One of the goals of this work is to support that idea.

However, the promise of visual pre-training is to amortize the cost of collecting visually diverse demonstrations. The ideal pre-trained visual representation would enable a researcher to collect a handful of demos in the lab without having to vary the appearance of the background or the direction of the lighting. To use the reviewer's example, we would hope that the pre-trained representation would be able to produce features that are invariant to the feature A that is spuriously correlated with actions. Motivated by this, we chose to _evaluate_ on visually diverse scenes instead of training on them so that we can see how well these models are holding this promise.

To address the reviewer’s concern about the number of demonstrations, we added additional experiments with 25 demonstrations for ResNet50-insup, ViT-DINO, and MVP on FrankaKitchen and find that the same trend holds strongly: ViT models with attention heads that have a high Jaccard index outperform models designed for manipulation. Please find these new results in Appendix Section A.12.

Reviewer comment (Experiments):
It seems possible that if two representations A and B are tested using the paper’s protocol and A encodes some extra variables that are spuriously correlated to actions in the limited demo data, then the policy could latch onto them to identify actions for prediction. However, these extra variables could be useful for prediction over a wider distribution of environments/conditions, and if the demos weren’t spuriously correlated with these extra encoded variables, then the learnt policy might have been better than the one learnt using B. Therefore, this trend might change if you went from 10 to, say, 40 demos where in the 40 demos the lighting or another condition is varied in a limited range [x-y] and the test time condition is a value outside this range. Including experiments of this format would make a much more convincing case, and would be much closer to an actual real-world use case.

Author response:
Our training protocol is identical to the protocol used to evaluate pre-trained models for manipulation. In particular, we borrow training environments, experiment parameters, and code directly from R3M and recover similar performance within the train distribution as the original R3M does.  This evaluation protocol closely resembles the protocol typically used to evaluate many of the pre-trained models for manipulation.

With that said, we completely agree with the reviewer that the old way of evaluating models (that is, training with demonstrations in a single visual environment and evaluating performance in that environment) is suboptimal. One of the goals of this work is to support that idea.

However, the promise of visual pre-training is to amortize the cost of collecting visually diverse demonstrations. The ideal pre-trained visual representation would enable a researcher to collect a handful of demos in the lab without having to vary the appearance of the background or the direction of the lighting. To use the reviewer's example, we would hope that the pre-trained representation would be able to produce features that are invariant to the feature A that is spuriously correlated with actions. Motivated by this, we chose to _evaluate_ on visually diverse scenes instead of training on them so that we can see how well these models are holding this promise.

To address the reviewer’s concern about the number of demonstrations, we added additional experiments with 25 demonstrations for ResNet50-insup, ViT-DINO, and MVP on FrankaKitchen and find that the same trend holds strongly: ViT models with attention heads that have a high Jaccard index outperform models designed for manipulation. Please find these new results in Appendix Section A.12.

Reviewer comment (Experiments):
Why does Section 5.2 say the probe (and shape bias value) is unavailable for MVP/R3M/VIP models? A linear probe just has to be trained on ImageNet on top of the representation, right?

Author response:
Yes. Our goal in this section is to show that these models are not predictive which is clear with the current evaluation set.