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
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.
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 ± 569.1
8677.1 ± 1585.6
3015.6 ± 370.4
5138.6 ± 96.7
9319.4 ± 0.4
SAC
1
4566.4 ± 1069.0
11349.7 ± 326.5
2220.6 ± 815.9
5350.5 ± 361.5
9352.0 ± 6.4
MaxEntDP
20×10
2888.6 ± 691.9
6794.8 ± 591.4
3000.6 ± 888.7
3084.2 ± 1358.8
8578.2 ± 1478.4
SDAC
20×32
977.5 ± 152.2
11931.2 ± 578.4
2591.3 ± 539.3
4087.9 ± 262.0
8794.3 ± 1210.9
TRFP(ours)
4×4
4802.1 ± 95.6
10867.5 ± 406.4
3507.5 ± 79.6
5370.8 ± 367.6
9304.1 ± 2.6
TRFP(one-step)
1×4
4792.8 ± 142.3
10332.7 ± 348.1
3440.6 ± 75.6
5332.5 ± 406.5
9304.1 ± 2.6
Table 1: Final return comparison on 10 MuJoCo-v5 benchmarks. Results are reported as mean ± 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×b indicate a rollout steps with b sampled action candidates.
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.
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/
Sungha Kim, Gawon Lee, Jusuk Lee +3
Seoul National University · Georgia Institute of Technology
Generative policies offer expressive action distributions for online reinforcement learning (RL), but many diffusion- and flow-based policies require costly iterative sampling. We propose DriftSpread, a one-step generative policy based on Drifting for online maximum-entropy (MaxEnt) RL. To incorporate entropy without a tractable policy density, we build on an energy-based view of Drifting and derive critic-guided Drift and entropy-induced Spread from the negative MaxEnt objective. Our method learns these updates sequentially in the same actor, using action regression for value improvement and distribution matching to capture the effect of entropy. The actor updates require neither direct samples from the critic-defined Boltzmann target nor policy-density or score estimation. Under exact fitting with a fixed critic, we prove convergence to the Boltzmann target in the small-step limit followed by the long-time limit. On MuJoCo, DriftSpread achieves the highest aggregate normalized return among all evaluated methods, while outperforming the diffusion- and flow-based baselines in both training and inference efficiency.
Chunyu Luo, Hongyu Zhao, Guangyi Ni +2
The University of Hong Kong · Beihang University · Harbin Institute of Technology, Shenzhen +2
Combining entropy regularisation with mirror descent in online off-policy reinforcement learning requires expressive policies with efficient sampling and explicit entropy control. Gaussian policies struggle to represent multimodal targets, whereas generative policies often require iterative sampling or lack tractable entropy estimates. We introduce Stochastic MeanFlow Policies (SMFP), a one-step generative policy class that combines MeanFlow-based noise-to-action mappings with Gaussian reparameterisation. The latent mapping captures diverse action modes, while a learnable conditional Gaussian scale provides an analytic entropy lower bound. We use this bound to regulate stochasticity alongside direct value guidance and advantage-weighted MeanFlow regression, enabling entropy-regularised mirror-descent learning with one-step policies. Across seven MuJoCo benchmarks, SMFP achieves strong performance against Gaussian and generative baselines with single-step inference. Additional experiments demonstrate multimodal action modelling and effective exploration in sparse-reward navigation.
Zeyuan Wang, Da Li, Yulin Chen +5
Laboratory for Big Data and Decision, National University of Defense Technology, China · Samsung AI Center Cambridge, UK · Queen Mary University of London, UK +3