Towards Balanced Behavior Cloning from Imbalanced Datasets

TL;DR

The paper addresses how imbalanced human demonstrations bias imitation learning policies toward frequently seen subtasks. It formalizes demonstrations as a mix of sub-policies and analyzes why equal weighting in behavior cloning favors dominant behaviors, proposing several data-balancing strategies including a novel meta-gradient rebalancing method. The authors show theoretically that dataset proportions bias learning and empirically demonstrate improvements in downstream imitation tasks when balancing offline data, with careful consideration of the limitations of each approach. They introduce a principled procedure to learn target losses per sub-policy, enabling balanced learning without extra data collection and highlighting practical implications for multi-task robotic learning. Overall, the work provides a framework and toolbox for balancing heterogeneous imitation datasets, improving generalization across behaviors while outlining avenues for future task-aware offline balancing.

Abstract

Robots should be able to learn complex behaviors from human demonstrations. In practice, these human-provided datasets are inevitably imbalanced: i.e., the human demonstrates some subtasks more frequently than others. State-of-the-art methods default to treating each element of the human's dataset as equally important. So if -- for instance -- the majority of the human's data focuses on reaching a goal, and only a few state-action pairs move to avoid an obstacle, the learning algorithm will place greater emphasis on goal reaching. More generally, misalignment between the relative amounts of data and the importance of that data causes fundamental problems for imitation learning approaches. In this paper we analyze and develop learning methods that automatically account for mixed datasets. We formally prove that imbalanced data leads to imbalanced policies when each state-action pair is weighted equally; these policies emulate the most represented behaviors, and not the human's complex, multi-task demonstrations. We next explore algorithms that rebalance offline datasets (i.e., reweight the importance of different state-action pairs) without human oversight. Reweighting the dataset can enhance the overall policy performance. However, there is no free lunch: each method for autonomously rebalancing brings its own pros and cons. We formulate these advantages and disadvantages, helping other researchers identify when each type of approach is most appropriate. We conclude by introducing a novel meta-gradient rebalancing algorithm that addresses the primary limitations behind existing approaches. Our experiments show that dataset rebalancing leads to better downstream learning, improving the performance of general imitation learning algorithms without requiring additional data collection. See our project website: https://collab.me.vt.edu/data_curation/.

Figures (9)

• Figure 1: Robot learning how to open a drawer and move a slider from offline human demonstrations. Standard imitation learning equally weights each state-action pair. This results in a policy that imitates the dominant movement patterns seen in the dataset. When the dataset represents one behavior more commonly, imitation learning will learn that behavior at the cost of the other. However, the underrepresented behaviors may not necessarily be unimportant. For instance, opening a drawer and moving the slider require the robot to execute very different motions in different parts of the state space. When the dataset has disproportionately more demonstrations for opening drawer, by imitating the dominant behavior the robot learns to open the drawer but can struggle with moving the slider. So, how do we learn a balanced policy from imbalanced datasets?
• Figure 2: The manipulation tasks we perform in our experiments. On the left, we see picking where the robot must learn to grasp and pick up a red block to a certain height. The red block is initialized at a position that is sampled randomly from three regions: left side of the table, middle of the table, and right side of the table. The possible initialization locations of the red block are shown by transparent overlays. On the right is opening which consists of two tasks conditioned on the state of the environment. If the bulb is off, the robot must move a slider to the left and if the bulb is on, the robot must open the drawer. All the objects irrelevant to the task are initialized randomly.
• Figure 3: Results of our experiments from Section \ref{['sec:sim_1']}. We compare the performance of a policy trained on a balanced dataset with one trained on an imbalanced dataset. The balanced dataset contains $21$ demonstrations for each of the three sub-policies in picking totaling $63$ demonstrations. In opening the balanced dataset contains $15$ demonstrations for each of the two sub-policies totaling $30$ demonstrations. We introduce imbalance by reducing the number of demonstrations for one sub-policy at a time. We train a separate policy for each of the imbalanced datasets as well as a policy that is trained on the balanced dataset. The plots compare the success rate of the policies across $100$ rollouts. The gray bars represent the success rate for the balanced policy and the orange bars represent the success rate of the imbalanced policy for each case. (Left) Results for the three behaviors in picking: lifting up the red block that is on the left side, the middle, and the right side of the table. (Right) Results for the two behaviors in opening: moving the slider to the left when the bulb is off and opening the drawer when the bulb is on. The vertical bars show the standard deviation and $*$ indicates statistical significance.
• Figure 4: Results of the experiment demonstrating the benefits of upsampling data. We compare the performance of a policy trained on an imbalanced dataset with a policy trained on an upsampled dataset. In the imbalanced dataset one of the sub-policies is underrepresented with fewer demonstrations that the other sub-policies. In the upsampled dataset, we resample the fewer demonstrations to make the proportions of each sub-policy equal such that all sub-policies are equally represented. The plots compare the success rate of the trained policies across $100$ rollouts. Th gray bars represent the success rate for the imbalanced policy and the orange bars represent the success rate for the upsampled policies. (Left) The results for the three sub-policies in picking. (Right) Results for the two sub-policies in opening. The vertical bars show the standard deviation and $*$ indicates statistical significance.
• Figure 5: Results of the experiment testing existing approaches for reweighting/balancing datasets. We train three policies. First policy is trained on the imbalanced dataset, In picking, the imbalanced dataset contains $27$ demonstrations for the left and right positions of the red block and $9$ demonstrations for the middle position of the block. In opening the imbalanced dataset contains $10$ demonstrations for moving the slider and $20$ demonstrations for opening the drawer. We train a second policy on a balanced dataset that contains equal proportions of demonstrations for all sub-policies, i.e., $21$ for all three block positions and $15$ for both opening sims. These two policies serve as baselines for comparing the performance of reweighting algorithm Remix hejna2024re. Finally, we use Remix to balance the dataset and train a third policy on it. We compare the success rates of the three trained policies across $100$ rollouts. For reliable results we perform the experiments for $10$ trials. The first three plots show the success rates in picking and the last two plot show the success rates of the three policies in opening. The vertical bars show the standard deviation and $*$ indicates statistical significance.

Theorems & Definitions (2)

• proposition thmcounterproposition
• proof