Extracting Reward Functions from Diffusion Models

Summary

The study focuses on extracting a reward function by comparing low-reward and high-reward decision-making diffusion models, sharing the similar goal with inverse reinforcement learning. The researchers define a relative reward function of two diffusion models and establish its existence and uniqueness under certain conditions. They develop a practical learning algorithm to extract the reward function by aligning neural network-parametrized reward function gradients with the difference in outputs from both diffusion models. The method successfully finds correct reward functions in navigation environments and improves performance in standard locomotion benchmarks when steering the base model. Furthermore, the approach generalizes beyond sequential decision-making by learning a reward-like function from two large-scale image generation diffusion models, effectively assigning lower rewards to harmful images.

Strengths

1. The research problem tackled in this study is both significant and clearly defined. Its importance and relevance are well-recognized within the research community. 2. The paper features a reader-friendly format, with technical details coherently presented and logically interconnected. 3. The theoretical findings are fundamental and grounded in reasonable assumptions. Intuitive examples are provided to facilitate practical implementation. 4. Promising empirical results are demonstrated across various tasks, including gridworld, locomotion benchmarks, and image generation.

Weaknesses

One of my primary concerns lies in the connection to inverse reinforcement learning (IRL). Although the paper uses IRL to motivate the task, the relationship to commonly studied IRL methods has not been sufficiently introduced. 1. IRL algorithms typically formulate the problem as a two-player minimax game, assuming bounded rationality of expert data through a maximum entropy framework. In stochastic environments, these methods often analyze causal entropy. However, it appears that such methods have not been thoroughly compared or investigated in this work. 2. In IRL, researchers typically have access to an environment for interaction, which allows for exploration and the generation of nominal trajectories. In the context of diffusion models, there seems to be no such interaction, or at least none observed within the algorithm. Consequently, this paper may be more akin to inverse optimal control, where preference learning is commonly implemented offline without any interaction. Despite the aforementioned concerns, the overall quality and contributions of this work remain substantial and noteworthy.

Questions

I am wondering whether the diffuision models can achieve comparable performance with popular RL methods like SAC, PPO or DDPG.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

None.

Summary

The authors propose a new problem setting, similar to IRL, of extracting reward functions from decision-making diffusion models: one expert and one base (non-expert), where the reward is defined as a function that can steer the base model to the expert through classifier guidance. Using this intuition, they derive a relative reward function of two diffusion models as an approximation to this steering function and provide theoretical evidence for its uniqueness. They test this method on maze navigation, locomotion, and image generation tasks and provide a variety of analysis that the learned reward function is reasonable. Furthermore, they demonstrate that guiding the non-expert model with the extracted reward function does indeed increase task performance in the locomotion tasks.

Strengths

- The authors propose a novel problem of extracting reward functions from decision-making diffusion models. - The method is simple and practical with sound derivation and some theoretical backing. - Good range of qualitative and quantitative results across different domains in control and image generation. Overall, the proposed method seems to be able to extract reasonable reward functions.

Weaknesses

- The motivation for the problem setting is not obvious. Since you have access to the expert model already, extracting a reward function is not useful for task learning like in inverse RL cases. Further discussion or demonstrations of use cases for such a reward function would be helpful. - The experiments lack proper evaluation on classifier guidance except for the locomotion tasks. As stated in the problem setting, the objective is to extract the reward function that can steer the base model to the expert model through classifier guidance, so this evaluation should be included in more domains.

Questions

- For the locomotion experiments, what do the expert policy rewards look like? Does your method represent a significant improvement to the policy? It would also be helpful to have an upper bound figure using a ground truth classifier. - The problem setting, motivation, and evaluation criteria are all a bit unclear. Defining the problem explicitly (what makes a good reward function?), with convincing use cases, and bringing the experiments in line would make this a strong paper.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

2 fair

Limitations

Yes, they stated no real world experiments.

Summary

