Authors: Maojiang Su, Qijie Zhu, Yu Xiao, Shuyang Yu, Minshuo Chen, Zhaoran Wang, Han Liu
Organizations: Department of Computer Science, Northwestern University · Center for Foundation Models and Generative AI, Northwestern University · Department of Statistics and Data Science, Northwestern University · Department of Industrial Engineering and Management Sciences, Northwestern University
Diffusion and flow models provide expressive policy classes for robotic and sequential decision-making systems. However, their action likelihoods are generally intractable, limiting the direct use of widely adopted policy optimization methods such as PPO and GRPO. We propose Doob Policy Optimization (DoobPO), a framework that leverages Doob's h-transform for likelihood-free reinforcement learning fine-tuning of generative policies. DoobPO separates policy improvement from generative policy learning. It first learns a lightweight ratio model that represents the desired policy update, and then distills the resulting ratio-tilted target policy into the generative policy via conditional drift matching. This formulation applies uniformly to both diffusion and flow policies and supports a broad family of policy improvement rules, including PPO and GRPO. We prove that the practical conditional drift matching objective is gradient-equivalent to the ideal Doob drift matching objective. We further show that the induced diffusion process solves an endpoint-constrained stochastic optimal control problem. This establishes that it achieves the minimum-energy modification required to realize the desired policy tilt. Across MuJoCo continuous-control benchmarks, DoobPO achieves strong performance against competitive diffusion- and flow-policy baselines, showing a principled and practical route to RL fine-tuning for generative policies.
Figures & tables
Method
Generative policy
Policy update method
Ratio
DPMD ( Ma et al., 2025 )
Diffusion
Policy Mirror Descent (KL-proximal)
Analytic
SDAC ( Ma et al., 2025 )
Diffusion
Soft Actor Critic
Analytic
DoobPMD (ours)
Flow
Policy Mirror Descent (KL-proximal)
Analytic
DoobPPO (ours)
Diffusion / Flow
Proximal Policy Optimization
Learned
DoobGRPO (ours)
Diffusion / Flow
Group Relative Policy Optimization
Learned
Table 1: Doob Policy Optimization
Figure 1: Learning curves on 10 MuJoCo continuous-control tasks, evaluated every 200k environment steps. Curves show mean undiscounted return over 3 training seeds, with shaded regions indicating one standard deviation. We compare DoobPO against strong diffusion-policy baselines DPPO, SDAC, DIPO, and DPMD.
Figure 2: Learning curves on 10 MuJoCo continuous-control tasks, evaluated every 250k environment steps. Curves show mean undiscounted return over 3 training seeds, with shaded regions indicating one standard deviation. We compare DoobPO variants against flow-policy baselines FPO, ReinFlow, and Flow-SDE.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Algorithm
GPU Memory (MB)
Training Time (min)
DIPO
5212
15.96
DPPO
1087
12.71
DPMD
1175
25.29
DACER
1312
29.12
SDAC
1206
23.52
DoobPPO
1578
25.28
Appendix
Table 4: GPU memory allocation and total compute time of 200K iterations and 1 million environment interactions.
Name
Value
Critic learning rate
3e-4
Policy learning rate
3e-4, linear annealing to 3e-5
Diffusion steps
20
Diffusion noise schedules
Linear
Policy network hidden layers
3
Policy network hidden neurons
256
Appendix
Table 5: Hyperparameters
Name
Value
Discount factor γ
0.999
Actor learning rate
2×10−5
Ratio-network learning rate
1×10−4
Reward scale
2.5
PPO clipping coefficient
0.2
Maximum ratio weight
2
Appendix
Table 6: DoobPPO hyperparameters on Robomimic ( Can and Square ).
Figure 3: Ablation study of DoobPPO on OpenAI Gym MuJoCo HalfCheetah-v4 . (a) Effect of the ratio network architecture on policy performance. (b) Different diffusion steps lead to significantly different performance, with moderate steps performing best.
Figure 4: Ablation study of DoobGRPO flow policy on OpenAI Gym MuJoCo HalfCheetah-v4 . (a) Only the lightweight MLP 3 × 256 ratio network learns effectively, while deeper and larger architectures fail to train. (b) Performance is robust to the number of flow steps, with T=20 achieving the best result and T=10 , T=30 remaining competitive.
Generative policies provide expressive and multimodal action distributions, making them attractive for reinforcement learning (RL) in complex continuous-control tasks. Among them, flow-based policies are especially appealing because they generate actions through deterministic transport maps. However, applying such generative policies to likelihood-based on-policy learning remains limited by the difficulty of evaluating the probability of executed actions. Existing flow RL methods either replace the true action-density ratio with approximate surrogates, which can introduce biased updates, or recover exact likelihoods through dummy-action augmentation, which enlarges the policy space and increases computation. In this work, we propose GenPO++, a reversible generative policy optimization framework that uses history states as auxiliary memory in a high-order reversible ODE solver, yielding exact inversion without changing the original action dimension. The resulting generative policy map has a log-determinant determined only by fixed solver coefficients, enabling exact and Jacobian-free likelihood-ratio computation. This design preserves the expressiveness of generative flow policies while avoiding both action ratio bias and dummy-action overhead. We evaluate GenPO++ on large-scale simulated control, fine-tuning, and real-world robotic manipulation tasks, where it achieves competitive or superior performance over state-of-the-art on-policy RL methods, while improving training stability and computational efficiency.
Policy representations in reinforcement learning have expanded beyond Gaussian policies to diffusion and flow-matching based generative policies. However, existing RL methods remain representation-specific, lacking a unified optimization interface for emerging generative policies. In particular, extending likelihood-based objectives to multi-step generative policies couples optimization to their internal generation process, where accumulated likelihood ratios can weaken the update signal. Meanwhile, generative-policy-specific methods often rely on multi-candidate generation and critic-based selection, improving overall return while providing limited improvement to the policy itself. We introduce DR-APO (Decoupled-Representation Advantage-Guided Policy Optimization), which decouples policy improvement from policy representation through advantage-guided target improvement and representation-compatible direct action fitting, providing a shared optimization principle for Gaussian, diffusion, and flow policies. Across five MuJoCo-v5 tasks, DR-APO remains competitive with SAC for Gaussian policies and achieves the best aggregate policy-only performance for both diffusion and flow-matching policies. NavDP post-training improves the Success Rate (SR) from 70.33% to 89.55% in simulation and from 25% to 65% in real-world navigation. These results demonstrate a unified policy-improvement principle across Gaussian, diffusion, and flow policies, from continuous-control learning to embodied-policy post-training.
Runhua Zhang, Yingjie Niu, Chao Li
DEEP Robotics, Hangzhou, China · Zhejiang University, Hangzhou, China
Diffusion models have recently emerged as expressive policy representations for online reinforcement learning (RL). However, their iterative generative processes introduce substantial training and inference overhead. To overcome this limitation, we propose to represent policies using MeanFlow models, a class of few-step flow-based generative models, to improve training and inference efficiency over diffusion-based RL approaches. To promote exploration, we optimize MeanFlow policies under the maximum entropy RL framework via soft policy iteration, and address two key challenges specific to MeanFlow policies: action likelihood evaluation and soft policy improvement. Experiments on MuJoCo, DeepMind Control Suite and HumanoidBench benchmarks demonstrate that our method, Mean Flow Policy Optimization (MFPO), achieves performance comparable to or exceeding current diffusion-based baselines while considerably reducing training and inference time. Our code is available at https://github.com/dongxiaoyi-xyz/MFPO.
Xiaoyi Dong, Xi Sheryl Zhang, Jian Cheng
C2DL, Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences · AiRiA +1