Forward KL Regularized Preference Optimization for Aligning Diffusion Policies
Zhao Shan, Chenyou Fan, Shuang Qiu, Jiyuan Shi, Chenjia Bai
Abstract
Diffusion models have achieved remarkable success in sequential decision-making by leveraging the highly expressive model capabilities in policy learning. A central problem for learning diffusion policies is to align the policy output with human intents in various tasks. To achieve this, previous methods conduct return-conditioned policy generation or Reinforcement Learning (RL)-based policy optimization, while they both rely on pre-defined reward functions. In this work, we propose a novel framework, Forward KL regularized Preference optimization for aligning Diffusion policies, to align the diffusion policy with preferences directly. We first train a diffusion policy from the offline dataset without considering the preference, and then align the policy to the preference data via direct preference optimization. During the alignment phase, we formulate direct preference learning in a diffusion policy, where the forward KL regularization is employed in preference optimization to avoid generating out-of-distribution actions. We conduct extensive experiments for MetaWorld manipulation and D4RL tasks. The results show our method exhibits superior alignment with preferences and outperforms previous state-of-the-art algorithms.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 063ea9aa-3974-4306-bfad-e8befd2fa89fCited by top-tier papers7
- Radiology Report Generation via Multi-objective Preference OptimizationTing Xiao, Lei Shi, Peng Liu, Zhe Wang et al.AAAI 2025 · 21 citations
- Regularized Conditional Diffusion Model for Multi-Task Preference AlignmentXudong Yu, Chenjia Bai, Haoran He, Changhong Wang et al.NeurIPS 2024 · 11 citations
- Uncertainty-aware Preference Alignment for Diffusion PoliciesRunqing Miao, Sheng Xu, Runyi Zhao, Wai Kin (Victor) Chan et al.NeurIPS 2025 · 2 citations
- -DPO: Robust Preference Alignment for Diffusion Models via DivergenceYang Li, Songlin Yang, Wei Wang, Xiaoxuan Han et al.ICLR 2026
- From Noise to Control: Parameterized Diffusion PoliciesRenhao Zhang, Haotian Fu, Mingxi Jia, George Konidaris et al.ICML 2026
Builds on33
- Denoising Diffusion Implicit ModelsJiaming Song, Chenlin Meng, Stefano ErmonICLR 2021 · 11,743 citations
- Direct Preference Optimization: Your Language Model is Secretly a Reward ModelRafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning et al.NeurIPS 2023 · 10,924 citations
- Improved Denoising Diffusion Probabilistic ModelsAlexander Quinn Nichol, Prafulla DhariwalICML 2021 · 5,234 citations
- DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 StepsCheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen et al.NeurIPS 2022 · 2,653 citations
- Planning with Diffusion for Flexible Behavior SynthesisMichael Janner, Yilun Du, Joshua B. Tenenbaum, Sergey LevineICML 2022 · 1,115 citations
Related papers
- Diffusion Blend: Inference-Time Multi-Preference Alignment for Diffusion ModelsMin Cheng, Fatemeh Doudi, Dileep Kalathil, Mohammad Ghavamzadeh et al.ICLR 2026 · 6 citations
- Aligning Diffusion Behaviors with Q-functions for Efficient Continuous ControlHuayu Chen, Kaiwen Zheng, Hang Su, Jun ZhuNeurIPS 2024 · 13 citations
- Reinforcement Learning for Fine-tuning Text-to-Image Diffusion ModelsYing Fan, Olivia Watkins, Yuqing Du, Hao Liu et al.NeurIPS 2023 · 372 citations
- Behavior-Regularized Diffusion Policy Optimization for Offline Reinforcement LearningChen-Xiao Gao, Chenyang Wu, Mingjun Cao, Chenjun Xiao et al.ICML 2025
- Rethinking DPO-Style Diffusion Aligning FrameworksXun Wu, Shaohan Huang, Lingjie Jiang, Furu WeiICCV 2025 · 4 citations
