BayesFP: Posterior Estimation for Flow-Based Policies via Feynman-Kac Sampling

2026-06-18

AI Summary

This work proposes a unified framework that integrates safety constraints and task objectives into pretrained diffusion models and flow-matching policies during inference without requiring retraining. By formulating constrained trajectory generation as a Bayesian posterior sampling problem—using expert demonstration distributions as the prior and cost functions to define the likelihood—the method extends the Feynman–Kac corrector, for the first time, to deterministic flow-matching strategies, enabling consistent correction across both classes of generative models. Evaluated on Diffusion Policy, GR00T-N1.6, and π₀.₅, the approach demonstrates zero-shot capability in handling complex, non-convex obstacle avoidance tasks and significantly outperforms the original π₀.₅ policy in terms of constraint satisfaction and task performance.

Abstract

Robots must generate trajectories that remain faithful to learned expert behavior while satisfying safety constraints and task-specific objectives specified only at inference time. We formulate constrained trajectory generation for pretrained diffusion and flow-matching policies as Bayesian posterior sampling, with the learned demonstration distribution as a prior and an inference-time, cost-derived likelihood tilting it toward feasible, optimal trajectories. To sample from this posterior without any retraining of the base policy, we leverage the Feynman--Kac corrector framework, originally formulated for diffusion models, and extend it to deterministic flow-matching policies. The result is a unified, inference-time, retraining-free sampler for diffusion and flow policies. We validate the approach on pretrained Diffusion Policy, GR00T-N1.6, and $π_{0.5}$ checkpoints across simulated and real-world manipulation tasks, including planning around non-convex obstacles introduced at inference time, and show improvements over the base $π_{0.5}$ on zero-shot tasks.

Problem

• trajectory generation
• safety constraints
• task-specific objectives
• inference-time adaptation
• pretrained policies

Innovation

• Bayesian posterior sampling
• Feynman-Kac corrector
• flow-matching policies
• inference-time constraint handling
• retraining-free trajectory generation