ReMix: Optimizing Data Mixtures for Large Scale Imitation Learning
Authors: Joey Hejna
Chethan Anand Bhateja
Yichen Jiang
Karl Pertsch
Dorsa Sadigh
Venue: CoRL 2024 (2024)
Keywords: Data Curation
Data Quality
Robot Imitation Learning
TL;DR: We use techniques from robust optimization to learn data mixture weights for Bridge and RT-X datasets, and show they improve downstream performance.

Abstract: Increasingly large robotics datasets are being collected to train larger foundation models in robotics. However, despite the fact that data selection has been of utmost importance to scaling in vision and natural language processing (NLP), 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 during pre-training to maximize worst-case performance across all possible downstream domains using distributionally robust optimization (DRO). Unlike in NLP, we find that these methods are hard to apply out of the box due to varying action spaces and dynamics across robots. Our method, ReMix, employs early stopping and action normalization and discretization to counteract these issues. Through extensive experimentation on both the Bridge and OpenX datasets, we demonstrate that data curation can have an outsized impact on downstream performance. Specifically, domain weights learned by ReMix outperform uniform weights by over 40\% on average and human-selected weights by over 20\% on datasets used to train the RT-X models.

ReMix: Optimizing Data Mixtures for Large Scale Imitation Learning
Joey Hejna Stanford
Chethan Bhateja Stanford
Yichen Jiang Stanford
Karl Pertsch Stanford, UC Berkeley
Dorsa Sadigh Stanford
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 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 roboties datasets for robot foundation model pre-training. Concretely, 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- open-source robot manipulation dataset, the Open X-Embodiment dataset, we demonstrate that data curation have an outsized impact on downstream performance. Specifically, domain weights leamedd by Re-Mix outperform uniform weights by 38% average and outperform human-selected weights by 32% on datasets used to train existing generalist robot policies, specifically the RT-X models.
Keywords: Data Curation, Data Quality, Robot Imitation Learning
Introduction
Many breakthroughs in machine learning can be attributed to "Internet-scale" datasets, from the develop- ment of vision models like CLIP [1] to advancements transformer- -based language modeling powered by the Common Crawl dataset [2]. Seeking to capitalize on this trend, several recent efforts in robotics focus on collecting [3-6] or pooling [7] large scale robotics datasets with the goal of training more performant imitation learning policies. Learning from this data, however, is particularly challenging: robotics datasets collected with different robots, environments, state spaces, action spaces, and dynamics [8]. For exam- ple, the commonly used Bridge V2 Dataset [4] uses a third person camera on a small WidowX robot and cartesian delta control, while many datasets [9-12] collected on the popular and much larger Franka Panda robot use wrist cameras [3] or joint-space actions [13]. While embracing such heterogeneity quickly scales the amount of available training data [7]. amplifies the importance of a fundamental question: how do we curate these raw, heterogeneous data sources into effective training datasets for generalist robot policies?
While early vision and language models were trained on highly- academic datasets such as ImageNet [14], questions surrounding data selection have shaped modern training pipelines that use Internet-scale data [15-17]. For example, training large language models involves numerous stages of data filtering [18). Simi- larly, large vision datasets, e.g. LAION [19], use pre-trained models to assess data quality. Thus as we scale robot data, curation weill become critical. Unfortunately, prior filtering techniques are often inadequate in robotics; we cannot use n-grams and visual embeddings do not capture the sequential nature of episodic data.
Even though aspects of demonstration data such as action quality [20] and visual diversity [3, 4, 21] have been shown to be of paramount importance to downstream performance, approaches for robotics data curation remain limited. In imitation learning (IL), the data selection problem has only been characterized theoretically [22, 23] or in simple small-scale settings [24]. Thus in practice we are left with ad hoc
'Corespondence to jhejna@cs. www.noord.ddeoo.o.. edu, Code: h/ght//b///tttbb.../. com/ hamm/jjejjnaeemei
8th Conference Robot Leaming (CoRL 2024), Munich, Germany.

---

