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

Authors: gunshi  gupta
karmesh  yadav
yarin  gal
zsolt  kira
dhruv  batra
cong  lu
tim g j rudner
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 00:09:55

Abstract: Embodied AI agents require a fine-grained understanding of the physical world mediated through visual and language inputs. Such capabilities are difficult to learn solely from task-specific data. This has led to the emergence of pre-trained vision-language models as a tool for transferring representations learned from internet-scale data to downstream tasks and new domains. However, commonly used contrastively trained representations such as in CLIP have been shown to fail at enabling embodied agents to gain a sufficiently fine-grained scene understanding—a capability vital for control. To address this shortcoming, we consider representations from pre-trained text-to-image diffusion models, which are explicitly optimized to generate images from text prompts and as such, contain text-conditioned representations that reflect highly fine-grained visuo-spatial information. Using pre-trained text-to-image diffusion models, we construct Stable Control Representations which allow learning downstream control policies that generalize to complex, open-ended environments. We show that policies learned using Stable Control Representations are competitive with state-of-the-art representation learning approaches across a broad range of simulated control settings, encompassing challenging manipulation and navigation tasks. Most notably, we show that Stable Control Representations enable learning policies that exhibit state-of-the-art performance on OVMM, a difficult open-vocabulary navigation benchmark.

Main paper: 
 introduction: 
 As general-purpose, pre-trained "foundation" models [2, 5, 6, 24, 31, 34, 47] are becoming widely available, a central question in the field of embodied AI has emerged: How can foundation models be used to construct model representations that improve generalization in challenging robotic control tasks [4, 40, 64] ? Robotic control tasks often employ pixel-based visual inputs paired with a language-based goal specification, making vision-language model representations particularly well-suited for this setting. However, while vision-language representations obtained via Contrastive Language-Image Pre-training [CLIP; 33]-a state-of-the-art method-have been successfully applied to a broad range of computer vision tasks, the use of CLIP representations has been shown to lead to poor downstream performance for robotic control. This shortcoming has prompted the development of alternative, control-specific representations for embodied AI [25, 30] but has left other sources of general-purpose pre-trained vision-language representations-such as text-to-image diffusion models-largely unexplored for control applications. In this paper, we propose Stable Control Representations (SCR): pre-trained vision-language representations from text-to-image diffusion models that can capture both high and low-level details of a scene [17, 34] . While diffusion representations have seen success in downstream vision-language tasks, for example, in semantic segmentation [3, 46, 50] , they have-to date-not been used for control. We perform a careful empirical analysis in which we deconstruct pre-trained text-to-image diffusion model representations to understand the impact of different design decisions.In our investigation, we find that diffusion representations can outperform general-purpose models like CLIP [33] across a wide variety of embodied control tasks despite not being trained for representation learning. This is the case for purely vision-based tasks as well as for settings that require task understanding through text prompts. A highlight of our results is the finding that diffusion model representations enable better generalization to unseen object categories in a challenging open-vocabulary navigation benchmark [58] and provide improved interpretability through attention maps [45] .Our key contributions are as follows:1. In Section 3, we introduce a multi-step approach for extracting vision-language representations for control from text-to-image diffusion models. We show that these representations are capable of capturing both the abstract high-level and fundamental low-level details of a scene, offering an alternative to models trained specifically for representation learning. 2. In Section 4, we evaluate the representation learning capabilities of diffusion models on a broad range of embodied control tasks, ranging from purely vision-based tasks to problems that require an understanding of tasks through text prompts, thereby showcasing the versatility of diffusion model representations. 3. In Section 5, we systematically deconstruct the key features of diffusion model representations for control, elucidating different aspects of the representation design space, such as the input selection, the aggregation of intermediate features, and the impact of fine-tuning on performance. We have demonstrated that diffusion models learn versatile representations for control and can help drive progress in embodied AI. Figure 1 presents a summary of our approach and results. 2 
related work: 
 We begin with a review of prior work on representation learning and diffusion models for control.Representation Learning with Diffusion Models. Diffusion models have received a lot of recent attention as flexible representation learners for computer vision tasks of varying granularity-ranging from key point detection and segmentation [46, 50] to image classification [48, 57] . Wang et al. [50] has shown that intermediate layers of a text-to-image diffusion model encode semantics and depth maps that are recoverable by training probes. These approaches similarly extract representations by considering a moderately noised input, and find that the choice of timestep can vary based on the granularity of prediction required for the task. Yang and Wang [57] train a policy to select an optimal diffusion timestep, we simply used a fixed timestep per class of task. Several works [45, 46, 50] observe that the cross-attention layers that attend over the text and image embeddings encode a lot of the spatial layout associated with an image and therefore focus their method around tuning, post-processing, or extracting information embedded within these layers.Visual Representation Learning for Control. Over the past decade, pre-trained representation learning approaches have been scaled for visual discrimination tasks first, and control tasks more recently. Contrastively pre-trained CLIP [33] representations were employed for embodied navigation tasks by EmbCLIP [21] . MAE representations have been used in control tasks by prior works like VC-1 [27] , MVP [54] and OVRL-v2 [56] . R3M [30] and Voltron [20] leverage language supervision to learn visual representations. In contrast, we investigate if powerful text-to-image diffusion models trained for image generation can provide effective representations for control.Diffusion Models for Control. Diffusion models have seen a wide range of uses in control aside from learning representations. These can broadly be categorized into three areas. First, diffusion models have been used as a class of expressive models for learning action distributions for policies [7, 14, 32] ; They can improve model multimodality and generate richer action distributions than Gaussians. Second, off-the-shelf diffusion models have been used to augment limited robot demonstration datasets by specifying randomizations for object categories seen in the data through inpainting [19, 28, 60] . Third, planning can be cast as sequence modeling through diffusion models [1, 11, 18] . 
