Flow-based Policy With Distributional Reinforcement Learning in Trajectory Optimization
Authors: Ruijie Hao, Longfei Zhang, Yang Dai, Yang Ma, Xingxing Liang, Guangquan Cheng
Organizations: College of Systems Engineering, National University of Defense Technology, Changsha 410073, China · Aviation University of Air Force, Changchun 130000, China
Reinforcement Learning (RL) has proven highly effective in addressing complex control and decision-making tasks. However, in most traditional RL algorithms, the policy is typically parameterized as a diagonal Gaussian distribution, which constrains the policy from capturing multimodal distributions, making it difficult to cover the full range of optimal solutions in multi-solution problems, and the return is reduced to a mean value, losing its multimodal nature and thus providing insufficient guidance for policy updates. In response to these problems, we propose a RL algorithm termed flow-based policy with distributional RL (FP-DRL). This algorithm models the policy using flow matching, which offers both computational efficiency and the capacity to fit complex distributions. Additionally, it employs a distributional RL approach to model and optimize the entire return distribution, thereby more effectively guiding multimodal policy updates and improving agent performance. Experimental trails on MuJoCo benchmarks demonstrate that the FP-DRL algorithm achieves state-of-the-art (SOTA) performance in most MuJoCo control tasks while exhibiting superior representation capability of the flow policy.
Figures & tables
Figure 1: FM learns a velocity fields that transports samples from a simple prior distribution to the target data distribution via an ordinary differential equation (ODE).
Figure 3: Training curves on benchmarks. The solid lines correspond to mean and shaded regions correspond to the standard error of the mean (SEM) over Three runs.
Task
FP-DRL
DSAC-T
SAC
TD3
SAC-FLOW
DACER
Ant
6665 ± 0
3499 ± 1081
5179 ± 1408
3023 ± 0
5312 ± 1622
5971 ± 410
Half.
15227 ± 512
11694 ± 171
13629 ± 994
8824 ± 0
12609 ± 0
10277 ± 399
Hopper.
3298 ± 608
1958 ± 599
2448 ± 769
3319 ± 0
3338 ± 367
1534 ± 619
Human.
11940 ± 0
77 ± 2
5184 ± 572
69 ± 0
10017 ± 0
4553 ± 1134
Invert.
1000 ± 5
860 ± 174
941 ± 164
778 ± 0
25 ± 5
5 ± 2
Reacher
-4 ± 0
-4 ± 1
-5 ± 1
-5 ± 0
-8 ± 2
-5 ± 1
Table 1: Comparison of FP-DRL with baseline algorithms on MuJoCo continuous control tasks. Each value reports the mean ± standard deviation of episodic return over the last 10% of training steps. The best value for each task is bolded.
Figure 4: Ablation studies of FP-DRL. (a) Performance comparison between a Gaussian policy and the proposed transformer-based flow policy. (b) Comparison of modeling returns using the mean versus a quantile-based distributional critic. (c) Effect of varying the number of quantiles (N=16, 32, and 64) for the distributional critic. (d) Training curves across different transformer sequence lengths (K=4, 7, 10, and 12).
Reinforcement learning control uses reward-driven learning to make decisions in dynamic systems, from robots and vehicles to energy and infrastructure.
Flow-based policies offer an expressive representation for online reinforcement learning, but conventional flow matching requires samples drawn from the distribution to be modeled. This poses a challenge when the desired action distribution is defined only implicitly by a Q-function, since directly sampling actions from the resulting distribution is generally intractable. We propose Refinement-Based Flow Policy Optimization (RFPO), a novel framework for training a flow policy in online reinforcement learning by alternating between Q-guided sample refinement and self-target flow matching. RFPO first generates actions from Gaussian noise using the current flow policy and then uses a finite-step stochastic refinement procedure to move them toward an energy-based distribution induced by the Q-function. Each refined action is then paired with its corresponding initial noise sample and used as a fixed target for flow-matching training. By repeatedly refining its own outputs and learning from the resulting targets, RFPO incorporates Q-guidance into the policy without requiring direct samples from the target distribution, while retaining the capacity to represent multiple action modes. We further provide a theoretical analysis of the distributional dynamics induced by RFPO. Across six continuous-control tasks, RFPO matches or outperforms a standard Gaussian-policy baseline on almost every task. Experiments on six synthetic two-dimensional target distributions with diverse geometries demonstrate that RFPO captures complex multimodal structure without mode collapse.
Bumgeun Park, Hyukjun Yang, Donghwan Lee
Korea Advanced Institute of Science and Technology (KAIST)
Offline reinforcement learning aims to learn a policy solely from fixed datasets, which often contain multimodal action distributions. Flow policies can naturally represent such multimodal behaviors, but learning an efficient one-step flow policy remains challenging: standard value guidance often leads to mode collapse or exploits overestimation bias in out-of-distribution regions. To address this, we introduce One-step Flow policy via Optimal Transport (OptiFlow), a framework for one-step flow policy learning as a structured sample-allocation problem. OptiFlow jointly trains a value-aware reference flow policy and an efficient one-step policy, coupling their action samples through state-wise entropic optimal transport. For each state, critic-estimated values define the priority of distillation target actions, while the action-distance cost ensures geometrically compatible pairings. By avoiding direct critic maximization, our transport-guided approach enables in-distribution exploitation by anchoring the one-step policy to high-value, dataset-supported modes without the risk of out-of-distribution divergence. Experimental results demonstrate that OptiFlow effectively captures optimal multimodal behaviors and achieves strong performance across diverse offline RL benchmarks. Our code is available at https://github.com/Yonsei-DILLab/OptiFlow.
Jaehun Shon, Jinha Choi, Jongwook Jeon +1
Department of Artificial Intelligence, Yonsei University
In complex continuous-control reinforcement learning tasks, multimodal optimal actions often coincide with uncertain, multimodal return distributions, making reliable value estimation and multimodal exploration challenging. Existing value estimation methods using unimodal Gaussians restrict expressiveness and yield biased estimates. Recent generative policies can represent multimodal actions but often collapse to a few modes and under-explore high-value areas of the action space. Motivated by these challenges, we propose Dual-Flow RL, a unified actor-critic framework that jointly models a continuous return distribution and a multimodal policy distribution using conditional flow matching (CFM). This design supports reliable value estimation and sustained multimodal exploration. To further enhance exploration, we introduce an Entropy-Covariance Exploration Regulator (ECER) that enables state-aware exploration regulation leveraging policy entropy and action-uncertainty covariance. Experiments on DeepMind Control Suite and Humanoid-Bench show that Dual-Flow RL achieves state-of-the-art performance on most tasks, significantly outperforming prior diffusion-based and flow-based methods.
Qijun Li, Zheng Fu, Qi Song +4
School of Vehicle and Mobility and the State Key Laboratory of Intelligent Green Vehicle and Mobility, Tsinghua University, Beijing 100084, China