Behavior Regularization with Flow Latent Policy for Offline Reinforcement Learning

Abstract

Expressive generative models have recently shown promise in offline reinforcement learning (RL) by capturing the complex, multimodal structure of dataset behavior. However, directly integrating these models into policy optimization introduces substantial computational and stability challenges due to the intricacies of their sampling processes. We introduce Flow Latent Policy (FLP), an offline RL framework that decouples expressivity from optimization by operating entirely in the latent space of a pre-trained, frozen flow-based behavior model. FLP learns a simple latent Gaussian policy whose samples are transformed through the flow to produce complex, behavior-aligned actions. This design enables closed-form behavior regularization via latent-space KL divergence and allows policy optimization without expensive backpropagation through the generative model. Experiments on the OGBench benchmark demonstrate that FLP achieves competitive or superior performance across diverse tasks, combining the benefits of expressive modeling and tractable optimization.