stable control representations: 
 In this paper, we investigate the use of language-guided visual representations from the open-source Stable Diffusion model (v1.5) and their application to language-conditioned visual control tasks. We present background on latent diffusion models and text-to-image diffusion models, along with the notation we adopt in this work, in Appendix B.To extract representations, we follow a similar protocol as Wang et al. [50] , Traub [48] , and Yang and Wang [57] : Given an image-text prompt, s = {s image , s text }, associated with a particular task, we use the SD VQ-VAE model as the encoder E(•) and partially noise the encoded latents z 0 = E(s image ) to some diffusion timestep t, to obtain the noised latent z t . We then extract a representation composed of the intermediate layer outputs of the U-Net ϵ θ as it produces a denoising estimate ϵ θ (z t , t, s text ). This process is illustrated in Figure 2 . We refer to the extracted representations as Stable Control Representations (SCR). In Sections 3.1, 3.2, 3.3, and 3.4, we describe the design space for extracting SCR, and in Sections 3.5 and 3.6, we explain how we use the representations for control.We are interested in evaluating the internal representations from the denoiser network, that is, the U-Net ϵ θ (•). The first design choice we consider is which layers of ϵ θ to aggregate intermediate outputs from. The U-Net does not have a representational bottleneck, and different layers potentially encode different levels of detail. Trading off size with fidelity, we concatenate the feature maps output from the mid and down-sampling blocks to construct the representation. This results in a representation size comparable to that of the other pre-trained models we study in Section 4. This is shown at the bottom of Figure 2 and we ablate this choice in Section 5.1. Since outputs from different layers may have different spatial dimensions, we bilinearly interpolate them so that they are of a common spatial dimension and can be stacked together. We then pass them through a learnable convolutional layer to reduce the channel dimension before feeding them to downstream policies. The method used to spatially aggregate pre-trained representations can significantly affect their efficacy in downstream tasks, as we will discuss in Section 5.4. We use the best-performing spatial aggregation method for all the baselines that we re-train in Section 4.Next, we consider the choice of extraction timestep t for the denoising network (shown on the left of Figure 2 ). Recall that the images we observe in control tasks are un-noised (i.e., corresponding to x 0 ), whereas the SD U-Net expects noised latents, corresponding to z t for t ∈ [0, 1000]. The choice of timestep t influences the fidelity of the encoded latents since a higher value means more noising of the inputs. Yang and Wang [57] have observed that there are task-dependent optimal timesteps and proposed adaptive selection of t during training, while Xu et al. [55] have used t = 0 to extract representations using un-noised inputs to do open-vocabulary segmentation. We hypothesize that control tasks that require a detailed spatial scene understanding would benefit from a lower diffusion timestep, corresponding to a later stage in the denoising process where the inputs have less noise. We provide evidence consistent with this hypothesis in Section 5.2. To illustrate the effect of the timestep, we display final denoised images for various t values in different domains in Figure 8 .Since text-to-image diffusion models allow conditioning on text, we investigate if we can influence the representations to be more task-specific via this conditioning mechanism. For tasks that come with a text specifier, for example, the sentence "go to object X", we simply encode this string and pass it to the U-Net. However, some tasks are purely vision-based and in these settings, we explore whether constructing reasonable text prompts affects downstream policy learning when using the U-Net's language-guided visual representations. We present this analysis in Section 5.3.Wang et al. [50] and Tang et al. [45] demonstrate that the Stable Diffusion model generates localized attention maps aligned with text during the combined processing of vision and language modalities. Wang et al. [50] leveraged these word-level attention maps to perform open-domain semantic segmentation. We hypothesize that these maps can also help downstream control policies to generalize to an open vocabulary of object categories by providing helpful intermediate outputs that are category-agnostic. Following Tang et al. [45] , we extract the cross-attention maps between the visual features and the CLIP text embeddings within the U-Net. We test our hypothesis on an open-domain navigation task in Section 4.3, where we fuse the cross-attention maps with the extracted feature maps from the U-Net. We refer to this attention-map-augmented representation as SCR-ATTN.To solve visual control tasks with states given by s = [s image , s text ], where s text may be used to specify the task, we wish to use pre-trained vision-language representations capable of encoding the state s as f ϕ (s image , s text ). This encoded state is then supplied to a downstream, task-specific policy network, which is trained to predict the action a t . Our evaluation encompasses both supervised learning and reinforcement learning regimes for training the downstream policies. We train agents through behavior cloning on a small set of demonstrations for the few-shot manipulation tasks we study in Section 4.1. For the indoor navigation tasks we study in Sections 4.2 and 4.3, we use a version of the Proximal Policy Optimization [PPO, 39] algorithm for reinforcement learning. Finally, we consider fine-tuning strategies to better align the base Stable Diffusion model towards generating representations for control. This serves to bridge the domain gap between the diffusion model's training data (e.g., LAION images) and robotics datasets' visual inputs (e.g., egocentric tabletop views in manipulation tasks or indoor settings for navigation). Crucially, we do not require any task-specific losses for fine-tuning. Instead, we adopt the same text-conditioned generation objective as that of the base model for the fine-tuning phase. We use a small subset of the collection of datasets used by prior works on representation learning for embodied AI [27, 54] : we use subsets of the EpicKitchens [9] , Something-Something-v2 [SS-v2; 13], and Bridge-v2 [49] datasets. As is standard, we fine-tune the denoiser U-Net ϵ θ but not the VAE encoder or decoder. Image-text pairs are uniformly sampled from the video-text pairs present in these datasets. A possible limitation of this strategy is that text-video aligned pairs (a sequence of frames that correspond to a single language instruction) may define a many-to-one relation for image-text pairs. However, as we see in experiments in which we compare to the base Stable Diffusion model in Section 4, this simple approach to robotics alignment is useful in most cases. Further details related to fine-tuning are provided in Appendix E.1. We refer to the representations from this fine-tuned model as SCR-FT. 
