Pre-trained Text-to-Image Diffusion Models Are Versatile Representation Learners for Control
Authors: Gunshi Gupta
Karmesh Yadav
Yarin Gal
Dhruv Batra
Zsolt Kira
Cong Lu
Tim G. J. Rudner
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Embodied AI
Representation Learning for Control
Diffusion Models
Foundation Models
Primary area: diffusion_based_models
TL;DR: We investigate representations from pre-trained text-to-image diffusion models for control tasks and showcase competitive performance across a wide range of tasks.

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.

Pre-trained Text-to-lmage Diffusion Models Are Versatile Representation Learners for Control
Gunshi Gupta"!
Karmesh Yaday
Yarin Gal
Zsolt Kira?
Dhruv Batra
Cong Lu
Tim G. J. Rudner
University of Oxford
Georgia Institute of Technology
New York University
Abstract
Embodied AI agents require a fine-grained understanding of the physical world mediated through visual and language inputs. Such capabilities are difficult learn solely from task-specific data. This has led to the emergence pre-trained vision-language vision- language models as 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 capability vital for control. To address this shortcoming, we consider representa- tions from pre-trained text-to-image diffusion models, which are explicitly opti- mized 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, construct Stable Control Repre- sentations 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-th ing approaches across a broad range of simulated control settings, encompassing challenging manipulation and navigation tasks. Most notably, we show that Sta- ble Control Representations enable learning policies that exhibit state-of-the-art performance on OVMM, difficult open open-vocabulary navigation benchmark.
representation Iearn-
Code: ott.a........aaaaaocaacc..ooooer
1 Introduction
As general-pure 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 language-based goal specification, making vision-language model representations particularly well-suited for this set- ting. However, while vision-language representations obtained via Contrastive Language-Image Pre-training [CLIP; 33] oftt-oofthhearr method-haad 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 AT [25, 30J but has left other sources of general-purpose pre-trained vision-language representations such text-to-image diffusion models largell unexplored for control applications.
"Equal Contribution.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Picking kettle"
CLI Eroudo
Stable
Overall Representation Comparison
VI SHOOOR
U-Net
SCR.FT (Ours)
Deooder
SCR (Ours)
VC-1
Original Image
Visual Reprasentation
Denclaed Image
R3M
Evaluation Suite
CLIP
VAE
0.s
Average Norm Sucouss
Meta World
Fizeaa Kitchen
CVMH
ImageNaay
Beter Expression
Groop Precittion
Figure 1: Left: Our paper proposes Stable Control Representations, which uses pre-trained text- to-image diffusion models as source of language-guig guided visual representations for downstream policy learning. Right: Stable Control Representations enable learning control policies that achieve all-round competitive performance on a wide range of embodied control tasks, including domains that require open-vocabulary generalization. Empirical results are provided in Section 4.
In this paper, we propose Stable Control Representations (SCR): pre-trained vision- vision-languaae representations from text-to-image diffusion models that can capture both high and low-level details of 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-no 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 ceneral-purposs models like CLIP [33] across a wide variety of embodied control tasks despite not being trained for represen- tation learning. This the case for purely bssion-based tasks well for settings that require task understanding through text prompts. A highlight of our results is the finding that diffusion model rep- resentations enable better generalization to unseen object categories in challenging open-vocabulary navigation benchmark [58] and provide improved interpretability through attention maps [45].
Our key contributions
follows:
1. In Section 3, introduce 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 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 broad range 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.
Section 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 AL. Figure 1 presents a summary of our approach and results.?
2 Related Work
We begin with 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-ranginn from key point detection and segmentation [46, 50) to image classification [48, 57]. Wang et al. [50] has shown that intermediate layers of 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 policy to select optimal
"Code link:

---

diffusion timestep, we simply used fixed timestep per class of task. Several works [45, 46, 50] observe that the coss-attention layers that attend over the text and image embeddings encode 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 learn visual representations. contrast, we investigate if powerful lext-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 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].
3 Stable Control Representations
In this paper, we investigate the use of language-guin guidedd visual representations from the open-source Stable Diffusion model (y1.5) and their application to language-conditing 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 similar protocol as Wang et al. [50]. Traub [48], and Yang and Wang [57]] Given an image-text prompt, 8 = {Simage: Ster), associated with particular task, we use the SD VQ-VAE model as the encoder E-) and partially noise the encoded latents zn - (simage) to some diffusion timestep f, to obtain the noised latent žい. We then extract a representation composed of the intermediate layer outputs of the U-Net Eg as it produces a denoising estimate c0(2,, t, Swa). 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.
3.1 Layer Selection and Aggregation
We 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 eg 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 representation size comparable to that of the other pre-trained models we study in Section 4. This 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 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 spatially aggregate pre- pre-trained representations can significantly affect their efficacy in downstream tasks, as we will discuss in Section 5.4. We use the best best-performing spatial aggregation method for all the baselines that we re-train Section 4.
3.2 Diffusion Timestep Selection
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 æp), whereas the SD U-Net expects noised latents, corresponding to 2, for 1 E [0, 1000]. The choice of timestep 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
3

---

Language Prompt
CLIP Language Encoder
Language Embedding
Noise level (ot)
VAE Encoder
Conv
Down
Down 2
Down 3
Down 4
Mid Block
Inter palple
Inmerpollle
Task-Specific Model
Compression Layer
Concatenated Representation
Figure 2: Extraction of Stable Control Representations from Stable Diffusion. Given an image-text prompt, {Simage ),exr, encode and noise the image and feed it into the U-Net together with the language prompt. We then aggregate feature maps from multiple layers within the U-Net, as described Section 3. Shown here features from the mid and downsampling blocks of the U-Net.
and proposed adaptive selection of t during training, while Xu et al. [55] have used t 0 extract representations using un-noised inputs to do open-vocabulary segmentation. We hypothesize that control tasks that require detailed spatial scene understanding would benefit from 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 Section 5.2. To illustrate the effect of the timestep, we display final denoised images for various values in different domains in Figure 8.
3.3 Prompt Specification
Since text-to-image diffusion models allow conditioning on text, investigate if can influence the representations to be more task-specific via this conditioning mechanism. For tasks that come with text specifier, for example, the sentence "go 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.
3.4 Intermediate Attention Map Selection
Wang al. [50] and Tang et [45] demonstrate that the Stable Diffusion model generates localized attention maps aligned with text during the combined processing of vision and language modali- ties. Wang et al. [50] leveraged these word-level attention maps to perform open-domain semantic segmentation. We hypothesize that these maps also help downstream control policies gener- alize to an open vocabulary object categories by providing helpful intermediate outputs that are aategory-agnostiee Following Tang al. [45], extract the antsnttttnnionn 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 fuse the cross-attention maps with the extracted feature maps from the U-Net. We refer to this attention-map-augmennn representation as SCR-ATTN.
3.5 Using Text-to-Imape Diffusion Model Representations to Learn Control Policies
To solve visual control tasks with states given by [Simager Steu], where Sixa may be used specify the task, we wish to use pre-trained as (Smagge S.ea). This encoded state then supplied a downstream, task mask-specificc policy network. which is trained to predict the action dz- Our evaluation encompasses both supervised learning and reinforcement learning regimes for training downstream policies. We train agents through behavior cloning on 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 and 4.3, we use version of the Proximal Policy Optimization [PPO, 39] algorithm for reinforcement learning.
representations capable of encoding the state

---

Train
Val
Figure 3: Sample scenes from the Habitat environments for the ImageNav (left) and OVMM (center) tasks. Instances from training and validation datasets of the OVMM object set are shown on the right.
3.6 Fine-Tuning on General Robotics Datasets
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 '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 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-Sometting-- [SS-v2; 13], and Bridge-v2 [49] datasets. As is standard, we fine-tune the denoiser U-Net Eg but not the VAE encoder decoder. Image-text pairs are uniformly sampled from the video-text pairs present these datasets. A possible limitation of this strategy is that text-video aligned pairs (a sequence of frames that correspond to 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 F. We refer to the representations from this fine-tuned model SCR-FT.
Empirical Evaluation
In this work, evaluate Stable Control Representations (SCR) extensive suite of tasks from 6 benchmarks covering few-shot imitation learning for manipulation in Section 4,1, reinforcement marning-based indoor navigation in Sections 4.2 and 4.3, and tasks related fine 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.
4.1 Few-shot Imitation Learning
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 Tranka-Kitched environments included in the RoboHive suite [23], which entail tasks such as turning knob or opening 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.
Results. We report the best results of SCR and baselines Table la. On Merl-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- Kiinke-,ttchen,... 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 Meta- Meta-World. This observation allows us to gauge the improved performance
5

