Flow-based Policy With Distributional Reinforcement Learning in Trajectory Optimization

Ruijie Hao; Longfei Zhang; Yang Dai; Yang Ma; Xingxing Liang; Guangquan Cheng

cs.LGApr 1, 2026

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

Abstract

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).

Panel of Figure 2 — Figure 2: Benchmarks. (a) Humanoid-v4: $(s\times a)\in\mathbb{R}^{376}\times\mathbb{R}^{17}$ . (b) Ant-v4: $(s\times a)\in\mathbb{R}^{111}\times\mathbb{R}^{8}$ . (c) Hopper-v4: $(s\times a)\in\mathbb{R}^{11}\times\mathbb{R}^{3}$ . (d) HalfCheetah-v4: $(s\times a)\in\mathbb{R}^{17}\times\mathbb{R}^{6}$ . (e) InvertedPendulum-v: $(s\times a)\in\mathbb{R}^{4}\times\mathbb{R}^{1}$ . (f) Reacher-v4: $(s\times a)\in\mathbb{R}^{11}\times\mathbb{R}^{2}$ .

Panel of Figure 3 — 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 $\pm$ 0	3499 $\pm$ 1081	5179 $\pm$ 1408	3023 $\pm$ 0	5312 $\pm$ 1622	5971 $\pm$ 410
Half.	15227 $\pm$ 512	11694 $\pm$ 171	13629 $\pm$ 994	8824 $\pm$ 0	12609 $\pm$ 0	10277 $\pm$ 399
Hopper.	3298 $\pm$ 608	1958 $\pm$ 599	2448 $\pm$ 769	3319 $\pm$ 0	3338 $\pm$ 367	1534 $\pm$ 619
Human.	11940 $\pm$ 0	77 $\pm$ 2	5184 $\pm$ 572	69 $\pm$ 0	10017 $\pm$ 0	4553 $\pm$ 1134
Invert.	1000 $\pm$ 5	860 $\pm$ 174	941 $\pm$ 164	778 $\pm$ 0	25 $\pm$ 5	5 $\pm$ 2
Reacher	-4 $\pm$ 0	-4 $\pm$ 1	-5 $\pm$ 1	-5 $\pm$ 0	-8 $\pm$ 2	-5 $\pm$ 1

Table 1: Comparison of FP-DRL with baseline algorithms on MuJoCo continuous control tasks. Each value reports the mean

\pm

standard deviation of episodic return over the last 10% of training steps. The best value for each task is bolded.

Panel of Figure 4 — 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).

RL Control

Reinforcement learning control uses reward-driven learning to make decisions in dynamic systems, from robots and vehicles to energy and infrastructure.

See all 272 papers

Refinement-Based Flow Policy Optimization
Sep 14, 2026Bumgeun Park, Hyukjun Yang, Donghwan LeeRL Control Flow Matching
Korea Advanced Institute of Science and Technology (KAIST)
Learning Multimodal One-step Flow Policy via Value-weighted Optimal Transport
Sep 14, 2026Jaehun Shon, Jinha Choi, Jongwook Jeon +1Offline RL Flow-Based Policy Learning
Department of Artificial Intelligence, Yonsei University
Dual-Flow Reinforcement Learning with State-Aware Exploration
Jun 29, 2026Qijun Li, Zheng Fu, Qi Song +4Distributional RL Reinforcement Learning
School of Vehicle and Mobility and the State Key Laboratory of Intelligent Green Vehicle and Mobility, Tsinghua University, Beijing 100084, China

Sep 14, 2026J/K move · Enter open · S save

Refinement-Based Flow Policy Optimization

Bumgeun Park, Hyukjun Yang, Donghwan Lee

Korea Advanced Institute of Science and Technology (KAIST)

RL Control Flow Matching Policy Optimization Continuous Control Flow-Based Policy Learning

Open paper AskPDF arXiv

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.

Abstract

RL Control

Explore similar work

Refinement-Based Flow Policy Optimization

Learning Multimodal One-step Flow Policy via Value-weighted Optimal Transport

Dual-Flow Reinforcement Learning with State-Aware Exploration