empirical evaluation: 
 In this work, we evaluate Stable Control Representations (SCR) on an extensive suite of tasks from 6 benchmarks covering few-shot imitation learning for manipulation in Section 4.1, reinforcement learning-based indoor navigation in Sections 4.2 and 4.3, and tasks related to fine-grained visual prediction in Appendices D.2 and D.3. Together, these tasks allow us to comprehensively evaluate whether our extracted representations can encode both high and low-level semantic understanding of a scene to aid downstream policy learning. We describe the common baselines used across tasks in Appendix C, and present the individual task setups and results in the following subsections.We start by evaluating SCR on commonly studied representation learning benchmarks in few-shot imitation learning. Specifically, our investigation incorporates five commonly studied tasks from Meta-World [59] (same as CORTEXBENCH [27] ), which includes bin picking, assembly, pick-place, drawer opening, and hammer usage; as well as five tasks from the Franka-Kitchen environments included in the RoboHive suite [23] , which entail tasks such as turning a knob or opening a door. We adhere to the training and evaluation protocols adopted in their respective prior works to ensure our results are directly comparable (detailed further in Appendix G.1).Results. We report the best results of SCR and baselines in Table 1a . On Meta-World, we see that SCR outperforms most prior works, achieving 94.9% success rate. In comparison, VC-1, the visual foundation model for embodied AI and CLIP achieved 92.3 and 90.1% respectively. On Franka-Kitchen, SCR obtains 49.9% success rate, which is much higher than CLIP (36.3%) and again outperforms all other baselines except for R3M. We note that R3M's sparse representations excel in few-shot manipulation with limited demos but struggle to transfer beyond this setting [27, 20] . We see that while the SD-VAE encoder performs competitively on Franka-Kitchen, it achieves a low success rate on Meta-World. This observation allows us to gauge the improved performance of SCR from the base performance gain we may get just from operating in the latent space of this VAE. Additionally, we see that the task-agnostic fine-tuning gives SCR-FT an advantage (4%) over SCR on Franka-Kitchen while making no difference on Meta-World. Note that the other high-performing baselines (R3M and Voltron) have been developed for downstream control usage with training objectives that take temporal information into account, while VC-1 has been trained on a diverse curation of robotics-relevant data. In this context, SCR's comparable performance shows that generative foundation models hold promise for providing useful representations for control, even with relatively minimal fine-tuning on non-task-specific data.We now assess SCR in more realistic visual environments, surpassing the simple table-top scenes in manipulation benchmarks. In these complex settings, the representations derived from pre-trained foundational models are particularly effective, benefiting from their large-scale training. We study Image-Goal Navigation (ImageNav), an indoor visual navigation task that evaluates an agent's ability to navigate to the viewpoint of a provided goal image [63] . The position reached by the agent must be within a 1-meter distance from the goal image's camera position. This requires the ability to differentiate between nearby or similar-looking views within a home environment. This task, along with the semantic object navigation task that we study in Section 4.3, allows for a comprehensive evaluation of a representation's ability to code both semantic and visual appearance-related features in completely novel evaluation environments. We follow the protocol for the ImageNav task used by Majumdar et al. [27] and input the pre-trained representations to an LSTM-based policy trained with DD-PPO [52] for 500 million steps on 16 A40 GPUs (further details in Appendix G.3). Given the large compute requirements for training, we directly compare SCR and SCR-FT to the results provided in Majumdar et al. [27] .Results. We evaluate our agent on 4200 episodes in 14 held-out scenes from the Gibson dataset and report the success rate in Table 1b . We find that SCR outperforms all other representations, while the fine-tuned version SCR-FT is almost on par with the second-best-performing VC-1 (69.5% vs 70.3%), the SOTA visual representation from prior work. This can be expected given that it was fine-tuned on images solely from table-top manipulation datasets. We also see that R3M, the best model for few-shot manipulation from Table 1a performs very poorly (30.6%) in this domain, showing its limited transferability to navigation tasks.We now shift our focus to evaluating how Stable Diffusion's web-scale training can enhance policy learning in open-ended domains. We consider the Open Vocabulary Mobile Manipulation (OVMM) benchmark [58] that requires an agent to find, pick up, and place objects in unfamiliar environments. One of the primary challenges here is locating previously unseen object categories in novel scenes (illustrated in Figure 3 (left)). To manage this complex sparse-reward task, existing solutions [58] divide the problem into sub-tasks and design modular pipelines that use open-vocabulary object detectors such as Detic [62] to enable generalization to novel objects. We study a modified version of the Gaze sub-task (detailed in Appendix G.2), which focuses on locating a specified object category for an abstracted grasping action. [21, 26] . We also incorporate a baseline that trains a policy with access to ground truth object masks, evaluated using either the ground truth or Detic-generated masks at test time (labeled as Oracle/Detic).Results. Table 1c shows that SCR-FT surpasses VC-1 by 1.3%, beating CLIP and SCR by 3.2%. It is surprising that VC-1's visual representation does better than CLIP's image encoder representation, given that the downstream policy has to use these with the CLIP text encoder's embedding of the target object category. Comparing these to SCR-FT-ATTN, we can see the benefit of providing intermediate outputs in the form of text-aligned attention maps to the downstream policy (+1.7%). Samples of attention maps overlaid on images from an evaluation episode can be found in Appendix G. These word-level cross-attention maps simultaneously improve policy performance and also aid explainability, allowing us to diagnose successes and failures. Interestingly, the foundation model representations (CLIP, VC-1, SCR) perform better than Detic. While object detections serve as a category-agnostic input for downstream pick-and-place policies, noisy detections can often lead to degraded downstream performance, as we see in this case. Nonetheless, there is still a sizeable gap to 'Oracle' which benefits from ground truth object masks at test-time. 
