Helping AI learn from demonstrations that leave an uneven picture

Imagine learning to navigate a building from recordings that show one room far more often than the others. Every demonstrated turn might be correct, yet the collection could leave you poorly prepared for the building as a whole. A preprint by Seokin Seo and colleagues tackles an analogous problem in artificial intelligence: how to imitate an expert when the demonstrations disproportionately represent certain situations.

Their method, DrilDICE, improved performance across many deliberately altered datasets in simulated navigation and movement tasks. The results support a targeted approach to uneven data collection, rather than a general solution to unreliable demonstrations. As preprint findings, they should be treated as provisional.

Correct examples, distorted coverage

Behavioral cloning trains a system to predict the action an expert would take in a given situation. Its appeal is simplicity: learning can happen entirely from recorded demonstrations, without interacting with the environment during training.

But matching the recorded examples does not necessarily produce reliable behavior. Once a learner makes a mistake, it may encounter situations unlike those emphasized in training. The paper examines an additional source of mismatch: the recordings themselves may not reflect how often an expert would naturally encounter different situations.

For example, a recording device might capture some locations more frequently than others, or a dataset might contain only short pieces of longer demonstrations. The expert actions can remain correct even though the distribution of situations is skewed. This is the particular form of “covariate shift” studied here. The researchers do not primarily address incorrect expert actions or changes in how the environment behaves.

Preparing for plausible difficulties

DrilDICE changes how much influence individual demonstrations have during training. Instead of minimizing only the average action-prediction error in the collected dataset, it considers distributions of situations that would make the learner’s errors more costly.

An unrestricted version of that strategy can be excessively pessimistic, concentrating attention on extreme cases. The researchers therefore constrain the alternatives using what is called Bellman flow: a consistency condition linking how frequently situations occur to the starting conditions and the environment’s transitions. In accessible terms, the alternative distributions must fit how behavior could actually unfold, rather than being arbitrary rearrangements of examples.

The method also penalizes distributions that stray too far from the available data. This matters because estimates based on finite demonstrations become less dependable when asked to represent something very different.

Using stationary distribution correction estimation, or DICE, the algorithm calculates weights for the recorded examples and trains the imitator with those weights. It uses recorded situations, expert actions and subsequent situations—not just isolated input-action pairs.

What the simulations showed

The researchers first tested a grid-based navigation environment called Four Rooms. They altered the proportions of examples associated with particular rooms or movement directions. DrilDICE achieved the highest average normalized score in all eight settings reported in the main comparison. It also performed best when scores were averaged over the worst-performing quarter of evaluation episodes.

The team then used expert datasets for three simulated movement tasks: hopper, walker2d and halfcheetah. They constructed uneven datasets by changing the balance of states or actions, sampling different portions of trajectories, and collecting short segments.

In one state-rebalanced walker2d setting, DrilDICE achieved a mean normalized score of 73.7, compared with 22.9 for ordinary behavioral cloning. These are benchmark scores, not percentages of successful real-world operation. The continuous-control comparisons were repeated five times, with standard errors reported.

For datasets sampled unevenly across time, DrilDICE had the highest mean score in 11 of 12 settings. It did not win universally: another method performed better in one hopper setting, and some halfcheetah settings remained difficult for every approach.

Where the evidence stops

All the reported tasks were simulated. Real-world deployment remains future work, and the study does not establish robustness to faulty demonstrations or altered physical dynamics.

The experiments also included measures to preserve data coverage. The Four Rooms dataset received an example from every state, while the continuous-control constructions included a complete trajectory alongside the sampled sub-trajectories. The results therefore should not be read as showing that reweighting can recover expert behavior from entirely missing information.

Finally, the paper reports only a weak relationship between action-prediction error and task score. Better imitation on a mathematical loss is not automatically better behavior. DrilDICE’s practical promise lies in improving performance under the tested collection biases—not in guaranteeing dependable imitation under every kind of data shortage or shift.