Guiding Data Collection via Factored Scaling Curves

TL;DR

This work tackles the data-efficiency challenge in training generalist imitation policies by introducing factored scaling curves (FSC), which quantify how policy performance scales with added demonstrations for individual environmental factors and their combinations. By fitting power-law curves and optionally using offline proxy metrics (FSC-Proxy) based on embedding similarity, FSC enables principled data allocation strategies (One Factor, Pairwise, Group) that prioritize the most informative factor variations under a fixed budget. Across extensive simulation and real-world robotic manipulation tasks, FSC achieves up to 26% improvements over baselines, with FSC-Proxy providing nearly equivalent guidance without hardware trials. The framework is general across tasks and policy backbones, offering a practical path to data-efficient generalization in robotics.

Abstract

Generalist imitation learning policies trained on large datasets show great promise for solving diverse manipulation tasks. However, to ensure generalization to different conditions, policies need to be trained with data collected across a large set of environmental factor variations (e.g., camera pose, table height, distractors) $-$ a prohibitively expensive undertaking, if done exhaustively. We introduce a principled method for deciding what data to collect and how much to collect for each factor by constructing factored scaling curves (FSC), which quantify how policy performance varies as data scales along individual or paired factors. These curves enable targeted data acquisition for the most influential factor combinations within a given budget. We evaluate the proposed method through extensive simulated and real-world experiments, across both training-from-scratch and fine-tuning settings, and show that it boosts success rates in real-world tasks in new environments by up to 26% over existing data-collection strategies. We further demonstrate how factored scaling curves can effectively guide data collection using an offline metric, without requiring real-world evaluation at scale.

Figures (15)

• Figure 1: To efficiently collect demonstrations so as to maximize policy performance under a fixed data budget, we propose factored scaling curves: a principled tool to quantify how policy performance changes with the quantity of factor data. Based on factored scaling curves, we can allocate the data budget to collecting demonstrations that vary different factors based on their importance.
• Figure 2: Illustration of factored scaling curves used to inform data allocation. For the distractor factor, points are used to construct the scaling curve, and is the predicted policy success rate at $K$ additional demos of the factor over the initial dataset.
• Figure 3: Evaluating FSC in the real world. We visualize the task rollouts and report the average policy success rate trained with additional collected data. For pick-place task, we train the policies with diffusion policy. For all other experiments, we obtain policies by fine-tuning $\pi_0$. FSC achieves the best performance in all tasks, achieving up to 26% more improvement over all baseline methods. Compared to the zero-shot setting, fine-tuning $\pi_0$ with FSC yields up to $30\%$ success rate improvement. FSC-Proxy achieves nearly the same high success rate as FSC while eliminating the need for any on-hardware policy execution.
• Figure 4: Visualizing factored scaling curves for real world fine-tuning $\pi_0$ experiments. Solid lines are factored scaling curves we construct based on the initial dataset, and dashed lines are the extrapolations that predicts how policy performance change with additional factor data. Based on the Top strategy, FSC suggests picking the curve with the highest slope, shown in blue (left), purple (middle) and purple (right). Factored scaling curves can accurately predict how policy performance changes with additional factor data, thus able to provide informed data collection strategies. We also visualize how different methods allocate data collection budget to the factors in the top pie charts.
• Figure 5: Expected improvement for $\pi_0$ on three task settings using the Attention Weights from the last denoising step: Camera Pose -- Distractor (CP-D), Table Texture -- Lighting (TT-L), and Robot Pose -- Object Pose (RP-OP). Cosine Similarity projections are normalized to have the same expected value as Expected Improvement. Cosine similarity predicts the top-ranked expected improvement for Fold Towel (CP-D) and Mouse in Drawer (TT-L).