Truncated Rectified Flow Policy for Reinforcement Learning with One-Step Sampling

Xubin Zhou; Yipeng Yang; Zhan Li

cs.LGApr 10, 2026

Truncated Rectified Flow Policy for Reinforcement Learning with One-Step Sampling

Authors: Xubin Zhou, Yipeng Yang, Zhan Li

Organizations: Harbin Institute of Technology

Abstract

Maximum entropy reinforcement learning (MaxEnt RL) has become a standard framework for sequential decision making, yet its standard Gaussian policy parameterization is inherently unimodal, limiting its ability to model complex multimodal action distributions. This limitation has motivated increasing interest in generative policies based on diffusion and flow matching as more expressive alternatives. However, incorporating such policies into MaxEnt RL is challenging for two main reasons: the likelihood and entropy of continuous-time generative policies are generally intractable, and multi-step sampling introduces both long-horizon backpropagation instability and substantial inference latency. To address these challenges, we propose Truncated Rectified Flow Policy (TRFP), a framework built on a hybrid deterministic-stochastic architecture. This design makes entropy-regularized optimization tractable while supporting stable training and effective one-step sampling through gradient truncation and flow straightening. Empirical results on a toy multigoal environment and 10 MuJoCo benchmarks show that TRFP captures multimodal behavior effectively, outperforms strong baselines on most benchmarks under standard sampling, and remains highly competitive under one-step sampling.

Figures & tables

Panel of Figure 1 — Figure 1: Trajectory visualization in the toy multigoal environment. The agent starts near the center and moves toward four equally rewarded goals. Each panel overlays 20 evaluation episodes collected after training.

Panel of Figure 2 — Figure 2: Learning curves on 10 MuJoCo benchmarks. Each subplot compares TRFP under standard sampling ( $K=4$ ) and one-step sampling ( $K=1$ ) with four baselines. The solid lines denote the mean episode return over five random seeds, and the shaded regions represent one standard deviation.

Method	NFE	Ant-v5	HalfCheetah-v5	Hopper-v5	Humanoid-v5	Inverted2Pendulum-v5
TD3	1	4708.5 $\pm$ 569.1	8677.1 $\pm$ 1585.6	3015.6 $\pm$ 370.4	5138.6 $\pm$ 96.7	9319.4 $\pm$ 0.4
SAC	1	4566.4 $\pm$ 1069.0	11349.7 $\pm$ 326.5	2220.6 $\pm$ 815.9	5350.5 $\pm$ 361.5	9352.0 $\pm$ 6.4
MaxEntDP	$20\times 10$	2888.6 $\pm$ 691.9	6794.8 $\pm$ 591.4	3000.6 $\pm$ 888.7	3084.2 $\pm$ 1358.8	8578.2 $\pm$ 1478.4
SDAC	$20\times 32$	977.5 $\pm$ 152.2	11931.2 $\pm$ 578.4	2591.3 $\pm$ 539.3	4087.9 $\pm$ 262.0	8794.3 $\pm$ 1210.9
TRFP(ours)	$4\times 4$	4802.1 $\pm$ 95.6	10867.5 $\pm$ 406.4	3507.5 $\pm$ 79.6	5370.8 $\pm$ 367.6	9304.1 $\pm$ 2.6
TRFP(one-step)	$1\times 4$	4792.8 $\pm$ 142.3	10332.7 $\pm$ 348.1	3440.6 $\pm$ 75.6	5332.5 $\pm$ 406.5	9304.1 $\pm$ 2.6

Table 1: Final return comparison on 10 MuJoCo-v5 benchmarks. Results are reported as mean

\pm

standard deviation. Best results are shown in bold , and second-best distinct results are underlined . NFE denotes the number of function evaluations per action. For iterative generative policies, expressions of the form

a\times b

indicate

a

rollout steps with

b

sampled action candidates.

Panel of Figure 3 — Figure 3: Ablation results of TRFP. From left to right: the effect of removing flow straightening regularization, the effect of removing Q-guided action selection, and the effect of removing the stochastic tail. The first two panels are evaluated on HalfCheetah-v5, and the third panel is evaluated on Humanoid-v5. Solid lines denote the mean episode return over five random seeds, and shaded regions represent one standard deviation.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Number of denoising steps $K$	4
Tail length $L$	1
One-step evaluation tail length	0
Discount factor $\gamma$	0.99
Replay buffer size	1e6
Batch size	256

Appendix

Table 2: Detailed Hyperparameter Settings.

Maximum Entropy RL

Maximum entropy reinforcement learning balances reward seeking with stochastic policies that preserve exploration and support diverse behavior.

See all 69 papers

FLAG: Flow Policy MaxEnt-RL by Latent Augmented Guidance
May 29, 2026Sungha Kim, Gawon Lee, Jusuk Lee +3RL Control Maximum Entropy RL
Seoul National University · Georgia Institute of Technology
Drift, Then Spread: Learning One-Step Generative Policies for Online Reinforcement Learning
Oct 11, 2026Chunyu Luo, Hongyu Zhao, Guangyi Ni +2Reinforcement Learning Maximum Entropy RL
The University of Hong Kong · Beihang University · Harbin Institute of Technology, Shenzhen +2
Stochastic MeanFlow Policies: One-Step Generative Control with Entropic Mirror Descent
May 20, 2026Zeyuan Wang, Da Li, Yulin Chen +5Reinforcement Learning Maximum Entropy RL
Laboratory for Big Data and Decision, National University of Defense Technology, China · Samsung AI Center Cambridge, UK · Queen Mary University of London, UK +3

May 29, 2026J/K move · Enter open · S save

FLAG: Flow Policy MaxEnt-RL by Latent Augmented Guidance

Sungha Kim, Gawon Lee, Jusuk Lee, Jonghae Park+2

Seoul National University · Georgia Institute of Technology

RL Control Maximum Entropy RL Policy Learning Flow-Based Policy Learning

Open paper AskPDF arXiv

Maximum entropy reinforcement learning (MaxEnt-RL) enables robust exploration, yet practical implementations often restrict policies to simple Gaussians. While recent approaches incorporate expressive generative policies via importance-weighted supervised learning, they are prone to importance weight collapse, which limits their scalability in high-dimensional action spaces. Our key insight is to mitigate this limitation by localizing the sampling region, avoiding the weight degeneracy induced by importance sampling over the entire action space. To instantiate this insight, we introduce FLAG (Flow policy with Latent-Augmented Guidance). FLAG augments the state space with a flow latent variable and optimizes a provably consistent proxy MaxEnt-RL objective. We empirically demonstrate that FLAG enables expressive policy optimization with limited importance samples and scales to high-dimensional control tasks. Furthermore, FLAG achieves state-of-the-art performance across challenging benchmarks. Our project webpage: https://flag-rl.github.io/

Abstract

Maximum Entropy RL

Explore similar work

FLAG: Flow Policy MaxEnt-RL by Latent Augmented Guidance

Drift, Then Spread: Learning One-Step Generative Policies for Online Reinforcement Learning

Stochastic MeanFlow Policies: One-Step Generative Control with Entropic Mirror Descent