Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Forward KL Regularized Preference Optimization for Aligning Diffusion Policies

Zhao Shan, Chenyou Fan, Shuang Qiu, Jiyuan Shi, Chenjia Bai

TL;DR

Forward KL Regularized Preference Optimization for Aligning Diffusion Policies (FKPD) tackles the challenge of aligning diffusion-based decision policies with human preferences without relying on explicit rewards. It introduces a two-stage approach: first learning a rich, preference-free diffusion policy from offline data, then aligning it to preferences via Direct Preference Optimization (DPO) with forward KL regularization to avoid out-of-distribution actions. The method derives a tractable alignment objective that blends a Bradley–Terry-based DPO term with a forward KL constraint to a behavior policy, and provides practical approximations using the forward diffusion process and Denoising Mean Squared Error (D-MSE). Experiments on MetaWorld and D4RL show FKPD achieving superior alignment and competitive task performance, highlighting the advantages of forward KL regularization for diffusion-policy alignment and its potential for reward-free preference learning in complex robotics tasks.

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.

Forward KL Regularized Preference Optimization for Aligning Diffusion Policies

TL;DR

Forward KL Regularized Preference Optimization for Aligning Diffusion Policies (FKPD) tackles the challenge of aligning diffusion-based decision policies with human preferences without relying on explicit rewards. It introduces a two-stage approach: first learning a rich, preference-free diffusion policy from offline data, then aligning it to preferences via Direct Preference Optimization (DPO) with forward KL regularization to avoid out-of-distribution actions. The method derives a tractable alignment objective that blends a Bradley–Terry-based DPO term with a forward KL constraint to a behavior policy, and provides practical approximations using the forward diffusion process and Denoising Mean Squared Error (D-MSE). Experiments on MetaWorld and D4RL show FKPD achieving superior alignment and competitive task performance, highlighting the advantages of forward KL regularization for diffusion-policy alignment and its potential for reward-free preference learning in complex robotics tasks.

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.
Paper Structure (16 sections, 18 equations, 3 figures, 3 tables)

This paper contains 16 sections, 18 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Preference-based Reinforcement Learning
  4. Diffusion Model for Decision-Making
  5. Method
  6. DPO with Forward KL Regularization
  7. Practical Approximation for Diffusion Policy
  8. Advantage over Reverse KL Regularization
  9. Related Work
  10. Diffusion Model for Decision Making.
  11. RLHF and Preference-based RL.
  12. Experiments
  13. Performance of FKPD
  14. Ablation of Regularization
  15. Conclusions
  16. ...and 1 more sections

Figures (3)

  • Figure 1: Overall scheme of FKPD: FKPD comprises two phases: the behavior cloning phase and the alignment phase. In the behavior clone phase, FKPD fits a diffusion policy to a preference-free dataset $\mathcal{D}$. Subsequently, in the alignment phase, FKPD employs a direct preference optimization method to align the diffusion policy with a preference dataset $\mathcal{D}_{\rm pref}$. Throughout this process, the diffusion policy is required to maintain a forward KL distance with respect to the distribution of $\mathcal{D}_{\rm pref}$
  • Figure 2: Samples generated by a toy diffusion model. The left sub-figure displays samples generated by the initial model. The middle and right sub-figures show samples generated by the aligned models with forward and reverse KL regularization, respectively.
  • Figure 3: Average D-MSE and implicit accuracy for FKPD, NKPD, and NRPD during the alignment phase of Drawer Open (bottom row) and Walker2d-medium-replay (top row). We also provide the same variables of the pretrained policy on the preference-free dataset $\mathcal{D}$ for reference.