Optimal Transport Q-Learning for Flow Policy Steering and Acceleration
OTQL integrates advantage-weighted conditional optimal transport with flow models, achieving few-step inference and efficient policy fine-tuning, boosting success rates from 36% to 86%.
Key Findings
Methodology
OTQL combines advantage-weighted conditional optimal transport (COT) with neural ODE flow matching, guiding flows along straight paths toward high-value samples. It employs a critic to steer the transport, using advantage functions to weight sample importance. The core components include:• Conditional OT plan for direct noise-action mapping;• Critic-guided flow training to prioritize high-reward samples;• Few-step inference (2-3 steps) achieved by straight-line paths, avoiding complex distillation. The approach enables rapid policy fine-tuning using robot experience, with only 50-60 interaction episodes, significantly improving success rates and reducing inference steps.
Key Results
- OTQL boosts single-task success from 36% to 86% and pre-trained VLA performance from 38% to 76%, with a 70% reduction in inference steps across simulation and real robot tasks.
- In offline and online settings, OTQL outperforms baselines like FQL and QAM-F, especially with limited interaction data, demonstrating high sample efficiency.
- The method maintains high inference speed while achieving comparable or superior policy quality, enabling real-time adaptation and acceleration.
Significance
This work addresses the bottleneck of slow inference in flow-based policies, crucial for real-time robotic control. By integrating optimal transport with advantage weighting, it offers a scalable solution for fast, high-quality policy adaptation. The approach bridges the gap between offline pretraining and online fine-tuning, facilitating robots to learn efficiently from limited experience. It advances the field by providing a theoretically grounded, computationally efficient framework that enhances the robustness and generalization of robotic policies in dynamic environments, with broad implications for autonomous systems and intelligent control.
Technical Contribution
OTQL introduces a novel synergy of advantage-weighted COT and neural ODE flows, enabling direct, straight-line transport paths that require only 2-3 inference steps. It innovates by incorporating a critic to influence the transport plan, prioritizing high-value samples, and avoiding costly distillation. The method achieves significant speed-ups in inference and training, with theoretical guarantees on the straightness of transport paths. These contributions collectively push the boundaries of flow model efficiency and adaptability in reinforcement learning contexts.
Novelty
This is the first application of advantage-weighted conditional optimal transport to neural flow policies for robotic control. Unlike prior diffusion acceleration methods, OTQL leverages a critic-guided, straight-line transport plan, enabling rapid, few-step inference. Its integration of advantage weighting with OT provides a new paradigm for policy fine-tuning, combining theoretical rigor with practical efficiency, marking a significant innovation in the field.
Limitations
- The effectiveness depends on the critic’s accuracy; poor critic training can impair transport quality and policy performance.
- In highly complex or high-dimensional tasks, the straight-line assumption may still require multiple steps, limiting speed gains.
- Current experiments focus on manipulation tasks; applicability to other domains remains to be validated.
Future Work
Future directions include enhancing critic robustness, extending OTQL to multi-agent systems, and integrating multi-modal sensory data for richer state representations. Exploring self-supervised signals to improve transport path learning and generalization is also promising. Additionally, scaling the approach to more complex, high-dimensional tasks and real-world scenarios will be crucial for broader deployment in autonomous robots and intelligent systems.
AI Executive Summary
Robotics research has long sought methods to enable fast, reliable decision-making in complex environments. Traditional flow models, while powerful in capturing multimodal action distributions, suffer from slow inference due to iterative sampling processes. This bottleneck hampers their deployment in real-time control scenarios. Addressing this challenge, the present work introduces OTQL, a novel reinforcement learning framework that leverages advantage-weighted conditional optimal transport (COT) combined with neural ODE flow matching. The core idea is to guide the flow along direct, straight-line paths from noise to high-value actions, drastically reducing the number of inference steps needed.
The method employs a critic network to evaluate action quality, influencing the transport plan to favor high-reward samples. This integration ensures that the flow learns to generate optimal actions efficiently, with only 2-3 integration steps, compared to dozens in traditional methods. Extensive experiments in both simulation and real-world robotic tasks demonstrate that OTQL significantly boosts success rates—from 36% to 86% in single-task settings—and improves pre-trained vision-language models from 38% to 76%. It also reduces inference steps by 70%, enabling real-time deployment.
This breakthrough not only enhances the efficiency of policy fine-tuning but also opens new avenues for robots to adapt rapidly with limited interaction. The approach outperforms existing baselines like FQL and QAM-F in offline, online, and hybrid training scenarios, showcasing robustness and scalability. Its theoretical grounding in optimal transport guarantees the straightness and optimality of the learned paths, providing interpretability and stability.
Looking ahead, future work aims to incorporate multi-modal data, improve critic robustness, and extend the framework to more complex, high-dimensional tasks. The potential to enable autonomous robots to learn and adapt swiftly in dynamic environments marks a significant step toward truly intelligent robotic systems, with broad implications for industry, service robots, and autonomous vehicles.
Deep Dive
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%.