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

Summary

The research question is whether representation from image generation models are superior to other pre-training paradigms (i.e. contrastive CLIP) for embodied/navigation-type of tasks. The main intuition given is that these tasks require a lot of fine-grained (e.g. spatial) understanding. While CLIP has been used as a foundation model in traditional vision or some vision+language tasks, it has not shown much promise for robotic control. Their method is similar to previous works and each design choice is ablated/explained : First a certain timestep (i.e. noise level) is chosen for denoising a given image with an associated text prompt. Then a concatenation of intermediate U-Net layer activations is used as the “Stable Control Representation” (SCR). They evaluate their method on three different types of robotic-related tasks and show that SCR is overall the strongest, albeit not the strongest on every single task. For example, R3M is slightly stronger on the few-shot manipulation tasks (Meta World, Franka Kitchen) but is drastically worse on navigation.

Strengths

This is the first paper to make Diffusion Model representations work for robotic control tasks (novelty). The work is also situated clearly in the previous literature context and the technical background is well explained. Similarly, the paper is well structured and very easy to follow. Each part of the paper is well motivated and answers a natural question or follow-up step. Regarding the technical details, it is rigorous work that actually advances science and makes it easy for others to use for their research, e.g. in the appendix it is mentioned: “We re-ran all the methods using the evaluation repository provided with the work, and obtained different results compared to the reported numbers in Karamcheti et al. [20], which we attribute to a bug that we fixed related to the computation of the precision metrics.”. Or similarly they note in the appendix that the original setup for the Gaze task leads to a trivial solution and modify the task setup. This shows the authors don’t just blindly adopt tasks but critically investigate details of all the data or tasks they chose. The authors use baselines, i.e. showing SD-VAE is an insightful baseline that is easily beat. Similarly, ensuring fair and scientific baselines comparison i.e. in l. 327 - 329 where baselines are made stronger with insights from the paper to make it fair. Lastly, it is great to see an anonymous code repo that seems to be done well and with the aim of reproducibility.

Weaknesses

Only minor improvements, hard to beat some baselines properly on first glance. But on second glance it seems like stronger general representations, some of the other baselines such as R3M are only useful on some tasks but drastically drop on others. Maybe for people not familiar with this exact domain, you could explain more how these baselines work and why they get the numbers they do? For example, I am a vision-and-language person who is also familiar with all the text-to-image modeling advances but not so much with the navigation/control details. It is briefly acknowledged in section 5.1. but imo requires a deeper investigation: is the success primarily from fusing multiple layers and not from the fact that it is a diffusion model vs. a CLIP-like model? The authors acknowledge that such an approach is less feasible with i.e. ViT since the feature maps are much larger but nonetheless this question seems important enough to nonetheless investigate further. Overall the paper is very well written but if I had to pick a confusing section, it is 5.3 where too many things are introduced at once. The last paragraph in the discussion about the simplicity of finetuning SCR compared to CLIP comes out of nowhere and should be discussed earlier, and the claim about sophisticated negative sampling only cites the original CLIP paper which afaik only has random negative samples? There is an experiment in the appendix that somewhat supports this but also does not discuss the negative sampling. Not a major issue but should be made clearer.

Questions

Why cite Imagen in l. 36? Reference to Table 2a/b in section 5.1 and 5.2 are mixed up, just a small fix. “bilinearly interpolate them so that they are of a common spatial dimension and can be stacked together”: this should be specified or shown in equations formally to understand exactly what is done. Is this method common for extracting/unifying internal representations?

Rating

8

Confidence

3

Soundness

4

Presentation

4

Contribution

3

Limitations

There are not many concerns to address and the authors wrote an appropriate short section in the Appendix.

Summary

This paper studies how diffusion models can be adapted to provide rich representations for training policies. They study how to extract features from a pre-trained Stable Diffusion model in terms of 3 questions: which layers, which diffusion timestep and which text prompts help for downstream task performance. They evaluate these learned representations on a number of robotics tasks like manipulation, navigation and a combination of the two tasks. Their proposed approach performs better than baselines for many tasks.

Strengths

1. Lots of ablation experiments to validate the final design choice. 2. Evaluating a single model for both navigation and manipulation tasks. This is an interesting setup that is becoming more and more relevant with full-body control and robots that need to do both tasks.

Weaknesses

1. It is difficult to ascertain what is the key reason why the performance difference exists between the baselines and proposed method. The baselines have different training data, model sizes, and training losses. While the presented approach has better performance, it is important to identify what is key cause of this: more data, bigger model or the diffusion loss? Furthermore, there seems to be evidence that for manipulation tasks the baselines (LIV, R3M etc) are good but the performance drops only for navigation tasks which is not the key focus of those baselines.

Questions

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. 2.What if the same feature pooling strategy is applied to baselines like CLIP?. For example if you aggregate spatial features of CLIP from different layers does that help performance? Seems some multi-layer feature aggregation was done for Tables 1(b) and 1(c) but not 1(a) ? 3. Why is CLIP-B used in the ImageNav experiments and not CLIP-L?

Rating

6

Confidence

4

Soundness

2

Presentation

4

Contribution

2

Limitations

Yes.

Summary

