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

Imitation learning has enabled remarkable progress in robotic manipulation, especially with diffusion- and flow-based policies that generate complex visuomotor behaviors directly from demonstrations. Despite their strong performance, these policies often struggle to generalize across tasks and environments. A key aspect that remains underexplored is the local motion structure inherent in manipulation behaviors, which is not explicitly captured by continuous action prediction. 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 exclusively at test time and can be seamlessly integrated with pretrained diffusion- and flow-based policies without retraining the backbone policies. Extensive experiments on LIBERO, CALVIN, SIMPLER, and real-world robotic tasks demonstrate that PriGo consistently improves robustness and generalization across both diffusion- and flow-based policies.