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

arXiv: https://arxiv.org/abs/2510.19373 · HTML: https://mlsanity.com/paper/2510.19373/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 26 of its 34 references that are in the map.
- First posted on arXiv 22 Oct 2025.

## Builds on (most cited first, 26 of 26)

- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [π0: A Vision-Language-Action Flow Model for General Robot Control](https://mlsanity.com/paper/2410.24164/) (2024)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [LIBERO: Benchmarking Knowledge Transfer for Lifelong Robot Learning](https://mlsanity.com/paper/2306.03310/) (2023)
- [DROID: A Large-Scale In-The-Wild Robot Manipulation Dataset](https://mlsanity.com/paper/2403.12945/) (2024)
- [BridgeData V2: A Dataset for Robot Learning at Scale](https://mlsanity.com/paper/2308.12952/) (2023)
- [RoboCasa: Large-Scale Simulation of Everyday Tasks for Generalist Robots](https://mlsanity.com/paper/2406.02523/) (2024)
- [Prismatic VLMs: Investigating the Design Space of Visually-Conditioned Language Models](https://mlsanity.com/paper/2402.07865/) (2024)
- [RoboNet: Large-Scale Multi-Robot Learning](https://mlsanity.com/paper/1910.11215/) (2019)
- [Bridge Data: Boosting Generalization of Robotic Skills with Cross-Domain Datasets](https://mlsanity.com/paper/2109.13396/) (2021)
- [Scaling Robot Learning with Semantically Imagined Experience](https://mlsanity.com/paper/2302.11550/) (2023)
- [DALL-E-Bot: Introducing Web-Scale Diffusion Models to Robotics](https://mlsanity.com/paper/2210.02438/) (2022)
- [GenAug: Retargeting behaviors to unseen situations via Generative Augmentation](https://mlsanity.com/paper/2302.06671/) (2023)
- [Balancing Training for Multilingual Neural Machine Translation](https://mlsanity.com/paper/2004.06748/) (2020)
- [CACTI: A Framework for Scalable Multi-Task Multi-Scene Visual Imitation Learning](https://mlsanity.com/paper/2212.05711/) (2022)
- [Learning Neuro-Symbolic Skills for Bilevel Planning](https://mlsanity.com/paper/2206.10680/) (2022)
- [SkillMimicGen: Automated Demonstration Generation for Efficient Skill Learning and Deployment](https://mlsanity.com/paper/2410.18907/) (2024)
- [Re-Mix: Optimizing Data Mixtures for Large Scale Imitation Learning](https://mlsanity.com/paper/2408.14037/) (2024)
- [Multi-task Learning for Multilingual Neural Machine Translation](https://mlsanity.com/paper/2010.02523/) (2020)
- [STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning](https://mlsanity.com/paper/2412.15182/) (2024)
- [FlowRetrieval: Flow-Guided Data Retrieval for Few-Shot Imitation Learning](https://mlsanity.com/paper/2408.16944/) (2024)
- [Adaptive Scheduling for Multi-Task Learning](https://mlsanity.com/paper/1909.06434/) (2019)
- [Bandits Don't Follow Rules: Balancing Multi-Facet Machine Translation with Multi-Armed Bandits](https://mlsanity.com/paper/2110.06997/) (2021)
- [EXTRACT: Efficient Policy Learning by Extracting Transferable Robot Skills from Offline Data](https://mlsanity.com/paper/2406.17768/) (2024)
- [Order Matters in the Presence of Dataset Imbalance for Multilingual Learning](https://mlsanity.com/paper/2312.06134/) (2023)
- [Upsample or Upweight? Balanced Training on Heavily Imbalanced Datasets](https://mlsanity.com/paper/2410.04579/) (2024)

## Abstract (arXiv)

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.