PPGuide: Steering Diffusion Policies with Performance Predictive Guidance
Authors: Zixing Wang, Devesh K. Jha, Ahmed H. Qureshi, Diego Romeres
Organizations: Department of Computer Science at Purdue University, IN 47907, USA. · Contribution was conducted while the author was at Mitsubishi Electric Research Laboratories. · Mitsubishi Electric Research Laboratories, Cambridge, MA 02139 USA.
Diffusion policies have shown to be very efficient at learning complex, multi-modal behaviors for robotic manipulation. However, errors in generated action sequences can compound over time which can potentially lead to failure. Some approaches mitigate this by augmenting datasets with expert demonstrations or learning predictive world models which might be computationally expensive. We introduce Performance Predictive Guidance (PPGuide), a lightweight, classifier-based framework that steers a pre-trained diffusion policy away from failure modes at inference time. PPGuide makes use of a novel self-supervised process: it uses attention-based multiple instance learning to automatically estimate which observation-action chunks from the policy's rollouts are relevant to success or failure. We then train a performance predictor on this self-labeled data. During inference, this predictor provides a real-time gradient to guide the policy toward more robust actions. We validated our proposed PPGuide across a diverse set of tasks from the Robomimic and MimicGen benchmarks, demonstrating consistent improvements in performance.
Figures & tables
Fig. 1: Overview of our PPGuide framework. First, a diverse dataset of trajectories is collected using policy checkpoints from different training stages. Then, (1) an attention-based MIL model analyzes these trajectories to automatically label observation-action chunks as Success-Relevant (SR), Failure-Relevant (FR), or Irrelevant (IR). (2) A lightweight classifier is trained on this labeled data to predict relevance from observation-action pairs. Finally, during inference, gradients from the classifier steer the diffusion sampling process away from failure modes while promoting success-relevant behaviors.
Fig. 2: Evaluation tasks from the Robomimic and MimicGen benchmarks. (The Stack D1 task uses two cubes, while Stack Three D1 uses three).
Challenge
Stack D1
Stack Three D1
Coffee D2
Coffee Prep. D1
Kitchen D1
Mug Cleanup D1
Square
Transport
Long-horizon
✗
✗
✗
✓
✓
✓
✗
✓
Precision
✗
✗
✓
✓
✗
✗
✓
✗
Articulated Obj.
✗
✗
✓
✓
✗
✓
✗
✗
TABLE I: Task specifications across different benchmark tasks.
Method
Stack D1
Stack Three D1
Coffee D2
Coffee Prep. D1
Policy Epochs
500
550
500
550
500
550
500
550
DP
92%
92%
28%
30%
54%
46%
16%
18%
DP-SS
88% (-4%)
90% (-2%)
24% (-4%)
26% (-4%)
44% (-10%)
60% (+14%)
14% (-2%)
14% (-4%)
PPGuide-SS
92% (+0%)
90% (-2%)
32% (+4%)
34% (+4%)
56% (+2%)
60% (+14%)
24% (+8%)
20% (+2%)
PPGuide-CG
94% (+2%)
94% (+2%)
36% (+8%)
32% (+2%)
58% (+4%)
58% (+12%)
20% (+4%)
24% (+6%)
PPGuide
94% (+2%)
94% (+2%)
34% (+6%)
28% (-2%)
58% (+4%)
58% (+12%)
20% (+4%)
22% (+4%)
TABLE II: Benchmark Evaluation Results.
Method
Square
Transport
Policy Epochs
1300
1400
1500
1600
1300
1400
1500
1600
DP
54%
60%
62%
62%
56%
68%
74%
58%
PPGuide (Ours)
70% (+16%)
60% (+0%)
68% (+6%)
70% (+8%)
74% (+18%)
72% (+4%)
70% (-4%)
70% (+12%)
TABLE III: Heterogeneous Evaluation Results.
Fig. 3: This example shows how PPGuide steers the base policy to avoid misalignment during square insertion.
Fig. 4: Averaged performance of PPGuide, with base policies of epoch 500 and 550, on the Square task across different guidance strength values. We are able to achieve performance gains over the base DP across a range of guidance strength.
Fig. 5: Effect of Z-score selection. We build four datasets to train the classifier using different z-scores for Coffee D2 task. The results are the averaged results of experiments with guidance strength of 0.25, 0.3 and 0.35.