RLDT 解决了流匹配策略在强化学习中难以微调的痛点,做连续控制或机器人操作的团队可以直接参考其密度传输思路,比蒸馏或近似分布的方法更高效。
本文提出 RLDT(Reinforcement Learning with Density Transport),一种在线强化学习算法,用于微调连续控制问题中的流匹配策略。核心思想是将策略改进视为动作密度向高奖励区域的传输,与流匹配模型的传输公式自然对齐。RLDT 使用 Stein 变分梯度下降(SVGD)从最大熵 RL 目标构建传输场,然后微调预训练的流匹配策略以对齐该场。通过预期目标估计近似中间去噪步骤的动作,避免了不稳定的反向传播。实验表明,RLDT 在奖励质量和收敛速度上优于基线,适用于密集/稀疏奖励及基于状态/视觉的长期机器人操作任务。
Reinforcement Learning for Flow-Matching Policies with Density Transport
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/}.