Pre-trained Text-to-Image Diffusion Models Are Versatile Representation Learners for Control

Thank you for taking the time to review our manuscript and for providing helpful and constructive feedback. We were happy to see your review mention the novelty of our work and the rigorousness and fairness of our evaluation. We address your concerns below.

---

1) The mutliple layer results for CLIP should be investigated further.

We agree and include the results for this in the general comment. On a high level, our conclusions from this experiment are that we do see a benefit from including additional outputs from certain layers along with the last layer which is always used (as we anticipated in the discussion in sec 5.1). We find these layers using a grid search, and they are located towards the middle (10-14th layers). We note that while the feature map sizes after including another layer for CLIP are very large and not comparable anymore to the other methods, the performance still doesn’t match that of SCR.

2) Could you explain more how these baselines work and why they get the numbers they do?

We agree with the reviewer and will add further contextualization of the performance of different baselines.

• In the manipulation space, representation learning works adopt a few shot imitation learning setting, with the backbone pretrained on video datasets of humans manipulating objects in indoor environments.
• In the navigation space, the setting is large scale reinforcement learning, and the evaluation includes generalization to new objects or scenes. The backbone is pretrained on datasets of indoor navigation videos.
• We select some of the best performing baselines used in prior works in manipulation (R3M, MVP) or navigation (VC-1) tasks. Different baselines often do well on disparate sets of tasks and settings. Specifically, R3M representations are tailored to reduce overfitting in few-shot IL settings by forcing sparsity in representations (see Table 1 in the R3M paper where performance drops 6% when the L1 sparsity penalty is removed), which doesn’t offer an advantage in the extended training regime of RL-based navigation tasks.

3) Simplicity of finetuning SCR compared to CLIP should be discussed earlier, claim about sophisticated negative sampling only cites the original CLIP paper

We agree that the discussion about simplicity of finetuning SD along with the CLIP fine-tuning experiment appears late and are happy to push it into the main paper. We cited the CLIP paper for the use of really large batches sizes in contrastive learning, but a more detailed list of references would be the following:

• References for "contrastive learning needs large batch sizes to avoid learning degenerate solutions" : the CLIP and SimCLR papers which rely on in-batch negatives, as well as momentum contrast methods like MoCo which cache mini-batch computations and use them in later batches.
• References for "Mining hard negatives can avoid having to use large batch sizes (but is itself non-trivial)" : “Contrastive Learning with Hard Negative Samples” [2], “Debiased Contrastive Learning” [3].

We did not use hard negative mining in our experiments and simply use the largest batch size which we could fit across 8x46GB A40 GPUs for CLIP finetuning.

4) Sec 5.3 where too many things are introduced at once.

We agree that section 5.3 abruptly introduces a lot of new things, and are happy to take the reviewer’s suggestion on the presentation of this section. We had wanted to present the following three results and had to choose a subset to retain in the main paper due to space constraints:

• We wanted to introduce the proxy visual grounding tasks used in Voltron [20] and present results for SCR on them.
• We wanted to ablate the method of spatial aggregation for CLIP and other baselines and show that retaining spatial information changes the conclusions of [20] completely on these tasks. We further wanted to show that a similar gain (or at least as considerable) is not seen on the actual control tasks.
• We further used the above visual grounding tasks to ablate how we provide the language inputs to SCR in the case of language guided tasks, and assess to what extent language modulates the representation.

5) Bilinear interpolation should be shown in equations formally to understand exactly what is done

We will modify figure 2 in the paper to include the equation along with the visual description of this concatenation operation. We primarily do this to enable stacking of the differently sized U-net feature maps so that we can use a single convolutional layer to process them without adding any other learnable parameters. A similar resizing operation was most recently used in [1], which looks at training a feature adapter module for control.

6) a) Why cite Imagen in L36? b) Reference to Table 2a/b are mixed up.

Thanks for pointing these out. a) We will move this reference (meant to be a citation for examples of foundation diffusion models) to the background instead, and b) fix the swapped references.

