Search papers, labs, and topics across Lattice.
This paper introduces PriGo, a test-time adaptive framework designed to enhance the robustness and generalization of diffusion and flow-based policies in robotic manipulation. By employing a lightweight primitive prediction module (PANet) that infers primitive distributions from observations and a differentiable guidance mechanism that refines actions during inference, PriGo effectively aligns generated actions with the underlying intent of tasks. Experimental results across various benchmarks and real-world scenarios show that PriGo significantly improves long-horizon execution and adaptability without the need for retraining existing models.
PriGo enhances robotic manipulation by refining actions in real-time, leading to improved robustness and generalization without retraining.
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.
