FLAG: Flow Policy MaxEnt-RL by Latent Augmented Guidance
This work provides a significant advancement for researchers and practitioners in reinforcement learning by enabling more expressive and scalable policy optimization in high-dimensional control tasks, overcoming a key limitation of existing MaxEnt-RL approaches.
This paper addresses the scalability issues of Maximum Entropy Reinforcement Learning (MaxEnt-RL) with expressive generative policies in high-dimensional action spaces, which are often limited by importance weight collapse. The authors introduce FLAG, a method that augments the state space with a flow latent variable and optimizes a consistent proxy MaxEnt-RL objective, achieving state-of-the-art performance on challenging benchmarks.
Maximum entropy reinforcement learning (MaxEnt-RL) enables robust exploration, yet practical implementations often restrict policies to simple Gaussians. While recent approaches incorporate expressive generative policies via importance-weighted supervised learning, they are prone to importance weight collapse, which limits their scalability in high-dimensional action spaces. Our key insight is to mitigate this limitation by localizing the sampling region, avoiding the weight degeneracy induced by importance sampling over the entire action space. To instantiate this insight, we introduce \textbf{FLAG} (\textbf{F}low policy with \textbf{L}atent-\textbf{A}ugmented \textbf{G}uidance). FLAG augments the state space with a flow latent variable and optimizes a provably consistent proxy MaxEnt-RL objective. We empirically demonstrate that FLAG enables expressive policy optimization with limited importance samples and scales to high-dimensional control tasks. Furthermore, FLAG achieves state-of-the-art performance across challenging benchmarks. Our project webpage: https://flag-rl.github.io/