# Re-Mix: Optimizing Data Mixtures for Large Scale Imitation Learning

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

## Key facts

- Cited by 47 papers in the map (data through 14 Oct 2026).
- Placed in Robot manipulation imitation learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 64 of its 77 references that are in the map.
- First posted on arXiv 26 Aug 2024.

## Placed in

- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 30 of 64)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [LAION-5B: An open large-scale dataset for training next generation image-text models](https://mlsanity.com/paper/2210.08402/) (2022)
- [Unsupervised Cross-lingual Representation Learning at Scale](https://mlsanity.com/paper/1911.02116/) (2019)
- [RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control](https://mlsanity.com/paper/2307.15818/) (2023)
- [OpenVLA: An Open-Source Vision-Language-Action Model](https://mlsanity.com/paper/2406.09246/) (2024)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasets](https://mlsanity.com/paper/2311.15127/) (2023)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [RT-1: Robotics Transformer for Real-World Control at Scale](https://mlsanity.com/paper/2212.06817/) (2022)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [Ego4D: Around the World in 3,000 Hours of Egocentric Video](https://mlsanity.com/paper/2110.07058/) (2021)
- [Theoretically Principled Trade-off between Robustness and Accuracy](https://mlsanity.com/paper/1901.08573/) (2019)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Robustness May Be at Odds with Accuracy](https://mlsanity.com/paper/1805.12152/) (2018)
- [Open X-Embodiment: Robotic Learning Datasets and RT-X Models](https://mlsanity.com/paper/2310.08864/) (2023)
- [DROID: A Large-Scale In-The-Wild Robot Manipulation Dataset](https://mlsanity.com/paper/2403.12945/) (2024)
- [QT-Opt: Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulation](https://mlsanity.com/paper/1806.10293/) (2018)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [BridgeData V2: A Dataset for Robot Learning at Scale](https://mlsanity.com/paper/2308.12952/) (2023)
- [CALVIN: A Benchmark for Language-Conditioned Policy Learning for Long-Horizon Robot Manipulation Tasks](https://mlsanity.com/paper/2112.03227/) (2021)
- [Deep Learning on a Data Diet: Finding Important Examples Early in Training](https://mlsanity.com/paper/2107.07075/) (2021)
- [Beyond neural scaling laws: beating power law scaling via data pruning](https://mlsanity.com/paper/2206.14486/) (2022)
- [BC-Z: Zero-Shot Task Generalization with Robotic Imitation Learning](https://mlsanity.com/paper/2202.02005/) (2022)
- [Evaluating Real-World Robot Manipulation Policies in Simulation](https://mlsanity.com/paper/2405.05941/) (2024)
- [Dolma: an Open Corpus of Three Trillion Tokens for Language Model Pretraining Research](https://mlsanity.com/paper/2402.00159/) (2024)
- [RoboCasa: Large-Scale Simulation of Everyday Tasks for Generalist Robots](https://mlsanity.com/paper/2406.02523/) (2024)

## Built on by (most cited first, 30 of 47)

- [EgoMimic: Scaling Imitation Learning via Egocentric Video](https://mlsanity.com/paper/2410.24221/) (2024)
- [OneTwoVLA: A Unified Vision-Language-Action Model with Adaptive Reasoning](https://mlsanity.com/paper/2505.11917/) (2025)
- [Vision Language Models are In-Context Value Learners](https://mlsanity.com/paper/2411.04549/) (2024)
- [Vision-Language-Action Models for Robotics: A Review Towards Real-World Applications](https://mlsanity.com/paper/2510.07077/) (2025)
- [RoboMonkey: Scaling Test-Time Sampling and Verification for Vision-Language-Action Models](https://mlsanity.com/paper/2506.17811/) (2025)
- [ReinboT: Amplifying Robot Visual-Language Manipulation with Reinforcement Learning](https://mlsanity.com/paper/2505.07395/) (2025)
- [Sim-and-Real Co-Training: A Simple Recipe for Vision-Based Robotic Manipulation](https://mlsanity.com/paper/2503.24361/) (2025)
- [Shortcut Learning in Generalist Robot Policies: The Role of Dataset Diversity and Fragmentation](https://mlsanity.com/paper/2508.06426/) (2025)
- [DataMIL: Selecting Data for Robot Imitation Learning with Datamodels](https://mlsanity.com/paper/2505.09603/) (2025)
- [CUPID: Curating Data your Robot Loves with Influence Functions](https://mlsanity.com/paper/2506.19121/) (2025)
- [Model Predictive Task Sampling for Efficient and Robust Adaptation](https://mlsanity.com/paper/2501.11039/) (2025)
- [What Matters in Learning from Large-Scale Datasets for Robot Manipulation](https://mlsanity.com/paper/2506.13536/) (2025)
- [Is Diversity All You Need for Scalable Robotic Manipulation?](https://mlsanity.com/paper/2507.06219/) (2025)
- [Curating Demonstrations using Online Experience](https://mlsanity.com/paper/2503.03707/) (2025)
- [The One RING: a Robotic Indoor Navigation Generalist](https://mlsanity.com/paper/2412.14401/) (2024)
- [A Taxonomy for Evaluating Generalist Robot Manipulation Policies](https://mlsanity.com/paper/2503.01238/) (2025)
- [Empirical Analysis of Sim-and-Real Cotraining of Diffusion Policies for Planar Pushing from Pixels](https://mlsanity.com/paper/2503.22634/) (2025)
- [Guiding Data Collection via Factored Scaling Curves](https://mlsanity.com/paper/2505.07728/) (2025)
- [How to Train Your Robots? The Impact of Demonstration Modality on Imitation Learning](https://mlsanity.com/paper/2503.07017/) (2025)
- [Contextual Rollout Bandits for Reinforcement Learning with Verifiable Rewards](https://mlsanity.com/paper/2602.08499/) (2026)
- [Real-Time Out-of-Distribution Failure Prevention via Multi-Modal Reasoning](https://mlsanity.com/paper/2505.10547/) (2025)
- [Data Retrieval with Importance Weights for Few-Shot Imitation Learning](https://mlsanity.com/paper/2509.01657/) (2025)
- [SCIZOR: A Self-Supervised Approach to Data Curation for Large-Scale Imitation Learning](https://mlsanity.com/paper/2505.22626/) (2025)
- [VAMOS: A Hierarchical Vision-Language-Action Model for Capability-Modulated and Steerable Navigation](https://mlsanity.com/paper/2510.20818/) (2025)
- [R2RGEN: Real-to-Real 3D Data Generation for Spatially Generalized Manipulation](https://mlsanity.com/paper/2510.08547/) (2025)
- [X-Diffusion: Training Diffusion Policies on Cross-Embodiment Human Demonstrations](https://mlsanity.com/paper/2511.04671/) (2025)
- [An Anatomy of Vision-Language-Action Models: From Modules to Milestones and Challenges](https://mlsanity.com/paper/2512.11362/) (2025)
- [Ambient Diffusion Policy: Imitation Learning from Suboptimal Data in Robotics](https://mlsanity.com/paper/2606.12365/) (2026)
- [Learning from the Best: Smoothness-Driven Metrics for Data Quality in Imitation Learning](https://mlsanity.com/paper/2604.23000/) (2026)
- [SiMDex: Mining Similar Egocentric Videos for Cross-Embodiment Dexterous Manipulation](https://mlsanity.com/paper/2608.04196/) (2026)

## Abstract (arXiv)

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. Concrete, 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 different 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.