Optimal Transport Q-Learning for Flow Policy Steering and Acceleration

2026-07-07

AI Summary

Existing flow-based robotic policies suffer from performance limitations due to their reliance on high-quality demonstrations, sensitivity to distributional shifts, and slow inference. To address these challenges, this work proposes OTQL, a novel approach that integrates optimal transport with Q-learning for the first time. By leveraging advantage-weighted conditional optimal transport within a post-training reinforcement learning framework, OTQL enables efficient fine-tuning and acceleration of suboptimal flow policies without requiring costly distillation. Evaluated under extremely low interaction budgets, the method substantially improves performance: success rates for single-task policies increase from 36% to 86%, and those for pretrained vision-language-action (VLA) models rise from 38% to 76%, while simultaneously reducing action generation inference steps by 70%.

Abstract

Diffusion and flow policies have recently demonstrated remarkable performance in robotic applications by accurately capturing multimodal robot trajectory distributions, especially in the context of vision language action (VLA) models. However, high quality policy performance also requires fast inference and high quality demonstrations, which are often hard to get. Lack of these leads to suboptimal policy behaviors and failure under distribution shifts. In this work we address the problem of fine-tuning and accelerating suboptimal flow-based policies using the robot's experience through RL post-training. We introduce Optimal Transport Q-Learning (OTQL), a new method for finetuning flow policies using advantage weighted conditional optimal transport flow matching. OTQL can finetune and accelerate flows with an interaction budget of 50-60 episodes while avoiding computationally expensive distillation in simulation and real-world robot tasks. Our results show that OTQL post-trains flow policies using the robot's own experience, increasing average success percentage of single-task policies from 36% to 86% and of a pre-trained VLA from 38% to 76% while reducing the number of inference steps per action generation by 70%.

Problem

• flow policy
• policy fine-tuning
• inference acceleration
• distribution shift
• robotic policy

Innovation

• Optimal Transport
• Q-Learning
• Flow Policy
• Policy Finetuning
• Robot Learning