Explicitly capturing demonstrator multi-modality for distillation into diffusion policies

Shao Yuan Chew Chia, Rebecca H. Jiang, Abhinav Agarwal, Andreas Hofmann, Leslie Pack Kaelbling, Russ Tedrake

November 2026

TL;DR Instead of recording one action per observation, vine sampling re-queries a stochastic demonstrator 10 times at a random subset of visited states. At matched dataset size, the resulting diffusion policy keeps the demonstrator's multi-modality and does better in best-of-N search and expert iteration on a simulated 2D truck unloading task.

What they do

Test-time search and policy improvement need a policy whose samples are diverse. Standard behavior cloning data pairs each observation history with a single action, so a diffusion policy that memorizes it collapses to one mode per observation even when the demonstrator is multi-modal. The authors fix this at data collection: at some visited states, query the demonstrator k times from exactly the same state.

How it works

Analysis. A triangle inequality splits policy error into a memorization term and a sampling term that says how well k recorded actions represent the demonstrator. A memorizing policy can only improve through the second, which needs larger k.
Synthetic task. 16-D observations and actions with four valid action modes; k is swept over 1, 10 and 100 at a fixed budget of 1000 samples.
Vine sampling. While the demonstrator runs an episode, each timestep is re-queried k times with probability p. Only these states enter the dataset.
Truck unloading. In Drake, a planar suction arm removes four boxes. A randomized BiRRT-based planner is the demonstrator. Vine (k=10, p=0.1) and sliding-window datasets are both cut to 370,000 samples and train the same diffusion policy.

Results

Synthetic: all policies memorize, but behavior entropy rises from 0 at k=1 to 1 at k=100, and mode recall is far higher at larger k. With 100,000 samples the k=1 policy catches up, so the gain is specific to sparse data.
Truck unloading: the vine policy has lower policy error and an SD ratio closer to the demonstrator’s than sliding window, early stopping, SIME and PostBC, and the best success rate and throughput under best-of-N search. Best-of-100 on the sliding-window policy still falls well short.
Under expert iteration, vine’s throughput keeps improving, while sliding window’s barely changes and SIME’s degrades.

Limitations

Simulation only, on a 2D task, and both search and expert iteration use the simulator to score candidates. Collecting the vine dataset cost about 30 times the compute of the sliding-window one (55 s vs 1.8 s per episode). p and k were not tuned, and the demonstrator must be re-queryable from the same state.

Figures

From Figure 1: Actions from the truck unloading demonstrator (top), the vine policy (middle), and the sliding window ... — From Figure 1: Actions from the truck unloading demonstrator (top), the vine policy (middle), and the sliding window 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.

From Figure 2: The choices that a demonstrator has the capacity to take branch along an episode. Sliding window datas... — From Figure 2: 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.

From Figure 5: Truck unloading experiment. (a), (b) Policy error (Wasserstein distance between policy and demonstrato... — From Figure 5: Truck unloading experiment. (a), (b) Policy error (Wasserstein distance between policy and demonstrator) and SD ratio per method at held-out conditions. (c), (d) Success rate and throughput in best-of-N search; N=10 for solid bars, except more search where N=100. (e), (f) The same two metrics over expert iteration for the sliding window, SIME and vine methods. Plotted values are means and error bars are 95% intervals.

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.

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.

Matt Strong