Stem-OB: Generalizable Visual Imitation Learning with Stem-Like Convergent Observation through Diffusion Inversion
This addresses the challenge of robust visual imitation learning for real-world robotics and AI systems, offering a plug-and-play solution without incremental reliance on data augmentation.
The paper tackles the problem of visual imitation learning's lack of generalization to visual perturbations like lighting and texture changes by proposing Stem-OB, which uses pretrained diffusion models to suppress low-level differences while preserving high-level structures, resulting in an average 22.2% increase in success rates compared to baselines in real-world applications.
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.