solutions. For example, though the Open-X-Embodinen E-bboddment dataset (OpenX) [7] is comprised of more than 60 individual datasets totalling over 2M robot trajectories, the RT-X models released alongside it were trained on a mixture of only 12 datasets, weighted based on expert intuition. The recently released VLA [26] generalist policies were similarly trained on a subset of OpenX, where the authors chose which datasets to include at what sampling weight based on a subjective notion of "interestingness". While the resulting data mixes are shown work well in practice, their curation requires extensive domain knowledge and manual data inspection. Such ad hoc selection strategies are unlikely to scale to the rapidly growing datasets used to train robot policies (.
Octo [25] and LL
In this work, we ask: how can we automatically curate large-scale robotics datasets to maximize the performance of generalist IL policies across domains? Though many filtering techniques are not directly applicable to robotics, we can borrow ideas from language modeling that systematically optimize data mixtures based on model performance. Specifically, DoReMi [28] uses group distributionally robust opti- mization [29] maximize performance across "domains" in given dataset. In the context of robotics, such "domains" can correspond to different scenes within a single dataset, e.g., different toy kitchens in Bridge V2 dataset [4], refer to full robot datasets in mixtures such the OpenX dataset. However, due to the het- erogeneity of robotics datasets we find that naively applying such techniques does not work. Distributionally robust optimization approaches minimize worst-case loss. Differences in action spaces and their distribu- tions can cause loss magnitudes to be imbalanced across domains, leading some domains to be weighed more heavily than they should be. Moreover, the smaller size of robotics datasets makes overfitting easy. Both of these issues result in poor estimates of model performance, and consequently bad mixture weights.
To address these problems, we propose Re-weighing Robotic Dataset Mixtures with Minimax Optimization (Re-Mix for short), which instantiates the data curation problem as a min-max optimization, where policy minimizes its excess behavior cloning loss over a reference model subject to learned domain mixture weights that try to maximize it. Intuitively, the excess loss measures how much room the policy has to improve on a given domain, and the data mixture is optimized to maximize such improvement potential. Crucially, we carefully control the loss magnitudes between domains via domain-indeependet action normalization and discretization, even if the final policies we train are continuous diffusion models [30, 3 31]. Moreover, we find that selecting a reference model that has not overfit to any domain prevents drastic skewing of the downstream domain weights.
We empirically evaluate Re-Mix by using it to automatically optimize the training data mixture for the Bridge V2 dataset [4] and the OpenX-based dataset used to train RT-X [7]. We show that policies trained with our data mix improve performance by 38% and 32% respectively over naive data balancing and human-expert exrret-eurateedddd data mixtures in evaluations using WidowX and Franka robot arms Additionally, we show that weights from Re-Mix can effectively sub-sample both datasets, achieving competitive performance when using only 25% of the original data, while using uniform or human curated weights significantly reduces performance. Our contributions are as follows:
We introduce Re-Mix, method for automatically curating large-scale robotics datasets using group distributionally robust optimization over the behavior cloning loss.
We demonstrate Re-Mix's 's ability to curate effective training data mixtures for the Bridge [4] datasets and the subset of the OpenX dataset [7] used to train the RT-X models.
We curate 25% subsets of the Bridge and OpenX datasets which can be used for training generalist
policies with minimal loss in performance, while reducing the required compute budget.
2 Related Work
In congruence with the rise of deep learning in various fields, data selection has become of increasing interest. Here we review the most relevant works, organized by area.
The Data Problem in Robotics. Several recent works in robotics have focused on collecting large demon- stration datasets for imitation learning in simulation [20, 32, 33] and the real world [3, 7, 34-38] to train large-scale robot policies [6. 25, 39, 40]. Generally, these works along with others that study the influence of data collection on compositional generalization [21, 41, 42] show that aspects of dataset construction such as scene and task diversity have direct impact downstream policy generalization. Though several studies focus on how data should be collected via specific hardware [43], collection procedures [11, 21, 44],

---

or provide theoretic insights about data collection [22], little work in robotics addresses the post-hoc dataset selection and analysis problem. This is particularly important as the number and diversity of robot datasets are increasing with less clear conclusions about how to train a policy that effectively consumes all the col- lected data [3, 7, 25]. Baker et al. [45] train a classifier to predict data quality, but require human annotations which are impractical to scale. Perhaps most related are retrieval aalaaaedddddee methods that subset datasets [12, 46], but do so based on priori target task specifications and are thus inapplicable to training generalist policies.
Data Curation in Computer Vision. Computer vision datasets were originally hand-crafted and manually labeled [14, 47J. However, scaling datasets to beyond what is possible to curate by hand, while retaining quality, has been critical to increasing performance [1, 48]. Notably, filtering techniques based on metadata-count balancing [49]. embeddings [19], optical flow [50], and clustering [51] have shown to greatly improve downstream performance despite filtering out large amounts of data. At the extreme, coreset selection methods use active learning [52, 53], but have prohibitive computational requirements [54, 55] Data curation techniques from computer vision can only filter state-action trajectories in an action-agnostic manner potentalll removing useful parts of a dataset.
Data Curation in Natural Language Processing. When training on large real-world sources of text, language modeling pipelines employ a number of text-specific preprocessing steps including metadata, dialect, de-duplication, and toxicity filtering [16, 18, 56, 57]. More advanced methods also consider sub-setting data to maximize downstream performance, as in this work, but use techniques such as k-means clustering over embedded text [58, 59]. While such clustering techniques are potentially visually informative in robotics - similar to curation works in computer vision - they do not provide information about actions. Mixture techniques, such as Domain Reweighting with Minimax Optimization (DoReMi) [28] balance text domains using robust optimization and build upon ideas from prioritized training [60-62]. Our work is inspired by DoReMi such robust optimization approaches can be applied to imitation learning as well. In this work, we discuss the challenges of applying these techniques in robotics, and propose solution that addresses their limitations for effective dataset curation for imitation learning.
3 Re-weighing Robotic Dataset Mixtures with Minimax Optimization
In this section, we first formalize the problem of re-weighting data mixtures for II.. We then discuss our approach which uses distributionally robust optimization to select domain weights.
Problem Setup. We consider the general imitation learning problem, where we are given a dataset of demonstrations D= (TI...... consisting of state-action trajectories = (Tr Our goal is to learn a parameterized policy my that lears a mapping from states to actions S A. In practice, this often done through standard imitation learning algorithms such as behavior cloning (BC) by minimizing the expected negative ligeliiod of the actions under the policy:
(BC (To, D) = =\ -logo))(aao
(1)
However, datasets often contain more information than just state action pairs. We assume that the overall
dataset DD can be split into i heterogeneous domains D Dr. This a general assumption: while
these domains could be larger groups, like different datasets from the Open X-Embodiment dataset [7] with different embodiments, they could also be small single trajectories. Moreover, each of the k domains can differ in state space S, action space A. transition dynamics, or their distributions. In fact when learning large behavior models, such heterogeneity becomes necessary to access more sources of diverse data. In this work, we use the Bridge dataset [4] with different environments as the domains, and the --in-----modd dataset [7]- - with different robot embodiments as the differing domains.
Our goal is to learn a weighting vector a E A that specifies a probability distribution over all domains such that any model, when trained on domain mixture weighted according to a, attains maximum performance across all domains. Unlike the data retrieval, where data is selected for a particular target task, our goal is to curate datasets for effective pre- co-training without any knowledge of target task.
Distributionally Robust Optimization. When pre-training on large amounts of robot data we want policies to generalize to new settings and tasks, not master specific target task. With that in mind, we want to optimize for data mixture that results in models that i) can perform well as possible on
3

---

each domain, but ii) do not overfit to any one domain at the expense of another. Distributionally robust optimization (DRO) techniques aim to solve the same problem: learn models that minimize the worst-case training loss [29] - BC loss in the case of imitation learning across domains Di...t- Specifically, naively applying group robust optimization techniques in robotics would result in the following objective:
mingmax,ea* Ei= (i
With this objective, up-weights domains that have higher loss value, emphasizing the hardest domains. However, in practice we might not be interested in just fitting the domains with higher losses. For example, a dataset with complex rotations may always have higher BC loss than simple pick-place datasets. Thus, standard robust optimization techniques could end up ignoring the latter domain. Instead, as in prior work [28, 52, 63]. we consider the difference in loss between our learned policy TTg and reference policy Teet which is trained to convergence on an initial guess of the domain weights, usually assumed to be proportional to the size of each domain, i.e. uniform sampling. In Eq. (2) this equates to replacing Luc with (RC( D:) Lac (Ter, Dt). We refer to this difference as the excess loss, and use it for robust optimization. Like before, this will down-weight domains that the policy fits well, as it can achieve a loss similar to that of the reference model. However, it crucially also dewn-weights domains which are difficult to fit (i.e. they have a high policy loss and a high reference loss) due to the relative nature of the excess loss. Therefore, only domains that have a high excess loss, meaning the policy can improve to match the reference model, will be up-weighted as a chosen to maximize the excess overall loss.
Unfortunately, models trained directly with robust optimization often exhibit worse overall performance [64, 65]. Alternatively, we can use the learned a for downstream training in Xie et al. [28]. This gives us a set of reusable weights that can be used to train different policies without the need for robust optimization.
3.1 The Challenges of Applying Robust Optimization in Robotics
While Group DRO has been applied in language modeling [28], robust optimization techniques face unique challenges in robotics which we highlight here. We then detail how we adapt distributionally robust optimization pipeline select domain weights for roboties datasets.
Orpoisc
Ottridee
Unbalanced Losses. Large robotics datasets are often highly hetero- geneous: many are collected across different embodiments, controllers, frequencies and even units (c.g., inches vs meters). Within the same dataset, different scenes or tasks require vastly different ranges and speeds of motion. As a result, some datasets may have an outsized effect on robust optimization. To address this issue, one needs to align action losses across domains. Action normal- ization is often applied in imitation learning to standardize datasets to common distribution [20, 25J. In our case, we specifically apply Gaussian normalization to each embodiment individually. We note that bounds normalization [30] applied to each domain, would be insufficient as it would not align the moments of each domains action distribution. To underline the importance of aligning actions to a common distribution, we construct a simple experiment by training a policy with Group DRO [29] (Eq. (2)) when the action distributions match versus when they differ. Specifically, we construct a noise domain where a subset of the Bridge V2 dataset [4] is assigned random Gaussian actions and bridge domain which uses the original ac- tions, either normalized also be unit gaussian or rescaled between -1 and 1 using "bounds" normalization. When Gaussian normalization is applied to the bridge domain, the action distribution matches the random noise unlike with bounds normalization. We show the learned domain weights for each scheme in Table 1. While one might expect that a assigns majority weight to the bridge domain, as the noise domain is impossible to fit, this is actually only true in the "Gaussian" case when the action distributions of are aligned.
Bounds
0.943
0.057
Gaussian 0.158
0.842
Table 1: Learned a from toy set- ting in Section 3.1
Continuous Losses. Robust optimization has largely been applied in discrete classification problems with cross-entropy losses, for example in language modeling [63]. Popular policy learning approaches, however, often predict continuous actions and use LI or L2 loss functions [20, 30, 66, 67J. Applying robust optimization in these settings can be problematic for two reasons. First, action distributions can be multi- modal, and expressive continuous policy classes such as diffusion models only optimize an upper bound on the true loss. However, without the expressiveness to fit multi-modee distributions, both the reference policy and DRO would be unable to effectively minimize BC loss on domains with multi-moda actions. Second,

---

compared to language datasets, robot datasets often have a large number of action outliers which can heavily sway the value of continuous action losses. With L.1 or 1.2 loss, these outliers can significantly increase the loss of a given domain, causing DRO to believe it can still make progress. To resolve these problems, when applying robust optimization in the robotics domain, we discretize each action dimension via binning.
Overfitting. Datasets in language modeling often contain billions of tokens. As a result, robust optimization techniques like Xie et al. [28] do not experience overfitting when applied to these large scale datasets. On the other hand, large robot datasets comparatively small (~100-10Kk demonstrations). In this regime, it is common for high-capacity policies to achieve near-zero training loss for every datapoint [6, 25, 26, 683. This is problematic when using the excess loss for robust optimization: if the reference model achieves near-zero training loss on every data point within a domain, the excess loss is equivalent to the regular loss (since the reference loss is always = and a no longer reflects the potential for improvement on each domain. To counteract this problem, we employ aggressive early stopping on both the reference model and robust optimization. Specifically, we select the latest checkpoint from the reference model that has not overfit to any of the domains Di, as measured by the difference in loss values between the training dataset and a held-out validation dataset for the respective domain. This ensures that the reference model does not overfit to any individual domain and the learned weights a are informative.
3.2 Re-weighing Robotic Dataset Mixtures with Minimax Optimization
Our approach, Re-Mix, uses group distributional robustness to determine the weights of a data mixture [28] that could then be used for policy training and incorporates the key design considerations from the previous section, addressing issues around unbalanced losses, continuous losses, and overfitting. We note that Re-Mix only returns the weights of the data mixture o, as opposed to the final policy. This is to decouple the data curation problem from the policy training problem. After running Re-Mix, the resulting weights can be used for learning policies of a different type (i.e. diffusion) or at a larger scale.
Stage 1: Action Preprocessing. Following Section 3.1, we apply Gaussian normalization separately to every domain Dr with different action spaces and dynamies, and then discretize actions via binning.
Stage 2: Reference Model Training. Next, train discrete reference model Ter on the uniform mixture of domains D Dkk where each domain is weighted in proportion to its size. We select the final reference model checkpoint by validation loss per Section 3.1, and use it to estimate the excess loss per domain.
Stage 3: Group Distributionally Robust Optimization. We learn the domain weights via the following robust optimization with a discrete policy TT²
min max DEAR
aN
Di
Σ +ogggr (s,a)ED,
(3)
which minimizes the worst case excess BC loss of the learned policy (ags) + Log))rr(als) over all possible weightings of the domains a E A To update following [29]. we perform one step of exponentiated gradient ascent on a followed by domain-weighted gradient descent on 0 at each training step. Our resulting values of a upweight domains that we can still improve on, while downweighting domains that are trivial or impossible to fit. This means that Re-Mix directly filters data based on actions, unlike other techniques in vision and language that solely filter based on embeddings [59, 69]. We optimized Eq. (3) for the same number of steps as the reference model.
Stage 4: Data Weighting for Policy Training, After robust optimization over the excess loss, take the average value of a over the course of training, which we denote by ā. We can then use this value of a to re-weight different domains, or even subset datasets for policy training. In practice, this means that we can re-use the weights for several training runs with different configurations. For example, Re-Mix uses discrete actions, but we train final policies with diffusion which has shown to perform well empirically [25, 30].
4 Experiments
We aim to answer the following questions: (1) Does Re-Mix effectively curate large robot datasets for down- stream policy learnin?? (2) Can use Re-Mix to heavily sub-san sub-sample robot datasets while retaining good performance? (3) Which design decisions matter for effective automatic curation of large robot datasets?
5

---

1.0
1.0
REX Mix
0.80
0.4
0..9
Sucoase Rate
9.8
0,42
0.1
2.0
Average
Carrut to Rack
Pork Rack
OOD Cup
Cube Plate
Flip Bowl
Pen
Human
Figure 2: Results for curating the RT-X training mix. We test policies trained on different weightings of the data WidowX (left) and two Franka (right) tabletop manipulation tasks. We find that data mix curated with Re-Mix achieves strongest performance, even outperforming the
mixture used by RT-X the policy trained human-expent expert-curateedd data mix from RT-X [7]. Mean 1 StdErr across 4 tasks, 10 evaluations each.
4.1 Experimental Setup
Datasets, We test Re-Mix curation on two widely- used, large-scale robot datasets: (1) the Bridge V2 Dataset [4], consisting of 50k diverse teleoperated demonstrations of single-arm manipulation tasks with WidowX 6 DoF robot arm, and (2) the datasets from the Open X-Embodiment dataset used train RT-1-X and RT-2-X models [7] which have third-person person cameras, consisting of a total of 350k demonstrations which span disparate embodiments and environments. We use "RT-X" to refer to this set of datasets. We partition the Bridge V2 dataset into 32 domains based on the scenes the data was collected in. For OpenX, we use each of the 1l datasets the RT-X training set as a domain for our curation experiments. The OpenX data mix is particularly challenging effective curation due its heterogeneous data sources. For detailed list of all datasets and partitions, Appendix B. For simulation experiments see Appendix A.
Training and Evaluation Details. We aim to assess the quality of vari- ous curated pre-training data mixtures for downstream policy learning, To that end, we co-train generalist goal-conditioned policies on the curated datasets. As we not have access the robot setups used to collect the datasets we train on, we construct our own WidowX and Franka robot evaluation setups. Unfortunately, policies trained only the pre-training data failed to zero-shot generalize
Bridge
1.0
1.0
Uniform
ReMix
Success Rate
0.11
0.5
0.8
0.4
out-of-dit titto--diitribution setups. To ad- dress this, we follow prior works [3, 12, 46, 70] and co-train our policies on a small amount of in-domain data (25 demonstrations each for 3 rep- resentative tasks), added to the final training mixture at small weight of 5%. We then evaluate policies on tasks that are outtof--dstribution with respect the co-training data to test generalization. As result of co-training, all policies achieve non-zero success rate. However, we note that the in-domain dataset is small enough that the quality of the pre- pre-traine training data mix still has significant impact the evaluation result, providing a good test bed for data curation approaches. All models are evaluated in the real world with 10 trials task totaling over 500 real-world trials cumulatively. For all policies we use ResNet 50 image encoder [71]. For the Re-Mix reference model and Group DRO optimization, we use discrete MLP action head. For final policies we use the diffusion head from [4, 25, 72] and train all models for 400,000 gradient steps.
Carrot Rack Fork to Rack
Figure 1: Bridge V2 [4] there notable difference between uniform sampling vs. Re-Mix when training the full dataset.
Comparisons. We compare the quality of Re-Mix's curated data mixes naive baseline: sampling uniformly from each domain according the total number of state-action pairs (Uniform). For evaluations on the OpenX datasets, we additionally compare to a wuwanuumeppppp -eppee--ppetteeeeeeeeeeeeee cureeeed data mix, using the hand-crafted weights from RT-X [7]. For Bridge there is no expert-curated data mix uniform sampling the norm.
4.2 How do Re-Mix weights impact performance?
In Fig. show results for weighing datasets from the RT-X mix according different methods. For the WidowX robot, consider four tasks that test generalization 1) unseen objects: "Carrot to Rack", "OOD Cup", 2) unseen initial conditions: "Fork to Rack", and 3) distractors the goal location "Cube to

---

Method OURs
Kenting
Claco
Orbec
Orboo
Oiviet
Umiton 1.01%
0.43%
22.7%
0.81%
24.9%
1.94%
40.9%
0.60%
2.46%
3.42%
0.80%
Human
1.22%
1.56%
27.5%
1.95%
25.1%
2.35%
26.8%
1.40%
5.94%
4.13%
1.90%
Re-Mix
2.37%
0.20%
19.9%
0.39%
12.1%
1.14%
42.5%
0.63%
3.04%
16.3%
1.51%
Table 2: Dataset mixture weights by different methods the RT-X dataset mix 6,
10, 37, 73-76].
We color
relative increases of more than 25% from uniform green and relative decreases more than 25% red.
RIX-Mix 25% Subset
1.0
0.80
0.9
Human
Uniform
0.8
0.6
0.6
0.38
0.5
0.4
0.S
0.2
Atrerage
Carrol to Rack
Fork to Rack
Flip Bowl
Pen in Cup
0.90
Bridge 25% Subset
1.0
1.0
Untform
0.50
0.7
ReMa
0.6
0.4
0.0
Carrot Kack
Cube
Cube Plate
0.0
Carrot Right
Figure 3: Results sub-setting datasets via different strategies until they reach 25% of their original size. We again use 10 evaluations per task, and show the Mean 1 StdEm.
Plate". Similarly, for the Franka Panda robot we consider two tasks that test generalization to unseen initial conditions "Pen in Cup" and 2) motions not seen in the RT-X data "Flip Bowl". Additionally, our Panda rohot uses a Robotiq 2F-85 gripper, which was not present in any of the RT--X-datasets. Note that for the RT-X mix, we co-train the same model on both the Widow and Franka data. As expected, we find that the domain weights selected by human experts for the RT-X models outperform the naïve uniform sampling baseline by 6% on average. More interestingly, we find that weighting datasets according to Re-Mix outperforms uniform weighting by 38% average, and surprisingly outperforms the human curated weights by 32% average. We find that Re-Mix generalizes better unseen objects and locations in Bridge. Re-Mix performs particularly performss well for the "Flip Bowl" tasks, which is potentially because up-weights the relatively small "Toto" dataset in OpenX that contains similar pouring motions to the flip blow task, though performance in Pen-in-cup is similar across all models. We hypothesize this because requires less generalization than the other tasks, and thus depends less on the pre-training mixes. Fig. shows results using Re-Mix weights versus uniform weighting over scenes in the Bridge dataset. We find that performance in this setting is similar across both models. We posit that in the presence of the full Bridge dataset, selecting weightings is less important as the model is able to fit every scene well.
4.3 Analyzing Re-Mix Weights
Table 2 shows the weights produced by different methods on the RT-X dataset mix in comparison to the uni- form mixture, which corresponds to sampling each datapoint with equal probability or equivalently weight- ing each domain by its total size (as fraction of the total number of datapoints). The human- numan-expert--ee -et-ee--ddessineee weights largely nown-weight RT-1 [6], while up-weighting some of the smaller datasets like Routing [73]. and Taco [9], perhaps ensure they were sampled often enough to not be ignored. On the other hand, Re- Mix largely down-weights the Kuka dataset [76]. This dataset was autonomously collected and then filtered by success, making it of potentially lower action quality. Re-Mix also down- down-weights some smaller domains that are easy to fit; for example, Cable Routing has gripper actions and Jaco [74] only has three possible actions. Surprisingly, Re-Mix up-weights the Toto dataset [77] by more than 4x. We posit that this is because Toto has particularly multi-modal modal action distribution which deviates far from standard Gaussian even after normalization and thus may be more challenging to fit. See Appendix A for a plot of its action distribution.
4.4 How well does Re-Mix subset datasets?
Though co-training on diverse data important for performance [3, 70], doing so often expensive given that modem robot datasets like the OpenX dataset encompass TBs of data. In this section, we evaluate how well Re-Mix be used subset datasets. The key idea: if Re-Mix weights proportional to the importance of the data in each domain, can use them to effectively sub-set the dataset by removing

---

data from domains that Re-Mix assigns low weight. We subset the base datasets according to Re-Mix and baselines by first computing the target size of the entire data mix after sub-setting, in our case 25% of (D). Then, we remove datapoints according to the mixture weights ã. If a small dataset is upweighted too much (i.e. if domain is 1% of [D] but is upweighted to d;; 5%), there might not be enough data to exactly match ā from subsetting alone. Thus, after subsetting we sample the remaining points uniformly until we reach the target size. During training still weight datasets according to the exact a.
The goal of subsetting is retain performance when data is removed. We compare performance of subsetting with the Re-Mix weights to using the na've uniform and human expert weights. For Bridge, where no expert weighting exists, we additionally compare to a vision subsetting method called "Self-Supervised Prototypes" (SSP) [69] which runs k-means on image embeddings and discards data closest to each centroid to encourage diversity. We average CLIP embeddings across each trajectory and use 32, matching the number of domains used by Re-Mix. To provide more extensive evaluation on Bridge, we add two additional tasks. See Appendix B for details.
Subsetting results are shown in Fig. Overall, we find that subsetting exacerbates the difference between methods, the weights directly affect dataset composition. On the RT-X datasets (Fig. 3 top row) with only 25% of the data Re-Mix retains performance while human weights drop over 10%. This is likely because as shown in Table 2, Re-Mix places higher weights on some of the smaller datasets and down-weights some of the larger datasets such as the Kuka dataset from [40]. For example, when using domain weights for subsetting Re-Mix retains 72% of the URS Dataset but only 12% of the Kuka dataset, while the expent human weights retain only 30% of the URS Dataset but 24% of Kuka which is from suboptimal learned policies. Generally, the human weights do not down- nown-weight datasets that are potentially uninteresting or not useful for training. On Bridge (Fig. bottom row), Re-Mix also outperforms baseline methods. Overall SSP performs poorly, likely since robot trajectories may be aattof-distributinn for vision models such as CLIP.
4.5
What matters in Re-Mix?
In this section, we ablate several design choices used in Re-Mix (see Section 3.1), including action discretization and early stopping. We run all ablations in the 25% subset setting (see Section 4.4), since subsetting further amplifies the effects of the domain weights. In Fig. 4, we first analyze the effects of choosing a reference model checkpoint for Group DRO that is overfit to the training dataset. Empirically, we find that choosing a checkpoint just 50K steps after early stopping decreases performance by over 15% on average, likely because the reference model baseline used to determine the domain weights is less meaningful once it overfits. On the right half of Fig. 4, we show performance on Bridge when using continuous (Cont.) actions in Re-Mix instead of discrete for estimating a. We find that continuous actions lead significantly worse performance, as their loss functions fail to fit outliers or multi-modal actions.
5 Limitations and Future Work
In this work we present Re-Mix, method for automatically curating robotics datasets using distributionally robust optimization.
Evaluation. While we train on large, diverse robot datasets, the need for real world trials makes it difficult to exhaustively evaluate trained generalist policies many robot embodiments and setups. While our evaluations capture two widely used robot arms from prior works [4, 7, 252,, WidowX and Franka, future work should extend to more embodiments, perhaps via simulated environments [68].
Abnormal Action Distributions. We have noticed that Re-Mix upweights datasets with abnormal action distributions such as the Toto dataset. While resulting data mixes performed well, such up-weighting is not necessarily desirable. We hope to achieve less sensitivity to such irregularities in future work.
Computational Cost. Using our pre-computed weights can significantly reduce the compute required to train generalist policies. However, our approach for computing Re-Mix weights requires training policies on the full data twice, once for the reference model and once for Group DRO optimization. Future work can instead strive curate datasets con-the-fly"withhe one run.
Scaling Up. While we have demonstrated improvements on two large datasets, Bridge V2 and RT-X, scaling up even larger ones such as the entire OpenX dataset [7] (>2M episodes) is exciting extension.
8

---

Acknowledgments
Compute for this rescarch was provided by a Google TPU Research Cloud Grant. This work was supported by NSF #1941722, ONR project #N00014-22-112295 DARPA grant IN922NF210214, Stanford HAI, and TRI. JH is supported by an NDSEG Fellowship.
References
[1] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 7738-766.. PMLR, 2021.
[2] A. Conneau, K. Khandelwal, N. Goyal, V. Chaudhary, G. Wenzek, F. Guzmán, E. Grave, M. Ott, L. Zettlemoyer, and V. Stoyanov. Unsupervised cross-lingual representation learning scale. arXiv preprint 1rxi::911.02116 2019.
[3] A. Khazatsky, K. Pertsch, S. Nair, A. Balakrishna, S. Dasari, S. Karamcheti, S. Nasiriany, M. K. Srirama, L. Y. Chen, K. Ellis, al. Droid: A large-scale in-the-wild robot manipulation dataset. arXiv preprint arXiv:2403. /2945 2024.
[4] H. Walke, K. Black, A. Lee, M. Kim, M. Du, C. Zheng, T. Zhao, P. Hansen-Fstruchh Q. Vuong, A. He, Myers, K. Fang, C. Finn, and Levine. Bridgedata v2: A dataset for robot learning at scale. In Conference on Robot Learning (CoRL), 2023.
[5] H.-S. Fang, H. Fang, Z. Tang, J. Liu, J. Wang, H. Zhu, and C. Lu. Rh20t: A robotic dataset for learning diverse skills in one-shot. In RSS 2023 Workshop on Learning for Task and Motion Planning, 2023.
[6] A. Brohan, N. Brown, J. Carbajal, Y. Chebotar, J. Dabis, C. Finn, K. Gopalakrishnan, K. Hausman, A. Herzog. J. Hsu, J. Ibarz, B. Ichter, A. Irpan, T. Jackson, S. Jesmonth, N. Joshi, R. Julian, D. Kalashnikov, Y. Kuang, I. Leal, K.-H. Lee, S. Levine, Y. Lu, U. Malla, D. Manjunath, I. Mordatch, O. Nachum, C. Parada, J. Peralta, E. Perez, K. Pertsch, J. Quiambao, K. Rao, M. Ryoo, G. Salazar, P. Sanketi, K. Sayed, J. Singh, S. Sontakke, A. Stone, C. Tan, H. Tran, V. Vanhoucke, S. Vega, Q. Vuong, F. Xia, T. Xiao, P. Xu, S. Xu, T. Yu, and B. Zitkovich. Rt-1: Robotics transformer for real-world control at scale. In arXiv preprint arXiv:22 2tt.:222...011..7 2022.
[7] Open X-Embodiment Collaboration, A. Padalkar, A. Pooley, A. Jain, A. Bewley, A. Herzog, A. Irpan, A. Khazatsky, A. Rai, A. Singh, A. Brohan, A. Raffin, A. Wahid, B. Burgess-Limerich B. Kim, B. Schölkopf. B. Ichter, C. Lu, C. Xu, C. Finn, C. Xu, C. Chi, C. Huang. C. Chan, C. Pan, C. Fu, C. Devin, D. Driess, D. Pathak, D. Shah, D. Büchler, D. Kalashnikov, D. Sadigh, E. Johns, F. Ceola, F. Xia, F. Stulp, G. Zhou, G. S. Sukhatme, G. Salhotra, G. Yan, G. Schiavi, H. Su, H.-S. Fang, H. Shi, H. B. Amor, H. I. Christensen, H. Furuta, H. Walke, H. Fang, I. Mordatch, I. Radosavovic, I. Leal, J. Liang, J. Kim, J. Schneider, J. Hsu, J. Bohg. J. Bingham, J. Wu, J. Wu, J. Luo, J. Gu, J. Tan, J. Oh, J. Malik, J. Tompson, J. Yang. J. J. Lim, J. Silvério, J. Han, K. Rao, K. Pertsch, K. Hausman, K. Go, K. Gopalakrishnan, K. Goldberg, K. Byrne, K. Oslund, K. Kawaharazuka, K. Zhang, K. Majd, K. Rana, K. Srinivasan, L. Y. Chen, Pinto, L. Tan, Oit, L. Lee, M. Tomizuka, M. Du, M. Ahn, M. Zhang, M. Ding, M. K. Srirama, M. Sharma, M. Kim, N. Kanazawa, N. Hansen, N. Heess, N. Joshi, N. Suenderhauf, N. D. Palo, N. M. M. Shafiullah, O. Mees, O. Kroemer, P. R. Sanketi, P. Wohlhart, P. Xu, P. Sermanet, P. Sundaresan, Q. Vuong, R. Rafailov, R. Tian, R. Doshi, R. Martin-Martin, R. Mendonca, R. Shah, R. Hoque, R. Julian, S. Bustamante, S. Kirmani, S. Levine, S. Moore, S. Bahl, S. Dass, S. Song, S. Xu, S. Haldar, S. Adebola, S. Guist, S. Nasiriany, S. Schaal, S. Welker, S. Tian, S. Dasari, S. Belkhale, T. Osa, T. Harada, T. Matsushima, T. Xiao, T. Yu, T. Ding, T. Davchev, T. Z. Zhao, T. Armstrong, T. Darrell, V. Jain, V. Vanhoucke, W. Zhan, W. Zhou, W. Burgard, X. Chen, X. Wang. X. Zhu, X. Li, Y. Lu, Y. Chebotar, Y. Zhou, Y. Zhu, Y. Xu, Y. Wang, Y. Bisk, Y. Cho, Y. Lee, Y. Cui, Y. hua Wu, Y. Tang, Y. Zhu, Y. Li, Y. Iwasawa, Y. Matsuo, Z. Xu, and Z. J. Cui. Open X-Embodiment: Robotic learning datasets and RT-X models. bttps://arzen org/ /2300 08864, 2023.

---

[8] J. H. Yang, D. Sadigh, and C. Finn. Polybot: Training one policy across robots while embracing variability. In 7th Annual Conference on Robot Learning, 2023. URL htt.httts///ppereeiiee net/ TorumPid=HEIRj511e..
[9] E. Rosete-Beas, O. Mees, G. Kalweit, J. Boedecker, and W. Burgard. Latent plans for task agnostic offline reinforcement learning. In Proceedings of the 6th Conference Robot Learning (CoRL), 2022.
[10] Y. Zhu, A. Joshi, P. Stone, and Y. Zhu. Viola: Imitation learning for vision- based manipulation with object proposal priors. 6th Amnual Conference on Robot Learning (CoRL.), 2022.
[11] S. Belkhale, Y. Cui, and D. Sadigh. Hydra: Hybrid robot actions for imitation learning, In Conference on Robot Learning, pages 2113-2133, PMLR, 2023.
[12] S. Nasiriany, T. Gao, A. Mandlekar, and Y. Zhu, Learning and retrieval from prior data for skill-based imitation leamingg In Conference on Robot Learning (CoRL), 2022.
[13] H. Bharadhwaj.. J. Vakil, M. Sharma, A. Gupta, S. Tulsiani, and V. Kumar. Roboagent: Generalization and efficiency in robot manipulation via semanntc augmentations and action chunking, 2023.
[14] J. Deng, W. Dong, Socher, L..J. Li, K. Li, and Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248 255, 2009.
[15] L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, S. Presser, and C. Leahy. The Pile: An 800gb dataset diverse text for language modeling. arXiv preprint arXiv:
2020.
[16] G. Penedo, H. Kydlfčck, L. Werra, and T. Wolf. Fineweb, April 2024. URL /ttps:// huggingface. co/ dotaaess
[17] K. Grauman, A. Westbury, E. Bymee Z. Chavis, A. Furnari, R. Girdhar, Hamburger, H. Jiang, M. Liu, X. Liu, al. Ego4d: Around the world in 3,000 hours of egocentric video. Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 18995-19012, 2022.
[18] A. Albalak, Y. Elazar, S. M. Xie. S. Longpre. N. Lambert, X. Wang, N. Muennighoff, B. Hou, L. Pan, H. Jeong, C. Raffel, S. Chang, T. Hashimoto, and W. Y. Wang. A survey on data selection for language models, 2024.
[19] C. Schuhmann, R. Beaumont, R. Vencu, C. Gordon, R. Wightman, M. Cherti, T. Coombes, A. Katta, C. Mullis, M. Wortsman, et al. Laion-5b: An open large scale dataset for training next generation image-text models. Advances in Neural Information Processing Systems, 35:25278 25294, 2022.
[20] A. Mandlekar, D. Xu, J. Wong, S. Nasiriany, C. Wang, R. Kulkamni, L. Fei-Fei, S. Savarese, Y. Zhu, and R. Martin-Martin. What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint arXiv:2108 03298, 2021.
[21] J. Gao, A. Xie, T. Xiao, C. Finn, and D. Sadigh. Efficient data collection for robotic manipulation via compositional generalization, 2024.
[22] S. Belkhale, Y. Cui, and D. Sadigh. Data quality in imitation learning. Advances Neural Information Processing Systems, 36, 2024.
[23] S. Ross, G. Gordon, and D. Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 627-635. JMLR Workshop and Conference Proceedings, 2011.
[24] K. Gandhi, S. Karamcheti, M. Liao, and D. Sadigh. Eliciting compatible demonstrations for multi- human imitation learning. In Conference on Robot Learning, pages 1981-1991 PMLR, 2023.
10

---

[25] Octo Model Team, D. Ghosh, H. Walke, K. Pertsch, K. Black, O. Mees, S. Dasari. J. Hejna, C. Xu, J. Luo, T. Kreiman, Y. Tan, I. Y. Chen, P. Sanketi, Q. Vuong, T. Xiao, D. Sadigh, C. Finn, and
S. Levine. Octo: An open-source generalist robot policy. In Proceedings of Robotics: Science and Systems, Delft, Netherlands, 2024.
[26] M. Kim, K. Pertsch, S. Karamcheti, T. Xiao, A. Balakrishna, Nair, R. Rafailov, E. Foster, G. Lam, P. Sanketi, Q. Vuong, T. Kollar, B. Burchfiel, R. Tedrake, D. Sadigh, S. Levine, P. Liang, and C. Finn. Openvla: An open-source vision-language-actiie action model. arXiv preprint 2406:2400000266, 2024.
[27] S. Nasiriany, A. Maddukuri, L. Zhang, A. Parikh, A. Lo, A. Joshi, A. Mandlekar, and Y. Zhu. Robocasa: Large-scale simulation of everyday tasks for generalist robots. In Robotics: Science and Systems (RSS), 2024.
[28] S. M. Xie, H. Pham, X. Dong, N. Du, H. Liu, Y. Lu, P. S. Liang, Q. V. Lc, T. Ma, and A. W. Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 69798-69818 Curran Asso- ciates, Inc., 2023. URL httpp::/prooceedings Conferrnooffeeccceeeoo pdf.
mewpiiinnnnnnnnuinseeeeee neurips wwterips.ccpppppee htc/paper_filee wallee/ppppr//022/
[29] S. Sagawa, P. W. Koh, T. B. Hashimoto, and P. Liang. Distributionallyy robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv:1911.0874 2019.
[30] C. Chi, S. Feng, Y. Du, Z. Xu, E. Cousineau, B. Burchfiel, and S. Song. Diffussoo policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
[31] M. Reuss, M. Li, X. Jia, and R. Lioutikov. Goal conditioned imitation learning using cosee-based diffusion policies. In Robotics: Science and Systems, 2023.
[32] O. Mees, L. Hermann, E. Rosete-Beat Beas, and W. Burgard. Calvin: A benchmark for language- conditioned policy learning for long-horizon robot manipulation tasks. IEEE Robotics and Automation Letters (RA-L), 7(3)773277734,, 2022.
[33] H. Ha, P. Florence, and S. Song. Scaling up and distilling down: Language Language-guided guided robot skill acquisition. In Proceedings of the 2023 Conference on Robot Learning, 2023.
[34] E. Jang, A. Irpan, M. Khansari, D. Kappler, F. Ebert, C. Lynch, S. Levine, and C. Finn. Be-z: Zero-shot task generalization with robotic imitation learnin.. In Conference on Robot Learning, pages 991-1002. PMLR, 2022.
[35] S. Dasari, F. Ebert, S. Tian, Nair, B. Bucher, K. Schmeckpeper, S. Singh, S. Levine, and C. Finn. Robonet: Large-scale multi-robot learning. arXiv preprint 121.11111. 2019.
[36] P. Sharma, L. Mohan, L. Pinto, and A. Gupta. Multiple interactions made easy (mime): Large scale demonstrations data for imitation. In Conference on robot learning, pages 906-915. PMLR, 2018.
[37] Mandlekar, Y. Zhu, A. Garg, J. Booher, M. Spero, A. Tung, Gao, J. Emmons, A. Gupta, F. Orbay, et al. Roboturk: A crowdsourcing platform for robotic skill leamingg through imitation. In Conference Robot Learning, pages 879-893. PMLR, 2018.
[38] L. Pinto and A. Gupta. Supersizing sulersuper Learning to grasp from 50k tries and 700 robot hours. In 2016 IEEE international conference on robotics and automation (ICRA), pages 3406-3413. IEEE, 2016.
[39] B. Zitkovich, T. Yu, Xu, P. Xu, T. Xiao, F. Xia, Wu, P. Wohlhart, S. Welker, A. Wahid, et al. Rt-2: Vision-language-action models transfer web knowledge to robotic control. In Conference Robot Learning, pages 2165-2183, PMLR, 2023.
11

---

[40] S. Levine, P. Pastor, A. Krizhevsky, J. Ibarz, and D. Quillen. Learning hand-eye coordination for robotic grasping with deep learning and large-scale data collection. The International journal of robotics research, 37(4-5):421 436, 2018.
[41] K. Burns, Z. Witzel, J. 1. Hamid, T. Yu, C. Finn, and K. Hausman. What makes pre-trained visual representations successful for robust manipulation? arXiv preprint arXtv:2312.12444 2023.
[42] A. Xie, L. Lee, T. Xiao, and C. Finn. Decomposing the generalization gap in imitation learning for visual robotic manipulation. arXiv preprint arXiv:2307.0383 0036999,,....... 2023.
[43] S. Young. D. Gandhi, S. Tulsiani, A. Gupta, P. Abbeel, and L. Pinto. Visual imitation made easy. In J. Kober, Ramos, and C. Tomlin, editors, Proceedings of the 2020 Conference on Robot Learning, volume 155 of Proceedings of Machine Learning Research, pages 1992-2005. PMLR, 16-18 Nov 2021. URL https: / petttps:/prroceeeee pttp://eeoc/egees.eeeoeeee/ mlr. www.mml....p.....mmm preas//115/yoougg222 httm..
[44] M. Laskey, Lee, R. Fox, A. Dragan, and K. Goldberg. Dart: Noise injection for robust imitation learnig. In Conference robot learning, pages 143 156. PMLR, 2017.
[45] B. Baker, Akkaya, Zhokov, Huizinga, J. Tang, A. Ecoffet, Houghton, R. Sampedro, and
J. Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances in Neural Information Processing Systems, 01122433-244655 2022.
[46] M. Du, S. Nair, D. Sadigh, and C. Finn. Behavior retrieval: Few-shot imitation learning by querying unlabeled datasets. arXiv preprint arXiv:2304.08742, 2023.
[47] A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[48] M. Oquab, T. Darcet, T. Moutakanni, H. Vo, M. Szaframiec, V. Khalidov, P. Fernandez, D. Haziza, F. Massa, A. El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint 3rX11:2304.071933 2023.
[49] H. Xu, S. Xie, X. Tan, P.-Y. Huang, R. Howes, V. Sharma, S.-W. Li, G. Ghosh, I.. Zettlemoyer, and C. Feichtenhofer. Demystifying CLIP data. In The Twelfth International Conference on Learning Representations, 2024. URL https: pttttp:://ppnsee /ttp///eeeeeeeeeee/ee/eeeeeeeeee wwt//errunrrieee htrmm/ormmm/// dBC.ddd
[50] A. Blattmann, T. Dockhorn, S. Kulal, D. Mendelevitch, M. Kilian, D. Lorenz, Y. Levi, Z. English, V. Voleti, A. Letts, et al. Stable video diffusion: Scaling latent video diffusion models to large datasets. arXiv preprint arXh: 2023.
[51] V. Vo, V. Khalidov, T. Darcet, T. Moutakanni, N. Smetanin, M. Szafraniec, H. Touvron, C. Couprie, M. Oquab, A. Joulin, al. Automatic data curation for self-supervised learning: A clustering- based approach. arXiv preprint a3X5v:2405.15613, 2024.
[52] Chitta, M. Álvarez, E. Haussmann, and C. Farabet. Training data subset search with ensemble active learning. IEEE Transactions on Intelligent Transportation Systems, 23(9):14741-1775 2021.
[53] O. Sener and S. Savarese. Active learning for convolutional neural networks: A core-set approach. arXiv preprint 178.::1708.0088, 2017.
[54] B. Mirzasoleiman, J. Bilmes, and J. Leskovec. Coresets for data-efficient training of machine learning models. In International Conference on Machine Learning, pages 6950-6960 PMLR, 2020.
[55] V. Birodkar, H. Mobahi, and S. Bengio. Semantic redundancies in image-classfficttion datasets: The 10% you don't need. arXiv preprint arXiv: 2019.
[56] T. Computer. Redpajama: open dataset for training large language models, October 2023. URL attps://giiie cow honm//oeehherrccmttteeeeeeeee
-aaa.
12

---

[57] L. Soldaini, R. Kinney, A. Bhagia. D. Schwenk, D. Atkinson, R. Authur, B. Bogin, K. Chandu, J. Dumas, Y. Flazar, V. Hofmann, A. H. Jha, S. Kumar, I., Lucy, X. Lyu, N. Lambert, I. Magnusson, J. Morrison, N. Muennighoff, A. Naik, C. Nam, M. E. Peters, A. Ravichander, K. Richardson, Z. Shen, E. Strubell, N. Subramani, O. Tafjord, P. Walsh, L. Zettlemoyer, N. A. Smith, H. Hajishirzi, I. Beltagy, Groeneveld. J. Dodge, and Lo. Dolma: An Open Corpus of Three Trillion Tokens for Language Model Pretraining Research. arXiv preprint, 2024. URL https:://aatt:ti org/ abs awwaaaa/////
[58] K. Tinumala, D. Simig, A. Aghajanyan, and A. Morcos. D4: Improving Ilm pre- training via document de-duplication and diversification. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neu- ral Information Processing Systems, volume 36, pages 53983-53998 Curran Associates, Inc., 2023. URL https: httpss:/prroceedings wwt.uris..ccpaaaaaeeeeeee httppppccceppppeeeee/eeeeeeeeeee /ttls//ppper2222///aas. Paaaaaaasttppeee---attttttsssss and__eechmarrss pdf.
[59] A. Abbas, K. Tinumala, D. Simig, S. Ganguli, and A. S. Morcos. Semdedup: Data-efficient learning at web-scale through semantic deduplication, 2023.
[60] S. Mindermann, J. M. Brauner, M. T. Razzak, M. Sharma, A. Kirsch, W. Xu, B. Höltgen, A. N. Gomez, A. Morisot, S. Farquhar, and Y. Gal. Prioritized training on points that are learnable, worth learning, and not yet learnt. In K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 15630-15649, PMLR, 17-23 Jul 2022. URL https: / http///ereeeee..ee///////o .mm.iiinnnnmm..m...mmmm...m....m httpa///1..///mmdddmmma2222.... ltml.
[61] A. H. Jiang, D. L.-K. Wong, G. Zhou, D. G. Andersen, J. Dean, G. R. Ganger, G. Joshi, M. Kaminksy, M. Kozuch, C. I Lipton, et al. Accelerating deep learning by focusing on the biggest losers. arXiv 2019.
preprint arXiv:1910.00770
[62] M. Paul, S. Ganguli, and G. K. Dziugaite. Deep learning on a data diet: Finding important examples early in training. Advances in Neural Information Processing Systems, 34:20596 20607,, 2021.
[63] Y. Oren, S. Sagawa, T. B. Hashimoto, and P. Liang. Distributionally robust language modeling. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 2019.
[64] H. Zhang, Y. Yu, Jiao, E. Xing, L. El Ghaoui, and M. Jordan. Theoretically principled trade-off between robustness and accuracy. In International conference on machine learning, pages 7472-7482. PMLR, 2019.
[65] D. Tsipras, S. Santurkar, I., Engstrom, A. Tumner and A. Madry. Robustness may be at odds with accuracy, arXiv preprint arXiv: 12152, 2018
[66] T. Z. Zhao, V. Kumar, S. Levine, and C. Finn. Learning fine-grained bimanual manipulation with low-cost hardware. arXiv preprint ar 2iv:22334. 137005 2023.
[67] F. Ebert, Y. Yang, K. Schmeckpeper, B. Bucher, G. Georgakis, K. Daniilidis, C. Finn, and S. Levine. Bridge data: Boosting generalization of robotic skills with cross-domain datasets. arXiv preprint arXtv:2109.13396, 2021.
[68] X. Li, K. Hsu, J. Gu, K. Pertsch, O. Mees, H. R. Walke, C. Fu, I. Lunawat, I. Sieh, S. Kirmani, S. Levine, J. Wu, C. Finn, H. Su, Q. Vuong, and T. Xiao. Evaluating real-world robot manipulation policies simulation. arXiv preprint arXiv:2405.0m 05944, 2024.
[69] B. Sorscher, R. Geirhos, S. Shekhar, S. Ganguli, and A. Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 19523- 19536. Curran Associates, Inc., 2022. URL https: https::/ nww.iis.i.... ctcpppscccpppppppp www.ripp..cppppr..aaeei.eeeeeee. påf.
13

---

[70] Z. Fu, T. Z. Zhao, and C. Finn. Mobile aloha: Learning bimanual mobile manipulation with low-cost whole-body teleoperation. In arXiv, 2024.
[71] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition (CVPR), June 2016.
[72] P. Hansen-Estruch, Kostrikov, M. Janner, G. Kuba, and S. Levine. Idql: Implicit q-learnngg as an actor-critic method with diffiusinn policies, 2023.
[73] J. Luo, C. Xu, X. Geng, G. Feng, K. Fang, L. Tan, S. Schaal, and S. Levine. Multi-stage cable routing through hierarchical imitation learning. arXiv pre-print, 2023. URL htttps://aaxiiii org/ abs/ 2307.08921
[74] S. Dass, J. Yapeter, J. Zhang, J. Zhang, K. Pertsch, S. Nikolaidis, and J. Lim. Clvr jaco play dataset, 2023. https::///gthibb www.cooocclloeccccccoooooooeeooe cttpaa///vvrrajoooaaee//
[75] L. Y. Chen, S. Adebola, and K. Goldberg.
Berkeley URS
demonstration dataset.
[76] D. Kalashmikov, A. Irpan, P. Pastor, J. Ibarz, A. Herzog, E. Jang, D. Quillen, E. Holly, M. Kalakrishnan, V. Vanhoucke, al. Scalable deep reinforcement learning for vision- robotic manipulation. In Conference on robot learning, pages 651 1-136 PMLR, 2018.
[77] G. Zhou, V. Dean, M. K. Srirama, A. Rajeswaran, Pari, K. Hatch, A. Jain, T. Yu, P. Abbeel, L. Pinto, et al. Train offline, test online: A real robot learning benchmark. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pages 9197-9203. IEFE, 2023.
[78] Tensorflow. TensorFlow Datasets, a collection of ready-to-use datasets. http://wwwwwtnsofooo...o org/ org/datasets.
[79] J. Pari, N. M. Shafiullah, S. P. Arunachalam, and L.. Pinto. The surprising effectiveness of representa- tion learning for visual imitation, 2021.
[80] J. Ho, Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 6851,, 2020.
14

---

Additional Results
A.1
Ablations
RTX Early Stopping
Bridge Discretization
1.0
200K
Cont. Diseggi
0.9
0.50
0.7
0.7
0.8
0.5
0.4
0.1
0.1
0.0
0.0
Average
Carrot to Rack Fork to Rack
Pen In Cup
Average
Carrot
Carrot Right
Cup
Figure Ablations design choices Re-Mix. We the effects of left: reference model overfitiing by selecting a checkpoint once validation loss starts increasing 150K steps and right: using continuous actions Re-Mix. For ablations, we remove the "Flip Bowl" and 'Cube to Plate" tasks as all Re-Mix variants achieved 100% success.
A.2 10% Bridge Sub-setting
Here we include results for 10% subsetting of the bridge dataset as described in Section 4.4. In the supplemental material we include videos of rollouts from our experiments.
Bridge Subset
1.0
Uniform
SSP
0.9
ReMix
0.4
0.37
0.4
0.4
0.4
0.15
0.3
0.2
Average
Carrot Rack
Cube to Plate
Carrot Right
Fig Figure 5: Bridge 10% subsetting.
A.3
Simulation Experiments
We additinally run simulation experiments on the Robomimic NutAssemblySquare task from images [20]. We chose Robomimic because was collected using human operators like real world datasets. We divided the 300 multi-human demonstrations into six domains by operator, which have "better", "okay", and "worre"" labels. We run Re-Mix with the same architecture described in all other experiments, but train Conditional UNet Diffusion Policies [30] since they performed far better this benchmark. We evaluate checkpoints for 100 episodes after 400K training steps. The results included in Table 3 and leamed Re-Mix weights are shown in Table 4. We see that the Re-Mix determined weights outperform uniform weights at both 50% and 25% subsetting. This is likely because Re-Mix up-weights the "better" operators and comparatively down-weights the "worse" ones. Note that the natural or uniform domain weights are not even across all operators. This is because some of the operators take longer to complete the task than others.
Method
50% Subsetting 25% Subsetting
ReMix 77/100
59/100
Uniform
53/100
39/100
Table 3: Performance on the RoboMimic NutAssembly task, divided by operator.
Method
Better
Better 2
Okay
Okay 2
Worse 1
Worse 2
ReMix 22.8%
20.0%
11.9%
14.6%
18.0%
12.7%
Uniform
9.6%
13.6%
18.7%
14.4%
20.0%
23.7%
Table 4: Domain weights used by Re-Mix in comparison to the natural uniform domain weights.
15

---

A.4 Action Distributions
In Fig. 7 we show the action distribution for the BridgeV2 d ataset and in Fig. 8 we show the action distribution for the ToTo dataset, both in log-scale. The BridgeV2 dataset's action distribution is far more normal and symmetric than the ToTo action distribution, which is heavily multi-modal and skew. Robust optimization appears to be more well-behaved on the more normally distributed datasets.
B
Dataset Details
B.1 OpenX RTX Subset
We use subset of the OpenX Embodiment datast similar to that used train the RT-X models [7]. First, we the RLDS dataset modification repository (https://8 thub. com http///tttttsttteeeeeeeeeeeeeeee bewerttsc/rddd__datttteet.t
used by Octo Model Team et al. [25] to preprocess the raw datasets downloaded from Tensor Flow Datasets [78]. Specifically, we resize all images to 256 x 256, and filter the Kuka dataset [76] by an included success key. Note that this does warp images. We use the updated version of the Bridge dataset, available at https: /tta:///riiii htt.ariileee........emmom eecs. , berkeley. edu/ datp//daasesss/aaaaaa /brigge.bbrree..eer httpdddereeeeeeeeeeeeeeeeeeeeeee rew.eeelleeeeeeeeeeeeeeerr datte///att/tt///////////// The specific composition of the dataset is listed in Table Note that we only train on the primary third-person camera in each dataset. For this reason, we omit the NYU Reacher-grabil dataset [79] which only inlcudes wrist cameras. We align all action spaces by converting them to delta cartesian and delta euller angle and binarize all gripper actions.
B.2 Bridge V2 Dataset
For experiments on bridge-only, we split the bridge dataset into 32 domains. First, we re-downloaded the raw bridge dataset and converted it to RLDS using the DLimp convertor (https: gi thub. com/ kvablack/ dlimp/). We then partitioned the bridge dataset by domain using the file path metadata field that lists which setting demonstrations were collected in e.g. "toy-kitchen 1" "toy-sink-3". We then manually group the domains into 32 categories. We omitted data that was collected by scripted policy, as it did not contain the scene information in the filepath metadata. This means we ended up with around 45,000 training trajectories, instead of the 60K used in the full bridge dataset. In Table 5 we list the natural weights of each of these domains and the learned weights by Re-Mix. We can see that Re-Mix down-weights some of the largest domains and places their weight on smaller domains.
B.3 Co-Training Datasets.
Below we describe our co-training data and evaluation procedure for the real-world tasks on the WidowX 250 and Franka Panda robots.
WidowX Tasks We evaluate on 6-DoF WidowX 250 robot on several new pick place tasks in a toy kitchen setting. Our setup is similar to Bridge V2 [4] with fixed side camera and blocking controller. Following Walke et al. [4] we use blocking controller during evaluation. We collect teleoperated demonstrations using an Oculus Quest Headset for motion tracking and co-train on 25 demonstrations for each of the three tasks "Move Cube out Sink", "Move Cup into Sink", and "Move Fork from Sink to Rack."
During evaluation, we examine generalization on various axes. The "Carrot Rack" task tests generaliza- tion to picking up a new type of target object, "Fork to Rack" tests new unseen object positions, "OOD Cup" tests object with different shape, "Cube to Plate" and "Cube to Cup" test generalization to new containers, and "Carrot to Right" tests generalization to both a new target object and a new motion. For each of these tasks, we first take goal image and then evaluate our policies with fixed object locations for up to 100 seconds, stopping early if the robot or objects reach unrecoverable states. For "Carrot to Rack" we do five trials with the carrot facing down and five trials with it facing upwards. For "Fork Rack" we use an unseen initial position to the right side of the sink and rotate the fork left 45 degrees for five episodes additional "Cube Distractor" task in which the robot has to move the cube to the sink with distractor object present. We ommited this task
and to the right 45 degrees for the other five. We also tested
16

---

Domain
Uniform Weight
ReMix Weight
0 toykitchen2
0.18728751
0.0961817
0.094527
0.04846529
2 toykitchenl
0.069307
0.07683
3 toykitchenó
0.06940527
0.0573625
4 datacol2 toykitchen7
0,07133783
0.06905
5 datacol2 toykitchen2
0.04329277
0.03651583
6 toykitchen7
0.032803
0.03538789
7 ddtacol22 2ttol/fooodnnnnnbaba
0.038522
0.0809778049
8 datacol Malacoll_toyle Looykitchen
0.03606622
0,037404168
9 Latacol2_rrboo robooodooossooo.ooooo..eeeroeere
0.025810027
0.034152
10 gatacol2.toykitcher
0.02394393
0.02740302
1l fooliii table
0.0272809
0.013906823
12 datacol2 Jaaunoll Jaundry.machine machine laundry mmachin
0.02582954
0.0396389
13 datacol2.toykinn tooyktchhe,n, toykitchen5
0,0337366
0.049943
14 deepthought toykitchen2
0,0253313
0.013434348
15 deepthought robot desk
0.01978364
0.032410502
16 tabletop dark wood
0.0219985
0.024691
17 datacol2 datacol2.tossink2 toysink2.bwx bww
0.0225748
0.0198516
18 toykitchen2 room8052
0.01083554
0.0295857
19 deepthought loykitchenl, Sanacoll_tokkichhtt
0.01868
0.04047
20 3_ww, toysink3
minsky woldtable..taayyyy dotacol2..oykiichenn .try
0.037856699
0.0484
21 toysink3
0.01235829
0.014877
22 datacol2_tow .ookkicen.t
0.01155453
0.02194
23 loysinkI_room@800 toysink1
0,00979455
0,01831014
24 tool chest
0.00471524
0.00878
25 toysink5
0.00405418
2.78E-05
26 whitehoard
0.006774
0.0129337
27 toykitchen4
0.00371938
0,00537445
28 toysink4
0.00289793
1,80E-05
29 toykitchen3
0.00124406
2.72E-05
30 realkitchenl _dishwasher
0.00202648
0.000541
31
tabletop_white, realkitchenl _counter
0.004647549
0.005079152
Table Learned weights by Re-Mix on Bridge dataset.
from Fig. 2 as baselines were able to complete the task, which heavily skewed performance, We include results for this task in Appendix B.3.
Method
ReMix
Human
Uniform 0/10
Success Rate
7/10
0/10
Figure 6: Performance on the cube Distractor task. Right: Depiction of the Cube Distractor task.
B.4 Franka Tasks
We evaluate on a Franka Panda robot on several pick place tasks on a tabletop. We use fixed over the shoulder camera We co-train on 25 teleoperated demonstrations for each of the tasks "Pen into Cup," where we put a pen into a cup from 5 different start locations, and "Flip Bowl," where bowl is flipped into a drying rack. For the "Pen into Cup" task we use a different pen than in co-training. However, because our franka embodiment with the Robotiq 2F-85 is not found in our pre-training datasets, we evaluate the same tasks we co-trained on. We evaluate each start location of the pen twice from new set of predifined positions. As in the WidowX evaluations, take goal image for each task and evaluate for up to 100 seconds using a 10Hz controller without blocking control.
17

---

RTX
Bridge
Batch Size
512
384
Action Chunk
4
2
Image Resolution
x
x
Table 6: Hyperparameteess
C
Training Details
Architecture. We borrow our architecture from [4] with few minor changes. Our policies takes as input history of two consecutive frames and single goal image and output sequence of actions via DDPM [80].
First, we preprocess all images to fit between 11 and 1. Then, we channel-wise concatenate both the goal image and a grid containing the position of each pixel in (2.3) space also normalized between -1 and 1. Images are then fed to ResNet 50 encoder, which employs global average pooling on the output to obtain a 5122 dimension representation for each image. Both image representations are then concatenated and fed diffusion action prediction head.
Hyperparameters. We use a cosine decay learning rate schedule with an initial learning rate of 0.0002. We train all models for 400K steps and evaluate the final checkpoint, except for Bridge 10% subsetting. which found perform better after 200K steps. More detailed hyperparameters Note that there some differences between bridge and RTX which were made for computational reasons we iterated faster on the bridge dataset before scaling to RTX. We also did maintained aspect ratio for bridge, hence the different image input size, but did not for RTX follow Octo Model Team et al. [25]. We apply data augmentatiin to all images consistently across the time horizon and goal image (meaning that the goal image and all past images of each example have the same augmentation applied). We use random resize cropping, brightmess, contrast, and hue randomization. For k-means in SSP for Bridge we set 3= equal to the number of domains used for Re-Mix.
found in Table 6.
18

---

bridge Dim GAUSSIAN
125 X
Frequency
1.15 111x x 10"
1.05 x
10
9.5x 10"
x 10'
-2
Value
bridge Dim 2 GAUSSIAN
10a
Frequency
10*
10a
Value
bridge Dim 3 GAUSSIAN
10.
Frqueeccy
10
10a
102
-2
2
Value
bridge Dim GAUSSIAN
10
Prequency
10a
101-
Value
bridge Dim 5 GAUSSIAN
Prequency
10
20
10a
-2
2
Value
bridge Dim 6 GAUSSIAN
Frequency
10
102
O
2
Value
bridge Dim GAUSSIAN
Freuuancy
10
202-
Value
Figure 7: Action distributions for Bridge.
19

---

toto Dim 1 GAUSSIAN
frequency
-2
Value
toto Dim 2 GAUSSIAN
frequency
10a
102-
10
Value
toto Dim 3 GAUSSIAN
Irequency
10
101
toto Dim GAUSSIAN
10a
Irequency
102
10
Value
toto Dim 5 GAUSSIAN
frequency
Value
toto Dim 6 GAUSSIAN
10a
Frequency
10a
10a
101
toto Dim 7 GAUSSIAN
10a
10a
107
101
Value
Figure 8: Action distributions for Toto.
20


Reviews:
[{"review_id": "9mY7sCwHYq", "reviewer": "Reviewer_aygU", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 4, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "Re-Mix is a min-max approach for optimizing training data mixtures in large-scale robot imitation learning. Using group distributionally robust optimization, Re-Mix dynamically adjusts domain weights within robotic datasets to maximize performance across various domains (i.e., different robot datasets). As opposed to current practices that rely on expert intuition for data curation, Re-Mix offers a systematic method to assess domain significance through a min-max optimization framework. This approach results in better performance compared to both uniform and expert-curated data mixtures.", "summary_of_recommendation": "This paper presents a principled approach to selecting data to scale up imitation learning from large robot datasets. It provides clear intuitions and is well written. This is particularly relevant to the current paradigm in robotics and could be widely applicable.", "technical_quality": 5, "clarity_of_presentation": 5, "potential_impact": 4, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "It would be interesting to see if more diverse tasks can be added to the evaluations. Intuitively, carrot to rack and fork to rack tasks potentially involve very similar action distributions. If not, adding plots of action distribution similarity could also clarify this.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IJAJD00qMA", "reviewer": "Reviewer_Q6JL", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 3, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "The paper introduces Re-Mix, a robust optimization method tailored to find the composition of large-scale imitation learning datasets that consist of different robot datasets. The method takes ideas used in NLP to find optimal dataset mixtures and adapts them for robotic datasets by solving several challenges related to robotic data. The experiments show that policies trained with Re-Mix outperform policies on default robot dataset mixes and mixes by robot experts. Further, by using the generated importance weights the authors show that policies trained on better datasets can reduce the dataset by 75% and still outperform default mixes. while I appreciate the various experiments in the real world. Could the authors try to add 1 or 2 sim environments to the dataset mixture and report results in this sims? This would allow better benchmarking of future work in this field by using available simulation.", "summary_of_recommendation": "The paper tackles an important and novel problem of optimizing the dataset of foundation policies and contributes relevant contributions for the Robotics community. Overall, the work is timely and I am excited for the possibilities, especially with the improved efficiency in training with smaller datasets optimized by re-mix. In addition, several experiments in the real world are conducted to verify various claims. Thus, I recommend accepting this paper.", "technical_quality": 4, "clarity_of_presentation": 4, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "- while I appreciate the various experiments in the real world. Could the authors try to add 1 or 2 sim environments to the dataset mixture and report results in this sims? This would allow better benchmarking of future work in this field by using available simulation.\n- given stronger BC policies that have been introduced recently such as VQ-BeT would it be possible to reduce the double training when with a more expressive policy?\n\n—\n\nThe authors were able to resolve my concerns and added my requested experiments for better reproduction using simulation. This improves the paper and I maintain my high score.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FLHf6G6DzM", "reviewer": "Reviewer_zvQ8", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 1, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "This paper applies DoReMi into the field of robotics dataset for data mixing.", "summary_of_recommendation": "Overall the originality is quite limited. I would not recommend this paper for acceptance.", "technical_quality": 1, "clarity_of_presentation": 4, "potential_impact": 2, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "I don't have specific questions in mind.", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: **Paper summary**

This paper adapts techniques from NLP to improve the generalizability of robot foundation models by re-weighting datapoints in large-scale datasets. The paper presents extensive experiments in real-world and simulated tasks, demonstrating how the method improves the data-efficiency and generalizability of the trained model compared to the original, uniform data weighting.

**Review summary**

Summary of strengths:
+ The paper addresses a key problem in robot learning: optimizing training data mixtures. The paper has the potential for high impact in the CoRL community.
+ The evaluation is thorough, involves real-world experiments, and demonstrates how the remixed dataset enables faster training and higher-performance in the trained model.
+ The paper is well-written and clear.

Summary of weaknesses:
- The proposed method borrows heavily from prior work (DoReMi). While the paper describes how the application is different (robot policies vs language models), it is unclear whether the proposed method itself is sufficiently novel. **[Edit: the rebuttal argues that the paper overcomes several practical barriers that otherwise prevent DoReMi from being simply reused in robotics.]**
- The evaluation could be improved with added variance in the environments and tasks. **[Edit: the rebuttal includes additional experiments in higher-variance tasks.]** 

**Response to rebuttal**

Following the rebuttal, the reviewers remain split about the originality of the paper's contribution; 1 reviewer advocates for a strong reject (on the basis of the paper being a "rebrand" of DoReMi for a robotics application) and 2 reviewers advocate for a strong accept (on the basis of the paper overcoming practical barriers to adapt DoReMi for robotics and thoroughly demonstrating its effectiveness for the robot learning community).

I lean toward the latter view. This paper provides compelling evidence to demonstrate the importance and practical effect of data mixtures for robot learning. I believe this will be a highly useful and practical contribution for the CoRL community.


**Recommendations for improvement**

The revised paper should include the rebuttal's argument for why the robotics domain requires unique considerations compared to language modeling with DoReMi.
Decision: Accept