Exploring Conditions for Diffusion models in Robotic Control

TL;DR

This work tackles the challenge of deriving task-adaptive visual representations for robotic control using pre-trained diffusion models without fine-tuning. It finds that naive text-based conditioning offers limited or negative benefits due to domain shift between diffusion training data and control environments, and proposes ORCA, which learns per-task, frame-specific prompts and leverages dense visual prompts from a vision encoder to condition diffusion models. Trained end-to-end with a behavior cloning objective, ORCA achieves state-of-the-art results across three major robotic benchmarks, demonstrating the importance of task-adaptive, frame-wise conditioning for control. The study provides extensive ablations and analysis of attention grounding, layer selection, and efficiency, underscoring the practical value of targeted diffusion conditioning in imitation learning for robotics.

Abstract

While pre-trained visual representations have significantly advanced imitation learning, they are often task-agnostic as they remain frozen during policy learning. In this work, we explore leveraging pre-trained text-to-image diffusion models to obtain task-adaptive visual representations for robotic control, without fine-tuning the model itself. However, we find that naively applying textual conditions - a successful strategy in other vision domains - yields minimal or even negative gains in control tasks. We attribute this to the domain gap between the diffusion model's training data and robotic control environments, leading us to argue for conditions that consider the specific, dynamic visual information required for control. To this end, we propose ORCA, which introduces learnable task prompts that adapt to the control environment and visual prompts that capture fine-grained, frame-specific details. Through facilitating task-adaptive representations with our newly devised conditions, our approach achieves state-of-the-art performance on various robotic control benchmarks, significantly surpassing prior methods.

Figures (10)

• Figure 1: How can we condition diffusion models in robotic control? We investigate methods for conditioning text-to-image diffusion models rombach2022high to perform control, aiming to address various tasks in a task-adaptive manner. We observe that text prompts, unlike in other vision tasks zhao2023unleashing, are ineffective for robotic control. Therefore, we propose to learn task prompts in control environments and further incorporate dynamic details through visual prompts for conditioning diffusion models.
• Figure 2: Motivation. We aim to overcome the limitations of existing task-agnostic approach (a) with frozen pre-trained visual representations parisi2022unsurprising, by leveraging conditions in diffusion models for robotic control tasks in a task-adaptive approach (b). In this regard, we explore text conditions(§ \ref{['text']}), more advanced methods(§ \ref{['discussion']},§ \ref{['method']}) as conditions.
• Figure 3: Case study.(a) We find that text conditions can be disadvantageous in some control tasks. (b) For Button-press, the cross-attention maps (e.g., for button, press) are well-grounded to relevant image regions. (c) In contrast, for Cheetah-run, the attention maps (e.g., for cheetah, run) are noisy, which presumably leads to a decline in performance. Nonetheless, our approach of using task and visual tokens (§ \ref{['method']}) achieves consistent gains across all tasks, with its cross-attention maps capturing diverse regions of the image relevant to the downstream task.
• Figure 4: Proposed framework. We propose ORCA, a framework for learning task and visual prompts to condition diffusion models in robotic control. Specifically, we utilize the features from the downsampling blocks and the bottleneck block of Stable Diffusion rombach2022high to extract visual representations conditioned on our input, which are then fed to the policy network for predicting the action.
• Figure 5: Visualization of evaluation tasks.