Title: Pre-trained Text-to-Image Diffusion Models Are Versatile Representation Learners for Control
Venue: NeurIPS 2024
Accepted: yes
Keywords: Embodied AI

Representation Learning for Control

Diffusion Models

Foundation Models
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-Image Diffusion Models Are Versatile Representation Learners for Control

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 understandinga 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.

Code: github.com/ykarmesh/stable-control-representations

1 Introduction

As general-purpose, pre-trained 'foundation' models [2, 5, 6, 24, 31, 34, 47] are becoming widely available, a central question in the field of embodied AI has emerged: How can foundation models be used to construct model representations that improve generalization in challenging robotic control tasks [4, 40, 64]?

Robotic control tasks often employ pixel-based visual inputs paired with a language-based goal specification, making vision-language model representations particularly well-suited for this setting. However, while vision-language representations obtained via Contrastive Language-Image Pre-training [CLIP; 33]-a state-of-the-art method-have been successfully applied to a broad range of computer vision tasks, the use of CLIP representations has been shown to lead to poor downstream performance for robotic control. This shortcoming has prompted the development of alternative, control-specific representations for embodied AI [25, 30] but has left other sources of general-purpose pre-trained vision-language representations-such as text-to-image diffusion models-largely unexplored for control applications.

0.0

0.2

0.4

0.6

0.8

1.0

VAE

CLIP

R3M

VC-1

SCR (Ours)

SCR-FT (Ours)

Overall Representation Comparison

"Picking up a kettle"

CLIP Language Encoder

StableDiffusion

VAE

U-Net

VAE

Encoder

Decoder

Noise

Level

Original Image

Visual Representation

DenoisedImage

EvaluationSuite

Meta-World

FrankaKitchen

OVMM

ImageNav

ReferExpression

GraspPrediction

In this paper, we propose Stable Control Representations (SCR) : pre-trained vision-language representations from text-to-image diffusion models that can capture both high and low-level details of a scene [17, 34]. While diffusion representations have seen success in downstream vision-language tasks, for example, in semantic segmentation [3, 46, 50], they have-to date-not been used for control. We perform a careful empirical analysis in which we deconstruct pre-trained text-to-image diffusion model representations to understand the impact of different design decisions.

In our investigation, we find that diffusion representations can outperform general-purpose models like CLIP [33] across a wide variety of embodied control tasks despite not being trained for representation learning. This is the case for purely vision-based tasks as well as for settings that require task understanding through text prompts. A highlight of our results is the finding that diffusion model representations enable better generalization to unseen object categories in a challenging open-vocabulary navigation benchmark [58] and provide improved interpretability through attention maps [45].

Our key contributions are as follows:

In Section 3, we introduce a multi-step approach for extracting vision-language representations for control from text-to-image diffusion models. We show that these representations are capable of capturing both the abstract high-level and fundamental low-level details of a scene, offering an alternative to models trained specifically for representation learning.

In Section 4, we evaluate the representation learning capabilities of diffusion models on a broad range of embodied control tasks, ranging from purely vision-based tasks to problems that require an understanding of tasks through text prompts, thereby showcasing the versatility of diffusion model representations.

In Section 5, we systematically deconstruct the key features of diffusion model representations for control, elucidating different aspects of the representation design space, such as the input selection, the aggregation of intermediate features, and the impact of fine-tuning on performance.

We have demonstrated that diffusion models learn versatile representations for control and can help drive progress in embodied AI. Figure 1 presents a summary of our approach and results. 2

2 Related Work

We begin with a review of prior work on representation learning and diffusion models for control.

Representation Learning with Diffusion Models. Diffusion models have received a lot of recent attention as flexible representation learners for computer vision tasks of varying granularity-ranging from key point detection and segmentation [46, 50] to image classification [48, 57]. Wang et al. [50] has shown that intermediate layers of a text-to-image diffusion model encode semantics and depth maps that are recoverable by training probes. These approaches similarly extract representations by considering a moderately noised input, and find that the choice of timestep can vary based on the granularity of prediction required for the task. Yang and Wang [57] train a policy to select an optimal

diffusion timestep, we simply used a fixed timestep per class of task. Several works [45, 46, 50] observe that the cross-attention layers that attend over the text and image embeddings encode a lot of the spatial layout associated with an image and therefore focus their method around tuning, post-processing, or extracting information embedded within these layers.

Visual Representation Learning for Control. Over the past decade, pre-trained representation learning approaches have been scaled for visual discrimination tasks first, and control tasks more recently. Contrastively pre-trained CLIP [33] representations were employed for embodied navigation tasks by EmbCLIP [21]. MAE representations have been used in control tasks by prior works like VC-1 [27], MVP [54] and OVRL-v2 [56]. R3M [30] and Voltron [20] leverage language supervision to learn visual representations. In contrast, we investigate if powerful text-to-image diffusion models trained for image generation can provide effective representations for control.

Diffusion Models for Control. Diffusion models have seen a wide range of uses in control aside from learning representations. These can broadly be categorized into three areas. First, diffusion models have been used as a class of expressive models for learning action distributions for policies [7, 14, 32]; They can improve model multimodality and generate richer action distributions than Gaussians. Second, off-the-shelf diffusion models have been used to augment limited robot demonstration datasets by specifying randomizations for object categories seen in the data through inpainting [19, 28, 60]. Third, planning can be cast as sequence modeling through diffusion models [1, 11, 18].

3 Stable Control Representations

In this paper, we investigate the use of language-guided visual representations from the open-source Stable Diffusion model (v1.5) and their application to language-conditioned visual control tasks. We present background on latent diffusion models and text-to-image diffusion models, along with the notation we adopt in this work, in Appendix B.

To extract representations, we follow a similar protocol as Wang et al. [50], Traub [48], and Yang and Wang [57]: Given an image-text prompt, s = { s image , s text } , associated with a particular task, we use the SD VQ-VAE model as the encoder E ( · ) and partially noise the encoded latents z 0 ˙ = E ( s image ) to some diffusion timestep t , to obtain the noised latent z t . We then extract a representation composed of the intermediate layer outputs of the U-Net ϵ θ as it produces a denoising estimate ϵ θ ( z t , t, s text ) . This process is illustrated in Figure 2. We refer to the extracted representations as Stable Control Representations (SCR) . In Sections 3.1, 3.2, 3.3, and 3.4, we describe the design space for extracting SCR, and in Sections 3.5 and 3.6, we explain how we use the representations for control.

