{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/extracting-reward-functions-from-diffusion","title":"Extracting Reward Functions from Diffusion Models","arxiv_id":"2306.01804","date":"2023-06-01","proceeding":"NeurIPS 2023 11","authors":["Felipe Nuti","Tim Franzmeyer","João F. Henriques"],"abstract":"Diffusion models have achieved remarkable results in image generation, and have similarly been used to learn high-performing policies in sequential decision-making tasks. Decision-making diffusion models can be trained on lower-quality data, and then be steered with a reward function to generate near-optimal trajectories. We consider the problem of extracting a reward function by comparing a decision-making diffusion model that models low-reward behavior and one that models high-reward behavior; a setting related to inverse reinforcement learning. We first define the notion of a relative reward function of two diffusion models and show conditions under which it exists and is unique. We then devise a practical learning algorithm for extracting it by aligning the gradients of a reward function -- parametrized by a neural network -- to the difference in outputs of both diffusion models. Our method finds correct reward functions in navigation environments, and we demonstrate that steering the base model with the learned reward functions results in significantly increased performance in standard locomotion benchmarks. Finally, we demonstrate that our approach generalizes beyond sequential decision-making by learning a reward-like function from two large-scale image generation diffusion models. The extracted reward function successfully assigns lower rewards to harmful images.","url_abs":"https://arxiv.org/abs/2306.01804v2","url_pdf":"https://arxiv.org/pdf/2306.01804v2.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"extracting-reward-functions-from-diffusion","repo_url":"https://github.com/FelipeNuti/diffusion-relative-rewards","is_official":1,"mentioned_in_paper":0,"mentioned_in_github":0,"framework":"jax","reach":null}],"tasks":[{"task_slug":"decision-making","task_name":"Decision Making"},{"task_slug":"image-generation","task_name":"Image Generation"},{"task_slug":"sequential-decision-making","task_name":"Sequential Decision Making"}],"methods":[{"method_slug":"base","method_name":"BASE"},{"method_slug":"diffusion","method_name":"Diffusion"}],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2306.01804","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2306.01804"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/FelipeNuti/diffusion-relative-rewards","reach":null}],"summary":{"ran":1},"by_repo_kind":{"official":{"samples":1,"ran":1,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":1,"samples":[{"code_sha256_prefix":"a9464f634bca0d54","entry":"GradientRewardRegressor","repo":"FelipeNuti/diffusion-relative-rewards","repo_kind":"official","path":"safe-stable-diffusion/rrf_diffusion/gradient_matching.py","file_url":"https://github.com/FelipeNuti/diffusion-relative-rewards/blob/HEAD/safe-stable-diffusion/rrf_diffusion/gradient_matching.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"a9464f634bca0d54"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}