Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

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

About

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.

Xubin Zhou, Yipeng Yang, Zhan Li• 2026

Related benchmarks

TaskDatasetResultRank
Continuous ControlWalker2D v5
Avg Return5.99e+3
17
Continuous ControlHopper v5
Average Return3.51e+3
15
Continuous ControlHumanoid v5
Average Return5.37e+3
13
Continuous ControlSwimmer v5
Average Episodic Reward119.3
8
Continuous ControlInvertedPendulum v5
Average Episodic Reward1.00e+3
8
Continuous ControlReacher v5
Average Episodic Reward-4.9
8
Continuous ControlAnt v5
Final Return4.80e+3
6
Continuous ControlPusher v5
Final Return-30.4
6
Continuous ControlHalfcheetah v5
Final Return1.09e+4
6
Continuous ControlInverted2Pendulum v5
Final Return9.30e+3
6
Showing 10 of 10 rows

Other info

Follow for update