Reparameterization Flow Policy Optimization

Reparameterization Policy Gradient (RPG) has emerged as a powerful paradigm for model-based reinforcement learning, enabling high sample efficiency by backpropagating gradients through differentiable dynamics. However, prior RPG approaches have been predominantly restricted to Gaussian policies, limiting their performance and failing to leverage recent advances in generative models. In this work, we identify that flow policies, which generate actions via differentiable ODE integration, naturally align with the RPG framework, a connection not established in prior work. Yet, naively realizing this synergy is non-trivial: it often suffers from training instability and a lack of exploration. We propose Reparameterization Flow Policy Optimization (RFO). RFO computes policy gradients by backpropagating jointly through the flow generation process and system dynamics, unlocking high sample efficiency without requiring intractable log-likelihood calculations. RFO includes two tailored regularization terms for stability and exploration. We also propose a variant of RFO with action chunking. Extensive experiments on diverse locomotion and manipulation tasks—involving both rigid and soft bodies with state or visual inputs—demonstrate the effectiveness of RFO. Notably, on a challenging locomotion task controlling a soft-body quadruped, RFO achieves almost $2\times$ the reward of the state-of-the-art baseline.

Reparameterization Policy Gradient (RPG) achieves high sample efficiency by backpropagating gradients through differentiable dynamics, but it mostly relies on Gaussian policies. We observe that flow policies, which generate actions via differentiable ODE integration, are a natural structural match for RPG. We propose Reparameterization Flow Policy Optimization (RFO), which backpropagates jointly through the flow generation process and the environment dynamics, unlocking the sample efficiency of RPG for flow policies without ever computing intractable log-likelihoods. Two lightweight regularizers handle stability and exploration. Across diverse rigid- and soft-body locomotion and manipulation tasks, RFO consistently beats strong baselines, nearly doubling the reward of the state-of-the-art baseline on a challenging soft-body quadruped task.