deconstructing stable control representations: 
 In this section, we deconstruct Stable Control Representations to explain which design choices are most determinative of model robustness and downstream performance.We begin our investigation by examining how the performance of SCR is influenced by the selection of layers from which we extract feature maps. We previously chose outputs from the mid and downsampling layers of the U-Net (Figure 2 ), because their aggregate size closely matches the representation sizes from the ViT-based models (VC-1, MVP, and CLIP). Appendix E.2 details the feature map sizes obtained for all the models we study. Table 2a lists the success rates achieved on the Franka-Kitchen domain when we use different sets of block outputs in SCR. We present similar ablations for Meta-World in the top four rows of Table 3b .We observe that utilizing outputs from multiple layers is instrumental to SCR's high performance. This finding underscores a broader principle applicable to the design of representations across different models: Leveraging a richer set of features from multi-layer outputs should enhance performance on downstream tasks. However, it is important to acknowledge the practical challenges in applying this strategy to ViT-based models due to the high dimensionality of each layer's patch-wise embeddings (16×16×1024 for ViT-L for images of size 224×224). We present the success rates achieved on the four benchmarks when aggregating multi-layer embeddings from CLIP models in Tables 3a and 4 , alongside SCR (the representation size for which is now half in comparison). In Table 3a , we observe that moving towards middle layers leads to higher performance indicating that CLIP layers 10-14 encode some details useful to the Franka-Kitchen benchmark. While we see benefits from including the output of certain additional layers, it still underperforms SCR. Next, we characterize the sensitivity of task performance to the denoising step values chosen during representation extraction. We present results on the Franka-Kitchen tasks in Table 2b , and on the Meta-World tasks in the bottom three rows of Table 3b . We see that the performance across nearby timesteps (0 and 10 or 100 and 110) is similar, and that there is a benefit to doing a coarse grid search up to a reasonable noising level (0 vs 100 vs 200) to get the best value for a given task.Recall that in the OVMM experiments (Section 4.3), we concatenated the target object's CLIP text embedding to the visual representations before feeding it to the policy. For SCR and SCR-FT, we also provided the category as the text prompt to the U-Net, and additionally extracted the generated cross-attention maps for SCR-FT-ATTN. In this subsection, we seek to more closely understand how the text prompts impact the representations in SCR.We first consider the Franka-Kitchen setup from Section 4.1, which includes manipulation tasks that do not originally come with a language specification. We experiment with providing variations of task-relevant and irrelevant prompts during the representation extraction in SCR. Table 2c shows the downstream policy success rates for irrelevant ("an elephant in the jungle") and relevant ("a Franka robot arm opening a microwave door") prompts, compared to the default setting of not providing a text prompt We see that providing a prompt does not help with downstream policy performance and may even degrade performance as the prompt gets more irrelevant to the visual context of the input.We now move to the Referring Expressions Grounding task which requires predicting a bounding box for an object being referred to in a sentence, within an image depicting cluttered objects. We defer the main presentation of this task to Appendix D.2 and use it to probe the degree of language grounding in SCR in this section. To study the role of the U-Net in shaping the visual representations guided by the text, we examine different text integration methods to generate SCR representations in Table 5 .We compared the following approaches for providing the task's text specification to the task decoder (also depicted in Figure 4 ): Investigating the results of (a) and (b) in Table 5 , it is evident that incorporating the text prompt into the U-Net significantly enhances accuracy compared to ignoring the text altogether. The difference in scores between (b) and (c) indicates that directly providing text embeddings to the decoder improves performance, suggesting that certain crucial aspects of object localization are not fully captured by the representation alone. Comparing (c) to (d), we see that with concatenated text embeddings, further modulation of the visual representations does not provide significant benefits. Finally, the significant decrease in the score for (e) reveals the extent to which the task relies on text-based guesswork.These findings align with both intuition and recent research on controllable generation with diffusion models [61] that highlights the challenges associated with using long-form text guidance. There are ongoing research efforts, focused on training models with more detailed descriptions or leveraging approaches to encode and integrate sub-phrases of long texts, that seek to address these challenges.In this study, we refine the approach for extracting representations by integrating a convolutional layer that downsamples the spatial grid of pre-trained representations. This adjustment, referred to as a "compression layer" by Yadav et al. [56] , aims to reduce the high channel dimension of pre-trained model outputs without losing spatial details, facilitating more effective input processing by downstream task-specific decoders.We explore the effect of spatial aggregation methods by comparing the convolutional downsampling layer method to multi-headed attention pooling (MAP) used for CLIP embeddings in Karamcheti et al. [20] . We find that using a compression layer significantly improves performance on the fine-grained visual prediction tasks described in Appendix D as reported in Table 6 (columns 3-4). This result challenges the conjecture made in prior work that CLIP representations are limited in their ability to provide accurate low-level spatial information [20] and emphasizes the critical role of appropriate representation aggregation.Building on this result, we assess whether better spatial aggregation can improve the performance of CLIP representations on downstream control tasks. We present these results in Table 6 (columns 5-6) for VC-1 and CLIP on the MuJoCo tasks. We see that the compression layer often outperforms the use of CLS token embeddings (by 1-2%), but CLIP representations still fail to match the best-Table 6 : We ablate the spatial aggregation method for VC-1 and CLIP. On the fine-grained visual prediction tasks, we compare the average precision between using multi-head attention pooling (MAP) and the compression layer. On the Meta-World & Franka-Kitchen tasks, we compare the average success rates (± one standard error) between the CLS token and compression layer embeddings. performing models. This result provides evidence that the underperformance of CLIP representations on control tasks is unlikely due to a lack of sufficiently fine-grained visual information. Finally, we note that the compression layer aggregation technique was used for all baselines in Tables 1b and 1c to ensure a strong baseline comparison. We recommend that future studies adopt this methodology to enable a fairer comparison of representations. 