---

Table 1: Average Success Rate and standard error evaluated across different representations.
(a) Meta- World & Franka-Kitchen Kitchen.
(b) ImageNav
(c) OVMM
Model
Meta- Meta-World
Model
Success
Model
Success
R3M
96.0 + 1.1
57.6
3.3
R3M
30.6
Oracle
77.6
CLIP
90.1 - 3.6
36.3 + 3.2
CLIP-B
52.2
Detic
36.7
VC-1
92.3 + 2.5
4
VC-1
70.3
CLIP
38.7 + 1.7
Voltron
5.2
13.5
MVP
68.1
VC-1
40.6 + 2.2
SD-VAE
75.5 +.2
43.7
SD-VAE
46.6
SCR
38.7 + 1.2
SCR
94.4
45.0 3.3
SCR
73.9
SCR-FT
41.9 - 1.0
SCR-FT
94.9 + 2.0
49.9
3.4
SCR-FT
69.5
SCR-FT-ATTN
43.6 + 2.1
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 diverse curation of robotics-relevant data. In this context, SCR's (C''s comparable performance shows that generative foundation models hold promise for providing useful representations for control, even with relatively minimal fine-tuning non-task-specific data.
4.2 Image-Goal Navigation
We now assess SCR in more realistic visual environments, surpassing the simple table-top scenes in manipulation benchmarks. 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 's ability to navigate to the viewpoint of provided goal image [63]. The position reached by the agent must be within 1-meter distance from the goal image's s camera position. This requires the ability to differentiate between nearby similar-looking views within home environment. This task, along with the semantic object navigation task that we study in Section 4.3, allows for a comprehensive evaluation of representation's s ability code both semantic and visual appearance-related features in completely novel evaluation environments.
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 with the vs 70.3%), the SOTA visual representation from prior work. This can be expected given that 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 la performs very poorly (30.6%) this domain, showing its limited transferability to navigation tasks.
VC-1 (69.5%
4.3 Open Vocabulary Mobile Manipulation
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- apen-vocabblary object detectors such as Detic [62] to enable generalization novel objects. We study modified version of the Gaze sub-task (detailed in Appendix G.2), which focuses on locating a specified object category for abstracted grasping action.

---

Table 2: We analyze the impact of varying the denoising timestep, layers selection, and input text prompt for the performance of SCR on the Franka-Kitchen benchmark. standard error over random seeds.
We report the mean and
(a) Denoising timestep.
(b) Layers selection.
(c) Input text prompt.
Timestep Success Rate
Layers
Success Rate
Prompt Type Success Rate
0
49.9 + 3,4
Down[1-3] + Mid
49.9 + 3,4
None
49.9 + 3,4
10
48.2 + 3.1
Down[1-3]
43.0 + 3,4
Relevant Inrelevant
49.2 + 3,5
100
42.0 + 3.7
Mid
41.6 3.3
48.7 = 3,3
110
42.0 + 3.4
Mid Up[0]
42.1 + 3.6
200
35.1 + 3.2
Mid + Up[0-1]
48.1 + 3.6
The task's success is measured by the agent's ability to precisely focus on the target object category. This category is provided as an input to the policy through its CLIP text encoder embedding. The evaluation environments cover both novel instances of object categories seen during policy learning, as well as entirely unseen categories. We compare to VC-1, the best model from Section 4.2 and CLIP, since prior work has studied it for open-vocab navigation [21, 26). We also incorporate a baseline that trains policy with access to ground truth object masks, evaluated using either the ground truth or Detic-genan retie-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 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 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. (+1.7%) Samples of attention maps overlaid on images from an evaluation episode can 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 dategory-agnostie input for downstream pick-and placeplaceeee policies, noisy detections can often lead to degraded downstream performance, 'Oracle' which benefits from ground truth object masks at test-time.
in this case. Nonetheless, there is still sizeable gap to
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.
5.1 Layer Selection
We begin our investigation by examining how the performance of SCR 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-Kide Kitchen domain when we use different sets of block outputs in SCR. We present similar ablations for Meta-World in the top four rows Table 3b.
We observe that utilizing outputs from multiple layers is instrumental to SCR's 's high performance. This finding underscores broader principle applicable to the design of representations across different models: Leveraging richer set of features from multi-layer outputs should enhance performance on downstream tasks. However, it 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 (16x 16x 1024 for VIT-L for images of size 224x224). 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-K Manka-Kitchen benchmark. While we see benefits from including the output of certain additional layers, it still underperforms SCR.
7

---

Table 3: Layer-selection ablations across different benchmarks.
(a) Ablations for CLIP on Kitche..
(b) Ablations for SCR on Meta-
Meta-Worid. World.
Model
Layers
Success
Layers
Noise
Success
CLIP-L
23 (last layer)
36.3 1.7
Mid
200
94.7 2.8
CLIP-L
21+23
35,4 + 2.9
Down[3] Mid
200
97.3 + 1.4
CLIP-L
19+23
3
Down[I-3
200
94.1 1.9
CLIP-L
12+23
40.8 2.8
Down[ 1-3] + Mid
200
94.4 + 1.9
CLIP-L
10+23
40.2 + 3.2
Down[ 33] + Mid
100
94.4 1.9
SCR (ours)
Down[ + Mid 49.9 + 3.4
Down[ 1-3] + Mid
0
94.1 + 1.9
Table 4: Comparison of CLIP Layer Ablations on Meta- Moord,,r,d,,.. OVMM, and ImageNav
(a) Meta-World and OVMM
(b) ImageNav
Model
Layers
Meta-Wood
OVMM
Model
Layers
ImageNav
CLIP-L
23 (Last Layer)
90.1 3.6
38.7 1.7
CLIP-B
11 (Last Layer)
52.2
CLIP-L
21+23
2.3
CLIP-B
6+11
66.6
CLIP-L
12+23
91.7 2.6
38.6 1.6
SCR
Down[1-3] + Mid
73.9
SCR
Down[1-3] + Mid
94.9 + 2.0
43.6 2.1
5.2
Sensitivity to the Noising Timestep
Next, we characterize the sensitivity of task performance 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 the bottom three rows of Table 3b. We see that the performance across nearby timesteps (0 and 10 or 100 and I 110) is similar, and that there is benefit to doing a coarse grid search up reasonable noising level (0 100 vs 200) get the best value for given task.
5.3 How is Language Guiding the Representations?
Recall that 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, 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-Kitte Kitchen setup from Section 4.1, which includes manipulation tasks that do not originally come with 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 the jungle") ") and relevant ("a Franka robot opening microwave door") prompts, compared to the default setting of not providing text prompt We see that providing prompt does not help with downstream policy performance and may even degrade performance the prompt gets more irrelevant the visual context of the input.
We now move to the Referring Expressions Grounding task which requires predicting bounding box for object being referred to in sentence, within image depicting cluttered objects. We defer the main presentation this task to Appendix D.2 and use 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 Table
We compared the following approaches for providing the task's text specification to the task decoder (also depicted in Figure 4):
(a) No text input: Exclude text prompt from both SCR and the task decoder by passing an empty prompt to the U-Net and using only the resulting SCR output for the decoder.
(b) Prompt only: Pass text prompt only to the U-Net.
(c) Concat only: Concatenate the CLIP embedding of the text prompt with the visual representation,
feeding an empty prompt to the U-Net.
(d) Prompt + Concat: Combine "Prompt Only" and "Concat Only".
(c) Only text encoding: Ignore the visual representation and rely only CLIP text embeddings.

---

Stable
Table 5: Ablating text imput to
Image Diffusion
Text
CLIP
SCR on the referring expres- sions grounding task.
a) No Text Input
e) Text Only
GLIP
Text
CLIP
Configuration
Score
(a) No text input
14.8
Stabll
Stable Diffusion
Stable
Image-1 Diffusion
Image
Image
Diffusion
(b) Prompt only
82.7
(c) Concat only
92.2
b) Prompt Only
c) Concat Only
Prompt and Concat
(d) Prompt + Concat
92.9
Figure 4: Illustration of different approaches to providing rele-
(e) Only text encoding
37.5
vant vision-language inputs to downstream task-decoder.
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 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 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 leveraging approaches to encode and integrate sub-phrases of long texts, that seek to address these challenges.
5.4 The Effect of Spatial Aggregation
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 'comppession 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-heades attention pooling (MAP) used for CLIP embeddings Karamcheti et al. [20]. We find that using compression layer significantly improves performance the fine-grained visual prediction tasks described in Appendix 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 MuloCo tasks. We see that the compression layer often outperforms the 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, compare the average precision between using multi-head attention pooling (MAP) and the compression layer. On the Meta- World & Franka- Franka-Kitchen tasks, we compare the average success rates (#z standard error) between the CLS token and compression layer embeddings.
Model
Aggregation Method
Refer Exp. Grounding
Grasp Affordance Prediction
Meta-Wond World
Tranka-Kitchhe
VC-1
MAPICLS
93.2
24.7
88.8 + 2.2
52,0 + 3.4
VC-1
Compression
94.6
83.9
92.3 + 2.5
47.5 +: 3.4
CLIP CLIP
MAP/CLS
68.1
60.3
88.8 +: 3.9
35.3 4: 3,4
Compression
94.3
72.9
90.1 +: 3.6
36.3 +: 3.2

