FLAG: Flow Policy MaxEnt-RL by Latent Augmented Guidance

From Global to Local Importance Sampling

Prior methods train diffusion- and flow-based policies by matching a target action distribution, such as the unnormalized MaxEnt target $\exp(Q(s, a) / \alpha)$. Since this target cannot be sampled from directly, they use importance sampling (IS): sample actions from a proposal policy, reweight them toward the target, and perform weighted score or flow matching.

The limitation is that IS can only reweight actions that were already sampled. When the proposal-target discrepancy is large, especially in high-dimensional control, the proposal rarely samples target-relevant actions. This yields sparse supervision and lead to failure mode.

Our key insight is to change where IS is performed. Prior methods use Global IS, which reweights samples over the full action space. FLAG uses Local IS: it conditions both the proposal and target on the same flow latent variable $z$, so reweighting happens inside a shared local region.

We first illustrate this effect in a controlled multi-goal task then show the same pattern in high-dimensional continuous control tasks in the Result section.

Multi-goal Environment (Didactic Experiment)

Method

Q1 How do we construct local IS, and is it consistent with the original MDP?

Challenge Solution Bottom Line

Challenge

For local IS to be principled — not just a trick — two conditions must hold:

• Localization. The proposal and target distributions must share a region indexed by a latent $z$.
• Consistency. Optimizing inside that local region must be the same optimization problem as optimizing the RL objective in the original MDP.

Without Consistency, local IS optimizes the wrong objective and any gains are illusory.

Solution

Policy Parameterization Latent-Augmented MDP Consistency

The z-MDP preserves both state–action marginals and the Q-function — optimizing on the z-MDP is the same problem as on the original MDP. Local IS becomes principled, not a heuristic.

Corollary 4.2 — Marginal consistency

$$\rho_\pi(s, a) = \int \hat\rho_{\hat\pi}(s, z, a)\,dz$$

Corollary 4.4 — Q-function consistency

$$Q^{\hat\pi}(\hat s, a) = Q^{\pi}(s, a), \quad \forall\, s, z, a$$

Full statements (Corollaries 4.2, 4.4) are in the paper. Details are in Appendix B.

Bottom Line

We construct the z-MDP, and prove that optimizing the RL objective on the z-MDP is equivalent to optimizing it on the original MDP (Consistency). Local IS becomes principled, not a heuristic.

Q2 How do we incorporate this into MaxEnt-RL and update the policy?

Challenge Solution Bottom Line

Challenge

Two obstacles block plugging the z-MDP into MaxEnt-RL:

• Intractable Entropy. The composite policy’s log-probability $\log\pi(a \mid s)$ requires marginalizing over the latent $z$ — there is no closed form.
• RL via Supervised Learning. We want to cast MaxEnt-RL as an EM algorithm whose policy update reduces to a supervised learning problem — avoiding backprop through the flow ODE (BPTT).

Solution

Cross-entropy surrogate Localized EM

Proposition 4.3 — Cross-entropy surrogate

$$\begin{aligned} D_{\mathrm{TV}}(\pi, \tilde\pi) &= \mathcal{O}\big(\sqrt{\mathrm{tr}(\Sigma)}\big), \quad D_{\mathrm{KL}}(\pi \,\|\, \tilde\pi) = \mathcal{O}\big(\mathrm{tr}(\Sigma)^2\big) \\[4pt] \mathcal{H}(\pi, \tilde\pi) &= \mathcal{H}(\pi) + \mathcal{O}\big(\mathrm{tr}(\Sigma)^2\big) \approx \mathcal{H}(\pi) \end{aligned}$$

• Replace the entropy $\mathcal{H}(\pi)$ with the cross-entropy $\mathcal{H}(\pi, \tilde\pi)$ against the base flow $\tilde\pi$.
• Estimate $\log \tilde \pi(a\mid s)$ via the Hutchinson trace estimator — Use common random numbers (CRN) to reduce variance in the estimator (Appendix D.3).

Localized EM on the z-MDP

$$\begin{aligned} \textbf{E-step:}\;\; & q_k(a \mid \hat s) \propto \hat\pi(a \mid \hat s; \theta_k)\,\exp\!\big(f_{\hat s, k}(a)/\lambda\big) \\[6pt] \textbf{M-step:}\;\; & \theta_{k+1} \in \operatorname*{arg\,min}_{\theta}\; \mathbb{E}_{(s,z)\sim\hat\rho_{\hat\pi}}\big\|\, T_\theta(s,z) - \mu_k^*(\hat s) \big\|^2 \\[2pt] & \mu_k^*(\hat s) := \mathbb{E}_{q_k}[a] \end{aligned}$$

• E-step — a local Boltzmann tilt of the Gaussian proposal around each anchor.
• M-step — supervised regression of the flow onto improved action labels $\mu_k^*$, using Method of moments to estimate the labels with $N$ local samples.

Bottom Line

We use the cross-entropy as a tractable entropy surrogate and update the flow via an EM algorithm on the z-MDP — supervised distillation onto locally improved labels, not gradient-through-flow.

Q3 Does FLAG provably improve the policy, and how does it relate to SAC?

Challenge Solution Bottom Line

Challenge

FLAG updates the policy by supervised distillation, not by differentiating the objective through the flow. Two guarantees are therefore not obvious — and this section establishes both:

