Aligning Diffusion Behaviors with Q-functions for Efficient Continuous Control
Authors: Huayu Chen
Kaiwen Zheng
Hang Su
Jun Zhu
Venue: NeurIPS 2024 poster (2024)
Keywords: diffusion
offline
RL
alignment
continuous control
DPO
Primary area: robotics
TL;DR: We propose an alignment strategy for diffusion models in continuous control

Abstract: Drawing upon recent advances in language model alignment, we formulate offline Reinforcement Learning as a two-stage optimization problem: First pretraining expressive generative policies on reward-free behavior datasets, then finetuning these policies to align with task-specific annotations like Q-values. This strategy allows us to leverage abundant and diverse behavior data to enhance generalization and enable rapid adaptation to downstream tasks using minimal annotations. In particular, we introduce Efficient Diffusion Alignment (EDA) for solving continuous control problems. EDA utilizes diffusion models for behavior modeling. However, unlike previous approaches, we represent diffusion policies as the derivative of a scalar neural network with respect to action inputs. This representation is critical because it enables direct density calculation for diffusion models, making them compatible with existing LLM alignment theories. During policy fine-tuning, we extend preference-based alignment methods like Direct Preference Optimization (DPO) to align diffusion behaviors with continuous Q-functions. Our evaluation on the D4RL benchmark shows that EDA exceeds all baseline methods in overall performance. Notably, EDA maintains about 95\% of performance and still outperforms several baselines given only 1\% of Q-labelled data during fine-tuning.