3.1 Layer Selection and Aggregation

We are interested in evaluating the internal representations from the denoiser network, that is, the U-Net ϵ θ ( · ) . The first design choice we consider is which layers of ϵ θ to aggregate intermediate outputs from. The U-Net does not have a representational bottleneck, and different layers potentially encode different levels of detail. Trading off size with fidelity, we concatenate the feature maps output from the mid and down-sampling blocks to construct the representation. This results in a representation size comparable to that of the other pre-trained models we study in Section 4. This is shown at the bottom of Figure 2 and we ablate this choice in Section 5.1. Since outputs from different layers may have different spatial dimensions, we bilinearly interpolate them so that they are of a common spatial dimension and can be stacked together. We then pass them through a learnable convolutional layer to reduce the channel dimension before feeding them to downstream policies. The method used to spatially aggregate pre-trained representations can significantly affect their efficacy in downstream tasks, as we will discuss in Section 5.4. We use the best-performing spatial aggregation method for all the baselines that we re-train in Section 4.

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 x 0 ), whereas the SD U-Net expects noised latents, corresponding to z t for t ∈ [0 , 1000] . The choice of timestep t influences the fidelity of the encoded latents since a higher value means more noising of the inputs. Yang and Wang [57] have observed that there are task-dependent optimal timesteps

and proposed adaptive selection of t during training, while Xu et al. [55] have used t = 0 to extract representations using un-noised inputs to do open-vocabulary segmentation. We hypothesize that control tasks that require a detailed spatial scene understanding would benefit from a lower diffusion timestep, corresponding to a later stage in the denoising process where the inputs have less noise. We provide evidence consistent with this hypothesis in Section 5.2. To illustrate the effect of the timestep, we display final denoised images for various t values in different domains in Figure 8.

3.3 Prompt Specification

Since text-to-image diffusion models allow conditioning on text, we investigate if we can influence the representations to be more task-specific via this conditioning mechanism. For tasks that come with a text specifier, for example, the sentence 'go to object X', we simply encode this string and pass it to the U-Net. However, some tasks are purely vision-based and in these settings, we explore whether constructing reasonable text prompts affects downstream policy learning when using the U-Net's language-guided visual representations. We present this analysis in Section 5.3.

3.4 Intermediate Attention Map Selection

Wang et al. [50] and Tang et al. [45] demonstrate that the Stable Diffusion model generates localized attention maps aligned with text during the combined processing of vision and language modalities. Wang et al. [50] leveraged these word-level attention maps to perform open-domain semantic segmentation. We hypothesize that these maps can also help downstream control policies to generalize to an open vocabulary of object categories by providing helpful intermediate outputs that are category-agnostic. Following Tang et al. [45], we extract the cross-attention maps between the visual features and the CLIP text embeddings within the U-Net. We test our hypothesis on an open-domain navigation task in Section 4.3, where we fuse the cross-attention maps with the extracted feature maps from the U-Net. We refer to this attention-map-augmented representation as SCR-ATTN .

3.5 Using Text-to-Image Diffusion Model Representations to Learn Control Policies

To solve visual control tasks with states given by s = [ s image , s text ] , where s text may be used to specify the task, we wish to use pre-trained vision-language representations capable of encoding the state s as f ϕ ( s image , s text ) . This encoded state is then supplied to a downstream, task-specific policy network, which is trained to predict the action a t . Our evaluation encompasses both supervised learning and reinforcement learning regimes for training the downstream policies. We train agents through behavior cloning on a small set of demonstrations for the few-shot manipulation tasks we study in Section 4.1. For the indoor navigation tasks we study in Sections 4.2 and 4.3, we use a version of the Proximal Policy Optimization [PPO, 39] algorithm for reinforcement learning.

Train

Val

2U

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 training data (e.g., LAION images) and robotics datasets' visual inputs (e.g., egocentric tabletop views in manipulation tasks or indoor settings for navigation). Crucially, we do not require any task-specific losses for fine-tuning. Instead, we adopt the same text-conditioned generation objective as that of the base model for the fine-tuning phase. We use a small subset of the collection of datasets used by prior works on representation learning for embodied AI [27, 54]: we use subsets of the EpicKitchens [9], Something-Something-v2 [SS-v2; 13], and Bridge-v2 [49] datasets. As is standard, we fine-tune the denoiser U-Net ϵ θ but not the VAE encoder or decoder. Image-text pairs are uniformly sampled from the video-text pairs present in these datasets. A possible limitation of this strategy is that text-video aligned pairs (a sequence of frames that correspond to a single language instruction) may define a many-to-one relation for image-text pairs. However, as we see in experiments in which we compare to the base Stable Diffusion model in Section 4, this simple approach to robotics alignment is useful in most cases. Further details related to fine-tuning are provided in Appendix E.1. We refer to the representations from this fine-tuned model as SCR-FT .

4 Empirical Evaluation

In this work, we evaluate Stable Control Representations (SCR) on an extensive suite of tasks from 6 benchmarks covering few-shot imitation learning for manipulation in Section 4.1, reinforcement learning-based indoor navigation in Sections 4.2 and 4.3, and tasks related to fine-grained visual prediction in Appendices D.2 and D.3. Together, these tasks allow us to comprehensively evaluate whether our extracted representations can encode both high and low-level semantic understanding of a scene to aid downstream policy learning. We describe the common baselines used across tasks in Appendix C, and present the individual task setups and results in the following subsections.

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 Franka-Kitchen environments included in the RoboHive suite [23], which entail tasks such as turning a knob or opening a door. We adhere to the training and evaluation protocols adopted in their respective prior works to ensure our results are directly comparable (detailed further in Appendix G.1).

Results. We report the best results of SCR and baselines in Table 1a. On Meta-World, we see that SCR outperforms most prior works, achieving 94.9% success rate. In comparison, VC-1, the visual foundation model for embodied AI and CLIP achieved 92.3 and 90.1% respectively. On Franka-Kitchen, SCR obtains 49.9% success rate, which is much higher than CLIP (36.3%) and again outperforms all other baselines except for R3M. We note that R3M's sparse representations excel in few-shot manipulation with limited demos but struggle to transfer beyond this setting [27, 20]. We see that while the SD-VAE encoder performs competitively on Franka-Kitchen, it achieves a low success rate on Meta-World. This observation allows us to gauge the improved performance

