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

TL;DR

RoCoDA introduces a unified data augmentation framework that blends causality, $SE(3)$-equivariance, and visual invariances to enhance data-efficient imitation learning for robotics. By decomposing state into causally relevant/irrelevant parts, applying SE(3) pose transformations with corresponding action adjustments, and incorporating standard augmentations, RoCoDA generates diverse, causally consistent demonstrations. Empirical results across five manipulation tasks show improved policy performance, stronger generalization to unseen poses and textures, and greater sample efficiency compared to baselines like MimicGen. The approach provides a principled bridge between geometric symmetries and causal reasoning, enabling more robust and adaptable robotic policies.

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

Figures (4)

• Figure 1: RoCoDA is a data augmentation framework that leverages causality and several symmetry groups to expand BC datasets. Our pipeline consists of 5 augmentations applied in tandem -- causally invariant counterfactual data augmentation, $SE(3)$ equivariant object state and trajectory augmentation, translation and scale invariance (resize and crop), hue invariance (color jitter), and observation noise invariance (proprioception noise).
• Figure 2: Causal structures for various augmentation methods. (a) Visual augmentation performs transformations to images that do not affect state-action dynamics in the environment. (b) Counterfactual augmentation factors the state space into causally invariant subspaces ($s^i \text{ and } s^j$), and resamples these subspaces from the training data distribution to generate synthetic data. (c) Equivariant augmentation performs a transformation on the pose of a target object associated with a subtask, and uses that same transform to augment the associated actions. (d) RoCoDA unifies these augmentation schemas.
• Figure 3: RoCoDA consists of three stages: $SE(3)$ equivariant state-action augmentation, causal augmentation, and visual augmentation. We first apply rigid body transformations to object poses and identically transform the corresponding actions. Then, we resample subsets of the environment state that are causally invariant to the action. Lastly, we apply several standard augmentations, including random resize/crop, color jitter, and state noise, all of which leverage invariance aspects of the observation with respect to actions (see Section \ref{['sec:standardaug']}).
• Figure 4: RoCoDA learns to re-grasp cubes in the Three Block Stack task, despite this behavior not being present in the dataset. This suggests that policies trained with RoCoDA possess a deeper understanding of the causal structure underlying the task, enabling the robot to adapt to out-of-distribution inputs.