Truncated Rectified Flow Policy for Reinforcement Learning with One-Step Sampling
2026-04-10 • Machine Learning
Machine Learning
AI summaryⓘ
The authors focus on improving how computers learn decisions by using a method called maximum entropy reinforcement learning (MaxEnt RL), which usually struggles with representing actions that have multiple possible outcomes. To fix this, they created a new approach called Truncated Rectified Flow Policy (TRFP) that combines deterministic and random elements to better handle complex action choices. Their method makes training more stable and faster by simplifying certain calculations and using shorter sampling steps. Tests showed TRFP works well in both simple and more complex simulated tasks, often beating other strong methods.
Maximum entropy reinforcement learningGaussian policyMultimodal action distributionsDiffusion modelsFlow matchingEntropy regularizationGradient truncationFlow straighteningMuJoCo benchmarksSequential decision making
Authors
Xubin Zhou, Yipeng Yang, Zhan Li
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.