Using Temperature Sampling to Effectively Train Robot Learning Policies on Imbalanced Datasets
This addresses data imbalance issues in multi-task robot learning, enabling more effective use of model capacity, though it is incremental as it builds on existing sampling techniques.
The paper tackles the problem of imbalanced datasets in robot learning, where similar physical actions across tasks cause data imbalance, by proposing a simple temperature sampling strategy that improves generalization and shows substantial gains on low-resource tasks without harming high-resource performance.
Increasingly large datasets of robot actions and sensory observations are being collected to train ever-larger neural networks. These datasets are collected based on tasks and while these tasks may be distinct in their descriptions, many involve very similar physical action sequences (e.g., 'pick up an apple' versus 'pick up an orange'). As a result, many datasets of robotic tasks are substantially imbalanced in terms of the physical robotic actions they represent. In this work, we propose a simple sampling strategy for policy training that mitigates this imbalance. Our method requires only a few lines of code to integrate into existing codebases and improves generalization. We evaluate our method in both pre-training small models and fine-tuning large foundational models. Our results show substantial improvements on low-resource tasks compared to prior state-of-the-art methods, without degrading performance on high-resource tasks. This enables more effective use of model capacity for multi-task policies. We also further validate our approach in a real-world setup on a Franka Panda robot arm across a diverse set of tasks.