Source-linked AI summary

FLAC: Maximum Entropy RL via Kinetic Energy Regularized Bridge Matching

Lei Lv, Yunfei Li, Yu Luo, Fuchun Sun, Xiao Ma

arXiv:2602.12829v1cs.LGcs.AI

TL;DR

Iterative generative policies complicate Maximum-Entropy RL because their action log-densities are not directly accessible. FLAC addresses this by casting policy optimization as a Generalized Schrödinger Bridge and regulating stochasticity with kinetic energy relative to a high-entropy reference. It achieves competitive or superior performance on high-dimensional benchmarks while avoiding explicit density estimation.

  • Problem

    Iterative generative policies make Maximum-Entropy RL difficult because their action log-densities are not directly accessible for regulating stochasticity.

  • Method

    FLAC formulates policy optimization as a Generalized Schrödinger Bridge and penalizes the kinetic energy of its velocity field relative to a high-entropy reference.

  • Results

    FLAC achieves competitive or superior performance against strong baselines on high-dimensional continuous-control benchmarks including DMControl and HumanoidBench.

  • Takeaways & Limitations

    Kinetic energy provides a likelihood-free proxy for divergence from the reference while avoiding explicit action-density estimation.

  • Takeaways & Limitations

    The current framework applies isotropic regularization across action dimensions, treating distinct actuators uniformly.

Abstract

from arXiv · show

Iterative generative policies, such as diffusion models and flow matching, offer superior expressivity for continuous control but complicate Maximum Entropy Reinforcement Learning because their action log-densities are not directly accessible. To address this, we propose Field Least-Energy Actor-Critic (FLAC), a likelihood-free framework that regulates policy stochasticity by penalizing the kinetic energy of the velocity field. Our key insight is to formulate policy optimization as a Generalized Schrödinger Bridge (GSB) problem relative to a high-entropy reference process (e.g., uniform). Under this view, the maximum-entropy principle emerges naturally as staying close to a high-entropy reference while optimizing return, without requiring explicit action densities. In this framework, kinetic energy serves as a physically grounded proxy for divergence from the reference: minimizing path-space energy bounds the deviation of the induced terminal action distribution. Building on this view, we derive an energy-regularized policy iteration scheme and a practical off-policy algorithm that automatically tunes the kinetic energy via a Lagrangian dual mechanism. Empirically, FLAC achieves superior or comparable performance on high-dimensional benchmarks relative to strong baselines, while avoiding explicit density estimation.

1 Introduction

Iterative generative policies offer expressive modeling for continuous control but make Maximum-Entropy RL difficult because their action log-densities are inaccessible. FLAC reframes entropy regularization through GSB and kinetic energy, achieving competitive or superior benchmark performance.

  • Motivation: Iterative generative policies transport simple base noise into complex, state-conditioned action distributions, supporting rich multimodal behaviors.Their expressiveness is intended for high-dimensional continuous control.
  • Motivation: Maximum-Entropy RL is difficult to apply because iterative generators do not directly expose the policy log-density needed to regulate stochasticity.Existing solutions add auxiliary estimation or tractable distributional proxies, increasing complexity and computation.
  • FLAC: FLAC formulates entropy-regularized policy optimization as a Generalized Schrödinger Bridge relative to a high-entropy reference process.This makes maximum entropy a consequence of balancing terminal utility with path-space closeness to the reference.
  • FLAC: FLAC optimizes Q-values while minimizing velocity-field kinetic energy and automatically tunes the energy penalty to adapt exploration during training.The framework uses kinetic energy to preserve generation stochasticity without explicit terminal entropy estimation.
  • Empirical evaluation: FLAC achieves competitive or superior performance against state-of-the-art baselines on challenging DMControl and HumanoidBench continuous-control benchmarks.The reported evaluation covers high-dimensional control settings.

2 Related Work

Prior work develops flexible iterative policies for offline RL, imitation learning, and online actor-critic learning, while entropy regulation remains difficult when their action densities are implicit. Schrödinger Bridge methods provide the path-space foundation for FLAC's formulation.

  • Iterative Generative Policies: Diffusion and flow policies provide flexible behavior models or policy classes in offline RL and imitation learning, where mode coverage is central.Related methods also guide generators with values or learned energies while maintaining data support.
  • Iterative Generative Policies: Online RL has begun combining iterative policies with actor-critic updates and efficiency-oriented designs.This line of work addresses deployment of diffusion- and flow-based policies in online control.
  • Entropy Regulators for Generative Policies: Maximum-entropy or KL regularization supports exploration, but implicit action densities can make likelihood-based regulation expensive or fragile under limited solver budgets.ODE change-of-variables evaluation and SDE path marginalization are both nontrivial in practice.
  • Schrödinger Bridges: Schrödinger Bridges formulate stochastic evolution relative to a reference process, connecting entropy regularization, stochastic control, and optimal transport.Generalized Schrödinger Bridges relax hard terminal constraints into soft potentials or functional constraints suitable for utility-based decision settings.