discussion: 
 In Section 5, we deconstructed Stable Control Representations and highlighted techniques used in our approach that could be applied to extract representations from other foundation models. Our analysis in Sections 5.1 and 5.4 revealed that using multi-layer features and appropriate spatial aggregation significantly affects performance, and overlooking these factors can lead to misleading conclusions about the capabilities of previously used representations.Next, our investigation into how language prompts shape diffusion model representations uncovered nuanced results and showed that text influence on representations does not consistently increase downstream utility. This is particularly evident in tasks where text specification is not required and where training and test environments are congruent, minimizing the need for semantic generalization.In contrast, tasks like referring expressions grounding demonstrate the necessity of direct access to text embeddings for accurate object localization, even when representations are modulated to considerable success. For the OVMM task, we identified a scenario where multimodal alignment is essential and proposed a method to explicitly utilize the latent knowledge of the Stable Diffusion model through text-aligned attention maps, which is not straightforward to do for other multimodal models. Future research could design methods to derive precise text-associated attribution maps for other models.Finally, we contrasted the simplicity of fine-tuning diffusion models with that of the contrastive learning objective required to fine-tune CLIP representations. The former only requires image-text samples for the conditional generation objective, whereas the latter requires a sophisticated negative label sampling pipeline along with large batch sizes to prevent the model from collapsing to a degenerate solution [33] . We demonstrate this phenomenon empirically on the Franka-Kitchen environment by fine-tuning CLIP similarly to SCR-FT in Appendix D.1. 
