GR-MG: Leveraging Partially Annotated Data via Multi-Modal Goal-Conditioned Policy

TL;DR

GR-MG leverages partially-annotated data by conditioning on text and goal images, using diffusion-based image editing to improve robot manipulation generalization with fewer annotations.

Contribution

Introduces GR-MG, a novel method that combines text instructions and goal images, including generated images, to enhance robot learning from partially-annotated data.

Findings

Improves task completion rate from 3.35 to 4.04 in simulation.

Increases success rate from 68.7% to 78.1% in real robots.

Outperforms baseline methods in few-shot skill learning.

Abstract

The robotics community has consistently aimed to achieve generalizable robot manipulation with flexible natural language instructions. One primary challenge is that obtaining robot trajectories fully annotated with both actions and texts is time-consuming and labor-intensive. However, partially-annotated data, such as human activity videos without action labels and robot trajectories without text labels, are much easier to collect. Can we leverage these data to enhance the generalization capabilities of robots? In this paper, we propose GR-MG, a novel method which supports conditioning on a text instruction and a goal image. During training, GR-MG samples goal images from trajectories and conditions on both the text and the goal image or solely on the image when text is not available. During inference, where only the text is provided, GR-MG generates the goal image via a diffusion-based image-editing model and conditions on both the text and the generated image. This approach enables GR-MG to leverage large amounts of partially-annotated data while still using languages to flexibly specify tasks. To generate accurate goal images, we propose a novel progress-guided goal image generation model which injects task progress information into the generation process. In simulation experiments, GR-MG improves the average number of tasks completed in a row of 5 from 3.35 to 4.04. In real-robot experiments, GR-MG is able to perform 58 different tasks and improves the success rate from 68.7\% to 78.1\% and 44.4\% to 60.6\% in simple and generalization settings, respectively. It also outperforms comparing baseline methods in few-shot learning of novel skills. Video demos, code, and checkpoints are available on the project page: https://gr-mg.github.io/.