3 Preliminaries

The preliminaries define continuous-control RL and iterative policies whose terminal action densities are difficult to evaluate. They then motivate path-space kinetic energy as a likelihood-free regularizer that controls proximity to a reference action distribution.

  • 3.1 Entropy-Regularized RL: Continuous-control RL learns a policy over continuous states and actions to maximize discounted expected return, often adding entropy to prevent premature convergence and encourage exploration.Maximum-entropy RL uses policy stochasticity as an explicit exploration objective.
  • 3.1 Entropy-Regularized RL: Maximum entropy can be expressed through H(π) = −E_a∼π[log π(a | s)] and equivalently as minimizing KL divergence from the uniform action distribution.The resulting optimal policy has Boltzmann form π*(a | s) ∝ exp(Q(s, a)/α).
  • 3.2 Iterative Generative Policies: Iterative generative policies define actions through a state-conditioned SDE or transport process rather than directly outputting action samples.The latent process starts from a simple prior, and its learnable drift or velocity field produces the terminal action.
  • 3.2 Iterative Generative Policies: The terminal action density is not directly accessible, making log π(a | s) expensive and numerically unstable to evaluate during online training.This motivates likelihood-free regulation of stochasticity.
  • 3.3 The Schrödinger Bridge Problem: The Schrödinger Bridge seeks the most likely stochastic evolution between probability distributions relative to a reference process.The generalized version relaxes hard terminal matching into a soft potential or functional constraint.
  • 3.4 Kinetic Energy and Path Constraint: Kinetic energy lifts policy regularization from action space to path space and acts as a proxy for divergence from the reference measure.This proxy applies across stochastic and deterministic generation regimes.
  • 3.4 Kinetic Energy and Path Constraint: Path-space divergence upper-bounds divergence between the terminal policy distribution and the reference terminal marginal, making kinetic-energy minimization sufficient for the terminal constraint.The result connects a computable path objective to control of the induced action distribution.
  • 3.4 Kinetic Energy and Path Constraint: When the prior is uniform on a bounded action domain, minimizing kinetic energy discourages concentrated policies and encourages broadly supported stochastic actions.In the deterministic limit, this does not provide a strict entropy guarantee.

4 Reinforcement Learning as a Generalized Schr¨odinger Bridge Problem

The paper reframes policy optimization as a one-ended Generalized Schrödinger Bridge relative to a high-entropy reference, replacing explicit terminal-density control with path-space regularization. Kinetic-energy minimization balances reward-seeking transport with preservation of stochasticity and multimodal actions.

  • 4.1 The Generalized Schr¨odinger Bridge Formulation: FLAC formulates reinforcement learning as finding a trajectory distribution close to a high-entropy reference while optimizing a terminal reward potential.The formulation uses a fixed reference path measure, such as Brownian motion, initialized from a high-entropy prior.
  • 4.1 The Generalized Schr¨odinger Bridge Formulation: The one-ended bridge leaves the terminal distribution free and regulates it through a soft potential rather than imposing a hard target marginal.This differs from classical Schrödinger Bridges, which impose a terminal-marginal constraint.
  • 4.1 The Generalized Schr¨odinger Bridge Formulation: The optimal terminal marginal has an exponential-tilting form relative to the reference marginal, yielding a Boltzmann-style policy when the potential is value-informed.For an approximately uniform reference, p∗(X1) ∝ exp(−G(X1)/α).
  • 4.2 Energy-Regularized Policy Optimization: FLAC replaces the abstract divergence term with expected kinetic energy and combines it with negative discounted return in an energy-regularized actor objective.The kinetic term is a dynamic regularizer: the zero-drift Brownian reference encourages stochasticity while allowing effort toward high-value regions.
  • 4.2 Energy-Regularized Policy Optimization: In a 2D multi-goal toy task, unconstrained high-velocity transport collapses to one mode, whereas FLAC preserves stochasticity and recovers all 8 optimal modes.The result illustrates how limiting kinetic energy prevents premature elimination of diverse solutions.