---

performing models. This result provides evidence that the underperformance of CLIP representations on control tasks is unlikely due lack of sufficiently fine-grained visual information. Finally, we note that the compression layer aggregation technique was used for all baselines in Tables Ib and Ic to ensure a strong baseline comparison. We recommend that future studies adopt this methodology to enable fairer comparison of representations.
Discussion
In Section 5, we deconstructed Stable Control Representations and highlighted techniques used in our approach that could be applied 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 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, identified 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 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 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 the Franka-Ko Kitchen environment by fine-tuning CLIP similarly to SCR-FT in Appendix D.I.
7 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- -ottte-image diffusion models for policy learning can improve generalization across a wide range of tasks including manipulation, image-gont 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 other foundation models used for control. We hope that Stable Control Representations will help advance data-efficient control and enable open-vocabulary generalization challenging control domains as the capabilities of diffusion models continue to improve.
Acknowledgments
GG is funded by the EPSRC Centre for Doctoral Training in Autonomous Intelligent Machines and Systems (EP/S024050/1) and Toyota Europe. We gratefully acknowledge donations of computing resources by the Alan Turing Institute. The Georgia Tech effort was supported in part by ONR YIP and ARO PECASE. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies or endorsements, either expressed or implied, of the U.S. Government, or any sponsor.
10

---