Aligning Diffusion Behaviors with Q-functions for Efficient Continuous Control
Huayu Chen Kaiwen Zhengl Hang Su :2.3, Jun Zhu 1.233
Department of Computer Science and Technology, Tsinghua University
Institute for AI, BNRist Center, Singhua-Boseh Joint ML Center, THBI Lab, Tsinghua University Pazhoo Lab (Huangpu), Guangzhou, China
Abstract
Drawing upon recent advances in language model alignment, we formulate offline Reinforcement Learning two-stage optimization problem: First pretraining expressive generative policies on reward-fres behavior datasets, then fine-tuning these policies to align with task-specific annotations like Q-values. This strategy allows us to leverage abundant and diverse behavior data to enhance generalization and enable rapid adaptation downstream tasks using minimal annotations. In par- ticular, we introduce Efficient Diffusion Alignment (EDA) for solving continuous control problems. EDA utilizes diffusion models for behavior modeling. However, unlike previous approaches, we represent diffusion policies as the derivative of a scalar neural network with respect action inputs. This representation is critical because it enables direct density calculation for diffusion models, making them compatible with existing L.M alignment theories. During policy fine-tuning, we extend meference-based alignment methods like Direct Preference Optimization (DPO) to align diffusion behaviors with continuous Q-functions. Our evaluation the D4RL benchmark shows that EDA exceeds all baseline methods in overall performance. Notably, EDA maintains about 95% of performance and still outper- forms several baselines given only 1 1% of Q-labelled data during fine-tuning. Code: https: //git/gbitttt com/ thu-thuumit ml//fffffi /tfffffciientt... Difffusion Al illignment
1 Introduction
Learning diverse behaviors is generative modeling: transforming them into optimized policies is reinforcement learning. Recent studies have identified diffusion policies as powerful tool for representing heterogeneous behavior datasets [
38) However, these behavior policies incorporate suboptimal decisions in datasets, making them unsuitable for direct deployment downstream tasks. To get optimized policies, typical methods involve either augmenting the behavior policy with an or training a new evaluation policy supervised by the behavior policy 113 4 While functional, these methods fail to leverage the full potential of pretrained behaviors as they require constructing some new policy models from scratch. This raises the question: Can we directly fine-tune pretrained diffusion behaviors into optimized policies?
additional guidancc/cvallution network 21
32 [5]
Recent advances in Large Language Model (LLM) alignment techniques (57 43 offer valuable insights for fine-tuning diffusion behavior policies due to the fundamental similarity of the issues they aim to address (Fig. []. While pretrained LLMs accurately imitate language patterns from web-scale corpus, they also capture toxic or unwanted content within the dataset. Alignment algorithms, such as Direct Preference Optimization (DPO, 41), are designed remove harmful or useless content learned during pretraining. They enable quick adaptations of pretrained LLMs to human intentions
'The corresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Langaage
Help me solve this math
Enhanced Tluman Value Aligmment
Web-scale corpus
L.M Pretraining
Math is
:uure.
Preference Alignment
(oodinaaa Contirr
3.7 88
22
88
-0.1
Diffusion Modeling
Q-function Alignment
Optimizodd policy
Diverse behevior
State
Action
Qalue estimation
Figure Comparison between alignment strategies for LLMs and diffusion policies (ours).
by fine-tuning them on a small dataset annotated with human preference labels. These strategies, due their simplicity and effectiveness, have seen widespread applications in academia and industry.
Despite the high similarity in problem formulation and the immense potential of LLM alignment techniques, they cannot be directly applied to fine-tune diffusion policy in domains like continuous control. This is primarily because L.LMs employ Categorical models to deal with discrete actions (tokens). Their alignment relies computing data probabilities for maximum likelihood training (Sec. [S2c. However, diffusion models lack a tractable probability calculation method in continuous action space (5). Additionally, the data annotation method differs significantly between two areas: LM alignment uses binary preference labels for comparing responses, while continuous control uses scalar Q-functions for evaluating actions (Fig. D.
To allow aligning diffusion behavior models with O-functions for policy optimization, we introduce Efficient Diffusion Alignment (EDA). EDA consists of two stages: behavior pretraining and policy fine-tuning. During the pretraining stage, we learn a conditional diffusion behavior model on reward- free datasets. Different from previous work which constructs diffusion models network, we represent diffusion policies as the derivative of a scalar neural network with respect to action inputs. This representation is critical because it enables direct density calculation for diffusion policies. We demonstrate that the scalar network exactly outputs the unnormalized density of behavior distributions, making diffusion policies compatible with existing LLM alignment theories.
end-to-end
During the fine-tuning stage, we propose a novel algorithm that directly fine-tunes pretrained behavior models into optimized diffusion policies. The training objective is strictly derived by constructing classification task to predict the optimal action using log- bog-probability ratios between the policy and arbitrary
the behavior model. Our approach innovatively expands DPO by allowing fine-tuning number of actions annotated with explicit Q-values, beyond just the typical binary preference data.
One main advantage of EDA is that it enables fast and data-efficient adaptations of behavior models in downstream tasks. Our experiments on the D4RL benchmark [10] show that EDA maintains 95 % of its performance and still surpasses baselines like IQL [25] with just 1 1% of Q-labelled data relative to the pretraining phase. Besides, EDA exhibits rapid convergence during fine-tuning, requiring only about 20K gradient steps (about 2% of the typical IM policy training steps) to achieve convergence. Finally, EDA outperforms all reference baselines in overall performance with access to the full datasets. We attribute the high efficiency of EDA to its exploitation of the diffusion behavior models' generalization ability acquired during pretraining.
Our key contributions: 1. We represent diffusion policies as the derivative of a scalar value network to allow direct density estimation. This makes diffusion policies compatible with the existing alignment framework. 2. We extend preference-based alignment methods and propose EDA align diffusion behaviors with scalar Q-functions, showcasing its vast potential in continuous control.
2 Background
2.1 Offline Reinforcement Learning
Offline RL aims to tackle decision-making problems by solely utilizing a pre-collected behavior dataset. Consider typical Markov Decision Process (MDP) described by the tuple AS,A, A, 7) is the state space, A is the action space, P(s'|a ((s'|s,,0) a) is the transition function, r(s, a) is the reward function and y is the discount factor. Given a static dataset T == {s, a, s') representing
2

---

interaction history between an implicit policy and the MDP environment, our goal is to learn new policy that maximizes cumulative rewards in this MDP while staying close the behavior policy μa.
Offline RI. can be formalized as a constrained policy optimization problem 226 35 53
max Bi (1) (1 Q(s, a) 3DKL [
(1)
π
where (s,, a) is an action evaluation network that can learned from D 3 is a temperature coefficient, Previous work [40 39 proves that the optimal solution for solving Eq. 1 is:
(als) ((s) p ((als)e
In this paper, we focus on how to efficiently learn parameterized policies for modeling n²"
2.2 Direct Preference Optimization for Language Model Alignment
Direct Preference Optimization (DPO, 41 is fine-tuning technique for aligning pretrained LLMs with human feedback. Suppose we already have a pretrained LLM model {{(a(s) (), where s rep- resents instructions, and represents generated responses. The goal is to align (s with some implicit evaluation rewards p.MMss. (s, a) that reflect human preference. Our target model is (s.a)/3 (e.a""
x" (a(s)) & [*(aas)) 8) (eaaa)))"""
DPO assumes we only have access to some pairwise preference data {s - (aw > a:)} and the preference probability is influenced by LLM as) Formally,
(aM(a))
p(a_ > si == e (s.(a)))) (nlM() = a(pLM(M (s, aw) - (s,
(3)
where a is the sigmoid function.
In order to learn me x (a)s) x (als)e) (x.a)/3 (x.a)// DPO first parameterizes reward model using the log-probability ratio between my and pas and then optimizes this reward model through maximum likelihood training:
LDOO = - -B Jar log logo (s, aw) (s, a;)),
L. (s) = Blgg To
where
po (a
The key insight behind DPO's loss function is the equivalence and mutual convertibility between the policy model and the reward model. This offers a new perspective for solving generative policy optimization problems by applying discriminative classification loss.
2.3 Diffusion Modeling for Estimating Behavior Score Functions
excel representing heterogeneous behavior policies in continuous control [21 [5] 38] To train diffusion behavior model, we first gradually inject Gaussian noise into action points according the forward diffusion process:
Recent studies show that diffusion models
45 20
a, = a, + o,6,
(5)
where E [o, 1], and e standard Gaussian noise. ae o, E [0. 1] are manually defined so that at time 1= O, have a; a and time / policy ((a)s), the marginal distribution of a, at various time satisfies
have a, e. When is sampled from the behavior
(a / N(allaia of T t)da.
(6)
Intuitively, the diffusion training objective predicts noise added to the original behavior actions:
min (.e.ss.--p((aa)) [[cl(a))))t) c| a: tajatt"""
More formally, it can be proved that the learned "noise predictor" 6 actually represents the score function V log 14 (a s, t)) of the diffused behavior distribution f4 49
V log (a(az) s, t) = -e" (arls, t)/a, z (a, t)/o2.
(8)
With such a acore-function estimator, we can employ existing numerical solvers to reverse the diffusion process, and sample actions from the learned behavior policy pa.

---

(1) Behavior Pretraining
Diffusion loss
(2) Policy Fine-tuning
Cross entrayy loss
Inferencce
Back-propagttion
B Behavioral actions
fo (s,1s)
ab
Ma COE
(G,att
ST
of
Ala Col)
init
add coise
Sxutte Saer
Mip
of
Afe
a
State time
Bottleneck value
To (de1s,))
Q-modal
Figure 2:
Algorithm overview. Left: In behavior pretraining, the diffusion behavior model is represented as the derivative scalar neural network with respect to action inputs. The scalar outputs of the network can later be utilized to estimate behavior density, Right: In policy tie-ttunigg we predict the optimality of actions logit for each action is the density gap between the learned policy model and the frozen behavior model. We use cross-entropy loss to align prediction logits Afo := S0 with dataset Q-labels.
contrastive manner among K candidates. The prediction
3 Method
We decompose the policy optimization problem into two stages: behavior pretraining (Sec. 3.1 and policy alignment (Sec. 3.2
3.1
Bottleneck Diffusion Models for Efficient Behavior Density Estimation
Recent advances alignment techniques cannot be readily applied to continuous control tasks. Their successful applications in LLM fine-tuning require two essential prerequisites:
1. A powerful foundation model capable of capturing diverse behaviors within datasets.
2. A tractable probability calculation method that allows direct density estimation (Eq.
Language models primarily deal with discrete actions (tokens) defined by a vocabulary set V. and thus employ Categorical models. This modeling method enables casy calculation of data probability through a softmax operation and is capable of representing any discrete distribution. In contrast, for continuous action space, direct density estimation is not feasible. Diffusion policies only estimate the gradient field (a.k.a,, score) of data density instead of the density value itself [49], making impossible directly apply LLM alignment techniques 41 3]. Conventional Gaussian policies have a tractable probability formulation but lack enough expressivity and multimodality needed to accurately model behavior datasets 552 13 5. and catastrophically fail in initial experiments.
To address the above limitation diffusion models, we propose new diffusion modeling technique to enable direct density estimation. Normally, a conditional diffusion policy (a(alls, t) A x X - RIA maps noisy actions a, to predicted noises e R|AI In approach, redefine cø as the derivative of a scalar network fa t)
S x R R with respect to imput a,:
(s, == - oe V V. (arlss,)..
Given that fa a parameterized network, gradient computation can be conveniently performed by auto-differential libraries. The new training objective for fø then be reformulated from Eq. 7
min (o) = Er,e. [ + = a.=aeeto
(10)
As noted by 449 with unlimited model capacity, the optimal solution for solving Eq. 10
(a. 8, -a Va, log (a.
(arls,)) log (arls, C(s,t).
(11)
An illustration is provided in Figure 2 (left). Intuitively, proposed modeling method first compresses the imput action into scalar value with one single dimension. Then, this bottleneck value expanded back to R |A through back-propagation. We thus refer to fo as Boutleneck Diffusion
Models (BDMs).

---

Dataset
Preiniined f
fine-tuned fg
Qe:nf 5-
Dataset
Pretrained 9
fine-tuned fe
Q: 13-1
t= 0.0
t=
0.2
t= 1.0 t= 0.3
(a) Valuc-based optimization
(b) Preference- optimization
Q-value
Dispreferred
Preferre
Density
Figure 3: Experimental results of EDA in 2D bandit settings at different diffusion times. Column Visualization of diversified behavior datasets. Each dot represents two-dimensionall behavioral action. Its color reflects the action's 's Q-value. Column 2 & 3: Density maps of the action distribution as estimated by the pretrained or fine-tuned BDM models. The density for low- --w-Q-value actions has been effectively decreased after fine-tuning Column 4: The predicted action Q-values, calculated by Eq. align with dataset Q-values in Column See appendix for complete results.
The primary advantage of BDMs is their ability efficiently estimate behavior densities in single forward pass. Moreover, BDMs are fully compatible with existing diffusion-basee codebases regarding training and sampling procedures, inheriting their key benefits such training stability and model expressivity. BDMs can also be viewed as a diffused version of Energy- Based Models (EBMs, We refer interested readers Appendix A for detailed discussion.
(8).
3.2 Policy Optimization by Aligning Diffusion Behaviors with Q-functions
In this section, our goal to learn new policy my x use by fine-tuning previously pretrained behavior policy 14 on new dataset annotated by existing Q-function Q(s, a). We show that this policy optimization problem can actually be transformed into a simple classification task for predicting the optimal action among multiple candidates. We elaborate on method below.
Dataset construction. For any state 8, we draw K independent action samples !:K from He (-1s). Assume we already have an Q-function ((s. a) that evaluates input state-action pairs in scalar values, our dataset is formed as Dí = {s, a K Q(s, )
Action optimality. We first introduce formal notion of action optimality. We draw from the control- as-probabilistic-infeeene framework (29) and define a random optimality variable Ox. which is a one-hot vector of length K. The k-th index of Ox being 1 indicates that a* is the optimal action within K action candidates I:K We have
Qaa,a"
pOO als,al.h
ΣK eQ(
The optimality probability of a behavioral action a is proportional to the exponential of its Q-value, aligning with the optimal policy definition T" (als) x u(aa ((s.)/)
From policy optimization to action classification. We construct classification task by predicting the optimal action among K candidates. This requires learning Q-model termed Qe first. Drawing inspiration from DPO (Sec. 2.2 parameterize Qe the log probability ratio between Re and pş:
To (a(a)s)
a0 Bog
+ log 7(s),
Ho (als)
5

---

This parameterization allows us to directly optimize Te during training because To(a|s) = (aal)) (als))(0(...))/) constantly holds. Since Q-values in datasets define the probability of being the optimal action, the training objective can be derived by applying cross-entropy loss:
log Rala [+) tog Z(xe
K
e((aaa))
in
max C, (0) = ((s.aa))).
log
(12)
E elss,a)
@ log Meio Pala to
optimality probability
predicted prohability
The unknown partition function (s) automatically cancels out during division. is hyperparameter that can be tuned to control how far тy deviates from jAg.
Expanding to bottleneck diffusion behavior. A reliable behavior density estimation of p pa 8) is critical and is a main challenge for applying Eq. 12 Our initial experiments tested with Gaussian policies drastically failed and even underperformed vanilla behavior cloning. This highlights the necessity adopting a much more powerful generative policy, such as the BDM model (Sec. 3.1.
Diffusion policies define series of distributions Tz (a t) = (aa a, )oo t)dao To. Consequently,
different timesteps t E [0, I], rather than just single distribution instead of directly predicting the optimal action given raw actions a l:K we perturb all actions with K independent Gaussian noises according to the diffusion forward process by applying a
a,a o,e". Then predict action optimality given K noisy action a.
Qo (), = log (a, + 3 log ((s) ( 8,
(a, 1 (a.le s, () + Bloog ((8,)) (* (*,,
(13)
Similar to Eq. 12 all unknown terms above automatically cancel out in the training objective:
K
Q(s,a")
(aa((a))).))) (a
max (0) (
log
(14)
Lill ea(a.)
Li=1 (0) 12 (a)la.1)
optimality probability
predicted probability on noisy actions
Let få be the optimal solution of Problem 14d the eSt be the optimal diffusion policy. Assuming unlimited model capacity and data samples, we have the following results:
Proposition 3.1. (Proof in Appendix C
(a) Optimality Guarantee. At time l = 0, the learned policy T. converges to the optimal target policy.
(a {(ao,((ass) x pa(a 2(a.a)))
(b) Diffusion Consistency. At time L 0, T+20,0 models the diffused distribution of T
t / Narlara,, Tm (TT))(((a))))
(an))saan)))dl (aols)dam
asymptotically holds when K + oo and 3 = 1, satisfying the definition of diffusion process (Eq. 6
Fine-tuning Efficiency. In practice, the policy and the behavior model share the same architecture, so we can initialize 0 = ó to fully exploit the generalization capabilities acquired during the pretraining phase. This fine-tuning technique allows us to perform policy optimization with an incredibly small amount of Qllabelled data (e.g., 10k samples) and optimization steps (e.g., 20K gradient steps). These are less than 5% of previous approaches (Sec. (.2).
Related Work
4.1 Diffusion Modeling for Offline Continuous Control
Recent advancements in offline RI. have identified diffusion models as an effective approach for behavior modeling [38 [6), which excels at representing complex and multimodal distributions

---

Environment
Dataset
BCQ
CQL
IQL
DT
D-Diffuser
IDQL
Diffusion-QL.
QGPO
BDM (Ours)
HallCheetah
Medium Expert
64.7
91.6
86.7
86.8
90.6
95.9
96.8
93.5
93
1.2
HallChoetah
Mediam
40.7
44.0
47.4
12.6
19.1
51.0
51.1
64.1
57.0 0.5
HalfCheetah
Medium- Replay
38.2
45.5
44.2
36.6
39.3
45.9
47.8
47.6
51.6 1: 0.9
Hopper
Medium Expert
100.9
105.4
91.5
107.6
111.8
108.6
111.1
108.0
104.9 1: 7.4
Hopper
Medium
54.5
58.5
66.3
67.6
79.3
65.4
90.5
98.0
98.4 3.9
Hopper
Medium- Replay
85.0
82.7
100.0
92.1
92.7 10.0
Walker2d
Medium- Expert
57.5
108.8
109.6
108.1
108.8
112.7
110.1
110.7
11T.1
Medium
72.5
74.0
82.5
87
Walker2d
Medium Replay
15.0
77.2
73.9
66.6
75.0
85.1
95.5
84.4
89.2 J. 5.5
Average (D4RL. Locomotion)
50.9
77.6
76.9
74.7
81.8
88.0
86.0
AntMaze
Umaze
T3.0
74.0
50.2
96.4
+ 4.5
AntMaze
Umuze-Diverse
81.0
84.0
62.2
53.0
80.2
66.2
74.4
81.0 7.4
AntMaze
Mediam- Play
0.0
61.2
71.2
0.0
84.5
76.6
83.
79.0 + 4.2
Diverse
53.7
Average (D4RL Locomotion)
35.5
68.2
72.7
28.1
85.9
78.7
84.6
81.3
Kitchen
Complete
8.1
43.8
62.5
84.0
62.8
81 51.5-1 - 7.3
Kitchen Kitchea
Partial
18.9
19.8
46.3
37.0
60.5
66.0
69 3 :: 4.6
Mixed
51.0
51.0
65.0
6
45.5
65 + 2.2
Average (D4RL Locomotion)
11.7
48.2
69.0
Average (DRRI. Overall)
80.7
Table 1: Evaluation results of D4RL benchmarks (normalized according to (10). We report mean + standard deviation of algorithm performance across 5 random seeds the end of training. Numbers within 5 % of the maximum are highlighted.
compared with other modeling methods like Gaussians [54 55 However, optimizing diffusion models can be bit more tricky due to the unavailability of a tractable probability calculation method 55 [1]. Existing approaches include learning a separate guidance network to guide the sampling process during evaluation [2]
53 or VAES
32] applying classifier-free guidance 717, backpropagating sampled actions through the diffusion policy model to maximize Q-values 52 [22, distilling new policies from diffusion behaviors [13 4 using rejection sampling to filter out
behavioral actions with low Q-values [5] and applying planning-based techniques [30] 36 [6]. Our proposed method differs from all previous work in that it directly fine-tunes the pretrained behavior to align with specifcc annotations, instead of learning a new downstream policy.
4.2 Preference-basee Alignment in Offline Reinforcement Learning
Existing alignment algorithms are largely preference-batee based methods. Preference- basedd Reinforcement Learning (PbRL) assumes the reward function is unknown, and must be learned from data. Previous work usually applies inverse RL learn a reward model first and then
this reward model for
standard RL training [19] (7 28
34 This separate learning of a reward model adds complexity to algorithm implementation and thus limits its application. To address this issue, recent work like
44
OPPO (23) DPO [41), and CPL (17 respectively proposes methods to align Gaussian or Categorical policies directly with human preference. Despite their simplicity and effectiveness, these techniques require calculating model probability, and thus cannot be applied to diffusion policies. Existing diffusion alignment strategies 7 56 51 are incompatible with mainstream PbRL methods. Our work effectively closes this gap by introducing bottleneck diffusion models. We also extend existing PbRL methods to align with continuous Q-functions instead of just binary preference data.
5 Experiments
We conduct experiments to answer the following questions:
How does EDA perform compared with other baselines in standard benchmarks? (Sec. 5..1
Is the alignment stage data-efficient and training-efficient? How many training steps and
annotated data does EDA require for aligning pretrained behavior models? (Sec. 5.2
Does value-based alignment outperform preference-basee alignment? (Sec.
5.3
How do contrastive action number K and other choices affect the performance? (Sec. 5.4
7

---

5.1 D4RL Evaluation
Benchmark. In Table [] we evaluate the performance of EDA D4RL benchmarks T0 All evaluation tasks have continuous action space and can be broadly categorized into three types: MuJoCo locomotion on are tasks for controlling legged robots to move forward, where datasets are generated by variety of policies, including expert, medium, and mixed levels. Antmaze is about an ant robot navigating itself in a maze and requires both low-level control and high-level navigation. FrankaKi tchen are manipulation tasks containing real-world datasets. It is critical to faithfully imitate human behaviors in these tasks.
Experimental setup. Throughout our experiments, we set the contrastive action number K = 16. For each task, train an action evaluation model ( a) for annotating behavioral data during the fine-tuning stage. Implicit Q-learning [25] is employed for training Qs due to its simplicity and orthogonality policy training. We compare with other critic training methods in Figure The rest of the implementation details are in Appendix D
Baselines. We mainly consider diffusion-basse RL methods with various optimization techniques. Decision Diffuser [1 employs classifier-free guidance for optimizing behavior models. QCPO employs energy guidance. IDQL [15] does not optimize the behavior policy and simply uses rejection sampling during evaluation. Diffusion-Qu QL [52) has no explicit behavior model, but instead adopts diffusion regularization loss. We also reference well-studied conventional algorithms for different
classes of generative policies. BCQ | features a VAE-based policy, DT 6 has transformer-basse
architecture. CQL [27 and IQL 25 targets Gaussian policies.
Result analysis. From table I we find that EDA surpasses all referenced baselines regard- ing overall performance and provides a compet- itive number in each D4RL task. To ascertain whether this improvement stems from the align- ment algorithm rather than from superior critic model or policy class, we conduct additional controlled experiments. As outlined in Figure 4 we evaluate three variants of EDA using differ- ent Q-learning approaches and compare these against both diffusion and Gaussian baselines. The experimental results highlight the superior- ity of diffusion policies and further validate the effectiveness of our proposed method.
100
Loconotioo scores
Geussine
Diffasion haseline
90-
OGPO
IDA
TDS-BC
UBDA
IQL
70
Gazsiaz hawak
Online QL
-Solins QL
Implicit QL
Q-leaming method
Figure 4: Average performance of EDA combined with different Q-learning methods Locomotion tasks.
Lacomotion avenuge (9 tasks)
palichosal-modduu--olaaaass
hopee-mcdiim-oopettt2
Normaliand scare
80
Numaltad koaEs
50
Korzallodd scoces
100-
60
Ours
Ours
40
Ours
20
QGPO
QGPO
D-QL
IQL
JQL
JQL
100%
50%
20%
10%
25%
1%
500
1000
1560
0
S00
1000
1500
Fine tunning data numbes (% of pretriming)
Gradient (k)
Gradient steps
(a) Sample Efficiency
(b) Training Efficiency.
Figure 5: Aligning pretrained diffusion behaviors with task Q-functions is fast and data-efficient.
5.2 Fine-tuning Efficiency
The success of the pretraining, fine-tuning paradigm in language models is largely due to its high fine-tuning efficiency, allowing pretrained models to adapt quickly to various downstream tasks. Similarly, aim to explore data efficiency and training efficiency during EDA's fine-tuning phase.
To investigate EDA's data efficiency, we reduce the training data used for aligning with pretrained Q-functions by randomly excluding a portion of the available dataset (Figure (a)). We compare this with IQI., which uses the same Q-model as EDA but extracts a Gaussian policy via weighted regression. We also compare with QGPO, which shares our pretrained diffusion behavior models

---

but employs separate guidance network to augment the behavior model during evaluation, instead of directly fine-tuning the pretrained policy, Experimental results reveal that EDA is significantly more data-efficient than the QGPO and IQL baselines. Notably, EDA maintains about 95% of its performance in locomotion tasks when using only 1% of the Q-labeled data for policy fine-tuning. This even surpasses several baselines that use the full dataset for policy training.
In Figure (b), plot EDA's performance throughout the fine-tuning phase. EDA rapidly converges in roughly 20K gradient steps, a negligible count compared to the typical IM steps used for behavior pretraining. Note that behavior modeling and task task-oriented Q-definition are largely orthogonal in offline RL. The high fine-tuning efficiency of EDA demonstrates the vast potential of pretraining on large-scale diversified behavior data and quickly adapting individual downstream tasks.
5.3 Value Optimization Preference Optimization
significant difference between EDA and existing preference- RL methods such as DPO is that EDA is tailored for alignment with scalar Q- Q-values instead of just preference data.
For preference data without explicit Q-labels, we can similarly derive an alignment loss:
120) 315 2155 (0,
max C
(0) E ((aa
log K ( (ai
(15)
Here a represents the most preferred action among a I:K In practice, we select a as the action with the highest Q-value but abandon the absolute number. We'd d like to note that when K 2, the above objective becomes exactly the DPO objective (Eq. 4 in preference learning:
max coro (0) = (t,)( ",0"" log B[ f (ar = (a)ls, t) (a t)
(16)
We ablate different choices of K and compare our proposed value-based aligmment with prefer- ence methods in 9 D4RL Locomotion tasks (Fig- ure Results show that EDA generally outper- forms preference-based alignment approaches. We attribute this improvement to its ability to exploit the Q-value information provided by the pretrained Q-model. Besides, we notice the performance gap between the methods be- comes larger as increases. This expected because preference-ber based optimization greedily follows a single action with the highest Q-value. However, more action candidates are sampled from the behavior model, the final selected action will have a higher probability of being out-of- deitrior- data. This further hurts performance. In contrast, EDA is softer version of preference learning. This leads to greater tolerance for K.
Leconation scores
Preference optimization
90
Value-based optimization
80
DPO
K-2
K-4 Contrastive action numbers
K=128
Figure 6: Ablation of action numbers K and opti- mization methods.
5.4 Ablation Studies
We study the impact of varying temperature 3 algorithm performance in Appendix E We also illustratively compare our proposed diffusion models with other generative models for behavior modeling in 2D settings in Appendix A
6
Conclusion
We propose Efficient Diffusion Alignment (EDA) for solving offline continuous control tasks. EDA allows leveraging abundant and diverse behavior data to enhance generalization through behavior pretraining and enables rapid adaptation to downstream tasks using minimal annotations. Our experimental results show that EDA exceeds numerous baseline methods in D4RL tasks. It also demonstrates high sample efficiency during the fine-tuning stage. This indicates vast potential in learning from large-scale behavior datasets and efficiently adapting to individual downstream tasks.

---

Acknowledgments and Disclosure of Funding
We would like to thank Chengyang Ying and Cheng Lu for discussion. This work was supported by NSFC Projects (Nos. 62350080, 92248303, 92370124, 62276149, 62061 136001), BNRist (BNR2022RC01006), Tsinghua Institute for Guo Qiang, and the High Performance Computing Center, Tsinghua University. J. Zhu was also supported by the XPlorer Prize.
References
[1] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations, 2023.
[2] Kevin Black, Michael Janner, Yilun Du llya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv: 13301, 2023.
[3] Huayu Chen, Guande Hc, Hang Su, and Jun Zhu. Noisc contrastive alignment of language models with explicit rewards. arXiv preprint arXiv:2402.05336, 2024.
[4] Huayu Chen, Cheng Lu, Zhengyi Wang, Hang Su, and Jun Zhu. Score regularized policy optimization through diffusion behavior. arXiv preprint arXiv:2310.07297, 2023.
[5] Huayu Chen, Cheng Lu, Chengyang Ying. Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. In The Eleventh International Conference on Learning Representations, 2023.
[6] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems, 2021.
[7] Zibin Dong, Yifu Yuan, Jianye HAO, Fei Ni, Yao Mu, YAN ZHENG, Yujing Hu, Tangjie Lv, Changjic Fan, and Zhipeng Hu. Aligndifff: Aligning diverse human preferences via behavior-customisabbe diffusion model. In The Twelfth International Conference on Learning Representations, 2024.
[8] Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems, 32, 2019.
[9] Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela. Kto: Model alignment as prospect theoretic optimization. arXiv preprint arXiv:2402.01306, 2024.
[10] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.072199 2020.
[11] Scott Fujimoto, David Meger, and Doina Precup. Off- Off-policy deep reinforcement learning without exploration. In Proceedings of the 36th International Conference on Machine Learning. 2019.
[12] Seyed Kamyar Seyed Ghasemipour, Dale Schuurmans, and Shixiang Shane Gu. Emaq: Expected-max q-learning operator for simple yet effective offline and online rl. In International Conference Machine Learning, pages 3682-3691. PMLR, 2021.
[13] Wonjoon Goo and Scott Niekum. Know your boundaries: The necessity of explicit behavioral cloning in offline rl. arXiv preprint arXiv: 20.665, 2022.
[14] Ulf Grenander and Michael Miller. Representations of knowledge in complex systems. Journal of the Royal Statistical Society: Series B (Methodological), 56(4):549-582) 1994.
[15] Philippe Mansen-Estruuhh Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit g-learning as an actor-critic method with diffusion policies, arXiv preprint arXiv:2304. 10573, 2023.
[16] Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang. Weinan Zhang, Dong Wang, Bin Zhao, and Xuelong Li. Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning. Advances in neural information processing systems, 36, 2024.
[17] Joey Hejna and Dorsa Sadigh. Inverse preference learning: Preference-based without reward function. Advances in Neural Information P Processing Systems, 36, 2024.
10

---

[18] Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771-180,, 2002.
[19] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in neural information processing systems, volume 29, pages 4565-4573, 2016.
[20] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
[21] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022.
[22] Bingyi Kang. Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. arXiv preprint arXiv:2305.20000 2023.
[23] Yachen Kang. Diyuan Shi, Jinxin Liu, Li He, and Donglin Wang. Beyond reward: Offline preference-guided policy optimization. arXiv preprint arXiv:2305. 2023.
[24] Iya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofrr Nachum. Offline reinforcement learning with fisher divergence critic regularization. In International Conference Machine Learning, pages 5774-5783. PMLR, 2021.
[25] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit Q-learning. In International Conference on Learning Representations, 2022.
[26] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 2019.
[27] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 2020.
[28] Kimin Lee, Laura Smith, Anca Dragan, and Pieter Abbeel. B-pref: Benchmarking preference- based reinforcement learning. In Conference on Neural Information Processing Systems Datasets and Benchmarks Track (round 1), 2021.
[29] Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv: 1805.0099, 2018.
[30] Wenhao Li, Xiangfeng Wang, Bo Jin, and Hongyuan Zha. Hierarchical diffusion for offiine decision making. In Andreas Krause, Emma Brunskill, Kyunghyun Cho. Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research. PMLR. 23-29 Jul 2023.
[31] Zengyi Li, Yubei Chen, and Friedrich T Sommer. Learning energy-based models in high- dimensional spaces with multiscale denoising-score matching. Entropy, 25610) 2023.
[32] Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, and Jun Zhu. Contrastive energy prediction for exact energy-guided diffusion sampling in offline reinforcement learning. In Proceedings of the 40th International Conference on Machine Learning, 2023.
[33] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li. and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. arXiv preprins 0ttv:22006009227. 2022.
[34] Michaël Mathieu, Sherjil Ozair, Srivatsan Srinivasan, Caglar Gulcehre, Shangtong Zhang, Ray Jiang, Tom Le Paine, Richard Powell, Konrad Zoina, Julian Schrittwieser, et al. Alphastar unplugged: Large-scale offline reinforcement learning. arXiv preprint wxXiv:2308.03526 2023.
[35] Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09353 2020.
[36] Fei Jianye Hao, Yao Mu, Yifu Yuan, Yan Zheng, Bin Wang, and Zhixuan Liang. Metadiffuser: Diffusion model as conditional planner for offline meta-rl. In International Conference Machine Learning, pages 26087-26105, PMLR, 2023.
[37] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedhack. Advances Neural Information Processing Systems, 25:27730-27744, 2022.
11

---

[38] Tim Pearce, Tabish Rashid, Anssi Kanervisto, David Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, et al. Imitating human behaviour with diffusion models. In Deep Reinforcement Learning Workshop NeurIPS, 2022.
[39] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weiighted regresssson Simple and scalable off- reinforcement learning. arXiv preprint arXiv: 2019.
[40] Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. In Twenty- Fourth AAAI Conference Artificial Intelligence, 2010.
[41] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D) Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. In Thirty- Conference on Neural Information Processing Systems, 2023.
[42] Saced Saremi, Arash Mehrjou, Bernhard Schölkopf, and Aapo Hyvärinen. Deep energy estimator networks. arXiv preprint arXiv: wtw.:/88000080... 08306, 2018
[43] John Schulman, Barret Zoph, Christina Kim, Jacob Hilton, Jacob Menick, Jiayi Weng, Juan Felipe Ceron Uribe, Liam Fedus, Luke Metz, Michael Pokorny, et al. Chatgpt: Optimizing language models for dialogue. OpenAl blog, 2022.
[44] Daniel Shin and Daniel S Brown. Offline preference-based apprenticeship learning. arXiv preprint arXiv:2107.09251 2021.
[45] Jascha Dickseii,,,,,,, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2265,, 2015.
[46] Jammigg Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
[47] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[48] Yang Song and Diederik P Kingma. How to train your energy-based models. arXiv preprint arXiv:2101.0322 2021.
[49] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[50] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 2367):1661-1674, 2011.
[51] Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. arXiv preprint arXiv:2311. 12990,, 2023.
[52] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou Diffusion policies an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023.
[53] Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, and Nando de Freitas. Critic regularized regression. In Advances in Neural Information Processing Systems, 2020.
[54] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint 1rXii:1911.11360 2019.
[55] Haoran Xu, Xianyuan Zhan, Jianxiong Li, and Honglei Yin. Offline reinforcement learning with soft behavior regularization. arXiv preprint arXiv:2110.07333 2021.
[56] Zhilong Zhang, Yihao Sun, Junyin Ye, Tian-Shuo Liu, Jiaji Zhang, and Yang Yu. Flow to better: Offline preference-basee reinforcement learning via preferred trajectory generation. In The Twelfth International Conference on Learning Representations, 2024.
[57] Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. 2019.
arXiv preprint
12

---

A Comparing Bottleneck Diffusion Models with Energy-Based Models
Our proposed Bottleneck Diffusion Models (BDMs) can be viewed diffused variant Energy- Based Models (EBMs,
8 ). Both methods aim to model data distribution's unnormalized log
probability:
px(x) Ex(x)
Despite their conceptual similarity, the sampling and training approaches differ between diffusion models and EBMs. A prevalent sampling method for EBMs Langevin MCMC [4], which employs the energy gradient (x) to progressively transform Gaussian noise into data samples. Langevin MCMC generally necessitates hundreds or thousands of iterative steps to achieve convergence, significantly higher number compared with 15-50 steps required by diffusion models (Figure A)
Furthermore, the maximum-likelihond training EBMs (e.g., Contrastive Divergence 18] is more computationally expensive,
involves online data sampling from the model during the training process [48]. To avoid online data sampling, subsequent research [50, [47 has shifted away from directly modeling Ex(x) towards developing score-based models defined as sø(x) := Es
The score- matching objectives utilized in training these score- based models have subsequently been
adapted for training diffusion models [49]. Our work is inspired by prior work 42 47 that employed such score-matching objectives for training EBMs.
(a) Ground truth
(b) Caussians
(c) VAES
(d) EBMs (100 steps) (e) EBMs (1k steps)
(0 Ours (25 steps)
Figure 7: Comparison of various generative modeling methods in 2D modeling and sampling.
13

---

B
Complete 2D-bandit Experiment Results
Dataset
Pretrained f
fine-tuned fö
Qe:n -5
Dataset
Pretrained f
fine-tuned fg
Qa
0.0.0
0.0
t=
0.2
t=
== 0.3
1-00
f=
t= t= 0.15
1.10
0.0
0.3
t=
1==00 1.0
(a) Value-based optimization
(b) Preference-basee optimization
Figure 8: Illustration EDA's performance in 2D bandit settings different diffusion times.
14

---

C Theoretical Analysis
In this section, we present the theoretical proof for Proposition 3.1
Our proof is based on the theoretical framework of Contrastive Energy Prediction (CEP) for diffusion energy guidance [32]. For the case of readers, we incorporate the relevant theories from their work as lemmas below.
Lemma C.1. Let (() AxS - R. be a scalar function approximator. Consider the optimization problem
K
el(a,a
(l(aaa*)
max E p( p(a)(s)
log
(17)
Ei=1 (l(a0))
Σ ((aa'))
The solution for Problem satisfies
=
where ((s) can be arbitrary scalar functions conditioned s.
Proof. The proof is quite straightforward. Consider two discrete distributions
(a.a'))
Q(.aa²)
P ::=-
K
eleaaa))
K
Q(e,a))
e(0,0)))
2i=1
-i=
((s,aa)
((oaa)))
els,a*)
p:=
E cl(a,a))
> ((3,a')) ""
-iml el(,aa))
For any s and any a I:K we have
K
c (aaaa
elaaaa))
E p(s) II p(af(s)
log
2== e((aaat)
e((a,at)
Si=1
=Ep(s) IE ((a))s)
(PL(P| -
< u(a"la)
H(P)
According to the properties of KI. divergence, the equality holds if and only if P P for any and l:K This implies that
( C(s),
constantly holds.
Lemma C.2. Let ((s, a,,t) A R be scalar function approximator, p(acla, t) is any conditional transition probability. Consider the optimization problem
((aa)(1)))
Σ ((xaa)) log
max E Q
IT. (aaasa)
Σ 1=
((aratt)))
(18)
The solution for Problem 18 satisfies
(", ar, log a Maa C(s).
where pa(a at, 8, 8, p(als))(aala, (a. 1) the posterior action distribution, C(s) can be arbitrary scalar functions conditioned on
15

---

Proof.
((xaa))) log K
(leaff
B alae
>i=1 ellaaat))
ellaalii)
((a.a)) log
((ss) I7 (la
Li=1 c((oratit)
=B =Bo( p=(af(o)
log Pllaatat at an)
Eim1 R elliiati))
@(laatit)
11 ((())ln)
log ΣK Bmrta (
EK eliaaa)
According to Lemma C.I satisfies
for any state a and any diffused action a, at time the optimal solution
as, log a Maai) C(s).
Lemma C.3. Consider the behavior distribution p(a)s)) and the policy distribution "t ra (8) Xx
j(a le ((,,)) Their diffused distribution at time t are both defined by the forward diffusion process
(Eq. 5 Let p(a,la, ((alaa,,, of I), such that
pr(als,tt=
Niaalaaa t)
and
क: (a,ls,t)) / N(a.la,a Df (als, t)da.
Then the relationship between oं and μe can be derived as
(arls pelails, t)
where Qe(s, az, t) := log E,, (a (ala, e.) ((s.a).
Proof. According the definition, we have
/ N(ailoza, (als,t)de
म (aas)
/ Niailara N( (f P(s.a)
(alss) Z(s) da
/ plarla,t)ulals) e(s.a)
Z(s) da
= / palalan, 8, (ar(8) 1) M(sa))
7(x) da
1 Z(s) | (o(arls.t) /
s, ((s,a) da
s,
x44(a, ()) (t(0,a)))))))
Proposition C.4. Let f be the optimal solution of Problem 14 and nH xx elō be the optimal diffusion policy. Assuming unlimited model capacity and data samples, we have the following results: (a) Optimality Guarantee. At time t = 0, the learned policy T. converges to the optimal target policy.
क (ā(a)) Ti=o,((aas) 8) psla se (saa)/8
16

---

(b) Diffusion Consistency. At time t O, Rt>0,0 models the diffused distribution of mgz
Tia f N Naaaaa,,, )= (0ols
asymptotically holds when K + Do and 3 1, satisfying the definition of diffusion process (Eq. [
Proof. We first rewrite Problem 14 below:
e(s,a")
log -1
max (0) R p(x) II ps)aas) 6
Li== (l(,,a))
(a:aat))))
2i=1
(19)
(a) Optimality Guarantee. At time = 0, we have ( = aa, 0f) such that a; = a.
Define (s, = BIf s,t 0) f (a 0).. Since we assume unlimited model
capacity for for Q' can be arbitrary scalar functions. Lemma C.1 then be applied:
Tg 8) - (a)s
(aflaa
=
xpp(als)
(b) Diffusion Consistency. When K + oo, we ii=1 Σx Q(2,a') (a, 8) e( becomes constant and can be removed. Set 3 the optimization problem equation becomes
117 (17 (. 11) (a.)
max (0)
E ((*)) 11 u(a"ls)((aila'
Σ 9ls,a" log
(20)
D 16 (a) (a]ls.t]]
K
We can then similarly apply Lemma C.2 and get
(a.(s)))
lā (arls,)) (arls,)) log EgggEr (al Q(o.a)) + 2(8) log 11
ne T & pe(arls, (a (a Q(3.a)
According to Lemma
we have
(alas of
17

---

D Implementation Details for D4RL Tasks
We NVIDIA A40 GPU cards to run all experiments.
Behavior pretraining. For pretraining the bottleneck diffusion model, we extract a reward-free behavior dataset (s, a} from D == {s, a, r, s'). We adopt the model architecture used by IDQL [1] and SRPO [4] but sum up the final [Al-dimensional output to form a scalar network. The resulting model is basically 6-layer MLP with residual connections, layer normalizations, and dropout regularization. We train the behavior network for IM steps to ensure convergence. The batch size is 2048. The optimizer is Adam with learning rate of 3e-4. We adopt default VPSDE [49 hyperparameters as the diffusion data perturbation method.
Constructing alignment dataset. First, we leverage D and use existing methods such as IQL to learn a critic network that will later be utilized for data annotation. Then, for random portion of state in D", leverage the pretrained behavior model to generate K" 16 action samples. The original action in D is thrown away. We use the critic model to annotate each state-action pair and store them together as the alignment dataset DJ == {s, a !:K Q(ss, a [kgI:k). We use 2-layer MLPs with 256 hidden states. Critic training details is exactly the same with previous work [25]. Other critie learning methods used in ablation studies are also consistent with respective prior work 15 32..
25
Policy fine-tuning, The policy network is initialized to be the behavior network. Throughout the training, we fix the behavior model and only optimize the policy model. Behavior weights frozen. The optimizer is Adam and the learning rate is 5e-5. All policy models are trained for 200k gradient steps though we observe convergence at 20K steps in most tasks. We do not employ dropout regularization during fine-tuning because find it harms performance. For the temperature coefficient, we sweep over 3 E {0.1, 0.3, 0.5, 0.8, 1.0, 2.0} (Figure 9 10). Similarly be improved by adopting rejection sampling technique.
15 15), find the performance We select the action with the highest Q-value among 4 candidates during evaluation. We do not use such techniques in experimental plots (Figure TT) better reflect policy improvement. We use 20 test seeds in all experiments and report numbers at the end of training. We train all experiments independently with 5 seeds our main experiments and 3-5 seeds in ablation studies.
18

---

E Additional Experiment Results
Hall Hallcchetahh
Hopper
Walker2d
Modian- Exper
Nonalled Sonce
100
Notzalied Soarze
100
Nataaiial
24-
Mediam
Medium Expent Raddum-Roppay Medium
Replay
Medium
0.
6.819
0.5
0.2
1.0
Antmmaze
Kilchen
Nauntted Saro
Namaizzo
Mediar Direue
Modiar Rugly
Complex
Uhrazs-Dovre
Micel
03
0.5
0.5
1.0
2.0
30
50
03
05
0.91.
25
30
50
Figure 9: Ablation of the temperature coefficient D4RL benchmarks.
Holfcheetsh
Hopper (wio
Modian Exper Modian-Rejky Rejly Mofian
Nenakast Siotes
Notzaked Stota
50
Medkan-lipee
Medium--aae
Medim Heplay
Sedism--apple
Medaaz
0.5
0.2
0.7
0.5
0.0.0
2.0
6.1
0.2
0.3
05
ILI
0.1
0.2
0.3
Kitchen (wil
Nandited Stotes
Stotes
Stores
60
Medism-Repley Limare Diverse
12
Mixed
Complete
Perial
20
Figure 10: Ablation of the temperature coefficient / without rejection sampling.
Medium-expert
Medium
Mediuum-Replay
Halfcheetah
2.0
0.1
0.2
Medium-expert 2.0
Medium 0.2
Sediumm-Rppayy
Hopper
0.2
Medium-exper 0.3
Medium 0.5
Medium-Replay Replay
Medium 0.1
Walker2d
Umaze 0.5
Umanze-diverse 5.0
Medium (both)
AntMaze
1.0
Complete
Mixed
Partial
Kitchen
2.0
3,0
3.0
Table Temperature coefficient / for every individual task.
19

---

Additional Training Curves
halfcheetah--enn expert-v22
hopper-meding expert-v2
wallker expert-v2
Normllzze Nomalizzd Sonirdddoo
Normalzed scores
100
Ours
100
Ours
Ours
50
QCPO
QAPO
QCPO
D-QL
[QL
IQL
[QL
500
LUG0
1500
500
1000
1500
300
1000
1500
Crralient
(k)
Cenditett stops
Cirddient steps
malliheetah-mediinne
Hupper-mediumm-y
walker2d malker2d-medium-2
100-
Nomalizzd scons
40
60-
Norallzzd
50
Ours
Nomaal
Ours
Ours
20
QGIO
QGPO
30-
QGPO
D-QL
D-QL
20
D.QL
1QL
IQL
50
1000
1500
500
Inño
5000
Gradient stops
stops
stops (k)
hallichettah-melee replay-v2
hopper-mediun-roolay----
walker22ä-mcdiu---rppayy--
100
40
100-
Nomallzzd commallzzddooo scoess
75-
Normallaos
Ours
50-
20-
QGPO
QGPO
QGPO
THQQ
D-QI
25
D-QU
I0L
IQL
$0
50
1000
1000
Gradiont steps
Gradient stops
Gradient stops (k)
Figure Training curves of EDA (ours) and several baselines.
20

---

G
Reproducibility
To ensure that work is reproducible, we submit the source code as supplementary material. The code will open-source upon publication. Reported experimental numbers are averaged under multiple random seeds. We provide implementation details of our work in Appendix D
H Limitations and Broader Impacts
Limitations. EDA employs diffusion policy, which is natural fit for continuous control problems. However, cannot be directly applied in recision-mmakin tasks with a discrete action space. Addi- tionally, EDA places restrictions on the diffusion model design. The diffusion policy has to be the gradient a scalar neural network with respect to action inputs to enable direct density estimation. Lastly.as task reward is human-define' and varies from task to task, achieving optimal performance with EDA in downstream tasks requires tailored network architecture design and hyperparameter tuning. How to ensure that the same set of algorithm parameters can be applicable across various tasks remains a topic worthy further research.
Broader Impacts. EDA is rather theoretical paper and all experiments are done in simulation environments. However, its further application could assist the deployment of highly optimized AI systems in critical real-world applications, such autonomous driving or healthcare robots. This could have unintended consequences if the systems fail or produce incorrect results. Robust testing and validation are essential to mitigate these risks. There is also risk that EDA could be misused for malicious purposes. For instance, optimized AI policies could employed autonomous systems for surveillance other invasive activities without proper regulation and oversight.
NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer:
[Yes]
Justification: Appendix H
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors
21

---

should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-texl system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Appendix C
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: yes we have.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same
22

---

dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis-
sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (c.g.. a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-souree models, it may be that access to the model is limited in some way (e.g.. registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Source code is in supplemental material. We will open- spen-source it upon publica- tion.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co
swtamiiisssoooiiooiooooiii for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ //nips
www.oooooidissiooooooooooo sow.ooooiisioooooooccooooo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Appendix D
23

---

Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: See Sec. 5.1 and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (c.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: Appendix and Sec. 5.2
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http://////eeaepppcccccppttt//t/
Answer: [Yes]

---

Justification: Yes.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Appendix H
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation. generating fake profiles, surveillance), faireess considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: This paper is highly theoretical, there's little risk the released training code is misused.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
25

---

12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Yes, have.
Guidelines:
The answer NA means that the paper does not use existing assets. The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Yes, we have.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
26

---

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
27


Reviews:
[{"review_id": "ICR6zdAprO", "reviewer": "Reviewer_2HKR", "summary": "Diffusion models have shown impressive results in solving real-world problems. This paper builds on the success of large language models (LLMs) to enhance the development of diffusion models. Specifically, they introduced Efficient Diffusion Alignment (EDA), a pipeline to train diffusion models in two stages similar to LLM: pretraining and finetuning. In the pretraining stage, they proposed a bottleneck diffusion model (BDM), which modifies the score function from predicting control actions to predicting the scalar value of the control action. With this pre-trained score function, they suggest performing policy optimization by aligning the pre-trained diffusion model to Q-function values. The authors conducted experiments on various environments and compared the results to several relevant baselines. This paper's proposed EDA pipeline outperforms all the previously proposed algorithms for aligning diffusion models and provides empirical insights into why their algorithm works well. In particular, they demonstrated that their algorithm is very sample efficient, requiring very few samples to learn successfully, and they showed that their algorithm could be combined with various Q-learning methods successfully.", "questions": "- Is the Behavior density estimation model a value network instead of a policy? Because you are essentially predicting a scalar value given state and action at a particular time.\n- How are you minimizing equation (10) if your BDM model outputs a scalar but your Gaussian noise variable \\epsislon is a vector of dimension the number of actions?\n- What do you mean by the bottleneck value being expanded back to R^|A| through back-propagation? The input to the function f is an a_t, which is a scalar, not \\bold{a}_t, which is a vector.\n- Does f^*=\\nabla_{\\bold{a}_t} \\log\\mu_t(a_t|s,t)? If so, what is the difference between \\epsilon_\\phi and \\f_\\phi?\n- Is a_{t} = \\bold{a}_t * e_{i}, where e_i is a vector of the standard basis?\nShould the equation on line 154 be Q^* instead of Q?\n- Does f^\\pi and f^\\mu both output scalar values, not vector corresponding actions? If so, how are you in equation (13) outputting a vector value?\n- In table 1, how does vaniila BC and vanilla diffusion BC perform?\n- Line 242,  there is a typo; you meant to say figure 4, not table 4.\n- On line 254, you should reference the figure in the text so the reader can correlate what the text is saying with the results in the figure.\n- What is the sample efficiency of the diffusion-QL baseline? Because it is not included in Figure 5(a).", "limitations": "Yes", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper addresses an important problem regarding alignment with diffusion models.\n- The paper performs several ablation studies to showcase the importance of key algorithm design choices.\n- The policy optimization procedure to minimize the q-values between optimal and reparameterization q-function is interesting and works well in practice.", "weaknesses": "- The paper needs more clarity regarding the difference between behavior density estimation and the normal diffusion model.\n- The paper needs more clarity about how the pre-trained scalar network is trained.\n- The experiment results for other baselines are missing standard deviation bars, so it is hard to tell if the results are significant, given that some baselines are extremely close in performance.\n- The value and preference optimization experiments need to include the derivation of what you optimized when k>2 for the preference experiments.\nNo ablation experiments compare the proposed BDM method with traditional conditional diffusion models.\n- There are no vanilla BC experiments in the results presented.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Official Response to Reviewer 2HKR (1/2)\nWe thank the reviewer for the very detailed feedback! The reviewer's concerns are twofold: 1. confusion regarding what the BDM model is trying to model/learn and how to train it; 2. the comprehensiveness of the experiments. We hope our explanations and the newly added experimental results can help address these concerns.\n\n**Q1: Needs more clarity regarding \"the difference between behavior density estimation and the normal diffusion model\" and \"how the pre-trained scalar network is trained. \"**\n\n**A1:**\nSuppose our dataset distribution is $\\mu(a|s)$.\n\n| | Classic Gaussian Policy | Normal diffusion | Bottleneck Diffuson (ours)|\n|--|--|--|--|\n| network | $\\pi_\\theta(s):\\mathcal{S} \\rightarrow \\mathcal{A}$ | $\\epsilon_\\theta(a_t,s, t)=a_{t-1}:\\mathcal{A} \\times \\mathcal{S} \\times \\mathbb{R} \\rightarrow \\mathcal{A}$ |  $f_\\theta(a_t,s, t): \\mathcal{A} \\times \\mathcal{S} \\times \\mathbb{R} \\rightarrow \\mathbb{R}$ |\n| predicts| $\\pi_\\theta^*(s)=\\mu(s)$| $\\epsilon_\\theta^*(a_t,s, t) = \\nabla_{a_t}\\log\\mu_t(a_t\\|s)$ | $f_\\theta^*(a_t,s, t) = \\log\\mu_t(a_t\\|s) + C(s)$ |\n| predicts| average **action** | **score function** of action distribution | behavior action **density** (log probs) |\n| sample action | $a =\\pi_\\theta(s)$| $a_{t-1}=\\epsilon_\\theta(a_t,s, t) + \\text{noise};a=a_0$ | $a_{t-1}=\\nabla_{a_t}f_\\theta(a_t,s, t) + \\text{noise};a=a_0$ |\n| sampling type| direct| iterative | iterative |\n| Expressive model?| No| Yes | Yes|\n| Allow density calculation?| Yes | No | Yes|\n| training| $\\|\\pi_\\theta(s)-a\\|^2$| $\\|\\epsilon_\\theta(a_t,s, t)-\\epsilon\\|^2;a_t=a_0 + \\sigma_t\\epsilon;\\epsilon\\sim\\mathcal{N}(0, I)$ | $\\|\\nabla_{a_t}f_\\theta(a_t,s, t)-\\epsilon\\|^2;a_t=a_0 + \\sigma_t\\epsilon;\\epsilon\\sim\\mathcal{N}(0, I)$ |\n\nThe **only difference** between normal diffusion models and our BDM is as follows: the diffusion model $\\epsilon_\\theta$ directly estimates the score function (outputting a vector of dimension $\\mathcal{A}$), while BDM $f_\\theta$ uses its derivative $\\nabla_{a} f_\\theta$ to estimate the score function. Thus, $f_\\theta$ outputs a scalar, but $\\nabla_{a} f_\\theta$ is a vector of dimension $\\mathcal{A}$. **Therefore, we can think of BDMs as exactly normal diffusion models with a different model architecture.** (See Figure 2 in the paper). \n\nThe training methods remain exactly the same.\n\n\n\n**Q2: Is the Behavior density estimation model a value network instead of a policy?**\n\n**A2:** Following **A1** the BDM model $f_\\theta$ is a learned behavior policy, but it can also be used to estimate the behavior density for a given action.\n\n **Policy**: When considering $\\epsilon(a_t,s,t):= \\nabla_{a_t} f_\\theta(a_t, s, t)$, which is easily calculatable by Pytorch through back-probagation. Then BDM is exactly a diffusion policy that can be used to sample an action. **This feature is mainly used during pretraining and evaluation.** \n\n**Value**: Without the partial derivative, $f_\\theta(a_t, s, t) \\approx \\log \\mu_t(a_t|s)$ , the BDM model is like an energy-based model that can be used to estimate action likelihood. **This is mainly used during the alignment phase.**\n\n**Q3: Does $f^\\pi$ and $f^\\mu$ both output scalar values, not vector corresponding actions? If so, how are you in equation (13) outputting a vector value?**\n\n**A3:** Yes, $f^\\pi$ and $f^\\mu$ both output scalars. We may not fully understand the reviewer's question but the output of Eq 13, including $Q_{\\theta}$, $\\pi_{t, \\theta}(a_t|s, t)$ and $\\mu_{t, \\phi}(a_t|s, t)$ are indeed all scalars.\n\nOriginal Eq 13 in paper:\n\n$Q_{\\theta}(s, a_t, t) := \\beta \\log \\frac{\\pi_{t, \\theta}(a_t|s, t)}{\\mu_{t, \\phi}(a_t|s, t)} + \\beta \\log Z(s) = \\beta [f_\\theta^\\pi(a_t|s,t) - f_\\phi^\\mu(a_t|s,t)] + \\beta [\\log Z(s, t)- C^\\pi(s, t) + C^\\mu(s, t)]$\n\n**Q4:Does $f^ * =\\nabla_{a_t} \\log\\mu_t(a_t|s,t)?$ If so, what is the difference between $\\epsilon_\\phi$and $f_\\phi$?$**\n\n**A4:** **No.** $f_\\mu^* =\\log\\mu_t(a_t|s,t)$ while $\\epsilon_\\mu^* =\\nabla_{a_t}\\log\\mu_t(a_t|s,t)$. $\\epsilon_\\phi$ is actually a notation for **normal** diffusion model. You can see that $\\epsilon_\\phi$ is rarely used when describing BDM models.  \n\n**Q5: What do you mean by the bottleneck value being expanded back to R^|A| through back-propagation? The input to the function f is an a_t, which is a scalar, not a_t, which is a vector.**\n\n\n**A5:** The input $a_t \\in \\mathcal{A}$ of $f_\\theta$  is a **vector**.  (explained in **A1** and **A2**)\n\nStage 1: $f_\\theta(a_t,s, t): \\mathcal{A} \\times \\mathcal{S} \\times \\mathbb{R} \\rightarrow \\mathbb{R}$ takes in this vector and squeezes $a_t$ into a single scalar, which we called bottleneck energy.  \n \nStage 2:  When calculating $\\nabla_{a_t} f_\\theta(a_t,s, t)$, we perform back-propagation and the output is again a vector of dimension $\\mathcal{A}$. **(bottleneck value being expanded back to R^|A|)**\n\nVector -> Scalar->Vector.  \n\nIt is just like a U-Net. **Figure 2 (left) in the paper** is a more illustrative explanation.\n\n\n**Q6: No vanilla (diffusion) BC experiments in the results presented.**\n\n**A6:** We thank the reviewer for this suggestion and present additional experimental results below.\n||Average|HalfCheetah-ME|HalfCheetah-M|HalfCheetah-MR|Hopper-ME|Hopper-M|Hopper-MR|Walker2d-ME|Walker2d-M|Walker2d-MR|\n|--|---|--|----|--|--|--|---|---|-------|---|\n|**EDA(ours)**|**87.3**|93.2±1.2|57.0±0.5|51.6±0.9|104.9±7.4|98.4±3.9|92.7±10.0|111.1±0.7|87.4±1.1|89.2±5.5|\n|DiffusionBC|67.0|73.9±28.5|47.9±3.8|42.2±6.5|71.1±37.1|63.9±15.0|69.9±28.2|98.9±25.1|68.7±24.5|66.5±9.9|\n|BC|43.0|35.8|36.1|38.4|111.9|29.0|11.8|6.4|6.6|11.3|\n\n*BC results come from D4RL paper.\n\n***\n\n## Reminder\n**QA7-QA12:**\n**We refer the reviewer to the global rebuttal posted at the top of the webpage for the rest (part 2/2) of our response. This is due to the severe page limit**.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZbvByGYYmu", "reviewer": "Reviewer_jAma", "summary": "This paper introduces Efficient Diffusion Alignment (EDA) for solving offline reinforcement learning problems. The approach involves first training a behavior cloning model using only state-action data, without reward information. Subsequently, the model is fine-tuned with rewards using DPO. During the reward-free training phase, a diffusion model is used as the policy network. Since the diffusion model doesn't provide the likelihood of the predicted action, it is modified to predict a scalar representing the density/energy of the action. The noise predicted by the model is computed by backpropagating to the input actions. Experiments are done on D4RL dataset, the paper shows that the proposed method can greatly improve the training efficiency.", "questions": "1. Does Equation 10 require computing the gradient of the gradient for the proposed Bottleneck Diffusion Models? If so, this aspect should be experimentally evaluated and discussed.\n2. Why was the diffusion model chosen as the policy network? Could other energy-based generative models, or even simpler models like GMM, be used to obtain data likelihood?\n3. What are the benefits of using the proposed method over alternative methods that apply DPO on diffusion models in the image generation field, such as Diffusion-DPO?\n\nReference: Wallace, Bram, et al. \"Diffusion model alignment using direct preference optimization.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024", "limitations": "- The paper lacks comparison with other energy-based generative models and simpler methods for obtaining data likelihood.\n- The proposed method is only evaluated on the D4RL benchmark; additional benchmarks would strengthen the experimental validation.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and clearly explains the proposed method.\n- The approach appears to be novel.", "weaknesses": "1. If I understand correctly, the proposed method needs backpropagation to compute the predicted noise and then another backpropagation to update the network. Would this require computing higher-order gradients?\n\n2. If so, there is a lack of experiments and discussion about the efficiency and robustness of the method given the potential computational cost and sensitivity of computing higher-order gradients.\n\n3. The method is only evaluated on the D4RL benchmark.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Official Response to Reviewer jAma\n**Q1: Does EDA require backpropagation to compute the predicted noise and then another backpropagation to update the network? Would this involve computing higher-order gradients?**\n\n**A1:** Yes. The pretraining of BDM requires two backpropagation steps for a single training iteration. \n\n**However, we don't think it requires computing higher-order gradients** (such as the computationally expensive Jacobian Matrix).\n\nThe BDM loss gradient simplifies to:\n\n$\\frac{\\partial L_\\theta}{\\partial \\theta} = \\frac{\\partial }{\\partial \\theta} \\|\\frac{\\partial }{\\partial a_t}f_\\theta(a_t,s, t)-\\epsilon\\|^2$\n\nHigher order gradient $\\frac{\\partial^2  f_\\theta}{\\partial  \\theta \\partial a_t}$ is computationally expensive.  However, such **Jacobian calculations can be automatically avoided by PyTorch** in practice.\n\nThe basic logic is that: A computational graph can be constructed when calculating $\\frac{\\partial f_\\theta}{\\partial a_t}$. PyTorch will just regard $\\frac{\\partial f_\\theta}{\\partial a_t}$ as any normal neural network, except that it is 2-times larger.  **It is convenient to understand the proposed BDM models jus as a Normal network** with 2-times computation. (Figure 2 in paper)\n\n**Q2: Lack of experiments and discussion about the efficiency and robustness given the potential computational cost and sensitivity of computing higher-order gradients. \"**\n\n\n**A2:** \n**computational cost： EDA roughly matches other diffusion-based baselines**, and is somewhat more efficient because it has negligible policy alignment computation. The inference/training speed of BDM is less than 2 times of normal diffusion with the same model size.\n\n| Method | Behavior pretraining | Critic training | Policy Alignment | Dataset | Overall |\n|---|---|---|---|---|---|\n| Ours|3.4h (1M steps) | 1.4h (1M steps) | 0.07h (20k steps) | 0.6h | 5.5h |\n| IDQL|1.9h (1M steps) | 1.4h (1M steps) | - | - | 3.3h |\n| QGPO | 1.9h (1M steps) | 4.2h (1M steps) | 3.2h (1M steps) | 1.0h | 10.3h |\n| DiffusionQL | - | 4.1h (1M steps) | 2.9h (1M steps) | - | 7h |\n\n**robustness：** including derivative in loss function has **already been verified by various research fields**. Application in diffusion modeling for image generation: [1] Application in Training PINN : [2] \n\n[1] Should EBMs model the energy or the score? Tim Salimans, Jonathan Ho \n\n[2] Scientific machine learning through physics–informed neural networks.\n\n**Q3: Could other energy-based generative models, or even simpler models like GMM, be used to obtain data likelihood?  Why diffusion models exactly?\" Lacks comparison with other energy-based generative models**\n\n**A3:**\n1. We refer the reviewer to **Appendix C in the paper, where we have already compared various generative models**, including VAE and EBM.\n2. Simper generative models could potentially be also effective in continuous control. However, at present many research consistently indicate that **diffusion models outperform other generative methods.** ([2] (Figure 1), [3] (Table 4), and [4] (Figure 1)). They are also successfully deployed in real-world robots [1]. \n3. We conducted some preliminary experiments applying GMM to EDA, with the results below (mixture number = 100, EM training, 2 seeds):\n\n| |Average| Half-ME | Half-M | Half-MR | Hop-ME | Hop-M | Hop-MR | Walk-ME | Walk-M | Walk-MR | \n|---|--|---|--|--|--|---|--|--|---|---| \n| Diffusion BC| 67.0 | 73.9 | 47.9 | 42.2 | 71.1 | 63.9 | 69.9| 98.9 | 68.7 | 66.5 | \n| GMM BC| 63.7 | 68.2 | 45.3 | 44.7 | 56.6 | 64.3 | 67.4| 97.0 | 68.7 | 65.0 | \n| **Diffusion+EDA**| **87.3**|93.2|57.0|51.6| 104.9 | 98.4 | 92.7 | 111.1 | 87.4 | 89.2 | \n| **GMM+EDA** |28.3|1.8|40.2|36.6|3.0|33.1|2.2|47.5|62.0|28.0|\n\n[1] Octo: An Open-Source Generalist Robot Policy \n\n[2] DiffusionQL paper\n\n[3] Offline Reinforcement Learning via High-Fidelity Generative Behavior Modeling  ICLR 2023\n\n[4] Imitating Human Behaviour with Diffusion Models  Neurips 2023\n\n**Q4: The method is only evaluated on the D4RL benchmark. Additional benchmarks would strengthen the experimental validation. \"**\n\n**A4:** We definitely agree that additional results from other benchmarks could enhance the validation of our experiments. However, given the limited time for rebuttal, we face some practical dilemma:\n\n1. Our primary goal is to demonstrate the effectiveness and competitiveness of our proposed algorithm. This involves comparisons with state-of-the-art diffusion baselines in the field of offline reinforcement learning, where D4RL is the predominant benchmark. To our knowledge, there are **very few public benchmarks with a sufficient number of diffusion baselines** for meaningful comparisons besides D4RL.\n2. Our experiments already covers **16 tasks** within the D4RL benchmarks, covering **3 distinct fields**: Locomotion, Navigation, and Manipulation. Additionally, we provide illustrative 2D experimental results in **6 more tasks**. Considering that the core contribution of this paper is somewhat theoretical, we believe this is sufficient to validate the effectiveness and support the main claims of our method.\n\n**Q5: What are the benefits of using the proposed method over alternative methods that apply DPO on diffusion models in the image generation field, such as Diffusion-DPO? \"**\n\n**A5:**\n1. Most image diffusion alignment methods, including Diffusion-DPO, are preference-based. **In offline RL, the concept of rewards is essential, and such rewards are usually continuous instead rather than binary.** Aligning pretrained policies with continuous rewards is both theoretically and practically significant.\n2. Different from image generation, it is **more important in RL to have an efficient way for calculating data density**, which current diffusion models do not provide.  PPO/SAC/REINFORCE all require calculating data prob. The proposed BDM models introduce a novel approach to address this issue, potentially expanding the application of diffusion models in RL, as also noted by reviewer FQV2.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JS6nlzBsHM", "reviewer": "Reviewer_y7D4", "summary": "This paper introduces Efficient Diffusion Alignment (EDA) for offline continuous control, combining the preference alignment theories with reinforcement learning. Specifically, EDA bridges the alignment finetuning by representing the diffusion score as a derivative of a scalar neural network with respect to actions, which allows direct density estimation. During fine tuning, DPO is used for policy improvement. In their experiments, EDA has demonstrated both strong performance and good sample efficiency in fine tuning on D4RL.", "questions": "1/ I believe this is a general framework that converts the policy optimization problem to alignment. Have you tried other alignment methods and how they perform?\n\n2/ For a fixed environment but different offline datasets, have you tried pretraining on all data then performing alignment?", "limitations": "As discussed in the weaknesses, although the current method claimed to perform general pretraining, it seems to be performing single-task pretraining and alignment. Also, the overall pipeline is more complex and time consuming than the standard diffusion-based offline RL methods, and didn’t seem to provide much performance improvement. Nevertheless, I still think this is an interesting work, and I’ll vote for a weak acceptance.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This is overall a novel and interesting paper. The way it connects preference-based optimisation with diffusion is simple yet effective, and allows efficient training of the diffusion-based RL policies. Besides, the authors have also provided theoretical justifications and toy examples for intuitive explanations to help the understanding. EDA also achieves a strong performance on the D4RL benchmark.", "weaknesses": "1/ The presentation has a certain room for improvement. The introduction made me very confused about why we want to do alignment, rather than directly optimising the policy with standard RL. This is never explained in the paper as well. After reading the whole paper, it is clearer to me. I do think this should be improved to make the paper an acceptance. \n\n2/ Although EDA is simple by itself, the overall framework is quite complex as the training has been separated into multiple stages: pretraining of the diffusion policy, pretraining of the Q functions, and final alignment. This actually made the whole framework much more complex compared with the standard offline diffusion RL methods. It would be great if the authors could provide some explanations to the actual training time of EDA, compared with the baselines.\n\n3/ Normally for pretraining and alignment, we are referring to training general purpose agents, and performing task-specific alignment. However, the current form of EDA doesn’t seem to support training general purpose agents that solve tasks with different action dimensions. Instead, it seems the current experiments are all conducted in a single-task manner, rather than multi-task settings. This actually hinders the contribution of the work.\n\n4/ Considering the complexity of the framework, the proposed EDA doesn’t seem to provide much performance improvement compared with the baselines (83.7 overall return of EDA compared with 82.1 of DQL).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Official Response to Reviewer y7D4\n\nWe truly appreciate the reviewer for the valuable suggestions and comments. Concerns are mainly about the computational complexity of EDA and the experimentation/application of the algorithm.\n\n**Q1:  Could provide some explanations to the actual training time of EDA, compared with baselines.\"**\n\n**A1:**\n**Actual time tested on NVIDIA A40**:\n\n|Method|Behavior pretraining|Critic training|Policy Alignment|Dataset|Overall|\n|---|---|---|---|---|---|\n|Ours|3.4h (1M steps)|1.4h (1M steps)|0.07h (20k steps)|0.6h|5.5h|\n|IDQL|1.9h (1M steps)|1.4h (1M steps)|-|-|3.3h|\n|QGPO|1.9h (1M steps)|4.2h (1M steps)|3.2h (1M steps)|1.0h|10.3h|\n|DiffusionQL|-|4.1h (1M steps)|2.9h (1M steps)|-|7h|\n\nOverall the inference/training speed of BDM is less than 2 times of normal diffusion with the same model size. However, its advantage is\n\n**Density estimation**: **20+ times faster** than diffusion models. Similar to Gaussian and EBM models.\n**Alignment:** **Negligible computation** because EDA allows finetuning the model for very few steps instead of learning a brand-new model from scratch.\n\n**Q2:  The overall pipeline is more time consuming than the standard diffusion-based offline RL methods\"**\n\n**A2:** We respectfully disagree with this comment. From **A1** we can see that in computational demand **EDA roughly matches other diffusion-based baselines**, and is somewhat more efficient because it has negligible policy alignment computation.\n\n**Q3:  EDA should support training general-purpose agents instead of conducting all experiments in a single-task manner. \"**\n\n**A3:** We couldn't agree more with the reviewer on this comment.  This is actually the very first \"hidden\" motivation of our paper.  As a matter of fact, we considered applying EDA to multi-task settings and building general-purpose agents when we first came up with this idea, though we indeed faced some difficulties:\n\n1. Since EDA is a pretty new idea, before moving on to large-scale/real-world experiments, we have to convince the research field as well as ourselves that it is actually effective and competitive. This means comparing with existing SOTA diffusion baselines in well-recognized continuous control tasks, where D4RL is the predominant benchmark. To our knowledge, there are **very few public benchmarks with a sufficient number of diffusion baselines** for meaningful comparisons besides D4RL.\n\n2. We focus on the theoretical derivation and motivation of the EDA algorithm in this paper. **Applying EDA for general-purpose multi-task learning would  involve extensive engineering practices like data collection and task definition that are mostly orthogonal to the core focus of our original article.** This might divert the attention from the theoretical contributions of our paper, which we aim to emphasize.\n\nOverall, while we highly recognize the value of general-purpose experiments, we sincerely hope the reviewer finds it acceptable that we list this as a future work of our paper. **We believe this does not harm the key contribution of the current article.**\n\n**Q4:  Although EDA is simple by itself, the overall framework is quite complex (3 stages). This actually made the whole framework much more complex compared with the standard offline diffusion RL methods.\"**\n\n**A4:**  \nStandard offline diffusion RL methods also require at least 2 stages of training (IQL/AWR/DiffusionQL). Some work also requires 3 stages (BCQ/BEAR/BRAC/QGPO)\n\nLLM training pipeline also requires three stages for training (Pretrain, SFT, alignment).\n\n**There are multiple stages of training because we need to handle distinctive data type and data amount in real-world application**\n\nFor instance, in autonomous driving, we can collect all driving behaviors from drivers, which creates a feasible continuous action space and is helpful for learning a foundational end-to-end model. This requires large-scale pretraining.\n\nHowever, we may only have enough human resources to label all those very harmful actions that lead to car crashes and preferred actions that save drivers' lives. These data require efficient alignment.\n\n**Q5:  I believe this is a general framework that converts the policy optimization problem to alignment. Have you tried other alignment methods and how they perform?\"**\n\n**A5:** We compare other classic preference-based LLM alignment methods including  DPO, SimPO, and IPO in our initial experiments. There's no very specific number but basically, they all perform similarly and DPO should be chosen given its simplicity and theoretical elegancy.\nAlso, as our ablation results in Figure 6 have pointed out, we find that value-based alignment （EDA） outperforms preference-based approaches such as DPO. \n\n**Q6:  The improved performance seems marginal compared with the baselines.\"**\n\n**A6:**\n1. We'd like to stress that **D4RL is a highly competitive benchmark**. The **upper bound** score for D4RL is 100 and  most baselines have achieved 80+. However EDA still outperforms all baselines in average performance of 16 tasks.\n2. EDA demonstrates **obvious improvement in sample efficiency and convergence speed** (Figure 5). Given only 1% of data, EDA maintains 95% of performance while the best-performing baseline keeps only 80% of performance.  **We think this sample efficiency improvement is more critical\"**. The high sample efficiency is an essential part of the real-world application of the alignment algorithms.\n\n**Q7:  For a fixed environment but different offline datasets, have you tried pretraining on all data and then performing alignment?\"**\n\n**A7:** Unfortunately no, this is because for D4RL,  datasets for a single robot are largely overlapping/reused, for instance, the hopper-ME dataset is already a mixture of two datasets and hopper-M dataset is exactly its subset. Concatenating them together has not much empirical meaning.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oqkxpoHmPB", "reviewer": "Reviewer_FQV2", "summary": "This paper introduces Efficient Diffusion Alignment (EDA), which draws inspiration from the optimization paradigm of DPO to align the distribution of Diffusion policies with the optimal solutions of constrained policy optimization. Experimental results demonstrate the outstanding performance and fine-tuning efficiency of EDA.", "questions": "Could the authors provide any insights on the selection of the number $K$ of actions in the state-action pairs within the alignment dataset? What is a sufficient value for $K$?", "limitations": "The paper only conducts experiments on benchmarks from D4RL. Including more real-world application benchmarks, such as Robomimic, may enhance the impact of the paper.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to follow.\n- A standout aspect of the paper is the introduction of Bottleneck Diffusion Models (BDM), which utilize neural networks (NNs) to estimate the log-probability of noisy actions. In diffusion models, it is a classic paradigm to use NNs to estimate the score function (gradient of the log-probability), eliminating the need for gradient computations when solving reverse SDE/ODEs. This is particularly crucial in the application of large diffusion models, such as in image generation tasks. However, for smaller applications like action generation, the BDM approach, although requiring gradient computations on NNs during the denoising process, allows for flexible manipulation of likelihood. I believe this approach could lead to more diverse applications in future works beyond EDA.", "weaknesses": "- EDM shows only a marginal performance improvement compared to baselines in the experiments, being slightly better than DQL and QGPO.\n- Using EDM in practice may be cumbersome. Firstly, manual preparation of an alignment dataset is required. Additionally, EDA seems to be sensitive to the temperature coefficient $\\beta$, which is bound to the training process and cannot be adjusted during inference.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Official Response to Reviewer FQV2\nWe completely agree with the reviewer's assessment of our paper, regarding both the praise and the pointed-out limitations. We thank the reviewer for providing valuable feedback and the reviewer's expertise in delivering such an insightful review, even though some deep motivations of our work were not extensively described in the text.\n\n**Q1:  Any insights on the selection of the number K of actions in the state-action pairs within the alignment dataset?\"**\n\n**A1:**We would like to refer the reviewer to **Figure 8 in the paper, where we have already conducted ablation studies of K**.\nOur experience is that $K>8$ is already sufficient for reasonably good performance in D4RL tasks. For several environments like hopper,  $K>100$ will hurt performance (possibly because the Q-model is not good enough). \nA practical way to determine $K$ is to first try if rejection sampling can improve the performance of the pretrained behavior model: Selecting the best action with the highest Q-value among $K$ independently-sampled behavior actions. \n\n**Q2:  Using EDM in practice may be cumbersome. Firstly, manual preparation of an alignment dataset is required. Additionally, EDA seems to be sensitive to the temperature coefficient  β, which is bound to the training process and cannot be adjusted during inference.\"**\n\n**A2:** We agree these are critical limitations of the EDA method.\n1. The $\\beta$ needs to be predefined before training. This is indeed a bummer for EDA. There's currently a method (QGPO) where $\\beta$, as a diffusion guidance scale, can be tuned during inference, which is realized by training an independent diffusion guidance network. However, QGPO cannot be initialized from pretrained diffusion behavior and thus has to be trained from scratch. We currently do not know how to combine the strengths of QGPO and EDA elegantly and think this could be a very important and fundamental topic. Its meaning is beyond the scope of offline RL but could benefit the entire research field such as LLM alignment/diffusion alignment.\n2. How to remove the need for requiring multiple actions for one single state is another crucial topic that we are very interested in. Still, manual preparation of the dataset can be avoided by sampling actions during training in an online manner. However, this is a little bit computationally expensive and we find it not helpful for D4RL performance. \n3. We think the sensitivity to $\\beta$ is to some extent the problem of human-defined rewards. Since the reward signal is task-specific, it can be very difficult to apply the same hyperparameter to all tasks. We believe this is a common challenge for most RL methods (and perhaps a privilege for preference-based RL).\n\n**Q3: EDM shows only a marginal performance improvement compared to baselines in the experiments.\"**\n\n**A3:**\n5. We'd like to stress that **D4RL is a highly competitive benchmark**. The **upper bound** score for D4RL is 100 and most baselines have achieved 80+. However, EDA still outperforms all baselines in the average performance of 16 tasks.\n6. EDA demonstrates **obvious improvement in sample efficiency and convergence speed** (Figure 5). Given only 1% of data, EDA maintains 95% of performance while the best-performing baseline keeps only 80% of performance.  **We think this sample efficiency improvement is more critical\"**. The high sample efficiency is an essential part for real-world application of the alignment algorithms.\n\n\n**Q4:  Including more real-world application benchmarks, such as Robomimic, may enhance the impact of the paper.\"**\n\n**A4:** We definitely agree with the reviewer.  As a matter of fact, we have considered applying EDA to multi-task,  general-purpose, and real-world agents just like Robomimic when we first came up with this idea. The main difficulty is that other benchmarks lack enough diffusion baselines to compare with. Still, this could be a very meaningful research direction.\n\nWe hope the reviewer finds it acceptable that we list this as a future work and believe this does not harm the key contribution of the current article.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# Rebuttal Summary

We would like to thank all the reviewers for their valuable comments. We are encouraged to see all reviewers recognize the theoretical novelty of our work. Reviewers FQV2 and y7D4 highlight the critical importance of the problem with diffusion models that  we are trying to solve. They also point out the vast potential of the proposed BDM model. Concerns primarily relate to the method's computational efficiency, the clarity of the paper, and limitations in the D4RL experiments.

Below, we summarize the main actions taken during the rebuttal:
1. Provided detailed computational efficiency experiment results for EDA and compared these with several diffusion-based baselines.
2. Conducted additional experiments to offer more baseline results, such as BC and diffusion BC.
3. Additionally compare with Diffusion-QL for sample efficiency.
4. Conducted an extra ablation study on various generative models (switching from diffusion models to GMMs/EBMs).
5. Rigorously revised the manuscript for a clearer narrative.
6. Clarified several confusions or misunderstandings regarding our paper.

We look forward to further discussions with the reviewers!


---

|

|

|

***


# Official Response to Reviewer 2HKR (2/2)
continue due to page limit

**Q7: Experiment results for other baselines are missing standard deviation bars**

**A7:** We did not report std bars for other baselines because:
1. It is a **common practice for most previous work** (e.g., Diffuser, QGPO, DiffusionQL, IDQL) to not report std bars of other work.  This is mainly due to the severe page limit.
2. In Table 1. The experimental results for baselines are cited from previous work. Different work may use different metrics to calculate std [1] (std, max-min, 0.5 * std, etc.). Some work (IDQL) simply do not report std. 

Nonetheless, we present the results with std bar for the two highest-related work which share the same evaluation metric as ours:

|Algorithm |Half-ME|Half-M|Half-MR|Hop-ME|Hop-M|Hop-MR|Walk-ME|Walk-M|Walk-MR|Ant(U)|Ant(UD)|Ant(MP)|Ant(MD)|Kit-C|Kit-P|Kit-M|
|----|---|---|--|--|---|---|---|--|--|---|--|--|---|--|---|---|
|EDA(ours)|93.2±1.2|57.0±0.5|51.6±0.9|104.9±7.4|98.4±3.9|92.7±10.0|111.1±0.7|87.4±1.1|89.2±5.5|93.0±4.5|81.0±7.4|79.0±4.2|84.0±8.2|81.5±7.3|69.3±4.6|65.3±2.2|
|DiffusionQL|96.8±0.3|51.1±0.5|47.8±0.3|111.1±1.3|90.5±4.6|101.3±0.6|110.1±0.3|87.0±0.9|95.5±1.5|93.4±3.4|66.2±8.6|76.6±10.8|78.6±10.3|84.0±7.4|60.5±6.9|62.6±5.1|
|QGPO|93.5±0.3|54.1±0.4|47.6±1.4|108.0±2.5|98.0±2.6|96.9±2.6|110.7±0.6|86.0±0.7|84.4±4.1|96.4±1.4|74.4±9.7|83.6±4.4|83.8±3.5| | | | |

**Q8: No ablation experiments compare the proposed BDM method with traditional conditional diffusion models.**

**A8:** We note that 4 out of the 8 selected baselines in Table 1 (Diffuser, DiffusionQL, IDQL, QGPO) are based on traditional diffusion models. For instance, EDA(ours) and IDQL share the same Q pretraining method and behavior model architecture, making IDQL an appropriate ablation baseline for comparison with EDA.

Furthermore, we clarify that **traditional diffusion models cannot be used in our proposed alignment method**. These models estimate the score function but do not estimate behavior density, which is required by our loss function in Eq. 14 (see **A1**). Therefore, a direct ablation experiment is not feasible.  **This is the very motivation for us to propose the BDM model**.

**Q9: It is hard to tell if the results are significant, given that some baselines are extremely close in performance**
**A9:** We respectfully disagree with the reviewer's comment. 

1. **D4RL is a highly competitive benchmark**. The selected baselines like IDQL, Diffusion-QL, and QGPO are recognized state-of-the-art methods, and EDA outperforms each by at least 2% in overall performance. Given the inherent randomness in RL tasks and the diversity across 16 tasks, it is natural for performance between EDA and other SOTA methods to be close on some tasks.
2. EDA demonstrates **significant improvements in sample efficiency and convergence speed** (Figure 5). With just 1% of the data, EDA retains 95% of performance while the best-performing baseline maintains only 80%. **This improvement should not be considered as "close performance."** Such high sample efficiency is crucial for real-world applications of alignment algorithms.
3. Furthermore, EDA clearly surpasses preference-based methods (83 vs. 76), as detailed in Figure 6.


**Q10: Sample efficiency of the diffusion-QL baseline.**

**A10:**  We conduct additional experiments studying Sample efficiency of diffusion-QL. Results are averged across 9 Locomotion tasks and 5 random seeds each.



| Algorithm      | 100%       | 10%        | 1%         |
|----------------|------------|------------|------------|
| **EDA (ours)** | 87.3 ± 3.5 | 85.6 ± 3.7 | 87.3 ± 5.1 |
| **Diffusion-QL** | 87.0 ± 2.9 | 60.2 ± 6.6 | 53.0 ± 10.7 |
| IQL            | 75.7 ± 7.7 | 65.0 ± 7.5 | 52.8 ± 7.2  |
| QGPO           | 86.4 ± 3.5 | 74.5 ± 3.7 | 71.2 ± 5.1  |

For fair comparison, we do not use the original critic training method in Diffusion-QL, but switch to the IQL-critic training method, which is consistent with ours and the IQL baseline.


**Q11: Typo in Line 242 (Figure instead of Table).**

**A11:** We appreciate the reviewer's careful reading. The typo has been corrected (though we cannot update the official paper during the rebuttal). We have reviewed the paper to ensure no similar mistakes remain.

**Q12: On line 254, you should reference the figure in the text so the reader can correlate what the text is saying with the results in the figure.**

**A12:** We thank the reviewer for the detailed suggestion. We have now ensured that all figures and tables are referenced before they are discussed in the text.

Meta-review: ## Summary 
This paper proposes a novel two-stage approach for offline RL, called "Efficient Diffusion Alignment," inspired by alignment approaches for LLM. The paper treats the offline learning of policies as a two-stage process inspired by LLM alignment approaches:
1. Pre-training: The policy is pre-trained on large, reward-free behavior datasets over the state-action papers to estimate the behavior density.
2. Policy fine-tuning: The diffusion policy is fine-tuned with a method inspired by DPO, where they align the diffusion behavior policy with the Q-values.

The paper reports that EDA performs very well in D4RL environments and is much more sample-efficient than the baselines compared against.

## Decision
Overall, the paper is clear and easy to read. It studies an important problem: improving the sample efficiency of the offline RL algorithms. The proposed approach of using a two-stage algorithm similar to LLMs is not completely novel; a similar approach was used to train offline RL agents for Starcraft 2, which would be beneficial for this paper if the authors could cite it. However, the idea of pretraining a behavior density estimator with the diffusion policy first and then aligning them on the offline dataset with DPO is a novel and interesting approach. 

The authors did a good job addressing the reviewers' criticisms of the paper. Overall, the reviewers were unanimously positive about the paper's acceptance. There were some concerns raised by the reviewers during the rebuttal and some confusion about certain parts of the paper. I would recommend the authors go over the reviews carefully and incorporate the suggestions made by the reviewers to eliminate confusion. Overall, I think this paper would be of interest to the broader AI community, and the experiments are convincing. Thus I am suggesting accept for this paper.

[1] AlphaStar Unplugged: Large-Scale Offline Reinforcement Learning, Mathieu et al, 2023: https://arxiv.org/abs/2308.03526
Decision: Accept (poster)