Using Temperature Sampling to Effectively Train Robot Learning Policies on Imbalanced Datasets

TL;DR

This paper introduces a simple temperature sampling strategy to address dataset imbalance in robot learning, improving generalization and performance on low-resource tasks while maintaining high-resource task performance.

Contribution

The authors propose a novel, easy-to-implement sampling method that mitigates dataset imbalance in robotic policy training, enhancing multi-task learning capabilities.

Findings

Significant improvements on low-resource tasks compared to prior methods.

No degradation in performance on high-resource tasks.

Validated effectiveness on real-world robot experiments.

Abstract

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.