Shortcut Learning in Generalist Robot Policies: The Role of Dataset Diversity and Fragmentation
This addresses a critical problem for robotics researchers and practitioners by revealing dataset-related bottlenecks in generalization, offering practical solutions to enhance policy performance without new data collection.
The paper investigates the limited generalization of generalist robot policies trained on large-scale datasets like Open X-Embodiment, identifying shortcut learning as the key cause and attributing it to limited dataset diversity and fragmentation. It proposes dataset collection strategies and data augmentation methods to reduce shortcut learning and improve generalization, with demonstrated effectiveness in simulation and real-world environments.
Generalist robot policies trained on large-scale datasets such as Open X-Embodiment (OXE) demonstrate strong performance across a wide range of tasks. However, they often struggle to generalize beyond the distribution of their training data. In this paper, we investigate the underlying cause of this limited generalization capability. We identify shortcut learning -- the reliance on task-irrelevant features -- as a key impediment to generalization. Through comprehensive theoretical and empirical analysis, we uncover two primary contributors to shortcut learning: (1) limited diversity within individual sub-datasets, and (2) significant distributional disparities across sub-datasets, leading to dataset fragmentation. These issues arise from the inherent structure of large-scale datasets like OXE, which are typically composed of multiple sub-datasets collected independently across varied environments and embodiments. Our findings provide critical insights into dataset collection strategies that can reduce shortcut learning and enhance the generalization ability of generalist robot policies. Moreover, in scenarios where acquiring new large-scale data is impractical, we demonstrate that carefully selected robotic data augmentation strategies can effectively reduce shortcut learning in existing offline datasets, thereby improving generalization capabilities of generalist robot policies, e.g., $π_0$, in both simulation and real-world environments. More information at https://lucky-light-sun.github.io/proj/shortcut-learning-in-grps/.