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

Problem Statement

Robots operating in real-world environments must generate trajectories that are both expressive enough to capture multimodal expert behaviors and reactive enough to satisfy safety constraints and task-specific objectives. A critical challenge arises when these constraints (e.g., obstacle avoidance, joint limits, energy costs) are unknown at training time or change between deployments.

Existing approaches to constrained trajectory generation generally fall into two categories:

1. Post-hoc filtering/optimization: Methods that rectify unconstrained rollouts using gradient-based planners or safety filters. These often decouple generation from correction, leading to out-of-distribution trajectories, local minima, and high computational overhead.
2. Generative guidance: Methods that modify the sampling process directly (e.g., classifier-based guidance, projection, reflected dynamics). While effective, these are typically heuristic modifications to sampling dynamics. They lack a principled relationship to a well-defined posterior distribution that jointly respects the learned data prior and the inference-time constraints.

The central question addressed is: Given a pretrained diffusion or flow-matching policy, how can we sample trajectories that respect a constrained optimization problem specified only at inference time, while remaining faithful to the learned demonstration distribution?

Methodology

The authors formulate constrained trajectory generation as Bayesian posterior sampling. The learned demonstration distribution $p_{data}(x)$ serves as a prior, and a cost function $J(x)$ encoding the inference-time constraints serves as a likelihood. The target is to sample from the posterior:
$$p(x | O = 1) \propto p_{data}(x) \exp(\beta J(x))$$
where $O=1$ indicates optimality and $\beta < 0$ controls the strength of the cost-induced tilt.

To sample from this posterior without retraining the base policy, the paper leverages the Feynman-Kac Corrector (FKC) framework. The key methodological contributions are:

1. Feynman-Kac Weighted Dynamics

Instead of heuristically adding cost gradients to the policy's score or velocity field, the authors derive a weighted Stochastic Differential Equation (SDE). This SDE governs the evolution of a particle population where each particle carries a log-weight $w_t$. The dynamics are derived from the Fokker-Planck PDE of the target Gibbs-tilted distribution.

• Drift Correction: The drift term includes a cost-aware correction proportional to $\nabla J(x)$.
• Weight Evolution: The weights accumulate a reweighting term derived from the discrepancy between the simulated dynamics and the target density evolution. This ensures the particle population is an asymptotically unbiased estimator of the posterior, even if the drift correction is imperfect.
• Resampling: To mitigate weight degeneracy over long horizons, the method employs Sequential Monte Carlo (SMC) techniques, specifically systematic resampling over an active time interval.

2. Extension to Deterministic Flow-Matching

The FKC framework was originally formulated for stochastic diffusion processes. The authors extend it to deterministic flow-matching policies (which solve an ODE $dx = -v_t(x)dt$) by recasting the flow as a marginal-preserving SDE.

• Following Singh and Fischer [16], they introduce an auxiliary diffusion term $\tilde{\sigma}(t)$ and a corresponding score-based drift correction.
• Crucially, for rectified flows with a Gaussian source, the score $\nabla \log q_t(x)$ can be imputed in closed form from the learned velocity field $v_t(x)$, eliminating the need for a separate score network.
• This construction allows the FKC framework to be applied uniformly to both diffusion and flow-matching policies.

3. Theoretical Guarantees

The paper provides a theoretical guarantee (Corollary 4.2) that as the penalty weights and the magnitude of the inverse temperature $|\beta|$ increase, the Gibbs-tilted posterior concentrates arbitrarily tightly on trajectories that are feasible and near-optimal. This implies that the method can satisfy constraints without requiring external safety filters or backup planners.

Key Contributions

1. Principled Formulation: The work reframes constrained trajectory generation as Bayesian inference, offering a principled alternative to heuristic guidance and post-hoc projection.
2. Unified Sampler: It instantiates the posterior as a Feynman-Kac weighted SDE, extending the framework from stochastic diffusion models to deterministic flow-matching policies via marginal-preserving SDE recasting.
3. Theoretical Convergence: It proves that the resulting posterior concentrates on feasible, near-optimal solutions under sufficient penalty strength.
4. Empirical Validation: The approach is validated across simulated and real-world tasks using pretrained models (Diffusion Policy, GR00T-N1.6, and $\pi0.5$), demonstrating success in planning around non-convex obstacles introduced at inference time.

Experimental Results

The authors evaluated BayesFP on several benchmarks, comparing it against baselines including Linear-Combination Guidance, CASF (Constraint-Aware Streaming Flow), and JM2D (Joint Model-based Model-free Diffusion).

• Simulation Tasks:
  • Diffusion Policy: On RoboMimic tasks (Can, Transport) with added cylindrical obstacles, BayesFP achieved 0% collision rate compared to 30% (No Guidance) and 2% (Linear Combination).
  • GR00T-N1.6: On LIBERO-Object tasks with cylindrical and non-convex V-shaped obstacles, BayesFP achieved 84% and 75% success rates respectively, significantly outperforming baselines which struggled with non-convex geometries.
  • $\pi0.5$: On zero-shot RoboLab tasks (e.g., FoodPacking, BBQSauce), BayesFP improved success rates (e.g., 36% vs 28% for HammersInLeftBin) and reduced collision rates compared to linear combination baselines.
• Real-World Experiments:
  • Using an SO101 arm fine-tuned with only 20 demonstrations, BayesFP successfully handled a PickAndPlace task with a mug obstacle introduced at inference (reducing collision from 98% to 6%) and a BimodalMugSelection task where a language constraint was reinforced by a cost constraint to suppress an unwanted mode (achieving 100% correct-mug selection vs 72% for the baseline).
• Ablations:
  • Increasing the number of particles ($K$) improved performance but saturated around $K=4$ for simple obstacles.
  • Increasing the guidance strength $|\beta|$ led to an exponential decay in constraint violations, consistent with theoretical predictions.

Significance and Claims

The paper claims that BayesFP provides a unified, inference-time, retraining-free sampler for both diffusion and flow policies. Its primary significance lies in:

• Principled Guidance: Unlike linear-combination methods which sample from a sum of distributions, BayesFP samples from the true product (posterior) of the prior and the cost-likelihood, ensuring a mathematically consistent treatment of constraints.
• Handling Non-Convexity: By sampling whole trajectories from the posterior rather than applying local gradient corrections, the method naturally handles non-convex obstacles and multi-modal solution spaces.
• Generality: The framework is agnostic to the base policy architecture (diffusion or flow) and the specific nature of the constraints, provided they can be encoded in a differentiable cost function.

The authors acknowledge limitations, noting a computational slowdown (1.25–2$\times$) due to the sequential nature of Signed Distance Function (SDF) construction and the overhead of particle simulation, though this is mitigated by parallel GPU execution. They also note the requirement for the cost function to be differentiable. Future work is suggested to extend the framework to streaming flow policies for low-latency control and to incorporate richer constraint representations.