Invariance Co-training for Robot Visual Generalization

TL;DR

This work tackles robotic generalization by introducing invariance co-training, which learns viewpoint-, lighting-, and distractor-robust visual representations via a contrastive framework trained on diverse static data and robot demonstrations. It combines state similarity and relative observational disentanglement objectives with auxiliary cues (extrinsics and bounding boxes) to produce a vision encoder that generalizes across observational perturbations while preserving state-goal semantics. Empirical results show substantial gains over standard behavior cloning and baselines relying on simulation or generative augmentation, with notable benefits from real static imagery and scalable static data. The approach offers a scalable path to more robust, generalist robot policies without requiring high-fidelity physics-based simulation, and suggests directions for extending invariance co-training to broader tasks and vision-language-action models.

Abstract

Reasoning from diverse observations is a fundamental capability for generalist robot policies to operate in a wide range of environments. Despite recent advancements, many large-scale robotic policies still remain sensitive to key sources of observational variation such as changes in camera perspective, lighting, and the presence of distractor objects. We posit that the limited generalizability of these models arises from the substantial diversity required to robustly cover these quasistatic axes, coupled with the current scarcity of large-scale robotic datasets that exhibit rich variation across them. In this work, we propose to systematically examine what robots need to generalize across these challenging axes by introducing two key auxiliary tasks, state similarity and invariance to observational perturbations, applied to both demonstration data and static visual data. We then show that via these auxiliary tasks, leveraging both more-expensive robotic demonstration data and less-expensive, visually rich synthetic images generated from non-physics-based simulation (for example, Unreal Engine) can lead to substantial increases in generalization to unseen camera viewpoints, lighting configurations, and distractor conditions. Our results demonstrate that co-training on this diverse data improves performance by 18 percent over existing generative augmentation methods. For more information and videos, please visit https://invariance-cotraining.github.io

Figures (13)

• Figure 1: Invariance Co-training. Our method leverages diverse synthetic images, large-scale open-source datasets, and videos of static scenes to train a 2D vision encoder that generalizes to new camera viewpoints, lighting conditions, and background clutter. This encoder is co-trained with robot demonstration data to enable the policy to generalize effectively to novel visual conditions.
• Figure 2: Our Method. We co-train our policy with a contrastive loss using static data, two auxiliary losses using static images, and a behavior cloning loss using robot demonstration data.
• Figure 3: Generalization Results Across Observation Variations. Invariance Co-training achieves an average of $40\%$ higher success than standard behavior cloning across environments with perspective, distractor, and lighting variations.
• Figure 4: Ablation: Generalization. Invariance co-training achieves an average success rate increase of $30\%$ compared to standard co-training when leveraging the DROID dataset.
• Figure 5: Ablation: Auxiliary Tasks. Including auxiliary tasks yields an average absolute success rate increase of approximately $29\%$ compared to omitting them.