Aligning Diffusion Behaviors with Q-functions for Efficient Continuous Control

Authors: huayu  chen
kaiwen  zheng
hang  su
jun  zhu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-10 08:34:54

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.

Main paper: 
 introduction: 
 Learning diverse behaviors is generative modeling; transforming them into optimized policies is reinforcement learning. Recent studies have identified diffusion policies as a powerful tool for representing heterogeneous behavior datasets [21, 38] . However, these behavior policies incorporate suboptimal decisions in datasets, making them unsuitable for direct deployment in downstream tasks. To get optimized policies, typical methods involve either augmenting the behavior policy with an additional guidance/evaluation network [21, 32, 15] or training a new evaluation policy supervised by the behavior policy [13, 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?Recent advances in Large Language Model (LLM) alignment techniques [57, 37, 43] offer valuable insights for fine-tuning diffusion behavior policies due to the fundamental similarity of the issues they aim to address (Fig. 1 ). 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 to remove harmful or useless content learned during pretraining. They enable quick adaptations of pretrained LLMs to human intentions by fine-tuning them on a small dataset annotated with human preference labels. These strategies, due to 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 LLMs employ Categorical models to deal with discrete actions (tokens). Their alignment relies on computing data probabilities for maximum likelihood training (Sec. 2.2). 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. 1 ).To allow aligning diffusion behavior models with Q-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 rewardfree datasets. Different from previous work which constructs diffusion models as an end-to-end 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.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 a classification task to predict the optimal action using log-probability ratios between the policy and the behavior model. Our approach innovatively expands DPO by allowing fine-tuning on an arbitrary 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% 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 1M 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 to align diffusion behaviors with scalar Q-functions, showcasing its vast potential in continuous control. 
offline reinforcement learning: 
 Offline RL aims to tackle decision-making problems by solely utilizing a pre-collected behavior dataset. Consider a typical Markov Decision Process (MDP) described by the tuple ⟨S, A, P, r, γ⟩. S is the state space, A is the action space, P (s ′ |s, a) is the transition function, r(s, a) is the reward function and γ is the discount factor. Given a static dataset D µ := {s, a, r, s ′ } representing interaction history between an implicit policy µ and the MDP environment, our goal is to learn a new policy that maximizes cumulative rewards in this MDP while staying close to the behavior policy µ.Offline RL can be formalized as a constrained policy optimization problem [26, 35, 53] :EQUATIONwhere Q(s, a) is an action evaluation network that can be learned from D µ . β is a temperature coefficient. Previous work [40, 39] proves that the optimal solution for solving Eq. 1 is:EQUATIONIn this paper, we focus on how to efficiently learn parameterized policies for modeling π * .Direct Preference Optimization (DPO, [41] ) is a fine-tuning technique for aligning pretrained LLMs with human feedback. Suppose we already have a pretrained LLM model µ(a|s), where s represents user instructions, and a represents generated responses. The goal is to align µ ϕ with some implicit evaluation rewards r LM (s, a) that reflect human preference. Our target model is π * (a|s) ∝ µ ϕ (a|s)e r LM (s,a)/β .DPO assumes we only have access to some pairwise preference data {s → (a w > a l )} and the preference probability is influenced by r LM (s, a). Formally,EQUATIONwhere σ is the sigmoid function.In order to learn π θ ≈ π * (a|s) ∝ µ ϕ (a|s)e r LM (s,a)/β , DPO first parameterizes a reward model using the log-probability ratio between π θ and µ ϕ , and then optimizes this reward model through maximum likelihood training:EQUATIONwhere r LM θ (s, a) := β log π θ (a|s) µ ϕ (a|s) 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.Recent studies show that diffusion models [45, 20, 49] excel at representing heterogeneous behavior policies in continuous control [21, 5, 38] . To train a diffusion behavior model, we first gradually inject Gaussian noise into action points according to the forward diffusion process:EQUATIONwhere t ∈ [0, 1], and ϵ is standard Gaussian noise. α t , σ t ∈ [0, 1] are manually defined so that at time t = 0, we have a t = a and at time t = 1, we have a t ≈ ϵ. When a is sampled from the behavior policy µ(a|s), the marginal distribution of a t at various time t satisfiesEQUATIONIntuitively, the diffusion training objective predicts the noise added to the original behavior actions:EQUATIONMore formally, it can be proved that the learned "noise predictor" ϵ ϕ actually represents the score function ∇ at log µ t (a t |s, t) of the diffused behavior distribution µ t [49] :EQUATIONWith such a score-function estimator, we can employ existing numerical solvers [46, 33] to reverse the diffusion process, and sample actions from the learned behavior policy µ ϕ .  