5 Field Least-Energy Actor-Critic

FLAC turns the bridge formulation into an energy-regularized actor-critic method with policy-iteration guarantees and automatic control of the kinetic-energy budget. The algorithm optimizes value while keeping the generation process near the stochastic reference.

  • 5 Field Least-Energy Actor-Critic: FLAC optimizes a velocity field that transports prior noise toward high-reward regions while minimizing kinetic energy.The practical algorithm instantiates the GSB framework as an off-policy actor-critic method.
  • 5.1 Energy-Regularized Policy Iteration: The energy-regularized Bellman operator incorporates the kinetic cost of action generation into policy evaluation and improvement.This extends standard policy-iteration reasoning analogously to entropy-regularized soft Bellman backups.
  • 5.1 Energy-Regularized Policy Iteration: Under bounded rewards and finite energy, the evaluation operator is a γ-contraction, so iterative evaluation converges to a unique regularized value function.The result is stated as Proposition 2.
  • 5.1 Energy-Regularized Policy Iteration: The policy-improvement update guarantees monotonic improvement of the generalized objective toward a reward-and-exploration transport plan.This is stated as Proposition 3: JGSB(πnew) ≥ JGSB(π).
  • 5.3 Automatic Energy Tuning: Because fixed energy penalties can cause over-exploration or premature determinism, FLAC learns the penalty through a target energy budget and Lagrangian dual updates.The multiplier is parameterized as α = exp(log α), remaining positive during optimization.
  • 5.3 Automatic Energy Tuning: The adaptive regulator decreases α when stochasticity is sufficient and increases it when the policy becomes too deterministic.This allows aggressive high-reward trajectories when appropriate while pulling the policy toward the high-entropy prior when needed.

6 Experiment

Across DMControl and HumanoidBench, FLAC matches or exceeds strong model-free baselines and reaches comparable asymptotic returns to TD-MPC2 while remaining model-free. Ablations show robustness across energy budgets and benefits from adaptive tuning.

  • 6.1 Main Results: FLAC is evaluated on high-dimensional locomotion and human-like robot-control tasks from DMControl and HumanoidBench against model-free and model-based baselines.The comparison includes deterministic, Gaussian, diffusion/flow, and TD-MPC2 methods.
  • 6.1 Main Results: FLAC consistently matches or exceeds strong model-free baselines, including in DMC Dog with s ∈ R223 and a ∈ R38 and the contact-rich Unitree H1 task.It also attains comparable asymptotic returns to TD-MPC2 without world-model learning or online planning.
  • 6.1 Main Results: FLAC achieves superior or comparable asymptotic performance to diffusion- and flow-based baselines using N = 2 function evaluations per action.The cited comparison contrasts this with DIME using N = 16.
  • 6.2 Ablation Studies: FLAC maintains high performance across a broad range of target energy budgets, but tight budgets C ∈ {0, 0.1} cause significant degradation.At C = 0, the velocity field is suppressed and the policy becomes fully random.
  • 6.2 Ablation Studies: Automatic tuning consistently outperforms fixed regularization settings, whose failures arise from restrictive priors or insufficient regularization.The learned log α typically decreases early and increases later, shifting from value maximization toward tighter exploration control.

7 Conclusions

FLAC maps reinforcement learning to the Generalized Schrödinger Bridge framework, using kinetic energy to regularize policies toward a high-entropy reference without explicit density estimation. The paper establishes theoretical connections across stochastic and deterministic regimes, while noting that deterministic energy control does not guarantee high terminal entropy.

  • FLAC maps reinforcement learning onto the Generalized Schrödinger Bridge framework and interprets maximum entropy through kinetic-energy minimization.The framework treats kinetic energy as a computable geometric proxy for deviation from a reference process.
  • Under standard regularity, boundedness, and absolute-continuity assumptions, the analysis relates policy-path divergence and kinetic energy.The stochastic analysis assumes σ > 0, while the deterministic analysis considers the ODE limit.
  • In the ODE regime, kinetic energy bounds Wasserstein-2 distance from the reference through the transport cost of the learned flow.The learned flow is one feasible transport plan, so its energy upper-bounds the optimal transport cost.
  • When the reference prior is uniform, minimizing energy is related to maximum-entropy regularization by discouraging concentrated action distributions.The paper presents this as a geometric surrogate that encourages broadly supported stochastic action sampling.
  • In the deterministic limit, proximity to a uniform prior in Wasserstein-2 does not generally imply large terminal differential entropy.The paper therefore characterizes the energy/Wasserstein constraint as a practical geometric inductive bias rather than a strict entropy guarantee.
  • The path-space argument uses data processing to show that minimizing kinetic energy keeps the terminal policy close to the high-entropy reference.For a uniform reference, this terminal constraint is related to the maximum-entropy objective.

