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

2026-04-10

AI Summary

This work addresses the limitations of Gaussian policies in maximum entropy reinforcement learning, which struggle to model multimodal action distributions, as well as the intractable entropy computation, training instability, and high inference latency associated with diffusion- or flow-matching-based approaches due to their reliance on multi-step sampling. To overcome these challenges, the paper proposes the Truncated Rectified Flow Policy (TRFP), which introduces rectified flows and gradient truncation into policy design for the first time. By integrating flow straightening with a hybrid deterministic-stochastic architecture, TRFP enables analytically tractable entropy-regularized optimization and supports efficient single-step sampling. Experiments demonstrate that TRFP effectively captures multimodal behaviors on both a toy multimodal task and ten MuJoCo benchmarks, consistently matching or significantly outperforming strong baselines under both standard and single-step sampling evaluation settings.

Technology Category

• Machine Learning: Reinforcement Learning
• Humans and AI: Human-Aware Planning and Behavior Prediction
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy

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.

Problem

• maximum entropy reinforcement learning
• multimodal action distributions
• generative policies
• intractable likelihood and entropy
• multi-step sampling

Innovation

• Truncated Rectified Flow
• Maximum Entropy Reinforcement Learning
• Generative Policy
• One-Step Sampling
• Multimodal Action Distribution