conclusion: 
 In this paper, we proposed Stable Control Representations, a method for leveraging representations of general-purpose, pre-trained diffusion models for control. We showed that using representations extracted from text-to-image diffusion models for policy learning can improve generalization across a wide range of tasks including manipulation, image-goal and object-goal based navigation, grasp point prediction, and referring expressions grounding. We also demonstrated the interpretability benefits of incorporating attention maps extracted from pre-trained text-to-image diffusion models, which we showed can improve performance and help identify downstream failures of the policy during development. Finally, we discussed ways in which the insights presented in this paper, for example, regarding feature aggregation and fine-tuning, may be applicable to other foundation models used for control. We hope that  
 
 Appendix: This paper presents work whose goal is to advance the field of robotic control using pretrained foundation models. Since we primarily aim to learn policies for simulated control tasks, we do not anticipate any ethical concerns at this stage. However, when deploying foundation models in the real world, we should be wary that our policies do not inherit any structural biases from the pretraining stage or internet-scale data.Diffusion models [16, 42] are a class of generative models that learn to iteratively reverse a forward noising process and generate samples from a target data distribution p(x 0 ), starting from pure noise. Given p(x 0 ) and a set of noise levels σ t for t = 1, . . . , T , a denoising function ϵ θ (x t , t) is trained on the objectiveL DM (θ) = E x0,ϵ,t [∥ϵ -ϵ θ x t , t))∥ 2 2 = E x0,ϵ,t [∥ϵ -ϵ θ x 0 + σ t • ϵ, t))∥ 2 2 , (B.1)where ϵ ∼ N (0, 1), t ∼ Unif(1, T ), and x 0 ∼ p(x 0 ). To generate a sample x 0 during inference, we first sample an initial noise vector x T ∼ N (0, σ T ) and then iteratively denoise this sample for t = T, ..., 1 by sampling from p(x t-1 |x t ), which is a function of ϵ θ (x t , t).In some settings, we may want to generate samples with a particular property. For example, we may wish to draw samples from a conditional distribution over data points, p(x 0 |c), where c captures some property of the sample, such as classification label or a text description [34, 36] . In these settings, we may additionally train with labels to obtain a conditioned denoiser ϵ θ (x t , t, c) and generate samples using classifier-free guidance [15] .Latent diffusion models [34] reduce the computational cost of applying diffusion models to highdimensional data by instead diffusing low-dimensional representations of high-dimensional data.Given an encoder E(•) and decoder D(•), Equation (B.1) is modified to operate on latent representations, z 0 = E(x 0 ), yieldingL LDM (θ) = E x0,c,ϵ,t [∥ϵ -ϵ θ E(x 0 ) + σ t • ϵ, t, c)∥ 2 2 , (B.2)where ϵ ∼ N (0, 1), t ∼ Unif(1, T ), x 0 , c ∼ p(x 0 , c). After generating a denoised latent representation z 0 , it can be decoded as x 0 = D(z 0 ).. The SD model is trained on the LAION-2B dataset [38] and operates in the latent space of a pre-trained VQ-VAE image encoder [12] . The model architecture is shown at the top of Figure 1 and is based on a U-Net [35] , with the corresponding conditioning text prompts encoded using a CLIP language encoder [33] .We model our environments as Markov Decision Processes (MDP, Sutton and Barto [43] ), defined as a tuple M = (S, A, P, R, γ), where S and A denote the state and action spaces respectively, P (s ′ |s, a) the transition dynamics, R(s, a) the reward function, and γ ∈ (0, 1) the discount factor. Our goal is to optimize a policy π(a|s) that maximizes the expected discounted return E π,P [ ∞ t=0 γ t R(s t , a t )]. In this paper, we consider visual control tasks that may be language-conditioned, that is, states are given by s = [s image , s text ], where s text specifies the task. We are interested in pre-trained visionlanguage representations capable of encoding the state s as f ϕ (s image , s text ). This encoded state is then supplied to a downstream, task-specific policy network, which is trained to predict the action a t . Our evaluation encompasses both supervised learning and reinforcement learning regimes for training the downstream policies. We train agents through behavior cloning on a small set of demonstrations for the few-shot manipulation tasks we study in Section 4.1. For the indoor navigation tasks we study in Sections 4.2 and 4.3, we use a version of the Proximal Policy Optimization [PPO, 39] algorithm for reinforcement learning.We compare SCR and its variants (i.e., SCR-FT and SCR-FT-ATTN) to the following prior work in representation learning for control:1. R3M [30] pre-trains a ResNet50 encoder on video-language pairs from the Ego4D dataset using time-contrastive video-language alignment learning. 2. MVP [54] and VC-1 [27] both pre-train ViT-B/L models with the masked auto-encoding (MAE) objective on egocentric data from Ego4D, Epic-Kitchens, SS-v2, and ImageNet, with VC-1 additionally pre-training on indoor navigation videos. 3. CLIP [33] trains text and ViT-based image encoders using contrastive learning on web-scale data. 4. Voltron [20] is a language-driven representation learning method that involves pre-training a ViT-B using MAE and video-captioning objectives on aligned text-video pairs from SS-v2. 5. SD-VAE [34] is the base VAE encoder used by SD to encode images into latents.To assess how well the vision-only methods would do on tasks with language specification, we concatenate their visual representations with the CLIP text embeddings of the language prompts. While we are limited by the architecture designs of the released models we are studying, to ensure a more fair comparison we try to match parameter counts as much as we can. We use the ViT-Large (307M parameters) versions of CLIP, MVP, and VC-1 since extracting SCR involves a forward pass through 400M parameters.Table 7 : Performance on Franka-Kitchen after fine-tuning CLIP.Franka-Kitchen CLIP 36.9 ± 3.2 CLIP (FT) 34.2 ± 2.9We follow the same experimental constraints that we took into account while fine-tuning the diffusion model to get SCR-FT: we trained it on the same text-image pairs from the same datasets and used CLIP's contrastive loss to bring the visual embedding of the middle frames of a video closer to the video caption's text embedding. Specifically, for our experiment, we use the huggingface CLIP finetuning implementation and train the model with a batch size of 384 (the maximum number of samples we were able to fit on 8 A40 GPUs) with a learning rate of 5e-5 and a weight decay of 0.001 for 5000 update steps (same as SR-FT). We present the results in Table 7 for Franka-Kitchen, and note the lack of improvement on the task post-fine-tuning.In Sections 4.1 to 4.3, our analysis focused on the performance of various representations across an array of control tasks. We now turn our attention to two downstream tasks involving fine-grained visual prediction. The first task, Referring Expressions Grounding, is detailed within this section, while the second task, Grasp Affordance Prediction, is discussed in Appendix D.3. Karamcheti et al. [20] have previously examined the performance on these tasks as proxy measures to evaluate the efficacy of representations for control applications.The Referring Expressions Grounding task requires the identification and bounding box prediction of an object in an image based on its textual description. Similar to Karamcheti et al. [20] , we use the OCID-Ref Dataset [51] for our experiments. We show a sample image-text pair from the dataset to showcase the complexity of the task in Figure 5 . The frozen visual representation is concatenated with a text embedding and passed to a 4-layer MLP, which predicts the bounding box coordinates.Results. We report the bounding box accuracy at a 25% Intersection-over-Union (IoU) threshold across different scene clutter levels for SCR-variants and baselines in Table 8 . Firstly, we present the results from [20] for CLIP, R3M and Voltron and observe that Voltron outperforms the other two baselines. We see that SCR is tied with Voltron and that VC-1 and SD-VAE perform the best with a 1.5% lead. The better performance of these vision-encoder-only methods highlights that on this task, it is not a challenge for the downstream decoder to learn to associate the visual embeddings with the (CLIP) text encoding of the language specification. Since the training budget is fixed, we observed that some of the runs could potentially improve over extended training. However, we were primarily interested in this task not just to compare the downstream visual prediction performance, but to use it as a testbed for exploring the following two questions: (1) Do the performance differences between the representations we evaluated in Sections 4.1 to 4.3, stem from the absence of fine-grained spatial information encoded within the representations? We refute this claim in Section 5.4, where we present the impact of the representations' spatial aggregation method on prediction performance. ( 2) Additionally, we explore to what extent language prompting influences the representations from SCR on language-conditioned tasks in Section 5.3.The lemon on the rear left of the instant_noodles. In this section, we present our experiments on a second visual prediction task continuing from the experiments in Section 4.1. The Grasp Affordance Prediction task requires predicting per-pixel segmentation outputs for certain areas of objects in an RGB image. These areas correspond to parts of the surface that would be amenable to grasping by a suction gripper. The evaluation metric adopted in prior work is the precision of predictions corresponding to positive graspability at varying confidence levels (90, 95, and 99th percentile of the predicted perpixel probabilities, denoted as Top90, Top95, and Top99 in Table 9 ). We refer the reader to Karamcheti et al. [20] for the complete task setup details.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. The evaluation procedure for this task adopted in prior work involves a 5-fold cross-validation, and we observed a high variability in the results, with different runs of 5-fold cross-validation yielding different final test metrics. Our findings highlight that SCR and our adaptation of CLIP (in gray, detailed in Section 5.4) both excel at this task, achieving a Top99 score of 72.9. Interestingly, we see that fine-tuning did not enhance the performance of SCR on the visual prediction tasks explored in this section and Section 4.1, suggesting a potential disconnect between visual prediction and control task benchmarks. We include a comparison with LIV [25] on two tasks that involve manipulation and navigation. LIV is a visionlanguage representation learned through contrastive learning on the EpicKitchens dataset [9] . Similar to R3M results in the main paper, this representation does well on manipulation tasks but poorly on navigation tasks.In Table 11 , we present the consolidated scores across the four control benchmarks we study in Section 4, for all the representations we evaluate in this work. This is to give a higher-level view of the all-round performance of the different representations on the diverse set of tasks we consider. We see that VC-1, SCR, and SCR-FT emerge as the top three visual representations overall. While VC-1 is a representation-learning foundation model trained specifically for robotics tasks, SCR and SCR-FT are the diffusion model representations that we study in this paper, confirming the potential of large pre-trained foundation generative models across a wide array of downstream robotics tasks. We used the runwayml/stable-diffusion-v1-5 model weights provided by Huggingface to initialize our models and fine-tuned them using the diffusers library3 . As noted in Section 3.6, we used subset of the frames from the EpicKitchens, Something-Something-v2, and Bridge-v2 datasets.We extracted the middle-third of the video clips and sampled four frames randomly from this chunk to increase the chances of sampling frames where the text prompt associated with the video clip is most relevant for describing the scene. Since the Something-Something-v2 (SS-v2) dataset includes a human hand manipulating table-top objects while the Bridge-V2 dataset includes a robot gripper, we append the text "human hand" and "robot hand" respectively to the text caption associated eith the respective frames to reduce ambiguity in the image-text pairing. Using this procedure, we generated a paired image-language dataset with 1.3 million samples. Figure 6 shows samples of the images from the fine-tuning datasets. Since different embodiments (human and robot) are visible in the training images, we prepended the corresponding embodiment name to the text prompt for the associated image during training.We fine-tuned for only a single epoch (∼5,000 gradient steps) parallely on 2 Nvidia A100 GPUs with a mini-batch size of 512 and a learning rate of 1e -4 . Although the original Stable Diffusion model is trained on images of resolution 512 × 512, we fine-tuned the model on images downscaled to 256 × 256 since it aligned with the resolution requirements of the downstream application. We show sample generations from the diffusion model after fine-tuning in Figure 7 . We found that the model learns to associate the prompt with not just the human or robot hand but also with the style of the background and objects of the training datasets.Something-Something v2 For all baselines, we freeze the pre-trained vision model and train a policy using imitation learning on the provided set of 25 expert demonstrations. The results are then reported as the average of the best evaluation performance for evaluation runs over 3 seeds. All experiments are conducted on a single A100 GPU with 24 CPUs and 188 GBs of RAM.Meta-World. We follow Majumdar et al. [27] and use the hammer-v2, drawer-open-v2, bin-picking-v2, button-press-topdown-v2, assembly-v2 tasks from the Meta-World benchmark suite [59] . The mini-batch size is set at 128, with a learning rate of 10 -4 . We additionally correct a bug in the RoboHive implementation of the VC-1 baseline, specifically on input image normalization. Adjusting the image normalization to a 0-1 range resulted in a significant improvement in its performance.Open-Vocabulary Mobile Manipulation [OVMM; 58] is a recently proposed embodied AI benchmark that evaluates an agent's ability to find and manipulate objects of novel categories in unseen indoor environments. Specifically, the task requires an agent to "Find and pick an object on the start_receptacle and place it on the goal_recetacle", where object, start_receptacle and goal_recetacle are the object category names. Given the long-horizon and sparse-reward nature of this task, current baselines [58] divide the problem into sub-tasks, which include navigation to the start receptacle, precise camera re-orientation to focus on the object (an abstracted form of grasping), navigating to the goal receptacle, and object placement.Since our aim is to investigate the open-vocabulary capabilities of pre-trained representations, we choose to evaluate the models on only the precise camera re-orientation task (more commonly known as the Gaze task). In the original Gaze task, the agent is initialized within a distance of 1.5m and angle of 15 • from the object which is lying on top of the start_receptacle. The episode is deemed successful when the agent calls the Pick action with the camera's center pixel occupied by the target object and the robot's gripper less than 0.8m from the object center. In our initial experiments, we found the current initialization scheme would lead the agent to learn a biased policy. This policy would call the Pick action after orienting towards the closest object in the field of view. Therefore, we chose to instantiate a harder version of the gaze task, where the episode starts with the agent spawned facing any random direction within 2.0m of the object.We carry out our experiments using the HomeRobot GitHub repository 5 . HomeRobot uses the Habitat simulator [44] with the episode dataset provided by Our agent is designed to resemble the Stretch robot, characterized by a height of 1.41 meters and a radius of 0.3 meters. At a height of 1.31 meters from the base, a 640x480 resolution RGBD camera is mounted. This camera is equipped with motorized pan and tilt capabilities. The agent's action space is continuous, allowing it to move forward distances ranging from 5 to 25 centimeters and to turn left or right within angles ranging from 5 to 30 degrees. Additionally, the agent can adjust the head's pan and tilt by increments ranging 0.02 to 1 radian in a single step.In our experiments, we use a 2 layer LSTM policy and pass in the visual encoder representations after passing them through the compression layer. We initialize the LSTM weights with the LSTM weights of the Oracle model to get a slight boost in performance. We train our agents using the distributed version of PPO [52] with 152 environments spread across 4 80GB Nvidia A100 GPUs. Each run also has access to 96 CPUs and 754 GBs of RAM. We train for 100M environment steps while evaluating the agent every 5M steps and report the metrics based on the highest success rate observed on the validation set.We conduct our ImageNav experiments in the Habitat simulator [37] , using the episode dataset from Mezghani et al. [29] . These experiments are conducted using the VC-1 codebase6 [27] . The dataset uses 72 training and 14 validation scenes from the Gibson [53] scene dataset with evaluation conducted on a total of 4200 episodes. The agent is assumed to be in the shape of a cylinder of height 1.5m and radius 0.1m, with an RGB camera mounted at a height of 1.25m from the base. The RGB camera has a resolution of 128×128 and a 90 • field-of-view.At the start of each training episode, an agent is randomly initialized in a scene and is tasked to find the position from where the goal image was taken within 1000 simulation steps. At each step, the agent receives a new observation and is allowed to take one of the four discrete actions including MOVE_FORWARD (25cm), TURN_LEFT (30 • ), TURN_RIGHT (30 • ) and STOP. The episode is a success if the agent calls the STOP action within 1m of the goal viewpoint. Similar to [56, 27] we train our agents using a distributed version of DD-PPO [52] with 320 environments for 500M timesteps (25k updates). Our experiments are conducted using 2 nodes containing 8 A40 GPUs each with a total of 128 CPUs and 504 GBs of RAM. Each environment accumulates experience across up to 64 frames, succeeded by two PPO epochs using two mini-batches. While the pre-trained model is frozen, the policy is trained using the AdamW optimizer, with a learning rate of 2.5 × 10 -4 and weight decay of 10 -6 . We use the reward function proposed by [56] . Performance is assessed every 25M training steps, with reporting metrics based on the highest success rate observed on the validation set.Noised Generated OVMM Noised Generated Refer ExpressionNoised Generated t=100 t=200 t=300Figure 8 : Noising and denoising plots for images from 3 of our tasks using the fine-tuned Stable Diffusion model. For each image, we first add noise up to timestep t, where t ∈ {100, 200, 300}, and then denoise the image back to 0. We observe that inputs from different tasks are differently sensitive to the noising ranges, based on the amount of information the images contain. On Meta-World, SD is able to reconstruct the image correctly even at t = 200, while for the referring expressions grounding task, noising leads to information loss even at t = 100 with several small objects reconstructed differently to the original. This affects the range of the noise we could add to the ended latents at the time of representation extraction. It should be noted however, that using un-noised inputs (t = 0) always worked well in our experiments, and this hyper-parameter could be ignored in practice.Input Image Refer Expression "pear" "A" "beside" "book" "a" "A" "and" "rocket" Input Image OVMM "chair" "a" Figure 9 : The Stable Diffusion model allows us to extract word-level cross-attention maps for any given text prompt. We visualize these maps in a robotic manipulation environment and observe that they are accurate at localizing objects in a scene. Since these maps are category agnostic, downstream policies should become robust to unseen objects at test time.(1) Find Tomato Code link: https://github.com/ykarmesh/stable-control-representations.https://huggingface.co/docs/diffusers/v0.13.0/en/training/text2imagehttps://github.com/vikashplus/robohivehttps://github.com/facebookresearch/home-robothttps://github.com/facebookresearch/eai-vc 