• Monotonic improvement. Does FLAG’s update actually raise the objective, i.e. is $\mathcal{J}_{k+1} \geq \mathcal{J}_k$ guaranteed?
• Relation to SAC. We optimize a MaxEnt-RL objective — so how does FLAG’s update relate to Soft Actor–Critic (SAC), and does it inherit SAC’s soft policy improvement?

Solution

Monotonic Improvement Relation to SAC

Theorem 4.5 — Monotonic Improvement Guarantee

$$\mathcal{J}_{k+1} - \mathcal{J}_k \;\geq\; \underbrace{\lambda \beta \mathcal{G}_k}_{\substack{\text{frozen-reward}\\\text{MPO improvement}}} \;-\; \underbrace{\alpha\, C_\Sigma\, \sigma_k^2\, \beta \mathcal{G}_k}_{\substack{\text{cross-entropy}\\\text{reward drift}}} \;-\; \underbrace{\lambda\, \epsilon_k^{\mathrm{proj}}}_{\substack{\text{CFM projection}\\\text{error}}} \;+\; \mathcal{O}(\beta^2)$$

• One FLAG step splits into a positive MPO improvement term, a cross-entropy drift that scales with the local variance $\sigma_k^2$, and the CFM projection error $\epsilon_k^{\mathrm{proj}}$.
• Whenever $\lambda > \alpha\, C_\Sigma\, \sigma_k^2$ and the projection error is small, the improvement term dominates, so $\mathcal{J}_{k+1} \geq \mathcal{J}_k$ — monotonic improvement.

Proof Roadmap

SAC — Soft policy improvement

$$\pi_{\mathrm{new}} = \operatorname*{arg\,min}_{\pi \in \Pi} D_{\mathrm{KL}}\!\left( \pi(a \mid s) \,\Big\|\, \tfrac{1}{Z}\exp\!\big(Q^{\pi_{\mathrm{old}}}/\alpha\big) \right) \;\Longrightarrow\; Q^{\pi_{\mathrm{new}}} \geq Q^{\pi_{\mathrm{old}}}$$

FLAG — zeroth-order realization

$$\begin{aligned} g_k(\mu) &:= \underbrace{\log \mathbb{E}_{\delta \sim \mathcal{N}(0,\, \Sigma_k)}\big[\exp\!\big( f_{\hat s, k}(\mu + \delta)/\lambda \big)\big]}_{\text{log-sum-exp smoothing}} \\[12pt] \underbrace{\mu_k^*(\hat s) - \mu_k}_{\text{FLAG M-step}} &\overset{(\text{Stein's identity})}{=} \Sigma_k \nabla_\mu g_k \;\approx\; \Sigma_k \underbrace{\nabla_a\!\big( Q^{\tilde\pi}(s, a) - \alpha \log \tilde\pi_{\theta_k}(a \mid s) \big)}_{\text{SAC action gradient}} \end{aligned}$$

• SAC’s soft policy improvement theorem is agnostic to how the KL minimizer is obtained.
• SAC restricts $\Pi$ to a Gaussian and takes a first-order reparameterization gradient; FLAG solves the same KL minimization with a non-parametric, action-level update — a zeroth-order approximation via Stein’s identity and log-sum-exp smoothing. Full statements and details are in Appendix C.2.

Proof Roadmap

Bottom Line

Monotonic improvement is established from two complementary perspectives — MPO (variational) and SAC (soft policy improvement) — and FLAG realizes SAC-style improvement for expressive flow policies without ever differentiating through the flow ODE.

Experimental Results

We present the results around the three questions used in Section 5 of the paper.

Q1 Scale to high-dim actions?

Q1 Scale to high-dim actions?

A1: Scalable & robust

Q1.1 Scalable?

Click to expand

Expand each banner Tap any result figure

Expand each banner to reveal results. Q1/Q2/Q3 cards and nested sub-banners (e.g., Q1.1) hide figures and tables until opened.

Click images to view detailed figure/table with caption. Cards show clean caption-free previews; the enlarged view shows the original caption from the paper.

Q1 Does FLAG scale to high-dimensional action spaces under limited sample budgets?

A1

FLAG is Scalable to High-dimensional Action Spaces and Robust to Sample Budget

Q1.1Is FLAG scalable to high-dimensional action spaces?

FLAG stays in the high-return, low-runtime region as action dimensionality increases, while global-proposal baselines degrade or require larger best-of-P budgets.

With CrossQ Without CrossQ

Q1.2Is FLAG robust to the sample budget?

Local proposal matching keeps importance samples informative even when the update uses only a small number of samples.

Click to expand with caption

×

Learning curveLearning Curves

Without best-of-P With best-of-P (P=32)

Q2 How does FLAG compare to action-gradient and BPTT-based actor-critic methods?

A2

FLAG Outperforms BPTT-based Actor-Critic and Action Gradient Methods

With CrossQ Without CrossQ

BibTeX citation

Displaying the BibTeX entry for your paper in a code block makes it easy to copy and paste.

@misc{kim2025flag,
  title         = {{FLAG}: Flow Policy MaxEnt-RL by Latent Augmented Guidance},
  author        = {Kim, Sungha and Lee, Gawon and Lee, Jusuk and Park, Jonghae and Kim, H. Jin and Cho, Daesol},
  year          = {2026},
  eprint        = {2605.00000},
  archivePrefix = {arXiv},
  primaryClass  = {cs.RO},
  url           = {https://arxiv.org/abs/2506.00000},
}