Reparameterization Flow Policy Optimization

TL;DR

This paper addresses the limited expressiveness and training challenges of applying Reparameterization Policy Gradient (RPG) to flow-based policies by introducing Reparameterization Flow Policy Optimization (RFO). RFO learns flow policies by backpropagating through both the flow generation process and the environment dynamics, and it adds two targeted regularizers—Past Data CFM Regularization and Uniform Exploration CFM—to stabilize training and promote exploration, plus an optional action-chunking variant for temporal consistency. Across a diverse suite of locomotion and manipulation tasks in differentiable simulators, including a high-dimensional soft-body quadruped with action space $\mathbb{R}^{222}$, RFO consistently outperforms RPG and diffusion/flow baselines, achieving up to nearly 2× the reward on Soft Jumper and strong gains on rigid-body tasks. The work demonstrates that integrating flow models with RPG can deliver high sample efficiency without requiring intractable log-likelihoods, enabling scalable, on-policy training in complex, differentiable simulators and setting the stage for offline-to-online extensions.

Abstract

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. However, naively exploiting this synergy proves ineffective, often suffering 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.

Figures (7)

• Figure 1: RFO optimizes the flow policy by jointly backpropagating through the dynamics and the flow policy.
• Figure 2: Comparison of RFO training with and without past data CFM regularization on the Soft Jumper task. We monitored the Conditional Flow Matching (CFM) loss of the current policy evaluated on actions sampled from the immediately preceding iteration, as well as the KL divergence between consecutive policy updates.
• Figure 3: Training curves across seven tasks. Solid lines represent the mean return, while shaded regions indicate the standard deviation. Curves are smoothed using a 100-episode moving average.
• Figure 4: Ablation study for the effectiveness of past data CFM regularization and uniform exploration CFM regularization. The results show that the two proposed regularization terms are critical for RFO's performance.
• Figure 5: (a) Ablation study on different weight combinations for past data CFM regularization and uniform exploration CFM regularization. (b) Ablation study on the number of flow integration steps.