of SCR from the base performance gain we may get just from operating in the latent space of this V AE. Additionally, we see that the task-agnostic fine-tuning gives SCR-FT an advantage (4%) over SCR on Franka-Kitchen while making no difference on Meta-World. Note that the other high-performing baselines (R3M and Voltron) have been developed for downstream control usage with training objectives that take temporal information into account, while VC-1 has been trained on a diverse curation of robotics-relevant data. In this context, SCR's comparable performance shows that generative foundation models hold promise for providing useful representations for control, even with relatively minimal fine-tuning on non-task-specific data.

4.2 Image-Goal Navigation

We now assess SCR in more realistic visual environments, surpassing the simple table-top scenes in manipulation benchmarks. In these complex settings, the representations derived from pre-trained foundational models are particularly effective, benefiting from their large-scale training. We study Image-Goal Navigation (ImageNav), an indoor visual navigation task that evaluates an agent's ability to navigate to the viewpoint of a provided goal image [63]. The position reached by the agent must be within a 1-meter distance from the goal image's camera position. This requires the ability to differentiate between nearby or similar-looking views within a home environment. This task, along with the semantic object navigation task that we study in Section 4.3, allows for a comprehensive evaluation of a representation's ability to code both semantic and visual appearance-related features in completely novel evaluation environments. We follow the protocol for the ImageNav task used by Majumdar et al. [27] and input the pre-trained representations to an LSTM-based policy trained with DD-PPO [52] for 500 million steps on 16 A40 GPUs (further details in Appendix G.3). Given the large compute requirements for training, we directly compare SCR and SCR-FT to the results provided in Majumdar et al. [27].

Results. We evaluate our agent on 4200 episodes in 14 held-out scenes from the Gibson dataset and report the success rate in Table 1b. We find that SCR outperforms all other representations, while the fine-tuned version SCR-FT is almost on par with the second-best-performing VC-1 (69.5% vs 70.3%), the SOTA visual representation from prior work. This can be expected given that it was fine-tuned on images solely from table-top manipulation datasets. We also see that R3M, the best model for few-shot manipulation from Table 1a performs very poorly (30.6%) in this domain, showing its limited transferability to navigation tasks.

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-vocabulary object detectors such as Detic [62] to enable generalization to novel objects. We study a modified version of the Gaze sub-task (detailed in Appendix G.2), which focuses on locating a specified object category for an abstracted grasping action.

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 a policy with access to ground truth object masks, evaluated using either the ground truth or Detic-generated masks at test time (labeled as Oracle/Detic).

Results. Table 1c shows that SCR-FT surpasses VC-1 by 1.3%, beating CLIP and SCR by 3.2%. It is surprising that VC-1's visual representation does better than CLIP's image encoder representation, given that the downstream policy has to use these with the CLIP text encoder's embedding of the target object category. Comparing these to SCR-FT-ATTN, we can see the benefit of providing intermediate outputs in the form of text-aligned attention maps to the downstream policy (+1.7%). Samples of attention maps overlaid on images from an evaluation episode can be found in Appendix G. These word-level cross-attention maps simultaneously improve policy performance and also aid explainability, allowing us to diagnose successes and failures. Interestingly, the foundation model representations (CLIP, VC-1, SCR) perform better than Detic. While object detections serve as a category-agnostic input for downstream pick-and-place policies, noisy detections can often lead to degraded downstream performance, as we see in this case. Nonetheless, there is still a sizeable gap to 'Oracle' which benefits from ground truth object masks at test-time.

5 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 is influenced by the selection of layers from which we extract feature maps. We previously chose outputs from the mid and downsampling layers of the U-Net (Figure 2), because their aggregate size closely matches the representation sizes from the ViT-based models (VC-1, MVP, and CLIP). Appendix E.2 details the feature map sizes obtained for all the models we study. Table 2a lists the success rates achieved on the Franka-Kitchen domain when we use different sets of block outputs in SCR. We present similar ablations for Meta-World in the top four rows of Table 3b.

We observe that utilizing outputs from multiple layers is instrumental to SCR's high performance. This finding underscores a broader principle applicable to the design of representations across different models: Leveraging a richer set of features from multi-layer outputs should enhance performance on downstream tasks. However, it is important to acknowledge the practical challenges in applying this strategy to ViT-based models due to the high dimensionality of each layer's patch-wise embeddings (16 × 16 × 1024 for ViT-L for images of size 224 × 224). We present the success rates achieved on the four benchmarks when aggregating multi-layer embeddings from CLIP models in Tables 3a and 4, alongside SCR (the representation size for which is now half in comparison). In Table 3a, we observe that moving towards middle layers leads to higher performance indicating that CLIP layers 10-14 encode some details useful to the Franka-Kitchen benchmark. While we see benefits from including the output of certain additional layers, it still underperforms SCR.

5.2 Sensitivity to the Noising Timestep

Next, we characterize the sensitivity of task performance to the denoising step values chosen during representation extraction. We present results on the Franka-Kitchen tasks in Table 2b, and on the Meta-World tasks in the bottom three rows of Table 3b. We see that the performance across nearby timesteps (0 and 10 or 100 and 110) is similar, and that there is a benefit to doing a coarse grid search up to a reasonable noising level (0 vs 100 vs 200) to get the best value for a given task.

5.3 How is Language Guiding the Representations?

Recall that in the OVMM experiments (Section 4.3), we concatenated the target object's CLIP text embedding to the visual representations before feeding it to the policy. For SCR and SCR-FT, we also provided the category as the text prompt to the U-Net, and additionally extracted the generated cross-attention maps for SCR-FT-ATTN. In this subsection, we seek to more closely understand how the text prompts impact the representations in SCR.

We first consider the Franka-Kitchen setup from Section 4.1, which includes manipulation tasks that do not originally come with a language specification. We experiment with providing variations of task-relevant and irrelevant prompts during the representation extraction in SCR. Table 2c shows the downstream policy success rates for irrelevant ( 'an elephant in the jungle' ) and relevant ( 'a Franka robot arm opening a microwave door' ) prompts, compared to the default setting of not providing a text prompt We see that providing a prompt does not help with downstream policy performance and may even degrade performance as the prompt gets more irrelevant to the visual context of the input.

