Reinforcement Learning for Flow-Matching Policies with Density Transport

2026-06-07

AI Summary

Existing reinforcement learning methods struggle to simultaneously achieve unbiased policy updates and multimodal action modeling in continuous control. This work proposes RLDT, a novel approach that formulates policy optimization as the transport of an action density toward high-reward regions. By integrating flow matching with a maximum-entropy reinforcement learning objective, RLDT constructs a transport field via Stein variational gradient descent and stabilizes fine-tuning of multi-step denoising policies through expectation-based target estimation. The method circumvents biases introduced by conventional distillation or distributional approximation techniques, preserving multimodal expressiveness while significantly enhancing training stability and convergence speed. Empirical results demonstrate that RLDT consistently outperforms current baselines across a range of continuous control tasks, including those with sparse or dense rewards and long-horizon robotic manipulation scenarios based on either state or visual observations.

Abstract

We present an online reinforcement learning (RL) algorithm for fine-tuning flow-matching policies in continuous-control problems. Our key insight is to view RL-based policy improvement as a transport of action densities towards regions of high reward, which naturally aligns with the transport formulation of flow matching models. Prior methods either approximate the current or optimal policy distribution or resort to distillation, which introduces biased gradients or sacrifices multimodal modeling capacity. In contrast, our approach for RL with Density Transport, which we name \emph{RLDT}, constructs a transport field from a maximum-entropy RL objective using Stein Variational Gradient Descent (SVGD). Then, it finetunes a pretrained flow matching policy to align with this field. Training with this alignment objective is nontrivial because flow-matching policies generate actions via a multi-step process, making direct gradient-based optimization challenging. To overcome this challenge and stabilize training, we approximate policy actions from intermediate denoising steps via expected-target estimation. This allows the transport-field update to propagate into the network parameters without unstable backpropagation through time. Experimental results demonstrate that RLDT outperforms competitive baselines in reward quality and convergence speed. This performance holds across diverse continuous-control tasks, encompassing both dense and sparse rewards, as well as state- and vision-based long-horizon robot manipulation. The project webpage is \href{https://rpfey.github.io/rldt/}{https://rpfey.github.io/rldt/}.

Problem

• Reinforcement Learning
• Flow Matching
• Density Transport
• Continuous Control
• Policy Optimization

Innovation

• flow matching
• density transport
• reinforcement learning
• Stein Variational Gradient Descent
• expected-target estimation