Aligning Diffusion Behaviors with Q-functions for Efficient Continuous Control

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.

Strengths

- The paper is well-written and easy to follow. - 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. - 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

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

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

3

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.

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.

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

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?

Rating

7

Confidence

4

Soundness

3

Presentation

2

Contribution

3

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.

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.

Strengths

- The paper is well-written and clearly explains the proposed method. - 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? 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

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

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

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.

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.

Strengths

- The paper addresses an important problem regarding alignment with diffusion models. - The paper performs several ablation studies to showcase the importance of key algorithm design choices. - 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. - 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

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

Rating

6

Confidence

2

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes