Zero-Shot Robotic Manipulation with Pre-Trained Image-Editing Diffusion Models

Published: 07 Nov 2023, Last Modified: 08 Nov 2023FMDM@NeurIPS2023EveryoneRevisionsBibTeX
Keywords: robot learning, diffusion model
TL;DR: We present a method to leverage pre-trained image-editing diffusion models for zero-shot robotic manipulation
Abstract: If generalist robots are to operate in truly unstructured environments, they need to be able to recognize and reason about novel objects and scenarios. Such objects and scenarios might not be present in the robot's own training data. We propose SuSIE, a method that leverages an image editing diffusion model to act as a high-level planner by proposing intermediate subgoals that a low-level controller attains. Specifically, we fine-tune InstructPix2Pix on robot data such that it outputs a hypothetical future observation given the robot's current observation and a language command. We then use the same robot data to train a low-level goal-conditioned policy to reach a given image observation. We find that when these components are combined, the resulting system exhibits robust generalization capabilities. The high-level planner utilizes its Internet-scale pre-training and visual understanding to guide the low-level goal-conditioned policy, achieving significantly better generalization than conventional language-conditioned policies. We demonstrate that this approach solves real robot control tasks involving novel objects, distractors, and even environments, both in the real world and in simulation. The project website can be found at https://subgoal-image-editing.github.io
Submission Number: 65
Loading