OXE-AugE: A Large-Scale Robot Augmentation of OXE for Scaling Cross-Embodiment Policy Learning
Authors: Guanhua Ji, Harsha Polavaram, Lawrence Yunliang Chen, Sandeep Bajamahal, Zehan Ma, Simeon Adebola, Chenfeng Xu, Ken Goldberg
Abstract
Large and diverse datasets are needed for training generalist robot policies that can control a variety of robot embodiments--robot arm and gripper combinations--across diverse tasks and environments. As re-collecting demonstrations and retraining for each new embodiment are prohibitively costly, we study whether existing robot data can be augmented to improve transfer and generalization across embodiments. The Open X-Embodiment (OXE) dataset, which aggregates demonstrations from over 60 robot datasets, has been widely used for training generalist policies. However, it is highly imbalanced: the top four robot types account for over 85% of its real data, which risks overfitting to robot--scene combinations. We present AugE-Toolkit, a scalable robot augmentation pipeline, and OXE-AugE, a high-quality open-source dataset that augments OXE with 9 different robot embodiments. OXE-AugE provides over 4.4 million trajectories, more than triple the size of the original OXE. We conduct a systematic study of how scaling robot augmentation impacts cross-embodiment learning. Results suggest that augmenting datasets with diverse arms and grippers improves policy performance not only on the augmented robots, but also on unseen robots and even the original robots under distribution shifts. In physical experiments, fine-tuning generalist policies such as OpenVLA and $\pi_0$ on OXE-AugE improves success rates by 24-45% on unseen robot-gripper combinations across four real-world manipulation tasks. Project website: https://OXE-AugE.github.io/.