We now move to the Referring Expressions Grounding task which requires predicting a bounding box for an object being referred to in a sentence, within an image depicting cluttered objects. We defer the main presentation of this task to Appendix D.2 and use it to probe the degree of language grounding in SCR in this section. To study the role of the U-Net in shaping the visual representations guided by the text, we examine different text integration methods to generate SCR representations in Table 5.

We compared the following approaches for providing the task's text specification to the task decoder (also depicted in Figure 4):

(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'.

(e) Only text encoding: Ignore the visual representation and rely only on CLIP text embeddings.

Investigating the results of (a) and (b) in Table 5, it is evident that incorporating the text prompt into the U-Net significantly enhances accuracy compared to ignoring the text altogether. The difference in scores between (b) and (c) indicates that directly providing text embeddings to the decoder improves performance, suggesting that certain crucial aspects of object localization are not fully captured by the representation alone. Comparing (c) to (d), we see that with concatenated text embeddings, further modulation of the visual representations does not provide significant benefits. Finally, the significant decrease in the score for (e) reveals the extent to which the task relies on text-based guesswork.

These findings align with both intuition and recent research on controllable generation with diffusion models [61] that highlights the challenges associated with using long-form text guidance. There are ongoing research efforts, focused on training models with more detailed descriptions or leveraging approaches to encode and integrate sub-phrases of long texts, that seek to address these challenges.

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 'compression layer' by Yadav et al. [56], aims to reduce the high channel dimension of pre-trained model outputs without losing spatial details, facilitating more effective input processing by downstream task-specific decoders.

We explore the effect of spatial aggregation methods by comparing the convolutional downsampling layer method to multi-headed attention pooling (MAP) used for CLIP embeddings in Karamcheti et al. [20]. We find that using a compression layer significantly improves performance on the fine-grained visual prediction tasks described in Appendix D as reported in Table 6 (columns 3-4). This result challenges the conjecture made in prior work that CLIP representations are limited in their ability to provide accurate low-level spatial information [20] and emphasizes the critical role of appropriate representation aggregation.

Building on this result, we assess whether better spatial aggregation can improve the performance of CLIP representations on downstream control tasks. We present these results in Table 6 (columns 5-6) for VC-1 and CLIP on the MuJoCo tasks. We see that the compression layer often outperforms the use of CLS token embeddings (by 1-2%), but CLIP representations still fail to match the best-

performing models. This result provides evidence that the underperformance of CLIP representations on control tasks is unlikely due to a lack of sufficiently fine-grained visual information. Finally, we note that the compression layer aggregation technique was used for all baselines in Tables 1b and 1c to ensure a strong baseline comparison. We recommend that future studies adopt this methodology to enable a fairer comparison of representations.

6 Discussion

In Section 5, we deconstructed Stable Control Representations and highlighted techniques used in our approach that could be applied to extract representations from other foundation models. Our analysis in Sections 5.1 and 5.4 revealed that using multi-layer features and appropriate spatial aggregation significantly affects performance, and overlooking these factors can lead to misleading conclusions about the capabilities of previously used representations.

Next, our investigation into how language prompts shape diffusion model representations uncovered nuanced results and showed that text influence on representations does not consistently increase downstream utility. This is particularly evident in tasks where text specification is not required and where training and test environments are congruent, minimizing the need for semantic generalization. In contrast, tasks like referring expressions grounding demonstrate the necessity of direct access to text embeddings for accurate object localization, even when representations are modulated to considerable success. For the OVMM task, we identified a scenario where multimodal alignment is essential and proposed a method to explicitly utilize the latent knowledge of the Stable Diffusion model through text-aligned attention maps, which is not straightforward to do for other multimodal models. Future research could design methods to derive precise text-associated attribution maps for other models.

Finally, we contrasted the simplicity of fine-tuning diffusion models with that of the contrastive learning objective required to fine-tune CLIP representations. The former only requires image-text samples for the conditional generation objective, whereas the latter requires a sophisticated negative label sampling pipeline along with large batch sizes to prevent the model from collapsing to a degenerate solution [33]. We demonstrate this phenomenon empirically on the Franka-Kitchen environment by fine-tuning CLIP similarly to SCR-FT in Appendix D.1.

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-image diffusion models for policy learning can improve generalization across a wide range of tasks including manipulation, image-goal and object-goal based navigation, grasp point prediction, and referring expressions grounding. We also demonstrated the interpretability benefits of incorporating attention maps extracted from pre-trained text-to-image diffusion models, which we showed can improve performance and help identify downstream failures of the policy during development. Finally, we discussed ways in which the insights presented in this paper, for example, regarding feature aggregation and fine-tuning, may be applicable to other foundation models used for control. We hope that Stable Control Representations will help advance data-efficient control and enable open-vocabulary generalization in 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.

References

Anurag Ajay, Seungwook 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 arXiv:2309.08587 , 2023.

Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, Iain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: a visual language model for few-shot learning. Advances in Neural Information Processing Systems , 35: 23716-23736, 2022.

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.

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 J Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Isabel Leal, Kuang-Huei Lee, Sergey Levine, Yao Lu, Utsav Malla, Deeksha Manjunath, Igor Mordatch, Ofir 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:2212.06817, December 2022. doi: 10 . 48550/arXiv . 2212 . 06817.

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-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.

Xi Chen, Xiao Wang, Soravit Changpinyo, AJ Piergiovanni, Piotr Padlewski, Daniel Salz, Sebastian Goodman, Adam Grycner, Basil Mustafa, Lucas Beyer, et al. Pali: A jointly-scaled multilingual language-image model. arXiv preprint arXiv:2209.06794 , 2022.

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.

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.

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.

Laura Downs, Anthony Francis, Nate Koenig, Brandon Kinman, Ryan Hickman, Krista Reymann, Thomas B McHugh, and Vincent Vanhoucke. Google scanned objects: A high-quality dataset of 3d scanned household items. In 2022 International Conference on Robotics and Automation (ICRA) , pages 2553-2560. IEEE, 2022.

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.00111 , 2023.

Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 12873-12883, June 2021.

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 58425850, 2017.

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.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications , 2021.

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.

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.

Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning , 2022.

Ivan Kapelyukh, Vitalis Vosylius, and Edward Johns. Dall-e-bot: Introducing web-scale diffusion models to robotics. IEEE Robotics and Automation Letters , 2023.

Siddharth Karamcheti, Suraj Nair, Annie S. Chen, Thomas Kollar, Chelsea Finn, Dorsa Sadigh, and Percy Liang. Language-driven representation learning for robotics. In Robotics: Science and Systems (RSS) , 2023.

Apoorv Khandelwal, Luca Weihs, Roozbeh Mottaghi, and Aniruddha Kembhavi. Simple but effective: Clip embeddings for embodied ai. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , June 2022.

Mukul Khanna, Yongsen Mao, Hanxiao Jiang, Sanjay Haresh, Brennan Shacklett, Dhruv 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.

Vikash Kumar, Rutav Shah, Gaoyue Zhou, Vincent Moens, Vittorio Caggiano, Jay Vakil, Abhishek Gupta, and Aravind Rajeswaran. Robohive - a unified framework for robot learning. In NeurIPS: Conference on Neural Information Processing Systems , 2023.

Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. arXiv preprint arXiv:2304.08485 , 2023.

Yecheng Jason Ma, William Liang, Vaidehi Som, Vikash Kumar, Amy Zhang, Osbert Bastani, and Dinesh Jayaraman. LIV: Language-Image Representations and Rewards for Robotic Control. arXiv e-prints , art. arXiv:2306.00958, June 2023. doi: 10 . 48550/arXiv . 2306 . 00958.

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-32352, 2022.

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 arXiv:2303.18240 , 2023.

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 arXiv:2212.05711 , 2022.

Lina Mezghani, Sainbayar Sukhbaatar, Thibaut Lavril, Oleksandr Maksymets, Dhruv Batra, Piotr Bojanowski, and Karteek Alahari. Memory-augmented reinforcement learning for imagegoal navigation. arXiv preprint arXiv:2101.05181 , 2021.

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.

OpenAI. GPT-4 technical report, 2023.

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.10677 , 2023.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, 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.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10684-10695, 2022.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Nassir Navab, Joachim Hornegger, William M. Wells III, and Alejandro F. Frangi, editors, Medical Image Computing and Computer-Assisted 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.

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 , 35:36479-36494, 2022.

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 on Computer Vision (ICCV) , 2019.

Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et al. LAION5B: An open large-scale dataset for training next generation image-text models. Advances in Neural Information Processing Systems , 35:25278-25294, 2022.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal Policy Optimization Algorithms. arXiv e-prints , art. arXiv:1707.06347, July 2017. doi: 10 . 48550/arXiv . 1707 . 06347.

Dhruv Shah, Blazej Osi´ nski, 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.

Mohit Shridhar, Jesse Thomason, Daniel Gordon, Yonatan Bisk, Winson Han, Roozbeh Mottaghi, 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.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning , pages 2256-2265. PMLR, 2015.

Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction . The MIT Press, second edition, 2018.

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 (NeurIPS) , 2021.

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 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , 2023.

Junjiao Tian, Lavisha Aggarwal, Andrea Colaco, Zsolt Kira, and Mar Gonzalez-Franco. Diffuse, attend, and segment: Unsupervised zero-shot segmentation using stable diffusion. arXiv preprint arXiv:2308.12469 , 2023.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothé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.

Jeremias Traub. Representation learning with diffusion models. arXiv preprint arXiv:2210.11058 , 2022.

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.

Jinglong Wang, Xiawei Li, Jing Zhang, Qingyuan Xu, Qin Zhou, Qian Yu, Lu Sheng, and Dong Xu. Diffusion Model is Secretly a Training-free Open Vocabulary Semantic Segmenter. arXiv e-prints , art. arXiv:2309.02773, September 2023. doi: 10 . 48550/arXiv . 2309 . 02773.

Ke-Jyun Wang, Yun-Hsuan Liu, Hung-Ting Su, Jen-Wei 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 of the Association for Computational Linguistics: Human Language Technologies , pages 5333-5338, Online, June 2021. Association for Computational Linguistics. doi: 10 . 18653/v1/2021 . naaclmain . 419.

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 preprint arXiv:1911.00357 , 2019.

Fei Xia, Amir R Zamir, Zhiyang He, Alexander Sax, Jitendra Malik, and Silvio Savarese. Gibson env: Real-world perception for embodied agents. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 9068-9079, 2018.

Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. arXiv:2203.06173 , 2022.

Jiarui Xu, Sifei Liu, Arash Vahdat, Wonmin Byeon, Xiaolong Wang, and Shalini De Mello. ODISE: Open-Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models. arXiv preprint arXiv: 2303.04803 , 2023.

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.

Xingyi Yang and Xinchao Wang. Diffusion model as representation learner. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 18938-18949, 2023.

Sriram Yenamandra, Arun Ramachandran, Karmesh Yadav, Austin Wang, Mukul Khanna, Theophile Gervet, Tsung-Yen Yang, Vidhi Jain, Alexander William Clegg, John Turner, et al. Homerobot: Open-vocabulary mobile manipulation. arXiv preprint arXiv:2306.11565 , 2023.

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 on Robot Learning (CoRL) , 2019.

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.11550 , 2023.

Jing Zhao, Heliang Zheng, Chaoyue Wang, Long Lan, and Wenjing Yang. Magicfusion: Boosting text-to-image generation performance by fusing diffusion models. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) , pages 22592-22602, October 2023.

