Normalizing Flows Policy in Reinforcement Learning
- The paper defines normalizing flows policy as a series of invertible transformations applied to a base distribution, ensuring exact likelihood evaluation and reparameterization.
- It enables multimodal, non-Gaussian action modeling that enhances exploration and performance in high-dimensional and sparse reward environments.
- The integration into RL frameworks like SAC, TRPO, and PPO demonstrates architectural substitution with challenges in numerical stability addressed through regularization techniques.
Searching arXiv for recent and foundational papers on normalizing-flow policies in reinforcement learning. A normalizing flows policy is a stochastic policy in which the action distribution is defined by an invertible transformation of a simple source distribution through a sequence of expressive, learnable mappings, rather than by a factorized Gaussian. In reinforcement learning, this policy class is used to retain tractable or exact likelihood evaluation, differentiable sampling, and compatibility with objectives based on entropy, KL constraints, or maximum likelihood, while enlarging the policy family to represent multimodal, skewed, constrained, or otherwise non-Gaussian action distributions (Ward et al., 2019). Across Soft Actor-Critic, Trust Region Policy Optimization, Proximal Policy Optimization, constrained control, dexterous manipulation, and visuomotor policy learning, normalizing flows have been studied as a replacement for the standard Gaussian actor, as a mechanism for enforcing action constraints, and as a way to obtain expressive policies with exact likelihoods and fast sampling (Tang et al., 2018).
1. Definition and mathematical structure
A normalizing flow policy starts from a base random variable and maps it into action space through a composition of invertible transformations. In the trust-region setting, one starts with noise and applies a chain of invertible transformations,
$x = g_{\theta_K} \circ g_{\theta_{K-1} \circ \cdots \circ g_{\theta_1}(\epsilon),$
so that the resulting density is tractable through the change-of-variables formula,
where is the inverse Jacobian matrix of (Tang et al., 2018). In control, the flow is commonly conditioned on state by inserting a state embedding into the transformation, yielding a generative map of the form
with an explicit, tractable density (Tang et al., 2018).
Within Soft Actor-Critic, the standard squashed Gaussian actor
is generalized by inserting a sequence of invertible transformations before the final ,
If $x = g_{\theta_K} \circ g_{\theta_{K-1} \circ \cdots \circ g_{\theta_1}(\epsilon),$0, then the transformed density is
$x = g_{\theta_K} \circ g_{\theta_{K-1} \circ \cdots \circ g_{\theta_1}(\epsilon),$1
and, after the final $x = g_{\theta_K} \circ g_{\theta_{K-1} \circ \cdots \circ g_{\theta_1}(\epsilon),$2 squashing, the policy density becomes
$x = g_{\theta_K} \circ g_{\theta_{K-1} \circ \cdots \circ g_{\theta_1}(\epsilon),$3
This preserves reparameterization and keeps the entropy term usable in maximum-entropy RL (Ward et al., 2019).
Coupling-based parameterizations are common. RealNVP-style layers use transformations such as
$x = g_{\theta_K} \circ g_{\theta_{K-1} \circ \cdots \circ g_{\theta_1}(\epsilon),$4
with triangular Jacobian and cheap determinant,
$x = g_{\theta_K} \circ g_{\theta_{K-1} \circ \cdots \circ g_{\theta_1}(\epsilon),$5
Neural Spline Flows replace affine transforms with monotone piecewise-rational or piecewise-cubic maps, while continuous normalizing flows define policies through ODEs driven by a velocity field (Tang et al., 2018).
2. Expressivity, multimodality, and exploration
The principal motivation for normalizing flows policies is that the diagonal Gaussian used in many RL algorithms is chosen mostly because it is easy to reparameterize, not because it is a strong model of the optimal action distribution (Ward et al., 2019). A Gaussian policy is unimodal, symmetric, and only flexible through mean and variance; by contrast, a flow can warp a simple Gaussian into a multimodal or highly non-uniform distribution that can assign nontrivial probability to multiple useful action neighborhoods (Mazoure et al., 2019).
In Trust Region Policy Optimization, the KL constraint controls divergence between distributions but does not force a new flow policy to remain geometrically close to the old policy’s center in action space. The core claim is that, under the same KL ball, a normalizing flow policy can generate samples far from the “center” of the previous policy iterate, potentially enabling better exploration and helping avoid bad local optima (Tang et al., 2018). The paper illustrates this with correlated-action and bimodal bandit examples, and reports that the gains are especially strong on high-dimensional, complex-dynamics tasks such as Humanoid and Sim. Humanoid (Tang et al., 2018).
In Soft Actor-Critic, the exploration argument is tied to entropy regularization. Entropy encourages stochasticity, but a Gaussian still tends to explore locally around its mean and may struggle to represent disconnected or irregular action regions that lead to rewards. A normalizing flow can preserve the maximum-entropy RL objective while broadening the policy family. In sparse-reward continuous grid world tasks, the flow policy is reported to match Gaussian SAC on dense rewards while being substantially better at sparse exploration, often reaching the goal earlier than all other methods (Ward et al., 2019). Related off-policy work reports that richer, potentially multimodal policies can discover better local behaviors early and avoid premature collapse to suboptimal solutions, with SAC-NF improving performance on MuJoCo and PyBullet Roboschool domains and doing so with as few as 5.5% the parameters for an equivalent SAC model (Mazoure et al., 2019).
A recurring misconception is that a maximum-entropy objective alone is sufficient to guarantee effective exploration. The flow-policy literature argues instead that policy family matters: entropy over a restrictive family need not produce the same exploratory behavior as entropy over a richer family (Ward et al., 2019). This suggests that exploration quality depends not only on the scalar entropy term, but also on the geometry of the policy class.
3. Integration into reinforcement learning algorithms
A notable feature of normalizing flows policies is that they are usually integrated by architectural substitution rather than by a complete redesign of the RL algorithm. In SAC, the flow policy simply replaces the Gaussian actor in the same soft actor objective, and the actor is optimized to maximize expected Q-value plus entropy using action samples and $x = g_{\theta_K} \circ g_{\theta_{K-1} \circ \cdots \circ g_{\theta_1}(\epsilon),$6 from the flow distribution rather than from a diagonal Gaussian (Ward et al., 2019).
In TRPO and ACKTR, the same pattern appears. Because these methods need only the policy’s log density $x = g_{\theta_K} \circ g_{\theta_{K-1} \circ \cdots \circ g_{\theta_1}(\epsilon),$7 and its gradient, a normalizing flow policy can be plugged into the same algorithmic machinery used for Gaussian policies. The algorithm computes $x = g_{\theta_K} \circ g_{\theta_{K-1} \circ \cdots \circ g_{\theta_1}(\epsilon),$8 using the change-of-variables formula; entropy and KL terms are estimated from samples when needed; and the rest of TRPO remains the same, including conjugate gradient and line search to satisfy the KL constraint (Tang et al., 2018).
More recent on-policy work studies continuous normalizing flow policies inside PPO-style updates. PolicyFlow defines a conditional CNF via an ODE,
$x = g_{\theta_K} \circ g_{\theta_{K-1} \circ \cdots \circ g_{\theta_1}(\epsilon),$9
with terminal action
0
To avoid evaluating likelihoods along the full flow trajectory during PPO updates, PolicyFlow approximates importance ratios using velocity field variations along a simple interpolation path, and adds the Brownian Regularizer as an implicit policy entropy regularizer inspired by Brownian motion (Yang et al., 1 Feb 2026). NFPO instead keeps the PPO clipped surrogate but uses a RealNVP actor with stabilized scale output, computing exact log-probability by the flow’s density transformation (Shi et al., 12 Mar 2026).
Normalizing flows policies have also been used outside standard actor parameterization. FlowPG studies action-constrained policy optimization by inserting a pretrained conditional RealNVP after a DDPG-style latent actor, so that the executed action is
1
and actor gradients propagate through the learned flow without differentiating through a projection solver (Brahmanage et al., 2024). SOFT-FLOW uses a conditional normalizing-flow actor over action chunks in off-policy dexterous fine-tuning, pairing it with an action-chunked critic and exact action-chunk likelihoods (Yang et al., 10 Feb 2026). In visuomotor imitation, NF-P formulates policy learning as conditional density estimation over action sequences, using exact likelihood optimization rather than iterative denoising (Lind et al., 25 Sep 2025).
4. Stability, regularization, and architectural limits
The main practical difficulty in training normalizing flows policies is instability. Flows are more numerically fragile than simple Gaussians, and several papers identify stabilization as a central requirement rather than an implementation detail. In flow-based SAC, reported stabilization tricks include a reparameterized Gaussian base distribution for the state encoder, weight clipping inspired by Wasserstein GANs, and removal of BatchNorm from the flow layers because in that RL setting the batch log-standard-deviation can overflow (Ward et al., 2019).
A related analysis appears in NFPO. Naïvely replacing PPO’s Gaussian policy with a RealNVP actor can cause the Jacobian determinant to grow very large and training to crash. The paper attributes this to overfitting due to expressiveness, the exponential scaling in RealNVP, and unbounded neural outputs. The proposed remedy is to normalize the RealNVP scale output with a 2-bounded parameterization,
3
using 4 in the final configuration (Shi et al., 12 Mar 2026). The paper argues that this is smoother than clipping and more stable across tasks (Shi et al., 12 Mar 2026).
Theoretical work also places limits on what flow architectures can express. A Bayesian-network reinterpretation shows that coupling and autoregressive layers encode explicit conditional independence assumptions, and that stacking multiple transformations relaxes these assumptions and entangles the model distribution (Wehenkel et al., 2020). The same paper states that a fundamental leap of capacity emerges when the depth of affine flows exceeds 3 transformation layers, but also proves the non-universality of the affine normalizing flow, regardless of depth (Wehenkel et al., 2020). For policy modeling, this means that shallow affine coupling policies may underfit complex correlated action distributions, while deeper stacks or richer scalar transforms are required for more difficult control distributions.
A complementary theoretical perspective views normalizing flows as local covariance learners, with
5
From this viewpoint, the singular values of the Jacobian determine how exploratory noise is stretched or compressed around each state, and Tikhonov regularization on the Jacobian Frobenius norm,
6
is motivated as a way to prevent pathological drive toward infinite Jacobian singular values when intrinsic dimension is smaller than ambient dimension (Feinman et al., 2019). Although this analysis is not framed as reinforcement learning, it is directly relevant to the numerical stability and geometric interpretation of flow policies.
5. Constraints, safety, geometry, and interpretability
Normalizing flows policies have also been used to encode structure beyond multimodality. In action-constrained RL, FlowPG learns an invertible, differentiable mapping between the feasible action space and the support of a simple latent distribution, using a conditional RealNVP and valid-action samples obtained through Hamiltonian Monte Carlo or probabilistic sentential decision diagrams (Brahmanage et al., 2024). The key practical claim is that a well-trained flow transforms policy output into a valid action without requiring an optimization solver, reducing constraint violations by up to an order of magnitude in Reacher and Half-Cheetah and being 2–3x faster than the Frank-Wolfe/QP-based baseline NFWPO (Brahmanage et al., 2024).
A different line of work constructs constrained normalizing flow policies analytically from known instantaneous constraints. The feasible action region in state 7 is written as
8
and the policy is built as a sequence of invertible transformations, each mapping actions into the subset that satisfies one specific constraint (Rietz et al., 2024). In the reported 2D navigation task, obstacle avoidance and battery constraints are handled by a rectangle squashing function and a circle squashing function, respectively, with priority order
9
Because each flow step corresponds to a specific constraint, the policy is described as interpretable and safe-by-construction, and it maintains quasi-perfect constraint satisfaction during the entire training process (Rietz et al., 2024).
Geometry can be incorporated more intrinsically when action spaces are curved or constrained manifolds. Riemannian continuous normalizing flows define the policy distribution directly on a smooth manifold by solving an ODE on the manifold,
0
with density evolution
1
The divergence depends on the Riemannian metric, and the model stays on the manifold through a manifold-aware ODE solver and tangent-space-valued vector field (Mathieu et al., 2020). This is relevant to policy learning whenever actions live on spheres, tori, hyperbolic spaces, or other smooth constrained domains.
Interpretability can also be imposed through restricted transformations. Restricted normalizing flow policies constrain the base distribution to be symmetric and the transform to be odd so that the transformed mean is analytically computable,
2
Because that restriction impairs expressiveness, Bit-RNF restores asymmetry and heavy tails with a bimodal student-t base distribution. In RL benchmarks and a real robot ball-plate task, Bit-RNF is reported to outperform previous models while allowing deterministic deployment through the analytic mean (Kobayashi et al., 2024).
6. Contemporary developments and comparative position
Recent work extends normalizing flows policies into domains where Gaussian policies are too simple and diffusion policies are expensive or awkward for likelihood-based optimization. SOFT-FLOW addresses real-world dexterous fine-tuning by modeling action chunks
3
using a conditional RealNVP-style actor with exact action-chunk likelihoods and a critic that scores whole chunks rather than single actions (Yang et al., 10 Feb 2026). The policy is trained with a conservative offline/online RL objective that combines Q maximization with imitation regularization through
4
The paper reports stable, sample-efficient adaptation on cutting tape with scissors retrieved from a case and in-hand cube rotation with a palm-down grasp, and states that this is the first demonstration of a likelihood-based, multimodal generative policy combined with chunk-level value learning on real robotic hardware (Yang et al., 10 Feb 2026).
In visuomotor policy learning, NF-P replaces diffusion policies with conditional normalizing flows over action sequences. The observation combines image features from a pre-trained ResNet18 and a vector representing the latest action; the model uses Neural Spline Flows with 10 coupling layers, alternating halves of the input vector, and a 3-layer feed-forward network of hidden size 2048 for spline parameterization (Lind et al., 25 Sep 2025). The policy can assign a probability score 5 to any candidate action sequence, which the paper interprets as a statistically grounded density estimate for confidence. On four simulated robotic tasks, NF-P is reported to be competitive with, and often better than, Diffusion Policy, with up to 30 times faster inference (Lind et al., 25 Sep 2025).
Broader RL evidence argues that normalizing flows are capable models not only for policies but also for Q-functions and occupancy measures. A single NF architecture based on repeated RealNVP-style coupling blocks and a linear flow inspired by Glow is used as a policy in imitation learning, as a policy in offline RL, and as a conditional density estimator for future-goal occupancy (Ghugare et al., 29 May 2025). The paper’s thesis is that NFs are unusually well matched to RL’s core algorithmic needs—exact likelihoods, fast sampling, and differentiable sampling/evaluation for variational updates—and that their expressivity has been underestimated (Ghugare et al., 29 May 2025).
At the same time, the literature does not present a single dominant verdict. Some works emphasize exact likelihood and one-pass sampling as decisive advantages over diffusion models (Lind et al., 25 Sep 2025), whereas CNF-based PolicyFlow keeps a diffusion- or flow-matching-style policy but avoids full likelihood evaluation along the flow path through a velocity-field-based approximation (Yang et al., 1 Feb 2026). This suggests an emerging division of design priorities rather than a settled consensus: exact density and computational simplicity on one side, or continuous-flow flexibility with approximate likelihood machinery on the other. Across these variants, the persistent theme is that normalizing flows policies occupy a middle position between restrictive Gaussian actors and more computationally intensive generative policies, combining expressive action modeling with tractable or approximately tractable policy optimization (Mazoure et al., 2019).