ReMix: Optimizing Data Mixtures for Large Scale Imitation Learning

Authors: Joey Hejna, Chethan Bhateja, Yichen Jiang, Karl Pertsch, Dorsa Sadigh · CoRL 2024

Licences

Open Access
Paper: CC BY 4.0

Everything below draws on the full paper under its CC BY 4.0 licence.

TL;DR

ReMix uses group distributionally robust optimization on a policy's excess loss to learn how much to sample each robot dataset, beating uniform and RT-X expert weights on real robots.

Why it matters

Teams pre-training on pooled robot data can replace hand-tuned sampling weights with weights computed from the data, and in the reported tests can drop 75% of the data with limited loss in performance. The learned weights also show which datasets to down-weight and why, for example the autonomously collected Kuka data, which helps decide what data to collect or keep.

Problem

Large robot datasets such as Open X-Embodiment pool data from different robots, cameras, action spaces and controllers, yet the mixtures used to train generalist policies such as RT-X, Octo and OpenVLA were chosen by hand. Data filtering methods from vision and language transfer poorly, because n-grams do not apply and visual embeddings ignore actions and the sequential nature of episodes.

Proposed solution

ReMix frames curation as a min-max problem: a discrete policy minimises its excess behaviour-cloning loss over a reference model, while domain weights are updated by exponentiated gradient ascent to maximise it, so weight goes to domains the policy can still improve on. To make losses comparable across robots, actions are Gaussian-normalised per domain and discretised into bins, and the reference checkpoint is chosen by early stopping before it overfits any domain. The weights averaged over training are then reused to train separate diffusion policies or to subsample the data.

Key insights

Using excess loss over a reference policy, rather than raw loss, down-weights both domains that are already fitted and domains that are too hard to fit.
Differences in action scale skew robust optimization: in a toy test with bounds normalisation, 94.3% of the weight went to a domain of pure random actions; with per-domain Gaussian normalisation the real data received 84.2%.
Continuous L1/L2 action losses are swayed by outliers and multimodal actions, so discretising actions gives more reliable weights even when the final policy is a continuous diffusion model.
Robot datasets (about 10k-100k demonstrations) are small enough for the reference model to overfit; using a checkpoint 50K steps past early stopping lowered performance by over 15% on average.

Main results

On the RT-X data mix, policies trained with ReMix weights outperformed uniform weights by 38% and the RT-X human-expert weights by 32% on average across four WidowX and two Franka real-world tasks; expert weights beat uniform by 6%.
On 25% subsets of the RT-X mix, ReMix retained performance while human weights dropped over 10%; ReMix kept 72% of the UR5 data and 12% of Kuka, against 30% and 24% for the human weights.
ReMix raised the small Toto dataset from 3.42% (uniform) to 16.3% and cut Kuka from 24.9% to 12.1%.
On the full Bridge V2 dataset there was no notable difference between ReMix and uniform weights; on 25% Bridge subsets ReMix outperformed uniform and the SSP image-embedding baseline.
On a WidowX cube task with a distractor object, ReMix weights succeeded 7/10 versus 0/10 for both human and uniform weights.
In simulation (Robomimic NutAssemblySquare split by operator), ReMix scored 77/100 vs 53/100 for uniform at 50% subsetting and 59/100 vs 39/100 at 25%.

Limitations

The authors note real-world evaluation covered only two arms, WidowX and Franka, and leave more embodiments to future work.
ReMix up-weights datasets with abnormal, multimodal action distributions such as Toto, which the authors say is not necessarily desirable.
Computing the weights requires training on the full data twice, once for the reference model and once for the robust optimisation.
Tested on Bridge V2 and the RT-X subset of OpenX, not on the full OpenX dataset of over 2M episodes. Section 4.1 describes Bridge V2 as 50k demonstrations, while Appendix B.2 says about 45,000 trajectories were used out of the 60K in the full dataset.
Policies trained only on the pre-training data did not generalise zero-shot to the authors' setups, so every evaluation relies on co-training with 25 in-domain demonstrations per task.

Field impact

Finalist for the CoRL 2024 Outstanding Paper Award and cited 74 times. The authors release code at github.com/jhejna/remix. The paper also publishes its learned domain weights for the RT-X mix (Table 2) and for the 32 Bridge V2 domains (Table 5).

Abstract

Increasingly large imitation learning datasets are being collected with the goal of training foundation models for robotics. However, despite the fact that data selection has been of utmost importance in vision and natural language processing, little work in robotics has questioned what data such models should actually be trained on. In this work we investigate how to weigh different subsets or “domains” of robotics datasets for robot foundation model pre-training. Concretely, we use distributionally robust optimization (DRO) to maximize worst-case performance across all possible downstream domains. Our method, Re-Mix, addresses the wide range of challenges that arise when applying DRO to robotics datasets including variability in action spaces and dynamics across datasets. Re-Mix employs early stopping, action normalization, and discretization to counteract these issues. Through extensive experimentation on the largest open-source robot manipulation dataset, the Open X-Embodiment dataset, we demonstrate that data curation can have an outsized impact on downstream performance. Specifically, domain weights learned by Re-Mix outperform uniform weights by 38% on average and outperform human-selected weights by 32% on datasets used to train existing generalist robot policies, specifically the RT-X models.

Open Source Robotics