This paper proposed a reward extraction method for reinforcement learning from diffusion models. Theoretical justification and experiments in diverse domains and settings were used to support the proposed method. I have no background in reinforcement learning and I think I'm not qualified to evaluate the contribution & novelty of this work. I will be stick with feedback from other reviewers and change my rating appropriately.

Strengths

* Theoretical analysis was provided to justify the proposed method; * Experiments under diverse settings and domains endorse the effectiveness of the proposed method.

Weaknesses

* What's the benefit of using diffusion models rather than other probabilistic generative models, e.g. VAE, FLOW and GAN, for reward function extraction? * The drawback of using diffusion models is not mentioned. Will diffusion models slowdown the reward function calculation because iterative denoising is super time-consuming. Can the proposed method be applied to realtime applications?

Questions

See weekness

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

NA

Summary

This paper proposed a method for extracting reward functions from diffusion models, which can be used for sequential decision-making. The authors defined the notion of a relative reward function of two diffusion models and show conditions under which it exists and is unique. They also propose a practical learning algorithm for finding the relative reward function by aligning the gradients of a neural network to the difference in outputs of the base and expert diffusion models. They demonstrate the results on several domains, such as navigation, locomotion, and image generation, and show that it can recover the correct reward functions and improve the performance of the base models.

Strengths

The paper provides a thorough analysis and empirical evaluation of its approach. It also provides a clear explanation of the proposed method and its relation to previous work. It proposes a new method for extracting reward functions from diffusion models that does not require environment access, simulators, or iterative policy optimization. It is also agnostic to the architecture of the diffusion models used, being applicable to continuous and discrete models, and making no assumption of whether the models are unguided or whether they use either classifier guidance or classifier-free guidance.

Weaknesses

1. The paper assumes that the two diffusion models have the same architecture and are trained on the same data distribution. This may limit the applicability of the method to scenarios where the diffusion models have different architectures or are trained on different data distributions. 2. The paper does not provide any ablation studies on the hyperparameters of the proposed method. The paper does not provide any qualitative analysis or visualization of the learned reward functions.

Questions

Can this method be used to reduce images with bad quality in the diffusion process? How?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

NA

Summary

The paper introduces the following problem: given two diffusion models, a base model and an expert model, extract a guidance function such that the output distribution of the base model steered with the guidance function would resemble the output distribution of the expert model. The authors first formalize this guidance function as a notion of relative reward between the two diffusion models. They provide some theoretical definitions and results: first, that under certain mild conditions, there exists a unique guidance function rho that solves the relative reward problem; second, they provide the closed-form unique reward guidance function in an ideal setting (if given the ground truth score functions); third, in the general setting, in a space of vector fields with a desirable property, there exists a unique minimizer of "distance" to the ground truth relative reward and they call it the optimal relative reward gradient; finally, they define epsilon-relative reward functions as a guidance function that is epsilon-close to the optimal relative reward gradient. The general setting is necessary because we cannot guarantee that diffusion models perfectly capture the ground truth score function or that they are conservative (due to the approximation error). Finally the authors give an algorithm for approximating the relative reward function via L2 loss assuming access to the diffusion models but not to a simulator. The authors find using their learned reward as a classifier can reliably separate base vs expert demonstrations in Maze2D, improve performance in locomotion environments, and mitigate generation of shocking images by StableDiffusion.

Strengths

The authors introduce and define the novel problem of extracting a relative reward function between two diffusion models. There is substantial and rigorous theoretical analysis on its existence and uniqueness, first in an idealized setting and then in the general setting as the minimizer of a distance quantity over a space with a desirable property. The content is presented in a clear organization, introducing complexity step by step starting from the simple but unrealistic ideal setting. The method does not require access to the environment or simulators. The experiments include a diverse set of environments (maze, locomotion, image generation) and a diverse set of evaluation methods (discriminating base vs test, performance improvement, manual examination of images, ...).

