RoCoDA: Counterfactual Data Augmentation for Data-Efficient Robot Learning from Demonstrations

TL;DR

RoCoDA introduces a unified framework leveraging invariance, equivariance, and causality to improve data augmentation in imitation learning for robotics, leading to better generalization and sample efficiency.

Contribution

It presents RoCoDA, a novel method that combines causal invariance and SE(3) equivariance for effective data augmentation in robotic imitation learning.

Findings

Enhanced policy performance and generalization across tasks.

Improved sample efficiency over existing augmentation methods.

Emergent behaviors like re-grasping indicate deeper task understanding.

Abstract

Imitation learning in robotics faces significant challenges in generalization due to the complexity of robotic environments and the high cost of data collection. We introduce RoCoDA, a novel method that unifies the concepts of invariance, equivariance, and causality within a single framework to enhance data augmentation for imitation learning. RoCoDA leverages causal invariance by modifying task-irrelevant subsets of the environment state without affecting the policy's output. Simultaneously, we exploit SE(3) equivariance by applying rigid body transformations to object poses and adjusting corresponding actions to generate synthetic demonstrations. We validate RoCoDA through extensive experiments on five robotic manipulation tasks, demonstrating improvements in policy performance, generalization, and sample efficiency compared to state-of-the-art data augmentation methods. Our policies exhibit robust generalization to unseen object poses, textures, and the presence of distractors. Furthermore, we observe emergent behavior such as re-grasping, indicating policies trained with RoCoDA possess a deeper understanding of task dynamics. By leveraging invariance, equivariance, and causality, RoCoDA provides a principled approach to data augmentation in imitation learning, bridging the gap between geometric symmetries and causal reasoning. Project Page: https://rocoda.github.io