Explicitly capturing demonstrator multi-modality for distillation into diffusion policies

Shao Yuan Chew Chia 1 , Rebecca H. Jiang 1 , Abhinav Agarwal 1 , Andreas Hofmann 2 , Leslie Pack Kaelbling 1 , Russ Tedrake 1
  • 1 Massachusetts Institute of Technology, Cambridge, MA, USA
  • 2 Pickle Robot Company, Charlestown, MA, USA

Pretrain to Adapt Workshop at the 10th Conference on Robot Learning (CoRL 2026)

Paper

TL;DR:

To make diffusion policies as multi-modal as their demonstrator, collect multiple expert actions from the same observation.

Abstract

The diversity of a distilled diffusion policy’s outputs matters beyond how it impacts the policy’s success rate. In particular, downstream applications like test-time search and policy improvement require exploration. Thus, we want to distill multi-modal demonstrators into diffusion policies which reproduce their conditional action distributions. Standard data collection for behavior cloning applies a sliding window to demonstration episodes, which pairs unique observation histories with actions. These datasets can be fit by a policy that is either multi-modal, or unimodal with high Lipschitz constant to capture a diversity of actions seen at nearby observation samples. Prior work finds that diffusion policies trained on sparse data learn the latter, memorizing one action per observation. We adapt vine sampling to behavior cloning data collection to resolve this ambiguity: at some states the demonstrator visits, query it k times from exactly the same observation. On a synthetic task, we study how the number of actions per observation, k, affects the learned policy, and find that larger k yields a policy that better matches the demonstrator’s conditional distribution, including its multi-modality and diversity. On a 2D truck unloading task, vine sampling likewise leads to a policy closer to the demonstrator and substantially improves both best-of-N and expert iteration performance.

Grid of truck unloading frames: four demonstrator actions, ten sliding window policy samples that all approach the same grasp, and ten vine policy samples that approach three different grasps.
Actions from the truck unloading demonstrator (top), the sliding window policy (middle), and the vine policy (bottom), queried at the same state of a held-out problem instance. Each panel shows the observation image with future arm poses overlaid; poses with suction on are shown in pink. Of 10 samples each, the vine policy's reach three of the four grasps, while the sliding window policy's all approach the same grasp.
A demonstrator's choices branching along an episode, with the actions recorded by sliding window and vine data collection.
The choices that a demonstrator has the capacity to take branch along an episode. Sliding window datasets record one action per observation history, while vine datasets record several counterfactual actions at a random subset of observation histories. Gray actions are recorded in neither dataset.

BibTeX citation

@inproceedings{chewchiaExplicitlyCapturingDemonstrator2026,
  author = {Chew Chia, Shao Yuan and Jiang, Rebecca H. and Agarwal, Abhinav and Hofmann, Andreas and Kaelbling, Leslie Pack and Tedrake, Russ},
  title = {Explicitly capturing demonstrator multi-modality for distillation into diffusion policies},
  booktitle = {Pretrain to Adapt Workshop at the 10th Conference on Robot Learning (CoRL)},
  year = {2026}
}