Xingyi Zhou, Rohit Girdhar, Armand Joulin, Philipp Krähenbühl, and Ishan Misra. Detecting twenty-thousand classes using image-level supervision. In ECCV , 2022.

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.

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-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-action models transfer web knowledge to robotic control. In 7th Annual Conference on Robot Learning , 2023.

Appendix

Table of Contents

A Wider Societal Impacts

This paper presents work whose goal is to advance the field of robotic control using pretrained foundation models. Since we primarily aim to learn policies for simulated control tasks, we do not anticipate any ethical concerns at this stage. However, when deploying foundation models in the real world, we should be wary that our policies do not inherit any structural biases from the pretraining stage or internet-scale data.

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 ( x 0 ) , starting from pure noise. Given p ( x 0 ) and a set of noise levels σ t for t = 1 , . . . , T , a denoising function ϵ θ ( x t , t ) is trained on the objective

where ϵ ∼ N (0 , 1) , t ∼ Unif(1 , T ) , and x 0 ∼ p ( x 0 ) . To generate a sample x 0 during inference, we first sample an initial noise vector x T ∼ N (0 , σ T ) and then iteratively denoise this sample for t = T, ..., 1 by sampling from p ( x t -1 | x t ) , which is a function of ϵ θ ( x t , t ) .