This paper introduces SCR, a method that extracts representations from a pre-trained text-to-image diffusion model for learning downstream robotics control policies. Given a noised input image and prompt text, the visual representation is extracted from selected layers of the denoising U-Net. The extraction method is meticulously designed with extensive experiments and analysis to support the design choices. Experiments on multiple benchmarks across three tasks—few-shot imitation learning, image-goal navigation, and open-vocabulary mobile manipulation—demonstrate competitive performance, verifying the method’s ability to generalize to open-ended environments.

Strengths

1. Although the idea of replacing CLIP with diffusion models for extracting representations for robotic control is straightforward, the technical details, supported by comprehensive ablation and analysis of the design choices, are insightful and convincing. 2. The paper has potential for future impact. The method's ability to handle complex and open-ended environments suggests it could be generalized and applied to many other tasks. 3. The paper is well-written and easy to read. 4. The comparison experiments are thorough across multiple tasks and benchmarks, demonstrating the proposed method's strong performance.

Weaknesses

1. As noted in the related work, [50] found that the optimal noising timestep selection can vary depending on the task due to the required granularity of predictions. However, the ablation study for fixed timestep selection is performed only on the Franka-Kitchen benchmark. Would the results differ if tested on different tasks or benchmarks? 2. Additionally, all ablation studies in Section 5 are conducted solely on the Franka-Kitchen benchmark. SCR-FT shows a performance drop compared to the non-finetuned version on the ImageNav benchmark, attributed to the finetuning being performed only with table-top manipulation datasets. This raises concerns that different tasks have notable domain gaps and may require different features from the inputs. More experiments, or at least reasonable explanations, on whether the optimal framework design varies across tasks would be beneficial.

Questions

Please kindly refer to the weaknesses.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes, the authors discussed the limitations that downstream performance depends on design choices for representation extraction and the higher run-time cost of using a pre-trained model with more parameters.

Summary

This paper introduces Stable Control Representations (SCR), a novel approach that aggregates intermediate outputs from diffusion models for robotic control tasks. The authors validate its effectiveness on various benchmarks, including manipulation, navigation, and grasp point prediction. The key design space has been thoroughly ablated, showing SCR's superiority over other visual representation methods.

Strengths

1. Innovative Exploration: Visual representation is a fundamental challenge in robotics. Exploring the effectiveness of diffusion models in this context is intriguing and novel. 2. Comprehensive Experiments: The authors conducted extensive experiments covering various robotic tasks, diffusion steps, intermediate layer selection, and aggregation methods. The design space has been thoroughly explored. 3. Clear Writing and Organization: The paper is well-written and organized, making the narrative easy to follow.

Weaknesses

1. Lack of Real-World Validation: While SCR performs well on simulation baselines, real-world validation is missing. Demonstrating its effectiveness on real robots would significantly enhance the credibility and practical applicability of SCR. 2. Insufficient Evidence: Although the authors provide extensive empirical evidence of SCR's effectiveness, the paper lacks theoretical and quantitative analysis explaining why diffusion-based visual representations outperform other baselines. This theoretical backing would strengthen the paper's overall argument.

Questions

None

Rating

6

Confidence

4

Soundness

2

Presentation

4

Contribution

3

Limitations

Refer to the weaknesses section.

Summary

The paper investigates whether latent representations from a pre-trained tex-to-image generation model are useful for object manipulation or navigation tasks in embodied agents. The paper considers a few design choices in how to extract latent features from Stable Diffusion, that could lead to more versatile and useful features for downstream control tasks. The method is referred to as Stable Control Representation (SCR) with variants like additionally using cross-attention maps (SCR-Attn) or a finetuned version (SCR-FT). Experimental results show that SCR series can achieve comparable or better performance than general VL encoders like CLIP or method designed for control tasks.

Strengths

- The paper studies an interesting problem of using text-to-image diffusion models for embodied agent control. The paper is well-written and easy to understand. - The method considers many design choices of SCR in a very comprehensive way. The choices of different factors, from denoising steps to layer selection, have been clearly presented and justified later in the experiment. - The experiment seems well-designed and comprehensive. Most of the results support the benefit of using SCR for control tasks.

Weaknesses

- The method is loosely connected with diffusion mechanism and simply adopt Stable Diffusion as a feature extractor with one pass if I understand correctly. 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 are not necessary for the task. In that sense, is it still necessary to use T2I diffusion models if one could find a visual encoder (e.g. ViT) trained with similar amount of images as Stable Diffusion? - The performance gap between SCR-FT and R3M in Meta-World and Franka-Kitchen indicates that the method may not generalize well to some scenarios. Could you give any insights about the performance gap? Is it related to domain shift or image style differences or any other reaons? - The order of Table 2 is different from descriptions in Sec. 5.1 and 5.2

Questions

See above

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Thank you for the response

Most of the responses are convincing and have addressed my concerns. But I am not sure if I understand the sentence, *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.*, correctly. Is it because that SD adopts a freezed CLIP text encoder, so that you hypothesize that the text encoder enforces encoding semantics? (Btw, what is enforced to encode semantics by the text-conditioning stage?) I guess I just don't see how the concept of "encoding" or "retention of visual details" could fit into the diffusion process or is it just a very high-level interpretation as the target is to extract representations from a generative model?