PolicyFlow: Policy Optimization with Continuous Normalizing Flow in Reinforcement Learning

TL;DR

PolicyFlow is a novel on-policy CNF-based reinforcement learning algorithm that integrates expressive CNF policies with PPO-style objectives without requiring likelihood evaluation along the full flow path, and the Brownian Regularizer is proposed, an implicit policy entropy regularizer inspired by Brownian motion.

Abstract

Among on-policy reinforcement learning algorithms, Proximal Policy Optimization (PPO) demonstrates is widely favored for its simplicity, numerical stability, and strong empirical performance. Standard PPO relies on surrogate objectives defined via importance ratios, which require evaluating policy likelihood that is typically straightforward when the policy is modeled as a Gaussian distribution. However, extending PPO to more expressive, high-capacity policy models such as continuous normalizing flows (CNFs), also known as flow-matching models, is challenging because likelihood evaluation along the full flow trajectory is computationally expensive and often numerically unstable. To resolve this issue, we propose PolicyFlow, a novel on-policy CNF-based reinforcement learning algorithm that integrates expressive CNF policies with PPO-style objectives without requiring likelihood evaluation along the full flow path. PolicyFlow approximates importance ratios using velocity field variations along a simple interpolation path, reducing computational overhead without compromising training stability. To further prevent mode collapse and further encourage diverse behaviors, we propose the Brownian Regularizer, an implicit policy entropy regularizer inspired by Brownian motion, which is conceptually elegant and computationally lightweight. Experiments on diverse tasks across various environments including MultiGoal, PointMaze, IsaacLab and MuJoCo Playground show that PolicyFlow achieves competitive or superior performance compared to PPO using Gaussian policies and flow-based baselines including FPO and DPPO. Notably, results on MultiGoal highlight PolicyFlow's ability to capture richer multimodal action distributions.

Figures (8)

• Figure 1: (PointMaze-Medium-Diverse-GDense-v3) Exploration Density Maps. (a) Environment overview: the agent is initialized at the green point for each episode, and the four red points indicate goal locations with equal rewards. (b) Exploration heatmap of PPO, showing limited coverage due to the simple Gaussian policy. (c) Exploration heatmap of PolicyFlow without the Brownian regularizer, which improves coverage but still leaves some regions under-explored. (d) Exploration heatmap of PolicyFlow with the Brownian regularizer, achieving near-complete coverage of all feasible locations.
• Figure 2: MultiGoal Test (Appendix \ref{['app: multi-goal']}): sample 1000 trajectories starting at the same original point. (a) PPO with Gaussian entropy regularization ($w_g=0.001$) covers only a limited set of goals. (b,c) DPPO and FPO collapse to a small number of modes, likely because neither method incorporates any form of entropy regularization. (d) PolicyFlow with uniform noise injection ding2024diffusion (weight 0.05) still suffers from mode collapse, concentrating on only a few modes. (e) PolicyFlow with only Gaussian entropy regularization ($w_g=0.001$) partially alleviates mode collapse. (f) PolicyFlow with the proposed Brownian regularizer ($w_b=0.25$) and Gaussian entropy regularization ($w_g=0.001$) achieves the most diverse and more balanced goal-reaching behaviors.
• Figure 3: Learning curves on MuJoCo Playground benchmarks. Plots show mean episodic reward with standard error (y-axis) over environment steps (x-axis, total 30M steps), averaged over 5 random seeds.
• Figure 4: Ablation studies on key components of PolicyFlow.
• Figure 5: Terminal training episodic rewards using different interpolation paths.