Official Response to Reviewer FQV2

We 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.

Q1: Any insights on the selection of the number K of actions in the state-action pairs within the alignment dataset?"

**A1:**We would like to refer the reviewer to Figure 8 in the paper, where we have already conducted ablation studies of K. Our 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). A 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.

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."

A2: We agree these are critical limitations of the EDA method.

1. 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.
2. 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.
3. 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).

Q3: EDM shows only a marginal performance improvement compared to baselines in the experiments."

A3: 5. 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. 6. 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.

Q4: Including more real-world application benchmarks, such as Robomimic, may enhance the impact of the paper."

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.

We 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.

Yes, though for the behavior pretraining dataset, reward labels are excluded. The diffusion model learns all kinds of behavior, regardless of whether they are good or bad.

We thank the reviewer for the prompt reply and the interest in our work.

We think the insight for being able to improve alignment efficiency here is that transforming $\mu$ into $\pi^*$ (alignment) is fundamentally much much easier than learning $\mu$ from scratch (pertaining). In alignment, we are simply trying to "suppress" some bad modes learned during pretraining instead of trying to find new meaningful modes.

The theoretical explanation could be $KL(\mu\|\pi^*)$ << $KL(\mu\|\text{uniform dist})$. In high-dimensional data space, the actual meaningful data is actually very scarce. Imagine if you are trying to sample an image from Gaussian distribution, you wouldn't get a visually realistic image from even $10^{100}$ candidates. However, if you already have a pretrained image-generation model like stable-diffusion, you can easily get a good-looking one from 4-16 image samples.

Why can EDA have high sample efficiency while other diffusion-based algorithms cannot? It is simply because EDA is completely initialized from the pretrained model. There are no network parameters that are required to be learned from scratch. Finetuning a model is much easier than learning a completely new one. Similar successes have already been proved by wide exploration in LLM alignment research.

Official Response to Reviewer y7D4

We 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.

Q1: Could provide some explanations to the actual training time of EDA, compared with baselines."

A1: Actual time tested on NVIDIA A40:

Overall the inference/training speed of BDM is less than 2 times of normal diffusion with the same model size. However, its advantage is

Density estimation: 20+ times faster than diffusion models. Similar to Gaussian and EBM models. Alignment: Negligible computation because EDA allows finetuning the model for very few steps instead of learning a brand-new model from scratch.

Q2: The overall pipeline is more time consuming than the standard diffusion-based offline RL methods"

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.

Q3: EDA should support training general-purpose agents instead of conducting all experiments in a single-task manner. "

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:

1. 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.

2. 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.

Overall, 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.

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."

A4:
Standard 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)

LLM training pipeline also requires three stages for training (Pretrain, SFT, alignment).

There are multiple stages of training because we need to handle distinctive data type and data amount in real-world application

For 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.

However, 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.

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?"

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. Also, as our ablation results in Figure 6 have pointed out, we find that value-based alignment （EDA） outperforms preference-based approaches such as DPO.

Q6: The improved performance seems marginal compared with the baselines."

A6:

1. 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.
2. 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.

Q7: For a fixed environment but different offline datasets, have you tried pretraining on all data and then performing alignment?"

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.

We are glad that the reviewer is happy with our responses! We also appreciate the reviewer for the prompt and positive feedback.

Official Response to Reviewer jAma

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?

A1: Yes. The pretraining of BDM requires two backpropagation steps for a single training iteration.

However, we don't think it requires computing higher-order gradients (such as the computationally expensive Jacobian Matrix).

The BDM loss gradient simplifies to:

$\frac{\partial L_\theta}{\partial \theta} = \frac{\partial }{\partial \theta} \|\frac{\partial }{\partial a_t}f_\theta(a_t,s, t)-\epsilon\|^2$

Higher 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.

The 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)

Q2: Lack of experiments and discussion about the efficiency and robustness given the potential computational cost and sensitivity of computing higher-order gradients. "

A2: 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.

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]

[1] Should EBMs model the energy or the score? Tim Salimans, Jonathan Ho

[2] Scientific machine learning through physics–informed neural networks.

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

A3:

1. We refer the reviewer to Appendix C in the paper, where we have already compared various generative models, including VAE and EBM.
2. 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].
3. We conducted some preliminary experiments applying GMM to EDA, with the results below (mixture number = 100, EM training, 2 seeds):

[1] Octo: An Open-Source Generalist Robot Policy

[2] DiffusionQL paper

[3] Offline Reinforcement Learning via High-Fidelity Generative Behavior Modeling ICLR 2023

[4] Imitating Human Behaviour with Diffusion Models Neurips 2023

Q4: The method is only evaluated on the D4RL benchmark. Additional benchmarks would strengthen the experimental validation. "

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:

1. 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.
2. 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.

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? "

A5:

1. 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.
2. 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.

Official Response to Reviewer 2HKR (1/2)

We 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.

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. "

A1: Suppose our dataset distribution is $\mu(a|s)$.

The 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).

The training methods remain exactly the same.

Q2: Is the Behavior density estimation model a value network instead of a policy?

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.

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.

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.

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?

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.

Original Eq 13 in paper:

$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)]$

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$?$

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.

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.

A5: The input $a_t \in \mathcal{A}$ of $f_\theta$ is a vector. (explained in A1 and A2)

Stage 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.

Stage 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|)

Vector -> Scalar->Vector.

It is just like a U-Net. Figure 2 (left) in the paper is a more illustrative explanation.

Q6: No vanilla (diffusion) BC experiments in the results presented.

A6: We thank the reviewer for this suggestion and present additional experimental results below.

*BC results come from D4RL paper.

Reminder

QA7-QA12: 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.