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

Gunshi Gupta,Karmesh Yadav,Yarin Gal,Dhruv Batra,Zsolt Kira,Cong Lu,Tim G. J. Rudner

2024-05-09

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.

Computer Vision and Pattern Recognition,Artificial Intelligence,Machine Learning,Robotics

What problem does this paper attempt to address?

The problem this paper attempts to address is: how to utilize pre-trained text-to-image diffusion models to generate visual-language representations suitable for control tasks, in order to enhance the generalization ability of robots in complex, open environments.

Specifically, the paper points out that the currently commonly used contrastive training visual-language models (such as CLIP) are insufficient in enabling robots to gain sufficiently detailed scene understanding, which limits their performance in control tasks. To address this issue, the authors propose a new method—Stable Control Representations (SCR), which extracts representations from pre-trained text-to-image diffusion models. These representations can capture both high-granularity and low-granularity details of the scene, thereby supporting the learning of downstream control strategies.

The main contributions of the paper include:
1. Proposing a multi-step method to extract visual-language representations from text-to-image diffusion models for control tasks.
2. Evaluating the representation learning capability of diffusion models in a wide range of robotic control tasks, demonstrating their competitiveness across various tasks.
3. Systematically analyzing the key features of diffusion model representations, explaining the impact of different design choices on performance.

Through this work, the authors demonstrate that diffusion models can provide effective control representations and contribute to the advancement of embodied AI.