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

• URL: http://arxiv.org/abs/2408.14037v1
• Date: Mon, 26 Aug 2024 06:14:25 GMT
• Title: Re-Mix: Optimizing Data Mixtures for Large Scale Imitation Learning
• Authors: Joey Hejna, Chethan Bhateja, Yichen Jian, Karl Pertsch, Dorsa Sadigh,
• Abstract summary: We investigate how to weigh different subsets or domains'' of robotics datasets for robot foundation model pre-training. Our method, Re-Mix, addresses the wide range of challenges that arise when applying DRO to robotics datasets.
• Score: 25.359270812682155
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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. 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.

Related papers

• Robots Pre-train Robots: Manipulation-Centric Robotic Representation from Large-Scale Robot Datasets [24.77850617214567]
  We propose a foundation representation learning framework capturing both visual features and the dynamics information such as actions and proprioceptions of manipulation tasks. Specifically, we pre-train a visual encoder on the DROID robotic dataset and leverage motion-relevant data such as robot proprioceptive states and actions. We introduce a novel contrastive loss that aligns visual observations with the robot's proprioceptive state-action dynamics, combined with a behavior cloning (BC)-like actor loss to predict actions during pre-training, along with a time contrastive loss.
  arXiv  Detail & Related papers  (2024-10-29T17:58:13Z)
• Scaling Cross-Embodied Learning: One Policy for Manipulation, Navigation, Locomotion and Aviation [49.03165169369552]
  By training a single policy across many different kinds of robots, a robot learning method can leverage much broader and more diverse datasets. We propose CrossFormer, a scalable and flexible transformer-based policy that can consume data from any embodiment. We demonstrate that the same network weights can control vastly different robots, including single and dual arm manipulation systems, wheeled robots, quadcopters, and quadrupeds.
  arXiv  Detail & Related papers  (2024-08-21T17:57:51Z)
• Mitigating the Human-Robot Domain Discrepancy in Visual Pre-training for Robotic Manipulation [16.809190349155525]
  Recent works have turned to large-scale pre-training using human data. morphological differences between humans and robots introduce a significant human-robot domain discrepancy. We propose a novel adaptation paradigm that utilizes readily available paired human-robot video data to bridge the discrepancy.
  arXiv  Detail & Related papers  (2024-06-20T11:57:46Z)
• Active Exploration in Bayesian Model-based Reinforcement Learning for Robot Manipulation [8.940998315746684]
  We propose a model-based reinforcement learning (RL) approach for robotic arm end-tasks. We employ Bayesian neural network models to represent, in a probabilistic way, both the belief and information encoded in the dynamic model during exploration. Our experiments show the advantages of our Bayesian model-based RL approach, with similar quality in the results than relevant alternatives.
  arXiv  Detail & Related papers  (2024-04-02T11:44:37Z)
• Efficient Data Collection for Robotic Manipulation via Compositional Generalization [70.76782930312746]
  We show that policies can compose environmental factors from their data to succeed when encountering unseen factor combinations. We propose better in-domain data collection strategies that exploit composition. We provide videos at http://iliad.stanford.edu/robot-data-comp/.
  arXiv  Detail & Related papers  (2024-03-08T07:15:38Z)
• BridgeData V2: A Dataset for Robot Learning at Scale [73.86688388408021]
  BridgeData V2 is a large and diverse dataset of robotic manipulation behaviors. It contains 60,096 trajectories collected across 24 environments on a publicly available low-cost robot.
  arXiv  Detail & Related papers  (2023-08-24T17:41:20Z)
• Scaling Robot Learning with Semantically Imagined Experience [21.361979238427722]
  Recent advances in robot learning have shown promise in enabling robots to perform manipulation tasks. One of the key contributing factors to this progress is the scale of robot data used to train the models. We propose an alternative route and leverage text-to-image foundation models widely used in computer vision and natural language processing.
  arXiv  Detail & Related papers  (2023-02-22T18:47:51Z)
• RT-1: Robotics Transformer for Real-World Control at Scale [98.09428483862165]
  We present a model class, dubbed Robotics Transformer, that exhibits promising scalable model properties. We verify our conclusions in a study of different model classes and their ability to generalize as a function of the data size, model size, and data diversity based on a large-scale data collection on real robots performing real-world tasks.
  arXiv  Detail & Related papers  (2022-12-13T18:55:15Z)
• MetaGraspNet: A Large-Scale Benchmark Dataset for Vision-driven Robotic Grasping via Physics-based Metaverse Synthesis [78.26022688167133]
  We present a large-scale benchmark dataset for vision-driven robotic grasping via physics-based metaverse synthesis. The proposed dataset contains 100,000 images and 25 different object types. We also propose a new layout-weighted performance metric alongside the dataset for evaluating object detection and segmentation performance.
  arXiv  Detail & Related papers  (2021-12-29T17:23:24Z)
• Multi-Robot Deep Reinforcement Learning for Mobile Navigation [82.62621210336881]
  We propose a deep reinforcement learning algorithm with hierarchically integrated models (HInt) At training time, HInt learns separate perception and dynamics models, and at test time, HInt integrates the two models in a hierarchical manner and plans actions with the integrated model. Our mobile navigation experiments show that HInt outperforms conventional hierarchical policies and single-source approaches.
  arXiv  Detail & Related papers  (2021-06-24T19:07:40Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.