In some settings, we may want to generate samples with a particular property. For example, we may wish to draw samples from a conditional distribution over data points, p ( x 0 | c ) , where c captures some property of the sample, such as classification label or a text description [34, 36]. In these settings, we may additionally train with labels to obtain a conditioned denoiser ϵ θ ( x t , t, c ) and generate samples using classifier-free guidance [15].

B.2 Latent Diffusion Models

Latent diffusion models [34] reduce the computational cost of applying diffusion models to highdimensional data by instead diffusing low-dimensional representations of high-dimensional data. Given an encoder E ( · ) and decoder D ( · ) , Equation (B.1) is modified to operate on latent representations, z 0 ˙ = E ( x 0 ) , yielding

where ϵ ∼ N (0 , 1) , t ∼ Unif(1 , T ) , x 0 , c ∼ p ( x 0 , c ) . After generating a denoised latent representation z 0 , it can be decoded as x 0 = D ( z 0 ) .

A popular instantiation of a conditioned latent diffusion model is the text-to-image Stable Diffusion model [SD; 34]. The SD model is trained on the LAION-2B dataset [38] and operates in the latent space of a pre-trained VQ-VAE image encoder [12]. The model architecture is shown at the top of Figure 1 and is based on a U-Net [35], with the corresponding conditioning text prompts encoded using a CLIP language encoder [33].

B.3 Policy Learning for Control

Wemodel our environments as Markov Decision Processes (MDP, Sutton and Barto [43]), defined as a tuple M = ( S , A , P, R, γ ) , where S and A denote the state and action spaces respectively, P ( s ′ | s, a ) the transition dynamics, R ( s, a ) the reward function, and γ ∈ (0 , 1) the discount factor. Our goal is to optimize a policy π ( a | s ) that maximizes the expected discounted return E π,P [ ∑ ∞ t =0 γ t R ( s t , a t )] .

In this paper, we consider visual control tasks that may be language-conditioned, that is, states are given by s = [ s image , s text ] , where s text specifies the task. We are interested in pre-trained visionlanguage representations capable of encoding the state s as f ϕ ( s image , s text ) . This encoded state is then supplied to a downstream, task-specific policy network, which is trained to predict the action a t . Our evaluation encompasses both supervised learning and reinforcement learning regimes for training the downstream policies. We train agents through behavior cloning on a small set of demonstrations for the few-shot manipulation tasks we study in Section 4.1. For the indoor navigation tasks we study in Sections 4.2 and 4.3, we use a version of the Proximal Policy Optimization [PPO, 39] algorithm for reinforcement learning.

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.

MVP [54] and VC-1 [27] both pre-train ViT-B/L models with the masked auto-encoding (MAE) objective on egocentric data from Ego4D, Epic-Kitchens, SS-v2, and ImageNet, with VC-1 additionally pre-training on indoor navigation videos.

CLIP [33] trains text and ViT-based image encoders using contrastive learning on web-scale data.

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 V AE encoder used by SD to encode images into latents.

To assess how well the vision-only methods would do on tasks with language specification, we concatenate their visual representations with the CLIP text embeddings of the language prompts. While we are limited by the architecture designs of the released models we are studying, to ensure a more fair comparison we try to match parameter counts as much as we can. We use the ViT-Large (307M parameters) versions of CLIP, MVP, and VC-1 since extracting SCR involves a forward pass through 400M parameters.

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 a video closer to the video caption's text embedding. Specifically, for our experiment, we use the huggingface CLIP finetuning implementation and train the model with a batch size

of 384 (the maximum number of samples we were able to fit on 8 A40 GPUs) with a learning rate of 5e-5 and a weight decay of 0.001 for 5000 update steps (same as SR-FT). We present the results in Table 7 for Franka-Kitchen, and note the lack of improvement on the task post-fine-tuning.

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 our attention to two downstream tasks involving fine-grained visual prediction. The first task, Referring Expressions Grounding, is detailed within this section, while the second task, Grasp Affordance Prediction, is discussed in Appendix D.3. Karamcheti et al. [20] have previously examined the performance on these tasks as proxy measures to evaluate the efficacy of representations for control applications.

The Referring Expressions Grounding task requires the identification and bounding box prediction of an object in an image based on its textual description. Similar to Karamcheti et al. [20], we use the OCID-Ref Dataset [51] for our experiments. We show a sample image-text pair from the dataset to showcase the complexity of the task in Figure 5. The frozen visual representation is concatenated with a text embedding and passed to a 4-layer MLP, which predicts the bounding box coordinates.

Results. We report the bounding box accuracy at a 25% Intersection-over-Union (IoU) threshold across different scene clutter levels for SCR-variants and baselines in Table 8. Firstly, we present the results from [20] for CLIP, R3M and Voltron and observe that Voltron outperforms the other two baselines. We see that SCR is tied with Voltron and that VC-1 and SD-VAE perform the best with a 1.5% lead. The better performance of these vision-encoder-only methods highlights that on this task, it is not a challenge for the downstream decoder to learn to associate the visual embeddings with the (CLIP) text encoding of the language specification. Since the training budget is fixed, we observed

that some of the runs could potentially improve over extended training. However, we were primarily interested in this task not just to compare the downstream visual prediction performance, but to use it as a testbed for exploring the following two questions: (1) Do the performance differences between the representations we evaluated in Sections 4.1 to 4.3, stem from the absence of fine-grained spatial information encoded within the representations? We refute this claim in Section 5.4, where we present the impact of the representations' spatial aggregation method on prediction performance. (2) Additionally, we explore to what extent language prompting influences the representations from SCR on language-conditioned tasks in Section 5.3.

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 certain areas of objects in an RGB image. These areas correspond to parts of the surface that would be amenable to grasping by a suction gripper. The evaluation metric adopted in prior work is the precision of predictions corresponding to positive graspability at varying confidence levels (90, 95, and 99th percentile of the predicted perpixel probabilities, denoted as Top90, Top95, and Top99 in Table 9). We refer the reader to Karamcheti et al. [20] for the complete task setup details.