---

We hope that these clarifications and additional experimental results have addressed any remaining questions and concerns.

If you find that we have addressed your questions and concerns, we kindly ask you to consider raising your score to reflect that. If issues still remain, then we would be more than happy to answer these in the discussion period. Thank you.

---

References

[1] Lin, Xingyu, et al. "Spawnnet: Learning generalizable visuomotor skills from pre-trained networks." arXiv preprint arXiv:2307.03567 (2023).

[2] Contrastive Learning with Hard Negative Samples : Joshua David Robinson, Ching-Yao Chuang, Suvrit Sra, Stefanie Jegelka

[3] Debiased Contrastive Learning: Ching-Yao Chuang, Joshua Robinson, Lin Yen-Chen, Antonio Torralba, Stefanie Jegelka

We thank you for engaging with our work and commending our presentation, comprehensive ablations and the key result of being able to employ a single representation model for navigation and manipulation tasks. We address your questions and concerns below.

---

1) How do the methods compare when a common model, feature pooling strategy and training dataset are used? It is okay to limit to one kind of task (manipulation/navigation) for this experiment.

Pre-training a foundational representation model is out of the scope of this work and would not be feasible given the limited compute resources available to us. We note that navigation and manipulation tasks are not the primary focus of diffusion model pre-training (similarly to how navigation is not the key focus of baselines that have been designed for manipulation tasks). This is what makes their all-round performance surprising.

2) What if the same feature pooling strategy is applied to baselines like CLIP?

We present multi-layer feature aggregation results for CLIP for all benchmarks in the general response of our rebuttal. Note that using multiple layers makes the representation size of CLIP twice as large as all the other methods, and is not feasible to do comprehensively for the bigger navigation domains, so we present a smaller ablation there. On a high level, our conclusions from this experiment are that we do see a benefit from including additional outputs from certain layers along with the last layer which is always used (as we anticipated in the discussion in sec 5.1). We find these layers using a grid search and they are located towards the middle (10-14th layers). We note that while the feature map sizes after including another layer for CLIP are very large and not comparable anymore to the other methods, the performance still doesn’t match that of SCR.

3) Seems some multi-layer feature aggregation was done for Tables 1(b) and 1(c) but not 1(a)

To clarify, we used the same representation extraction procedure for SCR in all of the tables (1a, 1b and 1c), i.e., using the mid+downsampling layers for SCR.

4) Why is CLIP-B used in the ImageNav experiments and not CLIP-L?

This is because we compare directly to the results from VC-1 paper [27] where they used the CLIP ViT-base model. We note that ImageNav training runs require 16 GPUs for each run over 6 days of training. We can aim to include a result with CLIP-L in the revised manuscript.

---

We hope that these clarifications and additional experimental results have addressed any remaining questions and concerns.

If you find that we have addressed your questions and concerns, we kindly ask you to consider raising your score to reflect that. If issues still remain, then we would be more than happy to answer these in the discussion period. Thank you.

Thank you for taking the time to review our manuscript and for providing helpful and constructive feedback. We were happy to see your review note our study’s potential for future impact, the thoroughness of our experiments and ablations which demonstrate the strong performance of the approach we study, and the clarity of presentation. We address your concerns below.

---

1) [50] found that the optimal noising timestep selection can vary depending on the task .. the ablation study for timestep is performed only on the Franka-Kitchen benchmark. Would results differ if tested on different tasks?

Yes, our observation is consistent with prior work. The optimal noising timestep depends on the task and we present the same ablation for timestep on the Meta-World environment below. SCR performs similarly across 0, 100 and 200 timestep on Meta-World because the images on this benchmark contain coarser details. We have provided further intuition and discussion around how to choose the best noising timestep in Figure 8 in the appendix, where we show reconstructions obtained from the diffusion model for different benchmark images starting at different denoising timesteps.

2) All ablation studies in Section 5 are conducted solely on the Franka-Kitchen benchmark.

