RoCoDA: Counterfactual Data Augmentation for Data-Efficient Robot Learning from Demonstrations
This addresses the challenge of costly data collection and limited generalization in robotics for researchers and practitioners, representing a novel integration of concepts rather than an incremental improvement.
The paper tackles the problem of data inefficiency and poor generalization in robot imitation learning by introducing RoCoDA, a data augmentation method that unifies invariance, equivariance, and causality, resulting in improved policy performance, generalization, and sample efficiency across five robotic manipulation tasks.
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