Decision: The work introduces Stable Control Representations, a technique that utilizes representations from general-purpose, pre-trained diffusion models for control tasks and are capable of capturing both the abstract high-level and fundamental low-level details of a scene. Experiments demonstrate improvements in generalization across various tasks.

Summary Of Reasons To Publish:

1) Novel approach proposing diffusion Model representations work for robotic control tasks.
2) Well-motivated and well-structured paper
3) Solid experiment setup and results

Summary Of Suggested Revisions:

No major revisions are suggested by the reviewers. I suggest authors to fix the presentation issues raised by reviewers in the final version.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: See above
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: None
limitations: limitations: Refer to the weaknesses section.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: Please kindly refer to the weaknesses.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: Yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: There are not many concerns to address and the authors wrote an appropriate short section in the Appendix.

Comments: title: Thank you for the response
comment: Dear Reviewer,

Thank you again for your engagement and valuable feedback that has helped us improve our paper, and for raising your score!

Best,

The authors
Thanks for the detailed response with references. It's interesting to learn about the connections and see how SD is also useful in embodied controls. I have raised my score.
Thank you for your response! We are happy to clarify our statement with respect to your question.

> “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?)”

Yes, more specifically, the signal from the text prompts is incorporated into the U-Net through the use of a (dot-product-based) cross-attention layer within each block. This incentivises the visual feature maps in each block to align with the text encodings selectively, thereby also aligning with language concepts. This is also evidenced by the fact that we can isolate word-aligned attention maps for each U-Net layer, which roughly correspond to the entity in the image referenced by the word (when projected onto the image). This is shown in Figures 9 and 10 of the appendix in our submission.