Thank you for suggesting this, we conducted additional ablations on the Meta-World benchmark and present these in the general response of the rebuttal. Note that ablating the layer + noise selection on Meta-World shows that we can beat the best method there (R3M[30]) if we use a smaller subset of the layers we used (thereby achieving a 97% success rate), however we recommend retaining the standardized settings which we used across benchmarks to retain simplicity while still getting overall good performance across tasks. Since the navigation benchmarks require a lot of compute resources, we can provide a full set of ablations on these for the camera ready version, if the reviewer suggests it.

3) SCR-FT shows a performance drop compared to the non-finetuned version on the ImageNav benchmark.

We agree that a visual domain gap exists between the manipulation and navigation benchmarks, which could be mitigated by co-finetuning with manipulation and navigation datasets (like Real Estate 10k[1]) as done in VC-1[27]. With SCR-FT, we aimed to present a proof-of-concept result for the simple non-task-specific fine-tuning that is possible when using an image generation representation model, using some of the commonly-used datasets in the representation learning literature - most of which focus on few-shot manipulation tasks. We also presented a negative result of trying to do this for the CLIP model in Appendix D.1. Finally, we have aimed to keep the representation selection procedure standardized for all tasks- by deciding a common subset of layer outputs which we use as the representation, as well as the denoising timestep value of 0 since that always does well (while higher denoising timesteps might also do well on some benchmarks). Beyond this there are no additional hyperparameters introduced for SCR.

---

We hope that these clarifications and additional experimental results have addressed any remaining questions and concerns.

If you find that we have addressed your questions and concerns, we kindly ask you to consider raising your score to reflect that. If issues still remain, then we would be more than happy to answer these in the discussion period. Thank you.

---

References

[1] Zhou, Tinghui, et al. "Stereo magnification: Learning view synthesis using multiplane images." arXiv preprint arXiv:1805.09817 (2018).

Thank you for taking the time to review our manuscript and for providing helpful and constructive feedback. We were happy to see your review note our study’s innovative exploration, thoroughness of our experiments which demonstrate the strong performance of the approach we study, and the clarity of presentation. We address your concerns below.

---

1) Real-world validation: We employed a broad range of tasks in our work to add evaluation signals from different domains (manipulation and highly photorealistic indoor navigation) and learning setups (few-shot vs RL). Real robot experiments require an expensive setup and largely feature in works that look at few-shot table-top manipulation tasks which can provide noisy performance signals due to low sample sizes. We believe this would not have added much additional value to our exploration of simply validating the usability of diffusion model representations as a first step, and is better suited for future work. Many influential works in this space also rely completely on evaluation in simulation environments [1, 2, 3, 4].

2) Theoretical evidence: Consistent with the embodied AI literature (e.g., VC-1[27], R3M[30], Voltron [20] etc.), our work focuses on displaying strong empirical evidence for the usefulness of pretrained representations for a wide range of control tasks. We agree that studying the properties of the representation manifold of denoising diffusion models is an interesting and impactful direction for future work.

---

We hope that these clarifications have addressed any remaining questions and concerns.

If you find that we have addressed your questions and concerns, we kindly ask you to consider raising your score to reflect that. If issues still remain, then we would be more than happy to answer these in the discussion period. Thank you.

---

References

[1] Khandelwal, Apoorv, et al. "Simple but effective: Clip embeddings for embodied ai." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

[2] Xiao, Tete, et al. "Masked visual pre-training for motor control." arXiv preprint arXiv:2203.06173 (2022).

[3] Yadav, Karmesh, et al. "Offline visual representation learning for embodied navigation." Workshop on Reincarnating Reinforcement Learning at ICLR 2023. 2023.

[4] Eftekhar, Ainaz, et al. "Selective visual representations improve convergence and generalization for embodied ai." arXiv preprint arXiv:2311.04193 (2023).

