ScoRe-Flow: Complete Distributional Control via Score-Based Reinforcement Learning for Flow Matching

2026-04-12

AI Summary

This work addresses the limitations of existing flow-matching approaches in imitation learning, which are constrained by suboptimal demonstration data, and the inefficiency of reinforcement learning (RL) fine-tuning that relies on noisy exploration. The authors propose a score-based RL fine-tuning framework that, for the first time, derives a closed-form expression of the score function directly from the velocity field without requiring auxiliary networks. By modulating the drift term and jointly learning variance prediction, the method enables decoupled control over both the mean and variance of stochastic transition distributions. Empirical results demonstrate that the approach achieves a 2.4× faster convergence on D4RL locomotion tasks and improves success rates by up to 5.4% on Robomimic and Franka Kitchen manipulation benchmarks.

Technology Category

• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Intelligent Robots: Learning & Optimization for ROB
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Flow Matching (FM) policies have emerged as an efficient backbone for robotic control, offering fast and expressive action generation that underpins recent large-scale embodied AI systems. However, FM policies trained via imitation learning inherit the limitations of demonstration data; surpassing suboptimal behaviors requires reinforcement learning (RL) fine-tuning. Recent methods convert deterministic flows into stochastic differential equations (SDEs) with learnable noise injection, enabling exploration and tractable likelihoods, but such noise-only control can compromise training efficiency when demonstrations already provide strong priors. We observe that modulating the drift via the score function, i.e., the gradient of log-density, steers exploration toward high-probability regions, improving stability. The score admits a closed-form expression from the velocity field, requiring no auxiliary networks. Based on this, we propose ScoRe-Flow, a score-based RL fine-tuning method that combines drift modulation with learned variance prediction to achieve decoupled control over the mean and variance of stochastic transitions. Experiments demonstrate that ScoRe-Flow achieves 2.4x faster convergence than flow-based SOTA on D4RL locomotion tasks and up to 5.4% higher success rates on Robomimic and Franka Kitchen manipulation tasks.

Problem

• Flow Matching
• Reinforcement Learning
• Score-Based Control
• Distributional Control
• Robotic Policy Fine-tuning

Innovation

• Score-based Reinforcement Learning
• Flow Matching
• Drift Modulation
• Distributional Control
• Stochastic Transitions