Fast Flow-based Visuomotor Policies via Conditional Optimal Transport Couplings

2025-05-02

AI Summary

Existing diffusion and flow-matching policies model multimodal robot trajectories but rely on ODE/SDE numerical integration, resulting in high inference latency and hindering real-time deployment. This paper proposes a fast flow policy based on conditional optimal transport (COT) coupling—introducing COT into the flow-matching framework for the first time. Our method explicitly models the joint distribution over noise, actions, and observations, enabling high-quality, diverse action generation in just 1–2 inference steps. It requires no knowledge distillation or additional training, maintaining training complexity comparable to Diffusion Policy and standard Flow Matching. In simulation benchmarks, our approach improves task success rate by 4% and accelerates inference by 10×. Furthermore, real-robot experiments demonstrate its capability for low-latency, robust real-time control.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Intelligent Robots: Learning & Optimization for ROB
• Search and Optimization: Learning to Search

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Diffusion and flow matching policies have recently demonstrated remarkable performance in robotic applications by accurately capturing multimodal robot trajectory distributions. However, their computationally expensive inference, due to the numerical integration of an ODE or SDE, limits their applicability as real-time controllers for robots. We introduce a methodology that utilizes conditional Optimal Transport couplings between noise and samples to enforce straight solutions in the flow ODE for robot action generation tasks. We show that naively coupling noise and samples fails in conditional tasks and propose incorporating condition variables into the coupling process to improve few-step performance. The proposed few-step policy achieves a 4% higher success rate with a 10x speed-up compared to Diffusion Policy on a diverse set of simulation tasks. Moreover, it produces high-quality and diverse action trajectories within 1-2 steps on a set of real-world robot tasks. Our method also retains the same training complexity as Diffusion Policy and vanilla Flow Matching, in contrast to distillation-based approaches.

Problem

• Reduce computational cost of flow-based robot controllers
• Improve few-step performance in conditional tasks
• Maintain training simplicity while enhancing speed

Innovation

• Uses conditional Optimal Transport couplings
• Enforces straight solutions in flow ODE
• Improves few-step performance with condition variables