We re-ran all the methods using the evaluation repository provided with the work, and obtained different results compared to the reported numbers in Karamcheti et al. [20], which we attribute to a bug that we fixed related to the computation of the precision metrics. The evaluation procedure for this task adopted in prior work involves a 5-fold cross-validation, and we observed a high variability in the results, with different runs of 5-fold cross-validation yielding different final test metrics. Our findings highlight that SCR and our adaptation of CLIP (in gray, detailed in Section 5.4) both excel at this task, achieving a Top99 score of 72.9. Interestingly, we see that fine-tuning did not enhance the performance of SCR on the visual prediction tasks explored in this section and Section 4.1, suggesting a potential disconnect between visual prediction and control task benchmarks.

D.4 Comparison with LIV

We include a comparison with LIV [25] on two tasks that involve manipulation and navigation. LIV is a visionlanguage representation learned through contrastive learning on the EpicKitchens dataset [9]. Similar to R3M results in the main paper, this representation does well on manipulation tasks but poorly on navigation tasks.

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 to give a higher-level view of the all-round performance of the different representations on the diverse set of tasks we consider. We see that VC-1, SCR, and SCR-FT emerge as the top three visual representations overall. While VC-1 is a representation-learning foundation model trained specifically for robotics tasks, SCR and SCR-FT are the diffusion model representations that we study in this paper, confirming the potential of large pre-trained foundation generative models across a wide array of downstream robotics tasks.

E Implementation Details

E.1 Fine-tuning Stable Diffusion

We used the runwayml/stable-diffusion-v1-5 model weights provided by Huggingface to initialize our models and fine-tuned them using the diffusers library 3 . As noted in Section 3.6, we used a subset of the frames from the EpicKitchens, Something-Something-v2, and Bridge-v2 datasets. We extracted the middle-third of the video clips and sampled four frames randomly from this chunk to increase the chances of sampling frames where the text prompt associated with the video clip is most relevant for describing the scene. Since the Something-Something-v2 (SS-v2) dataset includes a human hand manipulating table-top objects while the Bridge-V2 dataset includes a robot gripper, we append the text "human hand" and "robot hand" respectively to the text caption associated eith the respective frames to reduce ambiguity in the image-text pairing. Using this procedure, we generated a paired image-language dataset with 1.3 million samples. Figure 6 shows samples of the images from the fine-tuning datasets. Since different embodiments (human and robot) are visible in the training images, we prepended the corresponding embodiment name to the text prompt for the associated image during training.

We fine-tuned for only a single epoch ( ∼ 5,000 gradient steps) parallely on 2 Nvidia A100 GPUs with a mini-batch size of 512 and a learning rate of 1 e -4 . Although the original Stable Diffusion model is trained on images of resolution 512 × 512 , we fine-tuned the model on images downscaled to 256 × 256 since it aligned with the resolution requirements of the downstream application. We show sample generations from the diffusion model after fine-tuning in Figure 7. We found that the model learns to associate the prompt with not just the human or robot hand but also with the style of the background and objects of the training datasets.

Epic-Kitchens

Something-Somethingv2

Bridgev2

Human hand

Robothand

tossing aboxinthetrash

openingadrawer

tossingaboxinthetrash

openingadrawer

?

slicing a cucumberwith aknife

wipingthekitchencounter

slicing a cucumberwithaknife

wipingthekitchencounter

E.2 Representation Extraction Details

Here, we describe the representation extraction details for all our baselines assuming a 224 × 224 input image.:

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 × 256 . We pass the image to the VAE, which converts it into a 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 × 640 , 4 × 4 × 1 , 280 , 4 × 4 × 1 , 280 , and 4 × 4 × 1 , 280 , respectively. The total size is, therefore, 102 , 400 , and we linearly interpolated them to the same spatial dimension ( 8 × 8 ) before concatenating them channel-wise.

R3M [30]: For most of our experiments we use the original ResNet50 model, which outputs a 2048 dimensional vector. For the referring expressions and grasp affordance prediction tasks from the Voltron evaluation suite [20], a VIT-S is used, which outputs an embedding of size 14 × 14 × 384 = 75 , 264

MVP [54] and VC-1 [27]: The last layer ( 24 th ) outputs an embedding of size 16 × 16 × 1 , 024 = 262 , 144 .

CLIP [33]: For ViT-B, the last layer ( 12 th ) outputs an embedding of size 14 × 14 × 768 = 150 , 528 . For ViT-L, the last layer ( 24 th ) outputs an embedding of size 16 × 16 × 1024 = 262 , 144 .

Voltron [20]: We use the VCond-Base model which outputs a representation of size 14 × 14 × 768 = 150 , 528 .

SD-VAE [34]: Outputs a 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.

F Runtime Analysis

We include a runtime analysis for SCR versus VC-1 in Table 13 by reporting the time taken for a forward pass through both models over a single input image of size 256 × 256. We use an A100 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.

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 a single A100 GPU with 24 CPUs and 188 GBs of RAM.

Meta-World. We follow Majumdar et al. [27] and use the hammer-v2, drawer-open-v2, bin-pickingv2, button-press-topdown-v2, assembly-v2 tasks from the Meta-World benchmark suite [59]. Each task provides the model with the last three 256 × 256 RGB images, alongside a 4-dimensional gripper pose. The model consists of a 3-layer MLP with a hidden dimension of 256, utilizing ReLU as the activation function. It undergoes training for 100 epochs, with evaluations conducted every 10 epochs, following the approach of Majumdar et al. (2023). The training uses a mini-batch size of 256 and a learning rate of 10 -3 .

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 a 256 × 256 RGB image and a 24-dimensional vector representing the manipulator's proprioceptive state. For our experiments, we use the the RoboHive Kumar et al. [23] GitHub repository 4 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 a learning rate of 10 -4 . We additionally correct a bug in the RoboHive implementation of the VC-1 baseline, specifically on input image normalization. Adjusting the image normalization to a 0-1 range resulted in a significant improvement in its performance.