Thank you for taking the time to review our manuscript and for providing helpful and constructive feedback. We were happy to see your review mention the clarity of our presentation, the relevance of the problem we are studying, and the comprehensiveness of our evaluation. We address your concerns below.

---

1) Can a visual encoder (like ViT) trained on a similar dataset as Stable Diffusion suffice, as the generative aspect of T2I diffusion models does not seem to be essential for the task with t=0.

Tackling complex open-ended embodied tasks requires foundation vision-language model representations. We study the Stable Diffusion model since other existing foundation models like CLIP have their own shortcomings, we expand on this below.

• The foundation model representations which have been considered for control so far are trained with auto-encoding (MAE) or contrastive learning objectives (CLIP[33], LIV[25], R3M[30]) - and prior work [20] has noted that these seem to do well on disparate sets of tasks depending on the level of visual and semantic details they encode (with image-text contrastive approaches encoding semantics better and auto-encoding approaches encoding more finer details). We hypothesized that the text-conditioning of a T2I generative model might enforce encoding semantics and the generation objective might incentivise retention of more fine grained visual detail in the representations. The question was then whether it would be possible to localize a subset of representations that would sufficiently encode these details and do well across tasks. This is the main result we demonstrate.
• Moving to language-conditioned control tasks necessitates investigating vision-language model representations. Given the limited effectiveness of CLIP (which is the primary candidate for a vision language representation) as a backbone on control benchmarks, we intended our work to be an explorative study of a different kind of foundation model that could provide vision-language aligned representations for control. Since there is a lot of research momentum in the direction of image and video generation diffusion models, successfully leveraging them for control would be beneficial in the long run.

2) This loose connection is particularly amplified in Table 2a where timestep=0 ends up with the best performance. This basically indicates that the "generative" power of text-to-image models is not necessary for the task.

• We note that when using a noise value of 0, the U-net still forms a meaningful representation of the input. [2] also uses zero noising for input images, to do zero-shot open-set segmentation using intermediate outputs from the U-Net. [3] also follow the score estimates from the U-net at t=0 to continue to keep refining the data samples and to bootstrap a Langevin sampler from that. The fact that the score function gradient estimates are informative at t=0 implies that the intermediate representations in the model are informative too.
• We also note that the MAE (Masked AutoEncoders) [1] representation model (which is the backbone for the VC-1 model) follows a similar strategy. It is trained to reconstruct missing patches from the input, but when it is used as a representation backbone none of the patches are masked out. The uncorrupted input is passed through the model to derive the representation, thereby not exactly using the denoising power of the MAE in the forward pass.

3) Explain The performance gap between SCR-FT and R3M in Meta-World and Franka-Kitchen We are happy to provide more context about performance differences on the benchmarks in the revised manuscript. R3M representations are highly sparse and tailored specifically to reduce overfitting in few-shot IL settings (see Table 1 in the R3M paper where performance drops 4-6% when the L1 sparsity penalty is removed), which does not end up being an advantage on other kinds of tasks with large scale RL training. We also note that we standardized representation extraction settings across all benchmarks for SCR for simplicity and reproducibility, but if we had tuned the layer selection choice for each of the domains specifically, we would be able to beat R3M on Meta-World with the setting where we use the mid+down_3 layers with t=200 timestep, and get 97% on Meta-World (this result is presented in the general response of our rebuttal, where we replicate our Franka ablations for the layer and noise parameters on Meta-World).

4) Typos

Thank you for pointing out the swapped order of table column references in sections 5.1 and 5.2, we have fixed this in the manuscript.

---

We hope that these clarifications have addressed any remaining questions and concerns.

If you find that we have addressed your questions and concerns, we kindly ask you to consider raising your score to reflect that. If issues still remain, then we would be more than happy to answer these in the discussion period. Thank you.

---

References

[1] He, Kaiming, et al. "Masked autoencoders are scalable vision learners." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[2] Open-Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models. Xu et al., 2023.

[3] Iterated Denoising Energy Matching for Sampling from Boltzmann Densities. Akhound-Sadegh et al., 2024.