A.6 Proof of Proposition 3

The proof establishes monotone policy improvement for the energy-regularized Bellman operator by combining a policy improvement condition with contraction and monotonicity. The operator has a unique fixed point, and iterative application converges to the updated policy’s value function.

  • Fixed-point convergence: By the fixed-point theorem, each fixed policy has a unique regularized value function Qπ.The proof identifies this fixed point with the regularized return by unrolling the Bellman equation.
  • Policy improvement: The updated policy satisfies a statewise policy improvement condition involving expected Q-values and kinetic-energy penalties.The condition is used to compare the updated and current Bellman backups.
  • Policy improvement: The updated Bellman backup is at least as large as the current-policy backup under that condition.The comparison follows by applying the improvement condition at the next state and taking expectations.
  • Monotone improvement: Iterating the updated operator from the current value function converges to Qπnew, which proves monotonic improvement.Contraction gives convergence, while monotonicity preserves the pointwise ordering throughout the iteration.

D Toy Example Setup

The toy example studies least-action regularization in a two-dimensional multi-goal bandit with eight uniformly placed goals. Without regularization, the flow collapses to one mode, whereas FLAC maintains bounded energy and covers all goals.

  • Toy Example Setup: The toy environment is a 2D multi-goal bandit with eight goal positions uniformly placed on a circle of radius 4.The action space is R2 and the reward is defined by the maximum Gaussian bump over the goals.
  • Toy Example Setup: Both policies use a 2-layer MLP drift field, a Gaussian base distribution, and K = 24 Euler steps.The unregularized and regularized policies share this setup for the comparison.
  • Toy Example Results: 8/8 coverage: FLAC maintains bounded energy through dual ascent and discovers all eight goals.The result is presented as evidence that least-action regularization prevents mode collapse in this toy setting.
  • Adaptive Kinetic Budgeting: FLAC’s adaptive kinetic budgeting draws on SAC’s target-entropy heuristic and sets the target from the transport cost of traversing the action space.The analogy regulates the trade-off between control effort and stochasticity in the energy domain.

E.1 Geometric Derivation

FLAC derives an action-dimension-scaled kinetic-energy budget and uses adaptive Lagrangian tuning to regulate the generative policy. Experiments indicate that C=0.5 and NFE=2 provide robust performance without per-task tuning or sacrificing final quality.

  • Geometric Derivation: The minimum energy for shifting probability mass from the prior center to a boundary action is 0.5 units per dimension.This follows from the constant-velocity least-action trajectory.
  • Geometric Derivation: The total energy required to reach the action-space boundary in all dimensions scales as 0.5 × d.
  • The Energy Budget Formula: The target energy budget is set linearly as Etgt = C · dim(A), where C denotes allowable kinetic energy per action dimension.
  • The Energy Budget Formula: C=0.5 yields robust performance across tasks, while adaptive α reduces sensitivity to the precise energy-budget choice.Experiments found robust performance for C ∈[0.5, 2.5], with C=0.5 used without per-task tuning.
  • Sensitivity to NFE: NFE=2 is used throughout, with higher NFE accelerating early convergence but having little impact on final performance.This supports efficient training with a small fixed number of function evaluations.
  • Computational Efficiency: FLAC remains computationally efficient relative to DIME across seven challenging DMC-Hard tasks despite using a PyTorch implementation with torch.compile.The comparison uses wall-clock time as its horizontal axis.

F.3 Comprehensive Results

Across comprehensive evaluations, FLAC matches or exceeds baselines on most HumanoidBench tasks and across DMC-Hard tasks. It nevertheless underperforms a strong model-based baseline on a small HumanoidBench subset.

  • Comprehensive Results: Figures 7 and 8 report the complete HumanoidBench and DMC-Hard results, respectively.
  • HumanoidBench: On HumanoidBench, FLAC matches or outperforms all baselines on most tasks but underperforms a strong model-based baseline on a small subset.
  • DMC-Hard: On DMC-Hard, FLAC matches or outperforms all baselines across tasks.
Loading 2602.12829v1…