Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Boosting Maximum Entropy Reinforcement Learning via One-Step Flow Matching

Zeqiao Li, Yijing Wang, Haoyu Wang, Zheng Li, Zhiqiang Zuo

TL;DR

FLAME presents a principled integration of flow-based policy representations with Maximum Entropy RL by introducing a Q-Reweighted Flow Matching objective that cancels the partition function, and two entropy estimation strategies (FLAME-R and FLAME-M) to mitigate discretization bias while preserving one-step inference. The framework leverages MeanFlow to deliver expressive, low-latency action generation (NFE=1) and demonstrates state-of-the-art performance on MuJoCo benchmarks, matching or surpassing multi-step diffusion methods with significantly lower inference cost. Empirically, QRFM effectively aligns the policy with high-value regions, FLAME-R ensures unbiased entropy regularization, and FLAME-M provides a practical decoupled estimator that maintains performance under strict latency constraints. The results indicate that FLAME offers a scalable path to expressive, low-latency policies for real-time control, with broad applicability to multimodal action distributions and vision-based tasks.

Abstract

Diffusion policies are expressive yet incur high inference latency. Flow Matching (FM) enables one-step generation, but integrating it into Maximum Entropy Reinforcement Learning (MaxEnt RL) is challenging: the optimal policy is an intractable energy-based distribution, and the efficient log-likelihood estimation required to balance exploration and exploitation suffers from severe discretization bias. We propose \textbf{F}low-based \textbf{L}og-likelihood-\textbf{A}ware \textbf{M}aximum \textbf{E}ntropy RL (\textbf{FLAME}), a principled framework that addresses these challenges. First, we derive a Q-Reweighted FM objective that bypasses partition function estimation via importance reweighting. Second, we design a decoupled entropy estimator that rigorously corrects bias, which enables efficient exploration and brings the policy closer to the optimal MaxEnt policy. Third, we integrate the MeanFlow formulation to achieve expressive and efficient one-step control. Empirical results on MuJoCo show that FLAME outperforms Gaussian baselines and matches multi-step diffusion policies with significantly lower inference cost. Code is available at https://github.com/lzqw/FLAME.

Boosting Maximum Entropy Reinforcement Learning via One-Step Flow Matching

TL;DR

FLAME presents a principled integration of flow-based policy representations with Maximum Entropy RL by introducing a Q-Reweighted Flow Matching objective that cancels the partition function, and two entropy estimation strategies (FLAME-R and FLAME-M) to mitigate discretization bias while preserving one-step inference. The framework leverages MeanFlow to deliver expressive, low-latency action generation (NFE=1) and demonstrates state-of-the-art performance on MuJoCo benchmarks, matching or surpassing multi-step diffusion methods with significantly lower inference cost. Empirically, QRFM effectively aligns the policy with high-value regions, FLAME-R ensures unbiased entropy regularization, and FLAME-M provides a practical decoupled estimator that maintains performance under strict latency constraints. The results indicate that FLAME offers a scalable path to expressive, low-latency policies for real-time control, with broad applicability to multimodal action distributions and vision-based tasks.

Abstract

Diffusion policies are expressive yet incur high inference latency. Flow Matching (FM) enables one-step generation, but integrating it into Maximum Entropy Reinforcement Learning (MaxEnt RL) is challenging: the optimal policy is an intractable energy-based distribution, and the efficient log-likelihood estimation required to balance exploration and exploitation suffers from severe discretization bias. We propose \textbf{F}low-based \textbf{L}og-likelihood-\textbf{A}ware \textbf{M}aximum \textbf{E}ntropy RL (\textbf{FLAME}), a principled framework that addresses these challenges. First, we derive a Q-Reweighted FM objective that bypasses partition function estimation via importance reweighting. Second, we design a decoupled entropy estimator that rigorously corrects bias, which enables efficient exploration and brings the policy closer to the optimal MaxEnt policy. Third, we integrate the MeanFlow formulation to achieve expressive and efficient one-step control. Empirical results on MuJoCo show that FLAME outperforms Gaussian baselines and matches multi-step diffusion policies with significantly lower inference cost. Code is available at https://github.com/lzqw/FLAME.
Paper Structure (42 sections, 4 theorems, 59 equations, 11 figures, 6 tables, 1 algorithm)

This paper contains 42 sections, 4 theorems, 59 equations, 11 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Maximum Entropy Reinforcement Learning
  4. Flow Matching and Mean Flow
  5. Methodology
  6. Q-Reweighted Flow Policy Learning
  7. Entropy Regularization via Continuous Flows
  8. Practical Implementation and Algorithms
  9. Experiments
  10. Experimental Setup
  11. Performance on MuJoCo Benchmarks
  12. Bias-Variance Analysis of Entropy Estimation
  13. Policy Representation on Multimodal Environment
  14. Visual Analysis of Training Evolution
  15. Efficiency and Deployment
  16. ...and 27 more sections

Key Result

Proposition 3.1

Assuming the marginal probability density satisfies $p_t(a \mid s) > 0$ for all $a \in \mathcal{A}$ and $t \in [0, 1]$, the conditional objective $\mathcal{L}_{\mathrm{CFM}}$ (Eq. eq:cfm_loss) and the marginal objective $\mathcal{L}_{\mathrm{FM}}$ (Eq. eq:fm_loss) are equivalent up to a constant ind Consequently, their gradients satisfy $\nabla_{\!\theta}\mathcal{L}_{\text{FM}} \!=\! \nabla_{\!\th

Figures (11)

  • Figure 1: MultiGoal multimodality. Gaussian policies (SAC) and methods without entropy regulation collapse to one goal, while FLAME-R/M cover all four symmetric goals, demonstrating stable MaxEnt exploration.
  • Figure 2: Log-Likelihood Estimation Error vs. Integration Steps ($N_{\text{est}}$). The Mean Squared Error (MSE) drops significantly as $N_{\text{est}}$ increases. The convergence at $N_{\text{est}}=5$ validates our choice of multi-step estimation for the critic to reduce bias.
  • Figure 3: Evolution of sampling trajectories. From left to right: Policies at 5K (early) and 200K (converged) iterations. SDAC retains curvature, requiring multi-step sampling. FLAME-R gradually straightens the flow, enabling one-step inference at convergence. FLAME-M achieves straight trajectories almost immediately (5K), demonstrating superior training efficiency.
  • Figure 4: Training curves on 10 MuJoCo continuous control benchmarks. The x-axis represents environment steps ($\times 10^6$) and the y-axis shows the average episode return. Our methods, FLAME-R and FLAME-M, consistently match or outperform both Gaussian baselines (SAC) and multi-step diffusion policies (SDAC) while requiring only one-step inference. Shaded regions denote the standard deviation across 5 seeds.
  • Figure 5: Training curves on visual DMC tasks. FLAME-R (solid red) consistently matches or exceeds the performance of DPMD and FPMD baselines across all pixel-based environments.
  • ...and 6 more figures

Theorems & Definitions (8)

  • Proposition 3.1: Gradient Equivalence of FM and CFM
  • Proposition 3.2: Reweighting Invariance
  • Proposition 3.3: Discretization Error in Density Estimation
  • Corollary 3.4: Multi-Step Error Suppression
  • proof
  • proof
  • proof
  • proof