method: 
 We decompose the policy optimization problem into two stages: behavior pretraining (Sec. 3.1) and policy alignment (Sec. 3.2).Recent advances in 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. 4).Language models primarily deal with discrete actions (tokens) defined by a vocabulary set V, and thus employ Categorical models. This modeling method enables easy 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 so feasible. Diffusion policies only estimate the gradient field (a.k.a., score) of data density instead of the density value itself [49] , making it impossible to directly apply LLM alignment techniques [41, 9, 3] . Conventional Gaussian policies have a tractable probability formulation but lack enough expressivity and multimodality needed to accurately model behavior datasets [52, 13, 5] , and catastrophically fail in our initial experiments.To address the above limitation of diffusion models, we propose a new diffusion modeling technique to enable direct density estimation. Normally, a conditional diffusion policy ϵ ϕ (a t |s, t) : A×S ×R → R |A| maps noisy actions a t to predicted noises ϵ ∈ R |A| . In our approach, we redefine ϵ ϕ as the derivative of a scalar network f ϕ (a t |s, t) : A × S × R → R with respect to input a t :EQUATIONGiven that f ϕ is a parameterized network, its gradient computation can be conveniently performed by auto-differential libraries. The new training objective for f ϕ can then be reformulated from Eq. 7:EQUATIONAs noted by [49] , with unlimited model capacity, the optimal solution for solving Eq. 10 is:EQUATION)An illustration is provided in Figure 2 (left). Intuitively, our proposed modeling method first compresses the input action into a scalar value with one single dimension. Then, this bottleneck value is expanded back to R |A| through back-propagation. We thus refer to f ϕ as Bottleneck Diffusion Models (BDMs). The primary advantage of BDMs is their ability to efficiently estimate behavior densities in a single forward pass. Moreover, BDMs are fully compatible with existing diffusion-based codebases regarding training and sampling procedures, inheriting their key benefits such as training stability and model expressivity. BDMs can also be viewed as a diffused version of Energy-Based Models (EBMs, [8] ). We refer interested readers to Appendix A for a detailed discussion.In this section, our goal is to learn a new policy π θ ∝ µ ϕ e Q by fine-tuning the previously pretrained behavior policy µ ϕ on a new dataset annotated by an 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 our method below.Dataset construction. For any state s, we draw K > 1 independent action samples a 1:K from µ ϕ (•|s). Assume we already have an Q-function Q(s, a) that evaluates input state-action pairs in scalar values, our dataset is formed asD f := {s, a 1:K , Q(s, a k )| k∈1:K }.Action optimality. We first introduce a formal notion of action optimality. We draw from the controlas-probabilistic-inference framework [29] and define a random optimality variable O K , which is a one-hot vector of length K. The k-th index of O K being 1 indicates that a k is the optimal action within K action candidates a 1:K . We havep(O k K = 1|s, a 1:K ) = e Q(s,a k ) K i=1 e Q(s,a i ).The optimality probability of a behavioral action a is proportional to the exponential of its Q-value, aligning with the optimal policy definition π * (a|s) ∝ µ(a|s)e Q(s,a)/β .From policy optimization to action classification. We construct a classification task by predicting the optimal action among K candidates. This requires learning a Q-model termed Q θ first. Drawing inspiration from DPO (Sec. 2.2), we parameterize Q θ as the log probability ratio between π θ and µ ϕ :Q θ (s, a) := β log π θ (a|s) µ ϕ (a|s) + β log Z(s),This parameterization allows us to directly optimize π θ during training because π θ (a|s) = 1 Z(s) µ ϕ (a|s)e Q θ (s,a)/β 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:EQUATION)The unknown partition function Z(s) automatically cancels out during division. β is a hyperparameter that can be tuned to control how far π θ deviates from µ ϕ .Expanding to bottleneck diffusion behavior. A reliable behavior density estimation of µ ϕ (a|s) 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 of adopting a much more powerful generative policy, such as the BDM model (Sec. 3.1).Diffusion policies define a series of distributions π t (a t |s, t) = N (a t |α t a 0 , σ 2 t I)π 0 (a 0 |s, t)da 0 at different timesteps t ∈ [0, 1], rather than just a single distribution π = π 0 . Consequently, instead of directly predicting the optimal action given raw actions a 1:K , we perturb all actions with K independent Gaussian noises according to the diffusion forward process by applying a k t = α t a k + σ t ϵ k . Then we predict action optimality given K noisy action a 1:K t :EQUATIONSimilar to Eq. 12, all unknown terms above automatically cancel out in the training objective:max θ L f (θ) = E t,ϵ 1:K ,(s,a 1:K )∼D f K k=1 e Q(s,a k ) K i=1 e Q(s,a i ) optimality probability log e β[f π θ (a k t |s,t)-f µ ϕ (a k t |s,t)] K i=1 e β[f π θ (a i t |s,t)-f µ ϕ (a i t |s,t)]predicted probability on noisy actions .Proposition 3.1. (Proof in Appendix C) Let f * θ be the optimal solution of Problem 14 and π * t,θ ∝ e f * θ 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 π * θ converges to the optimal target policy.π * θ (a|s) = π * t=0,θ (a|s) ∝ µ ϕ (a|s)e Q(s,a)/β (b) Diffusion Consistency. At time t > 0, π t>0,θ models the diffused distribution of π * θ : π * t,θ (a|s, t) = N (a t |α t a, σ 2 t I)π * θ (a 0 |s)da 0asymptotically holds when K → ∞ and β = 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 θ = ϕ 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 Q-labelled data (e.g., 10k samples) and optimization steps (e.g., 20K gradient steps). These are less than 5% of previous approaches (Sec. 5.2).4 Related WorkQuestion: For each theoretical result, does the paper provide the full set 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 crossreferenced. • 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 a 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. 
