Fast Flow-based Visuomotor Policies via Conditional Optimal Transport Couplings

Diffusion and flow matching policies have recently shown remarkable
performance in robotic applications by accurately capturing multimodal robot
trajectory distributions. However, their computationally expensive inference, due
to iterative denoising or numerical integration of an ODE, limits their application as
real-time controllers for robots. We introduce a methodology that utilizes Optimal
Transport couplings between noise and samples, in order to force straight solutions
on the Probability Flow ODE. We show that naively coupling noise and samples
does not work well in conditional tasks and introduce the condition variables in
the coupling process to enhance few-step performance. The proposed few-step
policy achieves state of the art success rates in a diverse set of simulation and
real-world robot tasks while maintaining the same training complexity as vanilla
Flow Matching, in contrast to distillation methods.