Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Fast Flow-based Visuomotor Policies via Conditional Optimal Transport Couplings

Andreas Sochopoulos, Nikolay Malkin, Nikolaos Tsagkas, João Moura, Michael Gienger, Sethu Vijayakumar

TL;DR

This work addresses the real-time deployment gap of diffusion and flow-matching policies for robotic visuomotor control by introducing a conditional Optimal Transport coupling (COT) within a flow-matching framework. By discretizing continuous observations into condition labels and pairing noise-action samples with a conditional OT plan, the method produces straighter integration paths in the action-generating ODE, enabling accurate few-step sampling while preserving multimodality. Empirically, the 2-step COT Policy outperforms diffusion-based and OT-CFM baselines on simulated robot tasks and demonstrates robust multimodal behavior with low inference steps, including successful real-world experiments with minimal computation. The approach maintains training efficiency comparable to standard FM methods and offers a practical, scalable alternative to distillation-based acceleration for real-time robotic control.

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.

Fast Flow-based Visuomotor Policies via Conditional Optimal Transport Couplings

TL;DR

This work addresses the real-time deployment gap of diffusion and flow-matching policies for robotic visuomotor control by introducing a conditional Optimal Transport coupling (COT) within a flow-matching framework. By discretizing continuous observations into condition labels and pairing noise-action samples with a conditional OT plan, the method produces straighter integration paths in the action-generating ODE, enabling accurate few-step sampling while preserving multimodality. Empirically, the 2-step COT Policy outperforms diffusion-based and OT-CFM baselines on simulated robot tasks and demonstrates robust multimodal behavior with low inference steps, including successful real-world experiments with minimal computation. The approach maintains training efficiency comparable to standard FM methods and offers a practical, scalable alternative to distillation-based acceleration for real-time robotic control.

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.
Paper Structure (58 sections, 7 equations, 11 figures, 6 tables, 1 algorithm)

This paper contains 58 sections, 7 equations, 11 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. COT Policy
  3. Preliminaries: Conditional Flow Matching for Generative Modeling
  4. Neural ODE generative models.
  5. Conditional flow matching.
  6. Optimal transport couplings.
  7. Conditional generative modeling with CFM.
  8. Imitation learning as conditional generative modeling.
  9. Optimal Transport Couplings for Finite Conditional Tasks
  10. COT-CFM as a Robot Policy
  11. Continuous condition tasks.
  12. Robot policies.
  13. Training complexity.
  14. Experiments
  15. Baselines.
  16. ...and 43 more sections

Figures (11)

  • Figure 1: (I-)CFM, OT-CFM, and COT-CFM flows trained to generate the two moons distribution from the 8 Gaussians distribution. Generation with 100 (top row) and 1 (bottom row) euler integration steps is shown. CFM gives curved flows that cannot be integrated accurately in one step, while OT-CFM gives biased samples. Our proposed COT-CFM avoids both issues.
  • Figure 2: Tasks used for the evaluation of COT Policy.
  • Figure 3: Left: Evaluation of the ability of CFM and COT Policy to encode multiple modes in the maze navigation task for 1, 2, and 5 euler steps (NFE). Right: Trajectory Variance of CFM and COT Policy in the two Maze tasks and coffee_d1 task for $\text{NFE}=1$ on the top and $\text{NFE}=2$ on the bottom. The percentage of successful rollouts is written above each bar.
  • Figure 4: Success rates of COT Policy on the push-t task for $K=8, 32, 64, 128$ and $\text{NFE}=2,4,10$.
  • Figure 5: Success rates on the push-t task with varying NFE with the Left:euler solver and Middle:midpoint solver. Right: success rate over training epochs.
  • ...and 6 more figures