> “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?”

Diffusion-based image generation models are trained to generate a final image by producing progressively more denoised versions of an existing noised image. To be able to refine and produce the image corresponding to the last few steps of denoising (likely only needing a few pixels to be edited or denoised), the model needs to encode the same level of visual details that a reconstruction/auto-encoding model might also encode. The works [55, 57] which we cite within our related work section also provide similar evidence for representation learning within diffusion models (for computer vision tasks). We also refer the reviewer to a recent survey paper that explores the interplay between diffusion models and representation learning in more detail.

Thank you for giving us the opportunity to clarify the points above. Please let us know if you have any further questions!

[1] Diffusion Models and Representation Learning: A Survey, https://arxiv.org/abs/2407.00783
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?
Dear Reviewer 9EUL,


The discussion periods ends tomorrow (August 13).

We have made considerable effort to address the concerns and queries you have raised, and would be grateful for the opportunity to clear up any remaining queries you may have while we are still in the discussion period.


Thank you.
Thanks for the clarifications and the additional experiments. They have answered my questions. I believe the original score accurately reflects my positive assessment of the paper.

I noticed a small typo in the table for question 1—the first column says "0, 0, 0", which I assume should be "0, 100, 200" given the texts. This doesn't impact my overall assessment.
Dear Reviewer,

Thank you again for your valuable feedback that has helped us improve our paper, and for raising your score! 

Best,
The authors
Thanks for your response. I agree that real-world experiments are unnecessary at this stage and I'm willing to raise the score.
Dear Reviewer,

Thank you so much for your feedback which has helped us improve the paper and for recommending acceptance!

We will change the phrasing as you suggested and include a subsection to discuss these findings in the revised manuscript!

Best,
The authors
Thanks for adding the new experiments with CLIP and answering my questions in detail. Based on the replies, I am updating my score to weak accept. 

One final comment: given the new CLIP results it might be too harsh to say 'contrastively trained representations such as in CLIP have been shown to fail at enabling embodied agents' in the abstract.
Dear Reviewer,

Thank you so much again for all your feedback that has helped us improve our paper, and for raising your score to "strong accept"!

Best,
The authors
Thank you for running these crucial additional experiments and engaging with my comments in general! I am willing to raise my score.