References
[1] Anurag Ajay. Scungwook Han, Yilun Du, Shaung Li, Abhi Gupta, Tommi Jaakkola, Josh Tenenbaum, Leslie Kaelbling, Akash Srivastava, and Pulkit Agrawal. Compositional foundation models for hierarchical planning. arXiv preprint (xXiv23309008388 2023.
[2] Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, lain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: visual language model for few-shot learning. Advances in Neural Information Processing Systems, 35: 23716-23736, 2022.
[3] Dmitry Baranchuk, Andrey Voynov, Ivan Rubachev, Valentin Khrulkov, and Artem Babenko. Label-efficient semantic segmentation with diffusion models. In International Conference on Learning Representations, 2022.
[4]
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan. Karol Hausman, Alex Herzog, Jasmine Hsu, Julian Ibarz, Brian Ichter, Alex Irpan, Tomas Jackson, Sally Jesmonth, Nikhil Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Isabel Leal, Kuang-Huei Lee, Sergey Levine, Yao Lu, Utsav Malla, Deeksha Manjunath, Igor Mordatch, Ofrr Nachum, Carolina Parada, Jodilyn Peralta, Emily Perez, Karl Pertsch, Jornell Quiambao, Kanishka Rao, Michael Ryoo, Grecia Salazar, Pannag Sanketi, Kevin Sayed, Jaspiar Singh, Sumedh Sontakke, Austin Stone, Clayton Tan, Huong Tran, Vincent Vanhoucke, Steve Vega, Quan Vuong, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, Tianhe Yu, and Brianna Zitkovich. RT-1: Robotics Transformer for Real- World Control at Scale. arXiv e-prints, art arXiv:221 ww..x:2212.066:6 12.66817, December 2022. doi: 10.48550/arXiv 22....22222111 17.
[5] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Hersert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
[6] Xi Chen, Xiao Wang, Soravit Changpinyo, AJ Piergiovanni, Piotr Padlewski, Daniel Salz, Sebastian Goodman, Adam Gryener, Basil Mustafa, Lucas Beyer, al. Pali: A jointl--called multilingual language-ing -magg model. arXiv preprint 22Xiv:2009.06794, 2022.
[7] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137, 2023.
[8] Jasmine Collins, Shubham Goel, Kenan Deng, Achleshwar Luthra, Leon Xu, Erhan Gundogdu, Xi Zhang, Tomas F Yago Vicente, Thomas Dideriksen, Himanshu Arora, Matthieu Guillaumin, and Jitendra Malik. Abo: Dataset and benchmarks for real-world 3d object understanding. CVPR, 2022.
[9] Dima Damen, Hazel Doughty, Giovanni Maria Farinella, Sanja Fidler, Antonino Furnari, Evangelos Kazakos, Davide Moltisanti, Jonathan Munro, Toby Perrett, Will Price, and Michael Wray. Scaling egocentric vision: The epic-kitchens dataset. In European Conference on Computer Vision (ECCV), 2018
[10] Laura Downs, Anthony Francis, Nate Koenig, Brandon Kinman, Ryan Hickman, Krista Rey- mann, Thomas B McHugh, and Vincent Vanhoucke. Google scanned objects: A high high-quality dataset of 3d scanned houschold items. In 2022 International Conference on Robotics and Automation (ICRA), pages 2553-2560. IEEE, 2022.
[11] Yilun Du Mengjiao Yang. Bo Dai, Hanjun Dai, Ofir Nachum, Joshua B Tenenbaum, Dale Schuurmans, and Pieter Abbeel. Learning universal policies via text-guided video generation. URL https://arxiv.. org/abs/2302.000 /1, 2023.
[12] Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/VE Conference on Computer Vision and Pattern Recognition (CVPR), pages 12873- 12883, June 2021.

---

[13] Raghav Goyal, Samira Ebrahimi Kahou, Vincent Michalski, Joanna Materzynska, Susanne Westphal, Heuna Kim, Valentin Haenel, Ingo Fruend, Peter Yianilos, Moritz Mueller-Freitag, et al. The" something something" video database for learning and evaluating visual common sense. In Proceedings of the IEEE international conference on computer vision, pages 5842- 5850, 2017.
[14] Philippe Hansen-Estruch, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit q-learning as an actor-critic method with diffusion policies. arXiv preprint arXiv:2304. 10573, 2023.
[15] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurlPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.
[16] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 6840-6851. Curran Associates, Inc., 2020.
[17] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P. Kingma, Ben Poole. Mohammad Norouzi, David J. Fleet, and Tim Salimans. Imagen video: High definition video generation with diffusion models, 2022.
[18] Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022.
[19] Ivan Kapelyukh, Vitalis Vosylius, and Edward Johns. Dall-e-bot: Introducing web-scale diffusion models to robotics. IEEE Robotics and Automation Letters, 2023,
[20] Siddharth Karamcheti, Suraj Nair, Annie Chen, Thomas Kollar, Chelsea Finn, Dorsa Sadigh, and Percy Liang. Language-driven representation learning for robotics. In Robotics: Science and Systems (RSS), 2023.
[21] Apoorv Khandelwal, Luca Weihs, Roozbeh Mottaghi, and Aniruddha Kembhavi. Simple but effective: Clip embeddings for embodied ai. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), June 2022.
[22] Mukul Khanna, Yongsen Mao, Hanxiao Jiang, Sanjay Haresh, Brennan Shacklett, Dhruy Batra, Alexander Clegg, Eric Undersander, Angel X. Chang, and Manolis Savva. Habitat Synthetic Scenes Dataset (HSSD-200), An Analysis of 3D Scene Scale and Realism Tradeoffs for ObjectGoal Navigation. arXiv preprint, 2023.
[23] Vikash Kumar, Rutav Shah, Gaoyue Zhou, Vincent Moens, Vittorio Caggiano, Jay Vakil, Abhishek Gupta, and Aravind Rajeswaran. Robohive unified framework for robot learning. In NeurlPS: Conference on Neural Information Processing Systems, 2023.
[24] Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning, arXiv preprint arXiv:2304.08458 2023.
[25] Yecheng Jason Ma. William Liang, Vaidehi Som, Vikash Kumar, Amy Zhang, Osbert Bastani, and Dinesh Jayaraman. LIV: Language-Iman Representations and Rewards for Robotic Control. arXiv e-prints, art. arXiv:2306.00955. June 2023. doi: (0.48550/arxiv.22388 www..xw....3005...2.
[26] Arjun Majumdar, Gunjan Aggarwal, Bhavika Devnani, Judy Hoffman, and Dhruv Batra. Zson: Zero-shot object-goal navigation using multimodal goal embeddings. Advances in Neural Information Processing Systems, 35:32340-32355 2022.
[27] Arjun Majumdar, Karmesh Yadav, Sergio Arnaud, Yecheng Jason Ma, Claire Chen, Sneha Silwal, Aryan Jain, Vincent-Pierre Berges, Pieter Abbeel, Jitendra Malik, et al. Where are we in the search for an artificial visual cortex for embodied intelligence? arXiv preprint 2303. 18240, 2023.
[28] Zhao Mandi, Homanga Bharadhwaj. Vincent Moens, Shuran Song, Aravind Rajeswaran, and Vikash Kumar. Cacti: A framework for scalable multi-task multi-scene visual imitation learning. arXiv preprint 22.0::221. 2022.
12

---

[29] Lina Mezghani. Sainbayar Sukhbaatar, Thibaut Lavril, Oleksandr Maksymets, Dhruv Batra, Piotr Bojanowski, and Karteek Alahari. Memory-augmentt reinforcement learning for image- goal navigation. arXiv preprint arXiv:2 2021.
[30] Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A universal visual representation for robot manipulation. In 6th Annual Conference on Robot Learning, 2022.
[31] OpenAl. GPT-4 technical report, 2023.
[32] Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, et al. Imitating human behaviour with diffusion models. arXiv preprint arXiv:2301. 200::707..0677, 2023.
[33] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8748-8763. PMLR, 18-24 Jul 2021.
[34] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10695, 2022.
[35] Olaf Ronneherger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Nassir Navab, Joachim Homnegerr William M. Wells III, and Alejandro F. Frangi, editors, Medical Image Computing and Computer-Asiisted Intervention MICCAI 2015 18th International Conference Munich, Germany, October 5 9, 2015, Proceedings, Part III, volume 9351 of Lecture Notes in Computer Science, pages 234-241.
Springer, 2015.
[36] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 11136499364994 2022.
[37] Manolis Savva, Abhishek Kadian, Oleksandr Maksymets, Yili Zhao, Erik Wijmans, Bhavana Jain, Julian Straub. Jia Liu, Vladlen Koltun, Jitendra Malik, Devi Parikh, and Dhruv Batra. Habitat: A Platform for Embodied AI Research. In Proceedings of the IEEE/CVF International Conference Computer Vision (ICCV), 2019.
[38] Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et LAION- 5B: An open large-scale dataset for training next generation image-text models. Advances in Neural Information Processing Systems, 0125278-25294 2022.
[39] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal Policy Optimization Algorithms. arXiv e-prints, art. Stii:1707.063444 July 2017. doi:
10.48550/arXiv. 17w.a..wi............... 1707.06347.
[40] Dhruv Shah, Blazej Osiński, brian ichter, and Sergey Levine. Lm-nav: Robotic navigation with large pre-trained models of language, vision, and action. In Karen Liu, Dana Kulic, and Jeff Ichnowski, editors, Proceedings of The 6th Conference on Robot Learning, volume 205 of Proceedings of Machine Learning Research, pages 492-504. PMLR, 14-18 Dec 2023.
[41] Mohit Shridhar, Jesse Thomason, Daniel Gordon, Yonatan Bisk, Winson Han, Roozbeh Mot- taghi, Luke Zettlemoyer, and Dieter Fox. Alfred: A benchmark for interpreting grounded instructions for everyday tasks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10740 10749, 2020.
[42] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256-2266 PMLR, 2015.
13

---

[43] Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018.
[44] Andrew Szot, Alex Clegg, Eric Undersander, Erik Wijmans, Yili Zhao, John Turner, Noah Maestre, Mustafa Mukadam, Devendra Chaplot, Oleksandr Maksymets, Aaron Gokaslan, Vladimir Vondrus, Sameer Dharur, Franziska Meier, Wojciech Galuba, Angel Chang, Zsolt Kira, Vladlen Koltun, Jitendra Malik, Manolis Savva, and Dhruv Batra. Habitat 2.0: Training home assistants to rearrange their habitat. In Advances in Neural Information Processing Systems (NeurlPS), 2021.
[45] Raphael Tang, Linqing Liu, Akshat Pandey, Zhiying Jiang, Gefei Yang, Karun Kumar, Pontus Stenetorp, Jimmy Lin, and Ferhan Ture. What the DAAM: Interpreting stable diffusion using cross attention. In Proceedings of the 6 Annual Meeting of the Association for Computational Linguisties (Volume I: Long Papers), 2023.
[46] Junjiao Tian, Lavisha Aggarwal, Andrea Colaco, Zsolt Kira, and Mar Gonzalez-Franco. Diffuse, attend, and segment: Unsupervised zero-shot segmentation using stable diffusion. arXiv preprins 2308 /2469, 2023.
[47] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023.
[48] Jeremias Traub. Representation learning with diffusion models. arXiv:2210. 1058, 2022.
arXiv preprint
[49] Homer Walke, Kevin Black, Abraham Lee, Moo Jin Kim, Max Du, Chongyi Zheng, Tony Zhao, Philippe Hansen-Estruch, Quan Vuong, Andre He, et al. Bridgedata v2: A dataset for robot learning at scale. arXiv preprint arXiv:2308. 12952, 2023.
[50] Jinglong Wang, Xiawei Li, Jing Zhang, Qingyuan Xu, Qin Zhou, Qian Yu, Lu Sheng, and Dong Xu. Diffusion Model Secretly Training-free Open Vocabulary Semantic Segmenter. arXiv e-prints, art 2tt.:2330000773... September 2023. doi: 10.
[51] Ke-Jyun Wang, Yun-Hsuan Liu, Hung-Ting -ing Su, Jen Wang, Yu-Siang Wang. Winston Hsu, and Wen-Chin Chen. OCID-ref: A 3D robotic dataset with embodied language for clutter scene grounding. In Proceedings of the 2021 Conference of the North American Chapter the Association for Computational Linguistics: Human Language Technologies, pages 5333-5338, Online, June 2021, Association for Computational Linguistics. doi: 10. 186538v1/2222 .naccl- main.419.
[52] Erik Wijmans, Abhishek Kadian, Ari Morcos, Stefan Lee, Irfan Essa, Devi Parikh, Manolis Savva, and Dhruv Batra. Dd-ppo: Learning near-perfect pointgoal navigators from 2.5 billion frames. arXiv preprini arXiv: 1eXiv:1911.003333 2019.
[53] Fei Xia, Amir R Zamir, Zhiyang He, Alexander Sax, Jitendra Malik, and Silvio Savarese. Gibson env: worldworld perception for embodied agents. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 9068-9079, 2018
[54] Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. arXiv:2203.061/ 2022.
[55] Jiarui Xu, Sifei Liu, Arash Vahdat, Wonmin Byeon, Xiaolong Wang, and Shalini De Mello. ODISE: Open-Vone Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models, arXiv preprint arXiv: 2303. 04803, 2023.
[56] Karmesh Yadav, Arjun Majumdar, Ram Ramrakhya, Naoki Yokoyama. Alexei Baevski, Zsolt Kira, Oleksandr Maksymets, and Dhruv Batra. Ovrl-v2: A simple state-of-art baseline for imagenav and objectnav, 2023.
[57] Xingyi Yang and Xinchao Wang. Diffusion model as representation learner. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 18938-18949, 2023.
14

---

[58] Sriram Yenamandra, Arun Ramachandran. Karmesh Yadav, Austin Wang, Mukul Khanna, Theophile Gervet, Tsung- Yen Yang. Vidhi Jain, Alexander William Clegg, John Tumne. et al. Homerobot: Open-vocabulary mobile manipulation. arXiv preprint arXiv:2306. 1565, 2023.
[59] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference Robot Learning (CoRL), 2019.
[60] Tianhe Yu, Ted Xiao, Austin Stone, Jonathan Tompson, Anthony Brohan, Su Wang, Jaspiar Singh, Clayton Tan, Jodilyn Peralta, Brian Ichter, et al. Scaling robot learning with semantically imagined experience. arXiv preprint arXiv:2302. 1550, 2023.
[61] Jing Zhao, Heliang Zheng, Chaoyue Wang, Long Lan, and Wenjing Yang. Magicfusion: Boosting rext-to-image generation performance by fusing diffusion models. In Proceedings of the IEEE/CVP International Conference on Computer Vision (ICCV), pages 22592-22600 October 2023.
[62] Xingyi Zhou, Rohit Girdhar, Armand Joulin, Philipp Krähenbühl, and Ishan Misra. Detecting twenty-thousand classes using image-level supervision. In ECCV, 2022.
[63] Yuke Zhu, Roozbeh Mottaghi, Eric Kolve, Joseph J. Lim, Abhinav Kumar Gupta, Li Fei-Fei, and Ali Farhadi. Target-driven Visual Navigation in Indoor Scenes using Deep Reinforcement Learning. ICRA, 2017.
[64] Brianna Zitkovich, Tianhe Yu, Sichun Xu, Peng Xu, Ted Xiao, Fei Xia, Jialin Wu, Paul Wohlhart, Stefan Welker, Ayzaan Wahid, Quan Vuong. Vincent Vanhoucke, Huong Tran, Radu Soricut, Anikait Singh, Jaspiar Singh, Pierre Sermanet, Pannag R Sanketi, Grecia Salazar, Michael S Ryoo, Krista Reymann, Kanishka Rao, Karl Pertsch, Igor Mordatch, Henryk Michalewski, Yao Lu, Sergey Levine, Lisa Lee, Tsang-War Wei Edward Lee, Isabel Leal, Yuheng Kuang, Dmitry Kalashnikov, Ryan Julian, Nikhil J Joshi, Alex Irpan, brian ichter, Jasmine Hsu, Alexander Herzog, Karol Hausman, Keerthana Gopalakrishnan, Chuyuan Fu, Pete Florence, Chelsea Finn, Kumar Avinava Dubey, Danny Driess, Tianli Ding, Krzysztof Marcin Choromanski, Xi Chen, Yevgen Chebotar, Justice Carbajal, Noah Brown, Anthony Brohan, Montserrat Gonzalez Arenas, and Kehang Han. RT-2: Vision-language-accioo action models transfer web knowledge to robotic control. In 7th Annual Conference on Robot Learning, 2023.
15

---

Appendix
Table of Contents
Wider Societal Impacts
17
B Background
17
B.1 Diffusion Models
17
B.2 Latent Diffusion Models
17
B.3 Policy Learning for Control
17
Baselines
18
D Further Empirical Results
18
D.I
Fine-tuning CLIP
18
D.2 Referring Expressions Grounding
18
D.3 Grasp Affordance Prediction
19
D.4 Comparison with LIV
20
D.5
Overall Ranking of Representations
20
E
Implementation Details
21
E.1
Fine-tuning Stable Diffusion
E,2
Representation Extraction Details
22
E.3 Hyperparameters
22
Runtime Analysis
22
G Task Descriptions
23
G.1 Few-Shot Imitation Learning
23
G.2 Open Vocabulary Mobile Manipulation
23
G.3 ImageNav
24
16

---

A Wider Societal Impacts
This paper presents work whose goal to advance the field of robotic control using pretrained foundation models. Since primarily aim to learn policies for simulated control tasks, do not anticipate any ethical concerns at this stage. However, when deploying foundation models in the real world, should be wary that policies do not inherit any structural biases from the pretraining stage or internet-scale data.
B Background
B.1 Diffusion Models
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(æa), starting from pure noise. Given p(ar)) and set of noise levels o, for .,...,?7 T, a denoising function co c0(21, 1) is trained on the objective
(D) Emo.,, lll.lllllc - (zz, t)) [ = e.llcc co (zo + + c, +.)))
(B.1)
where (0, 1), t Un, and zu p(zo). To generate a sample By during inference, we first sample an initial noise vector Vy ~ NO0,, or) and then iteratively denoise this sample for l t= T .... by sampling from p( ( which is function of c,(2 ().
In some settings, we may want to generate samples with a partieular property. For example, we may wish to draw samples from conditional distribution over data points, (()o where 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 c0(x;, ł, c) and generate samples using classifier-free guidance [15].
B.2 Latent Diffusion Models
Latent diffusion models [34] reduce the computational cost of applying diffusion models to high- dimensional data by instead diffusing -in-innnsiooall representations high-dimensional data. Given an encoder E-) and decoder D.)), Equation (B.1) is modified to operate on latent representa-
tions, 20 E(xo). yielding
(LD = E.o.,,,. co (E(xo) + a. C. c,
(B.2)
where e ~ ((0, 1), t t~ (T) xo, p(zu, c). After generating a denoised latent representa- tion Zu, it can be decoded as xn D(zo).
popular instantiation of a conditioned latent diffusion model is the text-to-image Stable Diffusion model [SD; 34j. The SD model is trained the LAION-2B dataset [38] and operates in the latent space of a pre-trained VQ-VAE image encoder [12]. The model architecture shown at the top of
Figure and is based a U-Net [35], with the corresponding conditioning text prompts encoded using a CLIP language encoder [33].
B.3 Policy Learning for Control
We model our environments Markov Decision Processes (MDP, Sutton and Barto [43]), defined as tuple M = (S,A, P, R.),, where S and A denote state and action spaces respectively, P(s"|s, (s, a) the transition dynamics, R(ss a) the reward function, and E (0, 1) the discount factor. Our goal is to optimize policy (als) that maximizes the expected discounted return E,, IE y'R(s,, a:)].
In this paper, we consider visual control tasks that may be innguage-conditionn given by = [Simuge Stau) where Stea specifies the task. We are interested in pre-trained vision- language representations capable of encoding the state s as fa Simuge: Swar). This encoded state is then supplied downstream, task-specific policy network, which is trained to predict the action ar- 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.
that is, states are
17

---

C
Baselines
We compare SCR and its variants (i.e., SCR-FT and SCR-FT-ATTN) to the following prior work in representation learning for control:
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 (MAB) objective on egocentric data from Ego4D, Epic-Kitchens, SS-v2, and ImageNet, with VC-1 additionally pre-training on indoor navigation videos.
CLIP [33] trains text and ViT-hased 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.
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 forward pass through 400M parameters.
D Further Empirical Results
D.1 Fine-tuning CLIP
We 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 video closer to the video caption's 's text embedding. Specifically, for our experiment, we use the huggingface CLIP finetun- ing implementation and train the model with a batch size of 384 (the maximum number of samples we 5e-5 and a weight decay of 0.001 for 5000 update steps (same as SR-FT). We present the results in and note the lack of improvement the task post-fine-tuning.
Table 7: Performance on Franka-Kitchen after fine-tuning CLIP.
Model
Kinka-
CLIP CLIP (FT)
36.9 3.2
34.2 + 2.9
able fit on 8 A40 GPUs) with learning rate of
Table 7 for
D.2 Referring Expressions Grounding
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 attention to two downstream tasks involving fine-graines 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 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-ovee -nverunnoo- (IoU) threshold across different scene clutter levels for SCR-variane and baselines 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- SD-VAE perform the best with a 1 1.5% lead. The better performance of these vision-encoder-onll methods highlights that on this task, is not 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
18

---

that some of the runs could potentially improve over extended training. However, we were primarily interested in this task not just compare the downstream visual prediction performance, but use as testhed for exploring the following two questions: (1) Do the performance differences between the representations we evaluated in Sections 4.1 4.3, stem from the absence of fine- 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, explore what extent language prompting influences the representations from SCR on language-conditione tasks Section 5.3.
Maximum clutter
Medium clutter
Minimum
Model
Average
clutter
CLIP
68.1
60.3
76.6
67.0
R3M
63.3
55.3
68.3
63.3
Voltron
92.5
96.9
91.8
90.2
VC-1 94.6
93.7
96.5
93.7
SD-VAE
94.3
93.2
96.3
93.4
SCR
92.9
91.1
95.9
91.8
SCR-FT
91.8
90.1
94.8
90.8
The lemon on the rear left of the instant_noodles Figure 5: Sample from the OCID-Ref dataset used for the Referring Expressions task.
Table 8: Referring Expression Grounding (Accu- racy at threshold IoU of 0.25 with label.).
D.3 Grasp Affordance Prediction
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 cer- tain areas of objects in an RGB image. These areas cor- respond to parts of the surface that would be amenable to grasping by suction gripper. The evaluation metric adopted in prior work is the precision of predictions cor- responding to positive graspability at varying confidence
Table 9: Grasp Affordance Prediction:
Precision on pixels corresponding pos- itive graspability at varying probability threshold levels.
Model
Top99
Top95
Top90
CLIP
60.3
45.0
28.6
CLIP (Comp)
72.9
55.9
36.5
Voltron
62.5
42.8
32.1
levels (90, 95, and 99th percentile of the predicted per- pixel probabilities, denoted as Top90, Top95, and Top99 Table 9). We refer the reader Karamcheti al. [20] for the complete task setup details.
SD-VAE
55.6
41.3
33.8
SCR
72.8
55.9
54.5
SCR-FT
72.3
54.6
44.4
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 the computation of the precision metrics. The evaluation procedure for this task adopted in prior work involves 5-fold cross-validation and we observed 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 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.
19

---

D.4 Comparison with LIV
We include comparison with LIV [25] on two tasks that involve manipulation and navigation. LIV is vision- language representation learned through contrastive learn- ing on the EpicKitchens dataset [9]. Similar to R3M re- sults in the main paper, this representation does well on manipulation tasks but poorly on navigation tasks.
Table 10: Comparing to LIV on manipu- lation and navigation tasks.
Model
Franka Kinka-Kitcchn
OVMM
SCR
45.0
38.7
SCR-FT LIV
49.9
41.9
54.2
8.4
D.5 Overall Ranking of Representations
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 give higher-level view of the all-round performance 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 a representation-leamiin 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 wide array of downstream robotics tasks.
Table 11: Representation Performance Comparison: Numbers in the task columns (OVMM, Image- Nav, Meta World, Franka Kitchen) indicate relative scores of different representations (normalized by the highest score on that task), and the average normalized score column indicates the averaged scores across the task-wise relative scores where numbers are available.
Method
OVMM
ImageNav
MetaWorld
Franka Kitchen
Average Norm. Score
VAE
0.629
0.786
0.759
0.725
R3M
0.414
1,000
1,000
0,805
VC-1
0.969
0.951
0.961
0.825
0,927
CLIP
0.924
0.706
0.939
0.630
0.800
SR
0.924
1.000
0.983
0.781
0.922
SR-FT
1.000
0.942
0.989
0.866
0.949
20

---

E
Implementation Details
E.1
Fine-tuning Stable Diffusion
We used the runwayml/ http///aabeddfffiiieee initialize our models and fine-tuned them using the diffiusess library', As noted in Section 3.6, used subset of the frames from the EpicKitchens, 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 relevant for describing the scene. Since the Something- Something-Sometthngg human hand manipulating table-top objects while the Bridge-V2 V2 dataset includes 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, paired
model weights provided by Huggingface to
and Bridge-v2 datasets.
(SS-v2) dataset includes
generated 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 munet for only a single epoch (~5,000 gradient steps) parallely on Nvidia A100 GPUs with mini-batch size of 512 and learning rate of le Although the original Stable Diffusion model trained on images of resolution 512 512, we fine-tuned the model on images downscaled to 256 x 256 since it aligned with the resolution requirements of the downstream application. We show sample generations from the diffusion model after fine-tuning Figure 7. We found that the model learns to associate the prompt with not just the human robot hand but also with the style of the background and objects of the training datasets.
Figure 6: Snapshots from the datasets for fine-tuning the Stable Diffusion model.
Human hand
Robot hand
lossing box the trash
opening drawer
tossing the tresh
opening drawer
slicing a cucumber with knife
wiping the kitchen counter
slicing cucumber with knife
wiping the kitchen counter
Figure 7: Image generations from the fine-tuned Stable Diffusion model. We provided four different prompts, each prefixed with either "Human hand" or "Robot hand".
21

---

E.2 Representation Extraction Details
Here, we describe the representation extraction details for all our baselines assuming 224 224 input image.:
1. Stable Control Representations: The Stable Diffusion model downsamples the input images by a factor of 64. Therefore, we need to first resize the input image to a size of 256 x 256. We pass the image to the VAE, which converts it into latent vector of size 32 32 4 and passes it to the U-Net. We use the last three downsampling blocks' and the mid block's output feature map of sizes 8 8 x 640, 4 x 1, 280, 4 x 1, 280, and x 4 280, respectively. The total size is, therefore, 102, 400, and we linearly interpolated them to the same spatial dimension (8 x 8) before concatenating them channel-wise.
2. R3M [30]: For most of our experiments we use the original ResNet50 model, which outputs 2048 dimensional vector. For the referring expressions and grasp affordance prediction tasks from the Voltron evaluation suite [20], VIT-S is used, which outputs an embedding of size
14 14 384 75, 264
3. MVP [54] and VC-1 [27]: The last layer (24th) outputs an embedding of size 16 16 x 1, 0,024
262, 144.
4. CLIP [33]: For VīT-B, the last layer (12th) outputs embedding of size 14 14 768 = 150, 528. For ViT-L, the last layer (24") outputs an embedding of size x x 1024 262, 144.
5. Voltron [20]: We use the VCond-Base model which outputs representation of size 14 x 14 X
768 150, 528.
6. SD-VAE [34]: Outputs latent vector of size 32 32 4 4, 096.
E.3
Hyperparameters
We provide the hyperparameters used in Section 4 for Stable Control Representations in Table 12.
Table 12: Hyperparameters and configuration settings used across tasks and methods.
Benchmark
Timestep Prompt
Att
Layers
Post Compression Dim
Meta- Meta-World
0/100/200
No
No Mid Down [1-3]
3072
Franka-Kitchen
0
No
No Mid Down [1-3]
2048
ImageNav
0
No
No Mid Down [1-3]
2048
OVMM
100
Yes
Yes Mid Down [1-3]
2048
Referring Expression
0
Yes
No Mid Down [1-3]
8192
Grasp Prediction
0
No
No Mid Down [1-3] 8192
F
Runtime Analysis
We include a runtime analysis for SCR versus VC-1 in Table 13 by reporting the time taken for forward pass through both models over a single input image of size 256x256. We use an 100 GPU for our experiment, and use both the models at half precision. We see that SCR takes 0.021 seconds per inference step, being 1.5x slower than the forward pass through VC-1 Large.
Table 13: Runtimes for.
SCR
VC-1
Time per forward pass (Averaged over 1,000 passes)
0.021 seconds 0.014 seconds
22

---

G Task Descriptions
G.1 Few-Shot Imitation Learning
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 25 evaluation runs over 3 seeds. All experiments are conducted on single A00 GPU with 24 CPUs and 188 GBs of RAM.
Worad. We follow Majumdar et al. [27] and use the hammer-v2, drawer-open-v2. bin- ain-picking- v2, pren--prrsss--oooonn..
assembly-v2 tasks from the Meta-World benchmark suite [59]. Each task provides the model with the last three 256 x 256 RGB images, alongside a 4-dimensional gripper pose. The model consists of 3-layer MI.P with a hidden dimension of 256, utilizing ReLUU as the activation function. It undergoes training for 100 epochs, with evaluations conducted every 10 epochs, following the approach of Majumdar et (2023). The training uses mini-batch size of 256 and a learning rate of 10
Franka-Kitchen. The tasks involved here include Knob On, Knob Off, Microwave Door Open, Sliding Door Open, and L Door Open, each observed from three distinct camera angles. For each task, the model receives 256 256 RGB image and 24-dimensional vector representing the manipulator's proprioceptive state. For our experiments, we use the the RoboHive Kumar et al. [23] a
GitHub repository" and use a 2-layer MLP with a hidden dimension of 256 and train for 500 epochs. The mini-batch size is set at 128, with learning rate of 10-1. We additionally correct a bug in the RoboHive implementation of the VC-I baseline, specifically on input image normalization. Adjusting the image normalization to a 0-1 range resulted in a significant improvement in its performance.
G.2 Open Vocabulary Mobile Manipulation
vren-Vocabulary Mobile Manipulation [OVMM; 58] is a recently proposed embodied AI bench- mark that evaluates an agent's s ability to find and manipulate objects of novel categories in unseen indoor enviromments. Specifically, the task requires an agent to "Find and pick an object on the start_receptacee and place it the goal_recetacle" where object, Start_receptant and goal tall_reeetcclle are the object category names. Given the long-horizen 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 gpcabvocyyylary capabilities of pre-trained representations, we choose to evaluate the models 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_ paartrecepttcle. 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 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 HomeRobot uses the Habitat simulator [44] with the episode dataset provided by Yenamandra et al. [58]. This dataset uses 38 scenes for training and 12 scenes for validation, all originating from the Habitat Synthetic Scenes Dataset [HSSD; 22]. The episode dataset populates the scenes with extra objects from other datasets including Amazon Berkeley Objects [ABO: 8). Google Scanned Objects [ABO; 10] and Alfred [41]. The validation scenes are populated with objects not seen during training, spanning 106 seen and 22 unseen categories. The validation set consists of a total of 1199 episodes.
Our agent is designed to resemble the Stretch robot, characterized by height of 1.41 meters and radius of 0.3 meters. At a height of 1.31 I meters from the base, a 640x480 resolution RGBD camera is mounted. This camera is equipped with motorized pan and tilt capabilities. The agent's 's action space
23

---

is continuous, allowing it to move forward distances ranging from to 25 centimeters and to turn left or right within angles ranging from 30 degrees. Additionally, the agent can adjust the head's pan and tilt by increments ranging from 0.02 to radian in a single step.
In our experiments, passing them through the compression layer. We initialize the LSTM weights with the LSTM weights of the Oracle model get a slight boost in performance. We train our agents using the distributed version of PPO [52] with 152 environment spread across 4 80GB Nvidia A100 GPUs. Each run also has access 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.
2 layer LSTT policy and pass in the visual encoder representations after
G.3 ImageNav
We conduct our ImageNay experiments in the Habitat simulator [37], using the episode dataset from Mezghani et al. [29]. These experiments are conducted using the VC-I codebase" [27]. The dataset uses 72 training and 14 validation scenes from the Gibson [53] scene dataset with evaluation conducted total of 4200 episodes. The agent is assumed to be in the shape of cylinder of height 1.5m and radius O. 0.,, with an RGB camera mounted at height of 1.25m from the base. The RGB camera has resolution of 128x 128 and 90 field-of-view.
At the start of each training episode, an agent randomly initialized in a scene and is tasked to find the position from where the goal image 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 (25em), 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- pre-trained model is frozen, the policy is trained using the Adam W optimizer, with a learning rate of 2.5 x 10 and weight decay of 10-6. 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.
24

---

t=100
t=200
t=300
Noised
Generated
Meta-World
Noised
Generated
OVMM
Noised
Generated
Refer Expression
Figure 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 E {100, 200, 300}. and then denoise the image back to timestep O. We observe that inputs from different tasks are differently sensitive to the noising ranges, based on the amount of information the images contain. On Morld, SD is able to reconstruct the image correctly even at f 200, while for the referring expressions grounding task, noising leads to information loss even at 100 with several small objects reconstructed differently to the original. This affects the range of the noise could add to the ended latents at the time representation extraction. It should be noted however, that using un-noised inputs (1 0) always worked well in experiments, and this hyper-parameter could be ignored in practice.
25

---

Refer Exoression
Input Image 'A''
"pear"
"beside"
"a'
"book"
DVMM
Input Image
"rocket"
'and"
"chair"
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 robotic manipulation environment and observe that they are accurate at localizing objects a scene. Since these maps are category agnostic, downstream policies should become robust unseen objects time.
Find Tomaoo
(2) Find Gamnnn
Container
Spray
(6) Find Candle
Figure 10: Images from OVMM benchmark with their corresponding attention maps obtained from the fine-tuned Stable Diffusion (SD) model. The first 5 pairs of images correspond to failed episodes, with the bottom right pair corresponding
successful episode. The attention maps help interpret
the cause of failure: (1) Tomato SD wrongly attends strongly to an apple. (2) Gaming Console visible at the top of the image; however, SID attends multiple objects due low visual quality. (3) Plant Container SD instead focuses on the two glasses sees in the image. (4) Spray Bottle SD completely misses the spray bottles in the image and attends to the lava lamp. (5) Plant Container SD wrongly attends to the apple. (6) Candle Holder - SD correctly attends.
26

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction accurately summarise the two main contributions of the paper: (1) a procedure derive vision- lasion-lagguage representations from diffusion models for inputs in control tasks (2) a broad empirical comparative study of the representations on different domains where the visual representations from diffusion models match prior state-of-the-an representations in terms of downstream policy learning.
Justification:
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: In Section we discuss the interpretation of our experimental results, and highlight limitations of our work. We discuss the dependence downstream performance on design choices related to representation extraction. We also include run-time analysis in Table 13 which highlights the run-time cost of using a pre-trained representation model with more parameters.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification: This work does not include theoretical results.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide detailed description of experimental procedure in Appen- dices and G. Section 3 describes how to extract vision(and language) representations from the Stable Diffusion model. We only use publicly available datasets and benchmarks in this work. We will open-source the code that we provided as part of our supplementary material.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We will open-source the code provided as part of our supplementary material. This includes documented code for the representation extraction procedure well the code for re-running the experiments related to the Mujoco and ImageNav benchmarks used in Section 4. Additionally, all the pre-trained models, benchmarks and datasets we are publicly available, and we include the sources and links to these in the code documentation well as the appendices.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
27

---

Justification: We provide details related to the finetuning of Stable Diffusion model and extraction of representations Section 5 and appendix E, and details related to hyperpa- rameters, optimisation and train/test data and environments in Section 4 and appendices D and G. We follow the train/test splits adopted in prior work for each of the publicly available benchmarks we study in this work.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Our experimental results tables in Section 4 mention the standard error as- sociated with the reported numbers over 3 seeds. Only the experiments in Section 4.2 (ImageNav) report numbers that involve single random seed, since we were limited by compute resources on this benchmark as we specify in Section 4.2.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide the details related to the type of GPU resources used, as well as the batch sizes and number of optimisation steps for each task in Appendix G.
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: sttps://ntt
Answer: [Yes]
Justification: We have carefully looked at the guidelines and can confirm this work does not violate them.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer:
[Yes]
Justification:
We discuss the broder societal impact of this work in Appendix A.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We build on the already open-sourced image generation model, Stable Diffu- sion and finetune
images of human hand robots manipulating objects. This released model is then used for the task of predicting visual representations on simulated embodied AI tasks. We believe this model does not pose risks more than that posed by the original Stable Diffusion model.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
We have included the CC-by-NC4.0 license with our code which is based on the license used by the Visual Cortex repository. Dataset: The datasets used in the work are open-sourced by prior work with different licenses and we will point the users of our codebase to the relevant repository to download the datasets for their use. Model: Our work releases an updated Stable Diffusion checkpoint with the license matching that of the original license (creativemllopenrailln))
Justification: Code:
28

---

13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
We will open-source the code included in the supplementary material. The code is well-documented, and be used to reproduce our experimental results in Section as well as be used to extract diffusion representations for generic visual (Sangua-a tasks by independent researchers..
Justification:
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This work does not include research with human subjects.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This work does not involve research with human subjects.
29


Reviews:
[{"review_id": "jI2SAM7Ixy", "reviewer": "Reviewer_9EUL", "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.", "questions": "See above", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n- 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.\n- 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?\n- 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? \n\n- The order of Table 2 is different from descriptions in Sec. 5.1 and 5.2", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n\n**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.**\n\nTackling 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.\n- 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.\n- 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.\n\n**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.**\n\n- 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. \n- 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. \n\n**3) Explain The performance gap between SCR-FT and R3M in Meta-World and Franka-Kitchen**\nWe 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). \n\n**4) Typos**\n\nThank 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.\n\n---\n\nWe hope that these clarifications have addressed any remaining questions and concerns.\n\nIf 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.\n\n---\n\n**References**\n\n[1] He, Kaiming, et al. \"Masked autoencoders are scalable vision learners.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.\n\n[2] Open-Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models. Xu et al., 2023.\n\n[3] Iterated Denoising Energy Matching for Sampling from Boltzmann Densities. Akhound-Sadegh et al., 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7kkTB0FnaV", "reviewer": "Reviewer_QYDZ", "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.", "questions": "None", "limitations": "Refer to the weaknesses section.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "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.\n2. 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.\n3. 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n\n**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].\n\n**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.\n\n---\n\nWe hope that these clarifications have addressed any remaining questions and concerns.\n\nIf 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.\n\n---\n\n**References**\n\n[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.\n\n[2] Xiao, Tete, et al. \"Masked visual pre-training for motor control.\" arXiv preprint arXiv:2203.06173 (2022).\n\n[3] Yadav, Karmesh, et al. \"Offline visual representation learning for embodied navigation.\" Workshop on Reincarnating Reinforcement Learning at ICLR 2023. 2023.\n\n[4] Eftekhar, Ainaz, et al. \"Selective visual representations improve convergence and generalization for embodied ai.\" arXiv preprint arXiv:2311.04193 (2023).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AkMuacV9qC", "reviewer": "Reviewer_sdTb", "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.\n\nGiven 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.\n\nExperiments 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.", "questions": "Please kindly refer to the weaknesses.", "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.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n    \n2. 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.\n    \n3. The paper is well-written and easy to read.\n    \n4. 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?\n    \n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n\n**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?**\n\nYes, 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.\n\n| Timestep | Success       |\n| -------- | ------------- |\n| 0        | 94.1 &pm; 1.9 |\n| 0        | 94.4 &pm; 1.9 |\n| 0        | 94.4 &pm; 2.0 |\n\n**2) All ablation studies in Section 5 are conducted solely on the Franka-Kitchen benchmark.** \n\nThank 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.\n\n**3) SCR-FT shows a performance drop compared to the non-finetuned version on the ImageNav benchmark.**\n\nWe 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]. \nWith 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.\nFinally, **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.\n\n---\n\nWe hope that these clarifications and additional experimental results have addressed any remaining questions and concerns.\n\nIf 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.\n\n---\n\n**References**\n\n[1] Zhou, Tinghui, et al. \"Stereo magnification: Learning view synthesis using multiplane images.\" arXiv preprint arXiv:1805.09817 (2018).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uCjMJsjvcf", "reviewer": "Reviewer_ZVD8", "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.", "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.\n\n2.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) ?\n\n3. Why is CLIP-B used in the ImageNav experiments and not CLIP-L?", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "1. Lots of ablation experiments to validate the final design choice.\n2. 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?\n\nFurthermore, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n\n**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.**\n\nPre-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.\n\n**2) What if the same feature pooling strategy is applied to baselines like CLIP?**\n\n**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**.\n\n**3) Seems some multi-layer feature aggregation was done for Tables 1(b) and 1(c) but not 1(a)**\n\nTo 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.\n\n**4) Why is CLIP-B used in the ImageNav experiments and not CLIP-L?**\n\nThis 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.\n\n---\n\nWe hope that these clarifications and additional experimental results have addressed any remaining questions and concerns.\n\nIf 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4w6WWBnUti", "reviewer": "Reviewer_zJpU", "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.\nTheir 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).\nThey 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.", "questions": "Why cite Imagen in l. 36?\nReference to Table 2a/b in section 5.1 and 5.2 are mixed up, just a small fix.\n\n“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": "There are not many concerns to address and the authors wrote an appropriate short section in the Appendix.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This is the first paper to make Diffusion Model representations work for robotic control tasks (novelty).\nThe 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.\nEach part of the paper is well motivated and answers a natural question or follow-up step.\n\nRegarding 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.”.\nOr 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.\n\nThe authors use baselines, i.e. showing SD-VAE is an insightful baseline that is easily beat.\nSimilarly, 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.\n\nLastly, 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.\n\nIt 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. \n\nOverall 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.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n\n**1) The mutliple layer results for CLIP should be investigated further.**\n\n**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**. \n\n**2) Could you explain more how these baselines work and why they get the numbers they do?**\n\nWe agree with the reviewer and will add further contextualization of the performance of different baselines.\n- 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. \n- 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.\n- 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. \n\n**3) Simplicity of finetuning SCR compared to CLIP should be discussed earlier, claim about sophisticated negative sampling only cites the original CLIP paper**\n\nWe 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:\n- 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.\n- 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].\n\nWe 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.\n\n**4) Sec 5.3 where too many things are introduced at once.**\n\nWe 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:\n- We wanted to introduce the proxy visual grounding tasks used in Voltron [20] and present results for SCR on them.\n- 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.\n- 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.\n\n**5) Bilinear interpolation should be shown in equations formally to understand exactly what is done**\n\nWe 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.\n\n**6) a) Why cite Imagen in L36? b) Reference to Table 2a/b are mixed up.**\n\nThanks 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.\n\n---\n\nWe hope that these clarifications and additional experimental results have addressed any remaining questions and concerns.\n\nIf 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.\n\n---\n\n**References**\n\n[1] Lin, Xingyu, et al. \"Spawnnet: Learning generalizable visuomotor skills from pre-trained networks.\" arXiv preprint arXiv:2307.03567 (2023).\n\n[2] Contrastive Learning with Hard Negative Samples : Joshua David Robinson, Ching-Yao Chuang, Suvrit Sra, Stefanie Jegelka\n\n[3] Debiased Contrastive Learning: Ching-Yao Chuang, Joshua Robinson, Lin Yen-Chen, Antonio Torralba, Stefanie Jegelka", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for taking the time to review our manuscript and for providing thoughtful and constructive feedback. 

We were delighted that reviewers recognized that our work studies an “important” and “interesting” problem (**9EUL, QYDZ**) while being presented in a manner which is “easy to understand” (**9EUL**) and which “contextualizes the background and related work well” (**zJpU**). We are pleased to see the reviewers note the “strong performance” (**sdTb**) of the proposed approach and appreciate the “thorough ablations and analysis” (**QYDZ, sdTb**), along with a “fair treatment of the baselines” (**zJpU**) which we strived to ensure in our work. 

Reviewers also commended the “comprehensive evaluation” (**all**), "which becomes more important with paradigm shift towards multi task models" (**ZVD8**). We are grateful that they point out the work has “potential for future impact through the ability to handle complex open-ended environments” (**sdTb**). Finally, we thank the reviewers for acknowledging our open source code repository which will help reproducibility and accelerate future efforts in this direction (**zJpU**).

---

**Since reviewers ZVD8 and zJpU asked for results on multi-layer aggregation of CLIP representations, we present these results in the general response:**

As we note in Section 5.1 of the manuscript, we can see clear benefits from aggregating multi-layer features for SCR. The reason why we used these features is because of the smaller feature maps in the U-Net compared to ViTs (thus using multiple layer outputs for a fair comparison). This is an architecture-specific property, and not a property of diffusion models specifically, and we expect all models to benefit from this. However, as we mention in L262, this can lead to very large representation sizes for ViTs (each layer output is 16x16x1024).

Following the reviewer’s suggestion, we ran the same experiment for CLIP by concatenating outputs from the last layer + another layer in the range of layers 10-2.1 (We use two layers to keep GPU memory manageable, and note that this makes the CLIP representations twice as large compared to other methods.) We present the results below:

| Model     | Layers  		   | Success (&pm; std err)  |
| --------- | ---------------- | ----------------------- |
| CLIP-L    | 23 (last layer)  | 36.3 &pm; 1.7           |
| CLIP-L    | 21+23 		   | 35.4 &pm; 2.9           |
| CLIP-L    | 19+23 		   | 38.5 &pm; 3.2           |
| CLIP-L    | 17+23 		   | 39.0 &pm; 3.0           |
| CLIP-L    | 12+23 		   | 40.8 &pm; 2.8           |
| CLIP-L    | 10+23 		   | 40.2 &pm; 3.2           |
| SCR (ours)| Down[1-3] + Mid  | **49.9 &pm; 3.4**       |

We see an interesting trend, namely that moving towards middle layers leads to higher performance indicating that CLIP layers around layer 10-14 encode some details useful to the Franka benchmark. **While we do see a benefit from including the output of certain additional layers, performance CLIP performance still does not match SCR.**
We also wanted to see if the best pair of layers that we found here might be better on other tasks as well. To test this, we ran the same combination for Meta-World, OVMM, and Imagenav (adjusted for ViT-B). We present the results below.

| Model     | Layers  		  | Meta-World        | OVMM          |
| --------- |---------------- | -------------     | ------------- |
| CLIP-L    | 23 (Last Layer) | 90.1 &pm; 3.6     | 38.7 &pm; 1.7 |
| CLIP-L    | 12+23           | 91.7 &pm; 2.6     | 38.6          |
| CLIP-L    | 21+23           | 91.2 &pm; 2.3     | -             |
| SCR (ours)| Down[1-3] + Mid | **94.9 &pm; 2.0** | **43.6 &pm; 2.1** |

| Model     | Layers  		  | ImageNav |
| --------- |---------------- | -------- |
| CLIP-B    | 11 (Last Layer) | 52.2     |
| CLIP-B    | 6+11            | 66.6     |
| SCR (ours)| Down[1-3] + Mid | **73.9** |


We see a significant improvement on ImageNav and a slight improvement on Meta-World. (We did not see an improvement on OVMM but note that the results above are only for a single seed.)

---

**Following Reviewer sdTb's suggestion, we also replicated our ablations---which we had previously performed on the Franka kitchen benchmark---on Meta world. We present the results here**:

**Meta-World Ablations**
| Layers          | Noise | Success (&pm; std err) |
| --------------- | ----- | ---------------------- |
| Mid             | 200   | 94.7 ± 2.8             |
| Down[3] + Mid   | 200   | **97.3 ± 1.4**         |
| Down[1-3]       | 200   | 94.1 ± 1.9             |
| Down[1-3] + Mid | 200   | 94.4 ± 1.9             |
| Down[1-3] + Mid | 100   | 94.4 ± 1.9             |
| Down[1-3] + Mid | 0     | 94.1 ± 1.9             |

The top three rows ablate the layer used, and the bottom three rows ablate the noise values. We note that this ablation on Meta-World shows that it is possible to outperform the best method here (R3M[30]) if we use a smaller subset of the layers we used as the default (thereby achieving a 97% success rate). However, standardizing the representation extraction settings which we used across benchmarks helps us to retain simplicity while still achieving an overall good performance across tasks.

Meta-review: 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.
Decision: Accept (spotlight)