Doob-Based Generative Policy Optimization

Maojiang Su; Qijie Zhu; Yu Xiao; Shuyang Yu; Minshuo Chen; Zhaoran Wang; Han Liu

cs.LGOct 20, 2026

Doob-Based Generative Policy Optimization

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

Abstract

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 $\gamma$	0.999
Actor learning rate	$2\times 10^{-5}$
Ratio-network learning rate	$1\times 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 ).

Panel of Figure 3 — 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.

Panel of Figure 4 — Figure 4: Ablation study of DoobGRPO flow policy on OpenAI Gym MuJoCo HalfCheetah-v4 . (a) Only the lightweight MLP 3 $\times$ 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.

\displaystyle\beta_{k+1}\leftarrow\mathop{\mathrm{argmax}}_{\beta}\Big\{J_{\mathrm{RL}}(\beta)-\lambda_{\mathrm{reg}}\E_{o\sim\mathcal{D}}\Big[\big(\E_{a\sim\pi^{k}(\cdot|o)}[r_{\beta}(o,a)]-1\big)^{2}\Big]\Big\}.

Appendix

Algorithm 1 Diffusion Policy Optimization with DoobPO

\displaystyle\beta_{k+1}\leftarrow\mathop{\mathrm{argmax}}_{\beta}\Big\{J_{\mathrm{RL}}(\beta)-\lambda_{\mathrm{reg}}\E_{o\sim\mathcal{D}}\Big[\big(\E_{a\sim\pi^{k}(\cdot|o)}[r_{\beta}(o,a)]-1\big)^{2}\Big]\Big\}.

Appendix

Algorithm 2 Flow Policy Optimization with DoobPO

Policy Optimization

Policy optimization develops algorithms for improving decision rules in reinforcement learning, language models, and sequential decision systems.

See all 221 papers

GenPO++: Generative Policy Optimization with Jacobian-free Likelihood Ratios
Jun 5, 2026Ke Hu, Shutong Ding, Panxin Tao +2RL Control Policy Gradient Methods
ShanghaiTech University
DR-APO: Decoupled-Representation Advantage-Guided Policy Optimization from Continuous Control to Embodied Navigation
Oct 12, 2026Runhua Zhang, Yingjie Niu, Chao LiPolicy Optimization RL for Generative Models
DEEP Robotics, Hangzhou, China · Zhejiang University, Hangzhou, China
Mean Flow Policy Optimization
Apr 16, 2026Xiaoyi Dong, Xi Sheryl Zhang, Jian ChengMeanFlow Mean Flow
C2DL, Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences · AiRiA +1

Jun 5, 2026J/K move · Enter open · S save

GenPO++: Generative Policy Optimization with Jacobian-free Likelihood Ratios

Ke Hu, Shutong Ding, Panxin Tao, Jingya Wang+1

ShanghaiTech University

RL Control Policy Gradient Methods Flow-Based Generative Modeling RL for Generative Models

Open paper AskPDF arXiv

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.

Abstract

Policy Optimization

Explore similar work

GenPO++: Generative Policy Optimization with Jacobian-free Likelihood Ratios

DR-APO: Decoupled-Representation Advantage-Guided Policy Optimization from Continuous Control to Embodied Navigation

Mean Flow Policy Optimization