diffusion modeling for offline continuous control: 
 Recent advancements in offline RL have identified diffusion models as an effective approach for behavior modeling [38, 16] , which excels at representing complex and multimodal distributions Table 1 : Evaluation results of D4RL benchmarks (normalized according to [10] ). We report mean ± standard deviation of algorithm performance across 5 random seeds at the end of training. Numbers within 5 % of the maximum are highlighted.compared with other modeling methods like Gaussians [54, 55, 24, 35, 53] or VAEs [11, 26, 12] . However, optimizing diffusion models can be a bit more tricky due to the unavailability of a tractable probability calculation method [5, 51] . Existing approaches include learning a separate guidance network to guide the sampling process during evaluation [21, 32] , applying classifier-free guidance [1, 7] , 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, 15] and applying planning-based techniques [30, 36, 16] .Our proposed method differs from all previous work in that it directly fine-tunes the pretrained behavior to align with task-specific annotations, instead of learning a new downstream policy.Existing alignment algorithms are largely preference-based methods. Preference-based Reinforcement Learning (PbRL) assumes the reward function is unknown, and must be learned from data. Previous work usually applies inverse RL to learn a reward model first and then uses this reward model for standard RL training [19, 57, 28, 44, 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 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, 2, 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.Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)? 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 a dataset and/or model, the authors should describe the steps taken to 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 is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to 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 NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
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-based alignment? (Sec. 5.3)• How do contrastive action number K and other choices affect the performance? (Sec. 5.4)Benchmark. In Table 1 , we evaluate the performance of EDA in D4RL benchmarks [10] . All evaluation tasks have a continuous action space and can be broadly categorized into three types: MuJoCo locomotion are tasks for controlling legged robots to move forward, where datasets are generated by a 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.FrankaKitchen 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, we train an action evaluation model Q ψ (s, a) for annotating behavioral data during the fine-tuning stage. Implicit Q-learning [25] is employed for training Q ψ due to its simplicity and orthogonality to policy training. We compare with other critic training methods in Figure 4 . The rest of the implementation details are in Appendix D.Baselines. We mainly consider diffusion-based RL methods with various optimization techniques.Decision Diffuser [1] employs classifier-free guidance for optimizing behavior models. QGPO employs energy guidance. IDQL [15] does not optimize the behavior policy and simply uses rejection sampling during evaluation. Diffusion-QL [52] has no explicit behavior model, but instead adopts a diffusion regularization loss. We also reference well-studied conventional algorithms for different classes of generative policies. BCQ [11] features a VAE-based policy. DT [6] has a transformer-based architecture. CQL [27] and IQL [25] targets Gaussian policies. Result analysis. From table 1, we find that EDA surpasses all referenced baselines regarding overall performance and provides a competitive number in each D4RL task. To ascertain whether this improvement stems from the alignment algorithm rather than from a superior critic model or policy class, we conduct additional controlled experiments. As outlined in Figure 4 , we evaluate three variants of EDA using different Q-learning approaches and compare these against both diffusion and Gaussian baselines.The experimental results highlight the superiority of diffusion policies and further validate the effectiveness of our proposed method. 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, we 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 5 (a) ). We compare this with IQL, 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 a 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 5 (b), we plot EDA's performance throughout the fine-tuning phase. EDA rapidly converges in roughly 20K gradient steps, a negligible count compared to the typical 1M steps used for behavior pretraining. Note that behavior modeling and 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 to individual downstream tasks.A significant difference between EDA and existing preference-based RL methods such as DPO is that EDA is tailored for alignment with scalar Q-values instead of just preference data.For preference data without explicit Q-labels, we can similarly derive an alignment loss:EQUATIONHere a w represents the most preferred action among a 1:K . In practice, we select a w as the action with the highest Q-value but abandon the absolute number. We'd like to note that when K = 2, the above objective becomes exactly the DPO objective (Eq. 4) in preference learning: We ablate different choices of K and compare our proposed value-based alignment with preference methods in 9 D4RL Locomotion tasks (Figure 6 ). Results show that EDA generally outperforms 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 two methods becomes larger as K increases. This is expected because preference-based optimization greedily follows a single action with the highest Q-value. However, as more action candidates are sampled from the behavior model, the final selected action will have a higher probability of being out-ofbehavior-distribution data. This further hurts performance. In contrast, EDA is a softer version of preference learning. This leads to greater tolerance for K.EQUATIONWe study the impact of varying temperature β on 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.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: Source code is in supplemental material. We will open-source it upon publication.Guidelines:• The answer NA means that paper does not include experiments requiring code. Justification: Appendix D and Sec. 5.2Guidelines:• The answer NA means that the paper does not 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 in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
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 its vast potential in learning from large-scale behavior datasets and efficiently adapting to individual downstream tasks. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes]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 are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: Yes, we 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 or dataset.• The authors should state which version of the asset is used and, if possible, include a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 not available online, the authors are encouraged to reach out to the asset's creators. 
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 of 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 is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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. 
 
 Appendix: Our proposed Bottleneck Diffusion Models (BDMs) can be viewed as a diffused variant of Energy-Based Models (EBMs, [8] ). Both methods aim to model data distribution's unnormalized log probability:p θ (x) ∝ e -E θ (x) .Despite their conceptual similarity, the sampling and training approaches differ between diffusion models and EBMs. A prevalent sampling method for EBMs is Langevin MCMC [14] , which employs the energy gradient ∇ x E θ (x) to progressively transform Gaussian noise into data samples. Langevin MCMC generally necessitates hundreds or thousands of iterative steps to achieve convergence, a significantly higher number compared with 15-50 steps required by diffusion models (Figure A ).Furthermore, the maximum-likelihood training of EBMs (e.g., Contrastive Divergence [18] ) is more computationally expensive, as it 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 E θ (x) towards developing score-based models defined as s θ (x) := -∇ x E θ (x).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 some prior work [42, 47, 31 ] that employed such score-matching objectives for training EBMs. 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 ease of readers, we incorporate the relevant theories from their work as lemmas below.Lemma C.1. Let Q(s, a) : A×S → R be a scalar function approximator. Consider the optimization problemEQUATION)The solution for Problem 17 satisfiesQ * (s, a) = Q(s, a) + C(s),where C(s) can be arbitrary scalar functions conditioned on s.Proof. The proof is quite straightforward. Consider two discrete distributionsP := e Q(s,a 1 ) K i=1 e Q(s,a i ) , e Q(s,a 2 ) K i=1 e Q(s,a 2 ) ..., e Q(s,a K ) K i=1 e Q(s,a i ) P := e Q(s,a 1 ) K i=1 e Q(s,a i ) , e Q(s,a 2 ) K i=1 e Q(s,a i ) ..., e Q(s,a K ) K i=1 e Q(s,a i )For any s and any a 1:K , we haveE µ(s) K i=1 µ(a i |s) K k=1 e Q(s,a k ) K i=1 e Q(s,a i ) log e Q(s,a k ) K i=1 e Q(s,a i ) =E µ(s) K i=1 µ(a i |s) -D KL (P || P ) -H(P ) ≤E µ(s) K i=1 µ(a i |s) -H(P )According to the properties of KL divergence, the equality holds if and only if P = P for any s and a 1:K . This implies that Q * (s, a) = Q(s, a) + C(s), constantly holds. Lemma C.2. Let Q(s, a t , t) : A × S × R → REQUATION)The solution for Problem 18 satisfiesQ * (s, a t , t) = log E µt(a|at,s,t) e Q(s,a) + C(s),where µ t (a|a t , s, t) = µ(a|s)p(a t |a, t)/µ t (a t |s, t) is the posterior action distribution, C(s) can be arbitrary scalar functions conditioned on s.Proof.E µ(s) K i=1 µ(a i |s)p(a i t |a i ,t) K k=1 e Q(s,a k ) log e Q(s,a k t ,t) K i=1 e Q(s,a i t ,t) =E µ(s) K i=1 µt(a i t |s)µt(a i |a i t ,s,t) K k=1 e Q(s,a k ) log e Q(s,a k t ,t) K i=1 e Q(s,a i t ,t).=E µ(s) K i=1 µt(a i t |s) K k=1 E µt(a k |a k t ,s,t) e Q(s,a k ) log e Q(s,a k t ,t) K i=1 e Q(s,a i t ,t).=E µ(s) K i=1 µ ′ t (a i t |s) K k=1 E µt(a k |a k t ,s,t) e Q(s,a k ) K i=1 E µt(a i |a i t ,s,t) e Q(s,a i ) log e Q(s,a k t ,t) K i=1 e Q(s,a i t ,t).According to Lemma C. Proposition C.4. Let f * θ be the optimal solution of Problem 14 and π * t,θ ∝ e f * θ 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 π * θ converges to the optimal target policy. π * θ (a|s) = π * t=0,θ (a|s) ∝ µ ϕ (a|s)e Q(s,a)/β (b) Diffusion Consistency. At time t > 0, π t>0,θ models the diffused distribution of π * θ :π * t,θ (a|s, t) = N (a t |α t a, σ 2 t I)π * θ (a 0 |s)da 0 asymptotically holds when K → ∞ and β = 1, satisfying the definition of diffusion process (Eq. 6).Proof. We first rewrite Problem 14 below:EQUATION) (a) Optimality Guarantee. At time t = 0, we have p(a t |a, t) = N (a t |α t a, σ 2 t I) = N (a t |a, 0I) such that a t = a. Define Q * (s, a) := β[f π θ (a k t |s, t = 0) -f µ ϕ (a k t |s, t = 0)].Since we assume unlimited model capacity for f θ , Q * can be arbitrary scalar functions. Lemma C.1 can then be applied: ,a) becomes constant and can be removed. Set β = 1, the optimization problem equation 14 becomesπ * θ (a|s) =π * t=0,θ (a|s) ∝e f π θ (a i t |s,t=0) =e f µ ϕ (a i t |s,t=0)+ Q * (s,a)/β =e f µ ϕ (a i t |s,t=0) e [Q(s,a)+C(s)]/β ∝µ ϕ (a|s)e Q(s,a)/β (b) Diffusion Consistency. When K → ∞, we K i=1 e Q(s,a i ) = E µ (a|s)e Q(sEQUATIONWe can then similarly apply Lemma C. We use 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 [15] and SRPO [4] but sum up the final |A|-dimensional output to form a scalar network. The resulting model is basically a 6-layer MLP with residual connections, layer normalizations, and dropout regularization. We train the behavior network for 1M steps to ensure convergence. The batch size is 2048. The optimizer is Adam with a 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 [25] to learn a critic network that will later be utilized for data annotation. Then, for a random portion of state s in D µ , we 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 D f := {s, a 1:K , Q(s, a k )| k∈1:K }. We use 2-layer MLPs with 256 hidden states. Critic training details is exactly the same with previous work [25] .Other critic learning methods used in ablation studies are also consistent with respective prior work [15, 32] .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 are 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 we find it harms performance. For the temperature coefficient, we sweep over β ∈ {0.1, 0.2, 0.3, 0.5, 0.8, 1.0, 2.0} (Figure 9&10 ). Similarly to [11, 15, 5, 15] , we find the performance can be improved by adopting a rejection sampling technique. We select the action with the highest Q-value among 4 candidates during evaluation. We do not use such techniques in experimental plots (Figure 5&11 ) to 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 in our main experiments and 3-5 seeds in ablation studies. 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, e.g., if the approach was only tested on a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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, a 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 important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 


