Zero-Shot Visual Generalization in Robot Manipulation

TL;DR

This paper advances zero-shot visual generalization in robot manipulation by scaling disentangled representations and associative memory to complex tasks, enabling robustness to visual and camera perturbations in simulation and real-world settings.

Contribution

It introduces a scalable approach combining disentangled representations, associative memory, and model equivariance to improve visual robustness and zero-shot adaptability in manipulation policies.

Findings

Significant improvement in visual generalization over state-of-the-art methods.

Successful zero-shot adaptation to visual perturbations in simulation and hardware.

Enhanced robustness to camera rotations through a novel invariance technique.

Abstract

Training vision-based manipulation policies that are robust across diverse visual environments remains an important and unresolved challenge in robot learning. Current approaches often sidestep the problem by relying on invariant representations such as point clouds and depth, or by brute-forcing generalization through visual domain randomization and/or large, visually diverse datasets. Disentangled representation learning - especially when combined with principles of associative memory - has recently shown promise in enabling vision-based reinforcement learning policies to be robust to visual distribution shifts. However, these techniques have largely been constrained to simpler benchmarks and toy environments. In this work, we scale disentangled representation learning and associative memory to more visually and dynamically complex manipulation tasks and demonstrate zero-shot adaptability to visual perturbations in both simulation and on real hardware. We further extend this approach to imitation learning, specifically Diffusion Policy, and empirically show significant gains in visual generalization compared to state-of-the-art imitation learning methods. Finally, we introduce a novel technique adapted from the model equivariance literature that transforms any trained neural network policy into one invariant to 2D planar rotations, making our policy not only visually robust but also resilient to certain camera perturbations. We believe that this work marks a significant step towards manipulation policies that are not only adaptable out of the box, but also robust to the complexities and dynamical nature of real-world deployment. Supplementary videos are available at https://sites.google.com/view/vis-gen-robotics/home.