Using Temperature Sampling to Effectively Train Robot Learning Policies on Imbalanced Datasets
Cosine-warming Temperature Sampling raises Libero UniVLA success from 0.76 to 0.85 under severe task imbalance.
Key Findings
Methodology
The paper reshapes task sampling with p_i(τ)=|D_i|^(1/τ)/Σ_j|D_j|^(1/τ). τ=1 reproduces ordinary random sampling, whereas τ>1 relatively increases low-resource tasks. The default cosine-warming schedule increases τ from 1 to 5 during training, allowing the policy to learn broad representations from abundant data before specializing in scarce action primitives.
Key Results
- On imbalanced Libero fine-tuning with UniVLA, Random, fixed Upsample, and Temperature achieved overall success rates of 0.76, 0.77, and 0.85, respectively. Temperature obtained 0.96, 0.86, 0.84, and 0.73 on Libero-SPATIAL, GOAL, OBJECT, and Libero-10.
- On a Franka Panda, the policy was evaluated on Egg in Carton, Pen in Mug, Fold Towel, and Slide Cube. Cosine warming outperformed random sampling and cosine decay, with the clearest gains on the three low-resource tasks while preserving high-resource performance.
- RoboCasa ablations retained the advantage across 3.1M, 19M, and 56.7M models. Across 300:50, 1000:50, and 3000:50 imbalance ratios, the benefit over random sampling increased as imbalance became more severe.
Significance
The work shows that scaling robot data does not guarantee balanced skill coverage. When pick-and-place demonstrations dominate, optimization allocates model capacity to common behaviors while neglecting folding, sliding, and other operationally important skills. Temperature Sampling is lightweight, data-efficient, and compatible with existing behavior-cloning and foundation-model pipelines. It avoids the prohibitive augmentation cost of datasets such as Open-X Embodiment and offers a practical remedy for long-tailed robot datasets.
Technical Contribution
The main contribution is to formulate temperature sampling as an action-primitive balancing mechanism for robot policy learning, rather than merely transferring a classification heuristic. The paper introduces cosine warming from τ=1 to 5 and compares it with fixed upsampling, alternative schedules, and ReMix. Unlike ReMix, which requires reference and proxy models to estimate difficulty-based weights, the proposed method uses only task sample counts, requires minimal implementation changes, and avoids additional model training.
Novelty
Temperature sampling itself is not new, having appeared in multilingual learning and curriculum methods. The novelty lies in applying it to imbalanced robotic action distributions, defining the relevant domains around physical skills, and validating the mechanism across sparse parity, RoboCasa, Libero, and a physical Franka Panda. Compared with complexity-based or human-designed reweighting, the method directly corrects behavioral coverage and demonstrates that schedule direction and timing materially affect control success.
Limitations
- The method assumes reliable task or skill segmentation. It may be less effective when a trajectory contains several primitives, labels are noisy, or skills change continuously rather than forming clean task groups.
- The temperature range and schedule remain empirical. Although cosine warming performed best in the reported studies, the paper offers no general rule for selecting the final temperature from imbalance, model capacity, or task similarity.
Future Work
Future work should automatically discover primitives, adapt temperature online, and study segment-level sampling for multi-skill trajectories. Larger real-world datasets, more robot morphologies, long-horizon tasks, and cross-embodiment transfer are needed. A useful theoretical direction is to connect imbalance ratios, model capacity, temperature schedules, gradient variance, and generalization, enabling principled rather than searched hyperparameters.
AI Executive Summary
Robot learning is following the scaling logic of modern AI: collect more demonstrations and train larger networks. Yet task names often conceal repeated physical behaviors. “Pick up an apple” and “pick up an orange” may contribute nearly identical motions, while folding, sliding, or knob-turning remain rare. Random sampling therefore lets common behaviors dominate stochastic gradients. Data augmentation is expensive at the scale of Open-X Embodiment, and ReMix requires additional models to estimate task difficulty.
Patil and colleagues propose Temperature Sampling, which sets task probability to p_i(τ)=|D_i|^(1/τ)/Σ|D_j|^(1/τ). At τ=1, sampling is proportional to dataset size; increasing τ gives relatively more exposure to small tasks. Their practical schedule warms τ from 1 to 5 with a cosine curve: the policy first learns robust shared representations, then concentrates on under-represented skills. The method needs only task counts and can be inserted into behavior-cloning or foundation-model training.
Evidence spans sparse parity, RoboCasa, Libero, and a real Franka Panda. On Libero, UniVLA reached 0.85 overall success with Temperature, versus 0.76 for Random and 0.77 for fixed Upsampling; suite scores were 0.96, 0.86, 0.84, and 0.73. RoboCasa experiments preserved gains from 3.1M to 56.7M models and across increasingly severe imbalance. The broader lesson is that robot capability depends not only on how much data is collected, but also on which behaviors are sampled and when.
Deep Analysis
Background
Large collections such as Open-X Embodiment improve robot learning but inherit task-collection biases. RoboCasa contains 25 atomic tasks, 1,200 human trajectories, and 72,000 synthetic demonstrations; some pick-and-place behaviors are far more common than other skills. Foundation-model representations can partly absorb visual and linguistic variation, but they do not automatically correct the behavioral distribution presented during policy optimization.
Core Problem
For D=∪D_i with |D_1|≫|D_2|⋯, behavior cloning minimizes L(θ)=−(1/N)Σlogπθ(a_i|x_i,z_i). Uniform sampling over demonstrations makes large tasks dominate minibatches and gradients. The resulting policy may perform well on high-resource actions while failing on low-resource behaviors, making the challenge one of gradient allocation and model-capacity usage rather than merely label imbalance.
Innovation
The paper contributes three ideas. First, it defines sampling domains around task/action distributions rather than estimated task complexity. Second, it uses a Boltzmann-like temperature law to smoothly reshape the long tail. Third, it warms temperature from 1 to 5 with a cosine schedule, separating broad representation learning from late specialization. This is simpler than ReMix and more stable than permanently oversampling scarce tasks.
Methodology
- �� Input: task counts |D_i|, task condition z, and training progress t.
- �� Probability: compute p_i(τ)=|D_i|^(1/τ)/Σ|D_j|^(1/τ), then draw demonstrations by task.
- �� Schedule: τ(t)=τ_start+(τ_end−τ_start)(1−cos(πt/T))/2, normally 1→5.
- �� Optimization: update a behavior-cloning policy with SGD on expert action likelihood.
- �� From-scratch model: RoboCasa uses a 512-dimensional, 6-layer, 8-head BC-Transformer for 40k steps.
- �� Comparisons: Random, fixed-temperature upsampling, and ReMix, which needs reference and proxy models for difficulty-based weighting.
Experiments
The controlled experiment uses five sparse-parity tasks over binary strings of length 50, with four task-specific positions. Simulation uses RoboCasa and four Libero suites. RoboCasa is made imbalanced by giving seven pick-and-place tasks 3,000 demonstrations each and other atomic tasks 50 each. Libero uses 50, 15, or 20 demonstrations depending on suite. UniVLA is fine-tuned for 40k steps. The real study trains a ResNet-50 plus UNet diffusion policy from 588 demonstrations for eight Panda tasks.
Results
Temperature reaches 0.85 overall on Libero, exceeding Random at 0.76 and fixed Upsample at 0.77, while preserving the high-resource Spatial score of 0.96. RoboCasa results favor cosine warming across small, base, and large models. As imbalance increases from 300:50 to 1000:50 and 3000:50, the low-resource advantage grows. Schedule ablations report low-resource success of 0.31 for cosine warming, 0.28 for exponential warming, 0.24 for linear warming, and 0.24 for cosine decay.
Applications
The method suits multi-task pretraining, few-shot policy adaptation, and foundation-model fine-tuning. A practitioner needs task groups and sample counts, then modifies the data loader to apply temperature probabilities. It is particularly relevant when grasping dominates datasets while opening, pressing, folding, sliding, or turning are scarce, including warehouse, household, and laboratory robotics.
Limitations & Outlook
Reliable segmentation is a strong assumption, and the method may struggle with noisy labels, mixed-skill demonstrations, or continuous skill manifolds. The final temperature of 5 was selected empirically rather than theoretically. The study has limited coverage of very large real datasets, long-horizon tasks, different control frequencies, and cross-embodiment transfer. Repeated exposure to scarce demonstrations may also increase overfitting if temperature is too high.
Plain Language Accessible to non-experts
Imagine a school with several classes. The “picking things up” class has 3,000 practice sheets, while the “folding towels” class has only 50. If the teacher chooses worksheets at random, nearly every lesson comes from the large class. The students become excellent at picking things up but barely learn folding, even though both skills matter.
Temperature Sampling changes the timetable. At the beginning, the large class still appears often, because its many examples help everyone learn basic patterns. Later, the timetable deliberately gives the small classes more turns. The temperature is simply a control knob: raising it makes rare classes more likely to appear. The paper raises it smoothly from 1 to 5 rather than changing it suddenly.
Nothing new has to be filmed or invented. The system only changes how existing practice examples are selected. On Libero, this raised overall UniVLA success from 0.76 with random sampling to 0.85. Tests on a real Franka Panda also favored the gradual schedule, especially for rare skills. The lesson is simple: a learner’s ability depends not only on the size of the library, but on whether every important lesson receives enough attention.
ELI14 Explained like you're 14
Think of a game with lots of missions. The “pick up an item” mission has 3,000 walkthroughs, but “fold a towel” has only 50. If the game keeps showing walkthroughs randomly, your robot avatar becomes amazing at picking things up and terrible at almost everything else.
Temperature Sampling is like a smarter recommendation system. Early in training, it recommends popular missions because they teach useful basics. Later, it deliberately recommends rare missions more often. The temperature is not the robot’s body temperature—it is a dial controlling how strongly the system boosts the underrepresented missions. In the paper, the dial moves from 1 to 5 using a smooth cosine curve.
Researchers tested it on RoboCasa, Libero, and a real Franka Panda robot. On Libero, ordinary sampling scored 0.76 overall, fixed upsampling scored 0.77, and Temperature Sampling scored 0.85. It also worked with models containing 3.1M, 19M, and 56.7M parameters.
The coolest part is that timing matters! If you force rare examples too early, the learner may not have basic skills yet. If you build a foundation first and then practice weak areas, performance becomes more balanced. It is like studying for exams: learn the common ideas, then spend extra time on the subjects where you are struggling. Same robot, same data—smarter practice schedule!
Glossary
Temperature Sampling
A sampling method that changes task probabilities according to dataset size. Higher temperature relatively increases the probability of small tasks.
It is the paper’s main mechanism for balancing robot action distributions.
Behavior Cloning
An imitation-learning method that maximizes the likelihood of expert actions. Its standard objective is the negative log-likelihood of demonstrated actions.
BC-Transformer and diffusion policies are trained from demonstrations under this paradigm.
Low-Resource Task
A task with relatively few demonstrations and therefore fewer gradient contributions under ordinary sampling. Such tasks are the main beneficiaries of rebalancing.
Folding, sliding, and less common manipulation tasks are evaluated as low-resource cases.
ReMix
A group-distributionally robust optimization method that reweights domains using estimated difficulty. It trains reference and proxy models to derive task weights.
The paper compares it with the simpler count-based temperature method.
Cosine Warming
A schedule that gradually increases temperature according to a cosine curve. It produces stronger sampling of low-resource tasks late in training.
The default schedule moves from τ=1 to τ=5.
Action Primitive
A basic physical behavior composing a larger robot task, such as picking, pressing, turning, or folding. Its frequency shapes the policy’s behavioral coverage.
The paper focuses on imbalance in primitives rather than only language or visual domains.
Open Questions Unanswered questions from this research
- 1 How can primitives be discovered automatically without manual task labels, while preserving reliable counts and boundaries? Trajectory segmentation, latent representations, and online clustering are plausible directions but were not validated here.
- 2 Should the final temperature depend jointly on imbalance, model capacity, task similarity, and gradient variance? The reported search identifies a useful schedule, but not a universal selection rule.
- 3 Sampling by whole task may be inadequate when one trajectory contains multiple skills. Segment-level or hierarchical temperature control remains an open problem.
Applications
Immediate Applications
Multi-task robot pretraining
Robot-learning teams can count demonstrations by task and add cosine warming from τ=1 to 5 inside an existing data loader. No synthetic trajectories or auxiliary difficulty models are required, making the method suitable for datasets dominated by grasping but sparse in pushing, folding, opening, or pressing.
Foundation-policy adaptation
During UniVLA-like fine-tuning, teams can increase exposure to rare target tasks while retaining abundant demonstrations for common skills. The expected benefit is higher few-shot success without sacrificing established high-resource behavior, provided task labels and validation rollouts are available.
Long-term Vision
Balanced general-purpose robot policies
Combined with automatic skill discovery, temperature scheduling could support policies covering grasping, manipulation, folding, pushing, and tool use more evenly. The long-term goal is to reduce the common failure mode in which a large policy performs routine actions well but misses rare, deployment-critical skills.
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.