Zero-Shot Visual Generalization in Robot Manipulation

TL;DR

The paper tackles the problem of robust, zero-shot visual generalization for robot manipulation by combining disentangled representations with associative memory (ALDA) and extending them to diffusion-based imitation and SAC policies. It introduces a lightweight equivariant adaptation workflow to enforce discrete planar rotation invariance, enabling resilience to camera perturbations via learned canonicalization over the rotation group $C_n$ and invariances related to $SO(2)$. Empirically, ALDA-SAC and ALDA-DP outperform state-of-the-art baselines on ManiSkill3 with visual perturbations and demonstrate strong real-world performance on a Franka Panda, while revealing limitations tied to unmodeled visual variations such as table color changes and lighting extremes. Overall, the work presents a substantial advance toward generalist, robotics-ready policies that generalize across visual shifts without relying on heavy domain randomization or calibration.

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.

Figures (6)

• Figure 1: Behavior cloning with disentangled representations and associative latent dynamics achieves zero-shot generalization to various real world perturbations, such as changes in ambient lighting (left), object color (middle-left), directed lighting (middle-right), and the presence of distractor objects (right).
• Figure 2: Overview of ALDA + Diffusion Policy (ALDA-DP). ALDA-DP jointly learns a factorized representation of the image observation while training the policy. The diffusion model denoises actions conditioned on this representation.
• Figure 3: ManiSkill3 visual generalization tasks. Left to right: random lighting, random cube color, distracting background (DBG), DBG + random cube color, DBG + random lighting + random cube color.
• Figure 4: Average success rate of ALDA-SAC and ALDA-DP compared to various RL and BC baselines. The first two rows have both RL and BC results, while the 3rd row has only RL results and the fourth only BC results. ALDA-* methods overall perform the best, and with a large margin, especially on PickCube.
• Figure 5: Traversing the disentangled latent space of ALDA-DP trained on real demonstrations and visualizing their corresponding reconstructions. Rows correspond to latent traversals of a single latent dimension given a reference image. By editing specific latent dimensions, we can visualize what factor of variation they correspond to.