Weaknesses

The paper aims to extract a reward function from diffusion models but narrows its scope without explanation or justification to a *relative* reward function between a base model and expert model. It would be nice to see the authors' reasoning on this: do they claim this is an equivalent problem (perhaps if the base model is chosen cleverly)? Is it simply an easier problem than extracting an "absolute" reward function? Or perhaps we use the relative setup but interpret the result as an absolute reward function? (The experiments seem to support this interpretation since they evaluate whether the learned reward improves performance rather than reproduces the expert model distribution.) The algorithm for approximating the relative reward function is fairly obvious; if we want to guide model 1 to be like model 2, given the simple additive structure of guidance, we can simply "undo" model 1's drift term and add model 2's. (This weakness is balanced by the strength of the rigorous theoretical analysis!) It seems that the experiments primarily assess directional effect of the reward guidance, i.e. whether guiding with the reward increases performance compared to without guidance. But the paper claims to extract a relative reward; in the idealized setting using the reward guidance should reproduce the exact distributions of the expert model at every denoising step, and in the general setting we should approximate the expert distributions. This is a higher bar than simply showing improved performance and separation of base vs guided distributions. It would be ideal to see results showing that the base model guided with reward is *indistinguishable* from the expert model (i.e. fools a discriminator).

Questions

- It's not clear to me whether the learned per-step reward function can indeed be used as a reward function beyond guiding diffusion models, so it would help to see experiments supporting this claim. Have you explored whether it can be directly used as the reward signal for an RL agent? If so it would help motivate the setting, otherwise it may not be as compelling since we are learning a guidance function that turns one diffusion model into another when we already have access to the expert model. - Related to the last point under weaknesses, I would like to see results comparing the performance of expert model vs base guided with reward, not just base guided vs base without guidance. - It would be super interesting to see a comparison of reward learning from this method versus IRL or other classical methods given a common dataset of expert demonstrations / examples. Do you have a hypothesis about the relative strengths of reward learning via diffusion vs RL? - Is there any way to implement your method using classifier-free guidance? - What would happen if you guided the expert model with the relative reward? Or an intermediate model that captures some knowledge from expert but not all? Would the reward guidance echo (overweight) the effect of the already-learned features or would it have no effect? - In the first sentence under section 4 (methods) you mention you establish the connection between the cumulative reward of a trajectory and the value p(y|x) used for classifier guidance, but I either missed it or the connection wasn't made explicit. Is y the optimality variable of the entire trajectory?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

2 fair

Limitations

Yes

Summary

The authors explore the applications for diffusion models in IRL, and state that they provided the first method for extracting relative reward functions (RRF) from a pair of diffusion models, where the models differ in quality/performance. They propose a simple but intuitive idea: the difference in score function between base and expert models tells us the relative reward function \rho, and we can train a model to predict this \rho directly. They build off of the work of Janner et al’s Diffuser, which previously introduced the concept of training a separate model to predict the cumulative rewards of trajectory samples The authors introduce 3 sets of experiments across 5 domains: classification of expert vs. non-expert trajectories on Maze2D using their trained relative reward function (RRF), higher dimensional control problems (e.g. half cheetah) using classifier guidance (steered by RRF), and classification of synthesized images (with pretrained StableDiffusion network vs. SafeStableDiffusion.) They also provide proofs as to why the RRF is mathematically sound. I did not have time to check the proof in detail.

Strengths

Generally well written, grammatically correct. Impressive separability between predicted rewards for generic vs. safe image synthesis models (Figure 2). Shows 36% improvement for Walker2D with steering from RRF. Analysis on 5 different domains (4 RL domains, and I2P dataset).

Weaknesses

