PriGo: Test-Time Primitive Guidance to Diffusion and Flow Policies for Adaptive Robotic Manipulation

2026-07-08

AI Summary

Existing diffusion- and flow-based imitation learning approaches for robotics exhibit limited generalization across tasks and environments, often mimicking superficial actions while neglecting underlying behavioral intent. This work proposes PriGo, a framework that, for the first time, enables seamless integration of behavior primitive guidance into pretrained policies at test time without requiring retraining. PriGo employs a lightweight Primitive Attention Network (PANet) to infer the distribution over behavior primitives and combines it with a differentiable guidance mechanism to dynamically refine action generation, ensuring semantic consistency in executed trajectories. The approach is compatible with both diffusion models and flow-matching policies, significantly enhancing policy robustness, long-horizon execution capability, and cross-scenario generalization across diverse benchmarks including LIBERO, CALVIN, SIMPLER, and real-world robotic tasks.

Abstract

Imitation learning has enabled remarkable progress in robotic manipulation, especially with diffusion and flow-based policies that generate complex visuomotor behaviors directly from demonstrations. Yet, despite their strong performance, these policies often fail to generalize across tasks and environments. A key reason is that existing policies tend to imitate superficial action correlations rather than the underlying intent. Inspired by the compositional structure of human behaviors, we propose PriGo, a primitive-guided test-time adaptive framework for robust robotic manipulation. PriGo introduces PANet, a lightweight primitive prediction module that infers primitive distributions directly from observations. We further propose a differentiable primitive guidance mechanism that refines generated actions during inference, steering trajectories toward semantically consistent behaviors. Unlike prior primitive-conditioned approaches, PriGo operates entirely at test time and can be seamlessly integrated into pretrained diffusion and flow policies without retraining. Extensive experiments on LIBERO, CALVIN, SIMPLER, and real-world robotic tasks demonstrate that PriGo consistently improves robustness, long-horizon execution, and generalization ability across both diffusion and flow-based policies.

Problem

• generalization
• imitation learning
• robotic manipulation
• diffusion policies
• flow policies

Innovation

• primitive guidance
• test-time adaptation
• diffusion policies
• flow-based policies
• robotic manipulation