G.2 Open Vocabulary Mobile Manipulation

Open-Vocabulary Mobile Manipulation [OVMM; 58] is a recently proposed embodied AI benchmark that evaluates an agent's ability to find and manipulate objects of novel categories in unseen indoor environments. Specifically, the task requires an agent to ' Find and pick an object on the start_receptacle and place it on the goal_recetacle ', where object , start_receptacle and goal_recetacle are the object category names. Given the long-horizon and sparse-reward nature of this task, current baselines [58] divide the problem into sub-tasks, which include navigation to the start receptacle, precise camera re-orientation to focus on the object (an abstracted form of grasping), navigating to the goal receptacle, and object placement.

Since our aim is to investigate the open-vocabulary capabilities of pre-trained representations, we choose to evaluate the models on only the precise camera re-orientation task (more commonly known as the Gaze task). In the original Gaze task, the agent is initialized within a distance of 1.5m and angle of 15 ◦ from the object which is lying on top of the start_receptacle . The episode is deemed successful when the agent calls the Pick action with the camera's center pixel occupied by the target object and the robot's gripper less than 0.8m from the object center. In our initial experiments, we found the current initialization scheme would lead the agent to learn a biased policy. This policy would call the Pick action after orienting towards the closest object in the field of view. Therefore, we chose to instantiate a harder version of the gaze task, where the episode starts with the agent spawned facing any random direction within 2.0m of the object.

We carry out our experiments using the HomeRobot GitHub repository 5 . HomeRobot uses the Habitat simulator [44] with the episode dataset provided by 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 a height of 1.41 meters and a radius of 0.3 meters. At a height of 1.31 meters from the base, a 640x480 resolution RGBD camera is mounted. This camera is equipped with motorized pan and tilt capabilities. The agent's action space

is continuous, allowing it to move forward distances ranging from 5 to 25 centimeters and to turn left or right within angles ranging from 5 to 30 degrees. Additionally, the agent can adjust the head's pan and tilt by increments ranging from 0.02 to 1 radian in a single step.

In our experiments, we use a 2 layer LSTM policy and pass in the visual encoder representations after passing them through the compression layer. We initialize the LSTM weights with the LSTM weights of the Oracle model to get a slight boost in performance. We train our agents using the distributed version of PPO [52] with 152 environments spread across 4 80GB Nvidia A100 GPUs. Each run also has access to 96 CPUs and 754 GBs of RAM. We train for 100M environment steps while evaluating the agent every 5M steps and report the metrics based on the highest success rate observed on the validation set.

G.3 ImageNav

We conduct our ImageNav experiments in the Habitat simulator [37], using the episode dataset from Mezghani et al. [29]. These experiments are conducted using the VC-1 codebase 6 [27]. The dataset uses 72 training and 14 validation scenes from the Gibson [53] scene dataset with evaluation conducted on a total of 4200 episodes. The agent is assumed to be in the shape of a cylinder of height 1.5m and radius 0.1m, with an RGB camera mounted at a height of 1.25m from the base. The RGB camera has a resolution of 128×128 and a 90 ◦ field-of-view.

At the start of each training episode, an agent is randomly initialized in a scene and is tasked to find the position from where the goal image was taken within 1000 simulation steps. At each step, the agent receives a new observation and is allowed to take one of the four discrete actions including MOVE_FORWARD (25cm), TURN_LEFT ( 30 ◦ ), TURN_RIGHT ( 30 ◦ ) and STOP. The episode is a success if the agent calls the STOP action within 1m of the goal viewpoint. Similar to [56, 27] we train our agents using a distributed version of DD-PPO [52] with 320 environments for 500M timesteps (25k updates). Our experiments are conducted using 2 nodes containing 8 A40 GPUs each with a total of 128 CPUs and 504 GBs of RAM. Each environment accumulates experience across up to 64 frames, succeeded by two PPO epochs using two mini-batches. While the pre-trained model is frozen, the policy is trained using the AdamW optimizer, with a learning rate of 2 . 5 × 10 -4 and weight decay of 10 -6 . We use the reward function proposed by [56]. Performance is assessed every 25M training steps, with reporting metrics based on the highest success rate observed on the validation set.

t=200

t=300

Noised

Generated

Meta-World

Noised

Generated

OVMM

Noised

Generated

Refer Expression

Refer Expression

“A"

"pear"

"beside"

"book"

OVMM

“A”

"rocket"

"chair"

(1) Find Tomato

(2) Find Gaming Console

(3) Find Plant Container

(4)Find SprayBottle

(5)FindPlant Container

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]

Justification: The abstract and introduction accurately summarise the two main contributions of the paper: (1) a procedure to derive vision-language 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-art representations in terms of downstream policy learning.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: In Section 6 we discuss the interpretation of our experimental results, and highlight limitations of our work. We discuss the dependence of downstream performance on design choices related to representation extraction. We also include a 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 set 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 experimental results of the paper to the extent that it 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 a detailed description of our experimental procedure in Appendices E 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 to the data and code, with sufficient instructions 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 as well as 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 use are publicly available, and we include the sources and links to these in the code documentation as well as the appendices.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide details related to the finetuning of Stable Diffusion model and extraction of representations in Section 5 and appendix E, and details related to hyperparameters, 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 associated with the reported numbers over 3 seeds. Only the experiments in Section 4.2 (ImageNav) report numbers that involve a single random seed, since we were limited by compute resources on this benchmark as we specify in Section 4.2.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips . cc/public/EthicsGuidelines ?

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

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We build on the already open-sourced image generation model, Stable Diffusion and finetune it on images of human hand or 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]

Justification: Code: 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 all 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 (creativeml-openrail-m).

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We will open-source the code included in the supplementary material. The code is well-documented, and can be used to reproduce our experimental results in Section 4, as well as be used to extract diffusion representations for generic visual(-language) tasks by independent researchers..

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 as 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nOnly 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.\n\nQUESTIONS:\nWhy 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?\n\nLIMITATIONS:\nThere are not many concerns to address and the authors wrote an appropriate short section in the Appendix.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nPlease kindly refer to the weaknesses.\n\nLIMITATIONS:\nYes, 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

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

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

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

QUESTIONS:
None

LIMITATIONS:
Refer to the weaknesses section.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}