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

TL;DR

This work tackles the challenge of imbalanced action-primitive distributions in large robot datasets by introducing temperature-based sampling to rebalance multitask policy training. The method uses task-level sampling probabilities $p_i^{(\tau)} = { |D_i|^{1/\tau} } / { \sum_j |D_j|^{1/\tau} }$ and a cosine-warmup schedule that shifts emphasis toward low-resource tasks toward the end of training, enabling more effective use of model capacity. Validations span toy parity tasks, simulated robotic datasets RoboCasa and Libero, and real-world Franka Panda experiments, employing a Behavior Cloning Transformer and fine-tuning UniVLA, with 40k gradient steps in most settings. Results show substantial improvements for low-resource tasks while maintaining performance on high-resource tasks, demonstrating the practicality and generality of the approach for multi-task robotic policy learning and foundation-model fine-tuning.

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.

Figures (8)

• Figure 1: Temperature-based sampling rebalances the distribution over tasks, enabling more equitable training and improved generalization across both high- and low-resource domains.
• Figure 2: Total training loss over time for sparse parity tasks shows improvement in performance by over-sampling low-resource tasks in imbalanced dataset. While also highlighting the limitations of adaptive sampling methods like ReMix due to sensitivity to hyper-parameters.
• Figure 3: Real-world experiment setup as well as task distributions.
• Figure 4: Average subtask success rates in simulation for high- and low-resource tasks under different sampling strategies. Temperature sampling provides the highest performance on low-resource tasks while preserving high performance on high-resource ones.
• Figure 5: Real-world performance on a Franka Panda robot across four tasks, comparing random, cosine decay, and cosine warming schedules. Cosine warming consistently improves success rates, particularly for low-resource tasks.