- How useful is “extracting relative reward functions from two diffusion models”? Applications here don’t seem quite compelling, would be interesting to add a section discussing other potential applications. - How often would we be given a poor-performing diffusion policy and an expert diffusion policy? What are we trying to solve, and why? Would be helpful to discuss the applications further. Would “image appropriateness classification”, given two models, be a marquee application, and why? - Diffuser already introduce the idea of extracting the cumulative reward function. When is relative reward superior? No discussion of this. - Limited numbers of baselines – e.g. no comparison to other RRF techniques, or BC, or Janner et al numbers. Janner et al compare with many baselines (BC CQL IQL DT TT MOPO MOReL MBOP vs. their own). No quantifiable comparison with Diffuser? - No comparison to reward heatmaps learned in other ways, such as using method of Janner et al - Appendix is unclear about details given about which diffusion framework used, start/end of schedules, etc? Was Karras-style preconditioning used? - Requires access to the training dataset for marquee algorithm (Algorithm 1) Why is Algo 2 in the appendix, and not in main paper? Seems like the more useful version, than Algo 1? - Hard to tell whether it is actually unsafe content since it’s blurred out, overblurred to be not so convincing - Quantitative results and eval are very sparse. Figure 2: Histogram caption should say on which dataset . One qualitative histogram is not so easy to tell about quantitative performance – why no distributional distances measured? - No comparison with other baselines for the safe image generation experiment. - Missing an architecture figure with the two models? Need some sort of graphical illustration / system figure showing training vs. inference behavior and which models are being learned (too much text) - Bit too much background on diffusion, overpowers contribution. I recommend that the proof on page 5 can be moved to appendix. Main algorithm is only shown on page 7, seems it should be on page 3 - Content placement: Figure 2 is very far from where it is discussed (page 7 vs. page 9) - Line 278 – “peaks occurring at true goal” -> no discussion of within what neighborhood / distance? - Weird italicization and capitalization: “Physics informed neural networks” References seem not complete / unpolished / often incorrect attribution - Keeps citing Sohl-Dickstein [59] for classifier guidance L24, L47, L143, why not cite Dhariwal et al “Diffusion Models Beat GANs on Image Synthesis”? Where does Sohl-Dickstein mention anything about classifier guidance? - No author listed for “Artificial Neural Networks for Solving Ordinary and Partial Differential Equations” I. E. Lagaris, A. Likas, D. I. Fotiadis - No 2022 year listed for [31] Karras, T., Aittala, M., Aila, T. and Laine, S. [n.d.], Elucidating the design space of diffusion-based 418 generative models, in ‘Advances in Neural Information Processing Systems’. - No 2021 year listed for [63] “[63] Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S. and Poole, B. [n.d.], Score-based 489 generative modeling through stochastic differential equations, in ‘International Conference on Learning 490 Representations’.” - Janner et al bib entry is wrong — accepted to ICML 2022 (correct entry is here: https://diffusion-planning.github.io/files/bib.txt) - Not dated for [21] Ho, J. and Salimans, T. [n.d.], Classifier-free diffusion guidance, in ‘NeurIPS 2021 Workshop on Deep 398 Generative Models and Downstream Applications’. - L321 incorrect statement “Safe Stable Diffusion [58], a modified version of Stable Diffusion designed to mitigate the generation of shocking images.” -> See Table 1 of https://arxiv.org/pdf/2211.05105.pdf, shocking is just one of 7 categories of inappropriate images. Inappropriate images are the focus

Questions

- Is there any concept of “relative reward function” in existing literature? - L264 why are 32 expert models required? - L263 why are 8 expert models required? - Why is access to the original dataset required for Algorithm 1? In what real-world settings would this be available? - Is no steering occurring in Section 5.1, on Maze2D? - L334 why use “the dataset to generate sets of image embeddings rather than actual images”?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

No, the authors do not discuss any limitations of their work, and they do not discuss any potential negative societal impact. No Neurips checklist is provided. My understanding was that this checklist is a requirement.

Thank you for the authors’ responses

Thank you to the authors for their detailed replies and time spent providing clarifications. Thanks also for the clarification regarding the use of reward in Janner et al. While novel and of some interest theoretically, the method appears to be of limited value for the locomotion experiments, as one could simply use the expert model (and it would be preferred). In addition, assumptions introduce limitations, such as reliance on access to an expert model. I see more value on the interpretability side and content moderation side, but I don’t think the applications demonstrated in the paper indicate convincingly useful applications. The paper seems to be focusing on how one can train a new model with the RRF. If you already have the expert model, why would you want to do that, especially if the method doesn’t allow surpassing the performance of the expert/teacher model? In this case it appears the main value is in model interpretability, which I think should instead be the focus of the paper (main contribution and featured in abstract), whereas the locomotion experiments are presented as the marquee results. I don’t believe the authors have responded to my question regarding “ access to the original dataset required for Algorithm 1…In what real-world settings would this be available?” I also believe the authors have not responded to my question regarding: “Keeps citing Sohl-Dickstein [59] for classifier guidance L24, L47, L143, why not cite Dhariwal et al “Diffusion Models Beat GANs on Image Synthesis”? Where does Sohl-Dickstein mention anything about classifier guidance?” I was hoping the authors would have been willing to provide something based on my comment “Missing an architecture figure with the two models? Need some sort of graphical illustration / system figure showing training vs. inference behavior and which models are being learned (too much text)”, but they have not addressed this. I agree with Reviewer qcPp that “comparison of reward learning from this method versus IRL or other classical methods given a common dataset of expert demonstrations / examples” would strengthen the paper, and I appreciate the authors providing those experiments in their rebuttal. If included, that would have strengthened the paper. I think IRL should be discussed more in the related work section (just 1 sentence devoted to this, but it appears to be the conceptual foundation of the work — derive rewards from expert behavior). I also agree with their assessment that RRF is interpreted as an absolute reward function when training new models, but this is not made clear in the text. (Unclear what “medium-expert” is in the authors’ response). The work appears technically correct. Given the interesting theoretical approach, I believe the work and experiments could be revised and make a strong submission in a subsequent conference, but I keep my prior score. Thank you.

Thank you for the responses

Thank you for the explanation that if the base model represents unguided behavior (for example uniform random) then the learned reward is essentially absolute, this has addressed my question. This is a crucial detail that would greatly benefit readers if it were emphasized in the paper. I do not believe the authors adequately addressed my comment that while the paper claims to be able to match the denoising process at every step, the experiments only assess that using the learned reward improves performance. The experiments are a couple steps away from the theory. Thanks for the clarification that the method works for classifier-based and classifier-free guidance, and for pointing out that you already have experiments demonstrating your method on the classifier-free stable diffusion model. I think it would help to mention the fact that the method works for CG and CFG. Thanks also for the additional experiments demonstrating that 1) PPO agents can indeed be trained successfully using the learned reward, 2) that this outperforms AIRL in maze environments, and 3) that the reward can also improve medium-expert models. In 3) I think the small performance boost in the expert model likely is not statistically significant because the reward only contains information the expert model already has. The additional experiments certainly help me better understand the performance of the method, but I agree with reviewer CPFx that the primary concern is the unrealistic requirement of having an expert diffusion model. In order to obtain such a model, you would need to have the reward in hand already, in which case it doesn't make sense to extract a reward back out again. AIRL differs in that it only assumes access to optimal demonstrations, which is more realistic because we might have human experts provide demonstrations, but we don't have an accurate diffusion model of human demonstrations, and if we did then we would have already solved our problem. Overall I think the work is certainly novel, technically sound (the theory is solid), and the authors ran an impressive set of experiments, but the unrealistic requirement of having an expert diffusion model has not been motivated satisfactorily. If this were better addressed or the main contribution of the paper was reframed differently than reward learning (for example interpretability as reviewer CPFx suggests), the strength of this submission would certainly be improved. I will keep my score as is.