Stem-OB: Generalizable Visual Imitation Learning with Stem-Like Convergent Observation through Diffusion Inversion

TL;DR

Stem-OB leverages pretrained diffusion models to enhance visual imitation learning by removing low-level visual variations, significantly improving generalization and success rates in both simulated and real-world tasks.

Contribution

This paper introduces Stem-OB, a novel approach using diffusion inversion for robust visual imitation learning without additional training.

Findings

22.2% average success rate improvement in real-world tasks

Effective suppression of visual variations like lighting and textures

Plug-and-play method compatible with existing systems

Abstract

Visual imitation learning methods demonstrate strong performance, yet they lack generalization when faced with visual input perturbations, including variations in lighting and textures, impeding their real-world application. We propose Stem-OB that utilizes pretrained image diffusion models to suppress low-level visual differences while maintaining high-level scene structures. This image inversion process is akin to transforming the observation into a shared representation, from which other observations stem, with extraneous details removed. Stem-OB contrasts with data-augmentation approaches as it is robust to various unspecified appearance changes without the need for additional training. Our method is a simple yet highly effective plug-and-play solution. Empirical results confirm the effectiveness of our approach in simulated tasks and show an exceptionally significant improvement in real-world applications, with an average increase of 22.2% in success rates compared to the best baseline. See https://hukz18.github.io/Stem-Ob/ for more info.