Decision: ## 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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - The paper needs more clarity regarding the difference between behavior density estimation and the normal diffusion model.
- The paper needs more clarity about how the pre-trained scalar network is trained.
- 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.
- The value and preference optimization experiments need to include the derivation of what you optimized when k>2 for the preference experiments.
No ablation experiments compare the proposed BDM method with traditional conditional diffusion models.
- There are no vanilla BC experiments in the results presented.
questions: 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.
- 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?
- 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.
- 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?
- Is a_{t} = \bold{a}_t * e_{i}, where e_i is a vector of the standard basis?
Should the equation on line 154 be Q^* instead of Q?
- 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?
- In table 1, how does vaniila BC and vanilla diffusion BC perform?
- Line 242,  there is a typo; you meant to say figure 4, not table 4.
- 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.
- What is the sample efficiency of the diffusion-QL baseline? Because it is not included in Figure 5(a).
limitations: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: 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?

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

3. The method is only evaluated on the D4RL benchmark.
questions: 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.
2. 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?
3. 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?

Reference: 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: limitations: - The paper lacks comparison with other energy-based generative models and simpler methods for obtaining data likelihood.
- The proposed method is only evaluated on the D4RL benchmark; additional benchmarks would strengthen the experimental validation.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 

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

3/ 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.

4/ 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).
questions: 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?

2/ For a fixed environment but different offline datasets, have you tried pretraining on all data then performing alignment?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: - EDM shows only a marginal performance improvement compared to baselines in the experiments, being slightly better than DQL and QGPO.
- 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.
questions: 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: 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.

Comments: title: Additional Response
comment: Thank you for your positive feedback! We are glad that our responses help.
I have reviewed the rebuttal, and my concern has been addressed.
Thank you for your positive feedback! We are glad that our responses help.
Thank you for the detailed response and explanations. Since most of my concerns have been addressed, I will increase my rating to a 6.
We are glad that the reviewer is happy with our responses! We also appreciate the reviewer for the prompt and positive feedback.
I thank the reviewer for the detailed responses and explanations. My concerns are well addressed and I will increase my rating to 7.
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.
Thank you for the reply!

So, does the statement "It is simply because EDA is completely initialized from the pretrained model." implies that during the behavior pretraining phase, EDA trains a behavior policy using the entire dataset and then uses 1% of the dataset during the fine-tuning phase? If that's the case, then the high sample efficiency of EDA is indeed easy to understand.
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.
Thank you for your detailed response. Firstly, I want to apologize for my typo in writing "EDA" as "EDM."

I personally really like the experiments on sample efficiency. The fact that EDA shows almost no performance loss with only 1% of the data is quite surprising. I want to hear the authors' in-depth analysis and discussion of this result.  in the authors' view, what could be the main reasons behind this? Additionally, I strongly recommend that the authors consider adding a section in the paper for a detailed discussion on this topic. I believe this would bring many interesting insights.

And, I still stand by my judgment. I keep the rating unchanged.