Source-linked AI summary
Diffusion Reward: Learning Rewards via Conditional Video Diffusion
Tao Huang, Guangqi Jiang, Yanjie Ze, Huazhe Xu
TL;DR
Reward specification for complex visual RL is difficult because dense rewards are costly to design and sparse rewards provide limited supervision. Diffusion Reward pretrains a conditional video diffusion model on expert videos and uses negative history-conditioned entropy, together with exploration and sparse environmental rewards, to guide learning. Across 10 visual robotic manipulation tasks, it improves performance over the best baselines, generalizes to unseen tasks, and produces reasonable rewards for real robot settings.
Problem
Dense reward design is labor-intensive or infeasible, while sparse rewards provide insufficient supervision for efficient reinforcement learning.
Source: §1 · p. 1Method
Diffusion Reward uses conditional video diffusion on expert videos and negative history-conditioned entropy, combined with exploration and sparse environmental rewards, as a dense RL signal.
Source: §1 · p. 2, §4.2 · p. 8Results
Across 10 visual robotic manipulation tasks, Diffusion Reward improves performance by 38% on MetaWorld and 35% on Adroit over the best-performing baselines, while also showing unseen-task and real-robot applicability.
Source: §1 · p. 3, §6 · p. 14Takeaways & Limitations
Conditional video diffusion can provide informative dense rewards that guide visual RL toward expert behaviors and support generalization beyond pretrained tasks.
Source: §1 · pp. 2–3Computing the entropy is intractable without an explicit conditional distribution, so the method estimates a variational entropy bound.
Source: §4.2 · p. 7
Abstract
from arXiv · showhide
Learning rewards from expert videos offers an affordable and effective solution to specify the intended behaviors for reinforcement learning (RL) tasks. In this work, we propose Diffusion Reward, a novel framework that learns rewards from expert videos via conditional video diffusion models for solving complex visual RL problems. Our key insight is that lower generative diversity is exhibited when conditioning diffusion on expert trajectories. Diffusion Reward is accordingly formalized by the negative of conditional entropy that encourages productive exploration of expert behaviors. We show the efficacy of our method over robotic manipulation tasks in both simulation platforms and the real world with visual input. Moreover, Diffusion Reward can even solve unseen tasks successfully and effectively, largely surpassing baseline methods. Project page and code: https://diffusion-reward.github.io.
1 Introduction
Reward specification in visual RL is difficult because dense rewards are laborious to design while sparse rewards provide insufficient supervision. Diffusion Reward addresses this by using conditional video diffusion to learn dense rewards that guide exploration toward expert-like behaviors and outperform baselines across robotic manipulation tasks.
Motivation
- Learning rewards from expert videos offers dense task-execution information with lower collection effort than manually specifying rewards.This addresses the labor of dense reward design and the weak supervision of sparse rewards.Source: §1 · p. 1
- Video diffusion models are used because prior video-prediction reward methods struggle to model complex expert-video distributions and can assign lower returns to aligned out-of-distribution expert videos.Figure 2 motivates stronger video modeling and conditional-entropy-based rewards.Source: §1 · p. 2, Figure 2 · p. 3
Method
- Diffusion Reward estimates conditional entropy from video diffusion processes conditioned on historical frames, using its negative to encourage exploration of expert-like behaviors.The framework also combines standardized entropy with novelty-seeking and sparse environmental rewards.Source: §1 · p. 2, Figure 1 · p. 2
Reward analysis
- The learned reward distinguishes expert behaviors and varied-quality trajectories across seen and unseen MetaWorld tasks.The reward analysis aggregates results over 7 tasks and evaluates trajectories with different qualities.Source: Figure 3 · p. 3
Empirical validation
- 38% and 35% performance improvements over the best-performing baselines were achieved on 7 MetaWorld and 3 Adroit visual manipulation tasks, respectively.The evaluation covers 10 tasks across simulation domains.Source: §1 · p. 3
Generalization and real-world applicability
- The pretrained reward model can guide RL on unseen tasks and produces reasonable rewards for real robot tasks.The paper also reports associated favorable offline RL performance for real-world applicability.Source: §1 · p. 3
2 Related Work
Prior reward-learning methods use goals, adversarial discrimination, or video-model likelihoods, but can require additional information or struggle with complex dynamics. Diffusion Reward instead pretrains action-free rewards from expert videos using historical-frame-conditioned video diffusion.
Learning rewards from videos
- Goal-image distance methods can provide dense rewards, but obtaining goal images outside simulation limits their applicability to open-world tasks.Generative approaches avoid requiring future information when extracting rewards from videos.Source: §2 · p. 3
- Earlier video-prediction methods use predicted log-likelihood as rewards but struggle to model complex expert-video distributions, producing unproductive rewards.The limitation is especially relevant to tasks with intricate dynamics.Source: §2 · pp. 3–4
Diffusion models for RL
- Unlike diffusion-based planners, Diffusion Reward uses video diffusion conditioned on historical frames to learn rewards for online RL rather than directly generating policies or plans.It also differs from methods that fit separate diffusion models to different behaviors by requiring no action labels or extra task-agnostic data.Source: §2 · p. 4
Pretraining for RL
- Diffusion Reward pretrains rewards solely from action-free expert-only videos, avoiding additional rewards, action labels, or task-agnostic data.This distinguishes it from several policy-pretraining and behavioral-prior approaches.Source: §2 · p. 4
3 Preliminaries
The paper formulates visual RL with RGB observations and sparse binary rewards, then assumes access to unlabeled, action-free expert videos for reward pretraining. It also motivates video diffusion models by their stronger video-generation quality and lower diversity on seen expert trajectories.
Problem formulation
- The RL objective is to learn a policy maximizing discounted return over a finite-horizon Markov decision process.The MDP is specified by state space, action space, transition function, reward function, and discount factor.Source: §3 · p. 4
Video diffusion models
- Video diffusion models generate higher-quality videos than VideoGPT on real dexterous manipulation and simulation videos from seven MetaWorld tasks.The comparison uses real robotic manipulation videos and simulation videos, as summarized in Table 1.Source: Table 1 · p. 5
Key insight
- Seen expert trajectories exhibit the lowest generative diversity across seven MetaWorld tasks, with diversity proportional to distances among generated videos.This observation supports using conditional generative diversity to distinguish expert-like behavior.Source: Table 2 · p. 5
Problem formulation
- RGB images serve as the agent’s high-dimensional observations, while task success is represented by a binary sparse reward.This setting reflects robotics applications where reward engineering is difficult and visual RL can require substantial environment interaction.Source: §3 · pp. 4–5
Expert videos
- The method assumes unlabeled, action-free expert videos collected across multiple tasks without task identification, used to pretrain rewards before online exploration.Each task has a set of expert videos consisting of image trajectories.Source: §3 · p. 5
4 Method
Diffusion Reward models expert videos with conditional video diffusion and uses negative conditional entropy to reward expert-like behavior while supporting exploration. The method combines this learned reward with novelty-seeking and sparse task rewards for visual RL.
4 Method
- Diffusion Reward trains a conditional video diffusion model on expert videos and estimates history-conditioned entropy, using its negative to guide RL toward expert-like behaviors.The framework models video in a compressed latent space and conditions generation on historical frames.Source: §4 · p. 5, §4.1 · pp. 5–6, §4.2 · p. 7
4.3 Training Details
- The implementation uses VQ-GAN and VQ-Diffusion with two historical frames, ten denoising steps, one diffusion sample, and task-dependent reward coefficients.The default coefficient is 0.95 except for Pen, where it is reduced to 0 to avoid misleading exploration.Source: §4.1 · p. 6, §4.3 · p. 8
4.2 Conditional Entropy as Rewards
- Conditional diffusion exhibits lower generation diversity on expert-like trajectories, motivating entropy-based discrimination between expert-like and unexpert-like behaviors.The resulting reward favors trajectories whose conditioned future generations are less diverse.Source: §4.2 · p. 7
- Because exact conditional entropy is intractable, the method estimates a variational bound through randomized denoising samples and standardizes the resulting entropy reward across expert videos.The estimator uses a closed-form distribution compatible with VQ-Diffusion and repeats sampling M times.Source: §4.2 · pp. 7–8
- The final reward combines standardized conditional entropy, RND novelty, and sparse task reward as r^diff = (1-α)·r̄^ce + α·r^rnd + β·r^spar.The sparse reward remains important for complex tasks, while β is set to 1 across tasks.Source: §4.2 · p. 8
4.3 Training Details
- The method targets ten visual manipulation tasks with sparse rewards, using image observations from MetaWorld and Adroit.The experiments cover seven gripper tasks and three dexterous-hand tasks.Source: Figure 5 · p. 9
5 Experiments
Across simulation, unseen tasks, and real-robot settings, Diffusion Reward improves visual RL by combining expert-video guidance with exploration. Its strongest reported gains include 38% and 35% improvements over the best baselines on MetaWorld and Adroit.
5.1 Experimental Setup
- The evaluation covers 10 visual manipulation tasks with sparse rewards, including seven MetaWorld gripper tasks and three Adroit dexterous tasks.Baselines include sparse reward, RND, AMP, VIPER, and VIPER-std under matched RL settings.Source: §5.1 · p. 9
5.2 Main Results
- 38% and 35% performance improvements over the best-performing baselines are reported on MetaWorld and Adroit, respectively.The comparison uses the same training steps across the two simulation domains.Source: §5.2 · pp. 10–11, Figure 7 · p. 10
5.3 Zero-shot Reward Generalization
- On five unseen MetaWorld tasks, Diffusion Reward generalizes without additional tuning, produces reasonable rewards, and largely outperforms other baselines.Its retained discrimination between expert-like and expert-unlike behaviors also helps constrain RND exploration.Source: §5.3 · p. 11, Figure 8 · p. 11
5.4 Real Robot Evaluation
- The learned reward assigns higher returns to expert than random real-robot videos and supports a policy that presses a button at an unseen position.Offline RL on Button Press significantly outperforms VIPER in real-world applicability.Source: Figure 6 · p. 9, §5.4 · p. 12, Figure 10 · p. 12
5.5 Ablation Studies
- Ablations show that diffusion models paired with conditional entropy outperform transformer or log-likelihood alternatives, especially because they better capture complex expert-video distributions.Conditional entropy is also reported to provide more productive exploration and more generalizable reward inference.Source: §5.5 · pp. 12–13, Figure 11 · p. 13
- Approximately 10 denoising steps, stochastic reverse-process sampling, α around 0.95, and one or two context frames provide the strongest reported settings.More denoising slows inference, deterministic rewards degrade performance, extreme α values reduce performance, and longer contexts can hurt generalization.Source: §5.5 · pp. 13–14
6 Conclusion
Diffusion Reward extracts dense RL rewards from a pretrained conditional video diffusion model and shows improvements across simulation, real-robot, and unseen-task settings. The appendix specifies its architectures, entropy estimation, inference-speed choice, and baseline implementations.
6 Conclusion
- Diffusion Reward extracts dense rewards from a pretrained conditional video diffusion model and demonstrates improvements across MetaWorld and Adroit, real-robot tasks, and unseen tasks.The framework combines standardized entropy with exploration and sparse environmental rewards, and its unseen-task results surpass baseline methods.Source: §6 · p. 14
- Future work targets larger and broader diffusion datasets, additional modalities such as language, improved entropy–exploration balancing, better entropy estimation, and faster reward inference.These directions define the current scope boundary of the presented system.Source: §6 · p. 14
A Implementation Details
- All methods share the same RL backbone and sparse task reward, with other settings held identical except reward pretraining when applicable.The implementation uses DrQv2 for downstream RL and separately implements RND, VIPER, and AMP baselines.Source: A Implementation Details · p. 19, §A.1 · p. 19, §A.2 · p. 20
A.1 Diffusion Reward Implemenatation
- The conditional diffusion architecture tokenizes each historical frame into 8×8 tokens, concatenates them across frames, and produces a 128 × 1024 condition embedding.The embedding is passed to the subsequent conditional diffusion model.Source: §A.1 · p. 19
- Entropy estimation uses a variational bound with a closed-form distribution, while setting M = 1 retains discrimination between expert-like and unlike behaviors with faster inference.The appendix points to additional analysis of the sampling choice.Source: §A.1 · pp. 19–20
A.3 Real-world Offline RL Implementations
The real-world offline RL evaluation uses expert demonstrations from a Franka Button Press task to train reward models for relabeling offline data. Diffusion-process ablations further show that sampling choices affect exploration quality.
A.3 Real-world Offline RL Implementations
- Fifty expert videos from a Franka Button Press task train reward models that relabel offline data for vision-based DrQ+BC, with suboptimal data containing 50% random actions.Videos are resized to 84 × 84 resolution, and default DrQ+BC hyperparameters are used.Source: §A.3 · p. 20
- Using four denoised samples can hinder exploration through low estimated-entropy variance, while an appropriate sampling-noise scale produces more productive exploration.The ablation results average three seeds with shaded standard error, and red denotes the default.Source: Figure 13 · p. 21
B Task Descriptions
The benchmark contains ten diverse visual robotic manipulation tasks spanning MetaWorld gripper control and Adroit dexterous manipulation. Tasks use RGB observations and sparse environmental rewards with task-dependent expert-video collection.
B Task Descriptions
- The benchmark includes 7 MetaWorld gripper-manipulation tasks and 3 Adroit dexterous-manipulation tasks selected for diverse objects and skills.The task set includes examples such as Assembly, Coffee Push, Dial Turn, Hammer, Pen, and Door.Source: B Task Descriptions · p. 21
- All tasks provide 64 × 64 RGB observations and sparse environmental rewards for visual reinforcement learning.The task descriptions cover gripper actions in MetaWorld and higher-dimensional dexterous actions in Adroit.Source: B Task Descriptions · p. 21
- Expert-video collection varies by task complexity, with 20 videos per MetaWorld task and 50 videos per Adroit task.The demonstrations support learning from both gripper and dexterous manipulation behaviors.Source: B Task Descriptions · p. 21
C.1 Video Prediction Results
Video diffusion captures complex expert-video distributions and generalizes to unseen trajectories, while entropy-based rewards better discriminate behaviors than log-likelihood rewards. Qualitative errors identify a remaining exploration bottleneck.
C.1 Video Prediction Results
- The model captures complex expert-video distributions and generalizes to unseen expert videos, but sometimes mispredicts colored target points in Door Close and Reach.These errors may contribute to slower exploration during initial RL training and motivate more powerful diffusion models.Source: §C.1 · p. 22
- Video diffusion generalizes better than VideoGPT on unseen real-robot and simulation trajectories across SSIM, PSNR, and LPIPS, producing more informative rewards.The comparison is reported in Table 1.Source: §C.1 · p. 22
C.2 Reward Analysis
- Entropy-based rewards outperform log-likelihood rewards across video models, despite suboptimal Adroit behaviors sometimes receiving higher diffusion-model rewards than VideoGPT.The passage attributes this domain difference to possible variation in action spaces, scenes, backgrounds, and objects.Source: §C.2 · p. 22
D.1 More Ablations
Ablations show that sparse task-completion supervision and appropriately scaled entropy-related rewards are important for effective exploration, while additional denoising samples add cost without clear benefit. Comparisons further indicate that reward pretraining and generalization remain strong relative to alternatives.
D.1 More Ablations
- Performances decrease dramatically without sparse environmental rewards, showing that task-completion supervision is necessary for complex manipulation.Door performance is relatively more favorable because RND exploration and sparse reward overlap when opening the door.Source: §D.1 · p. 23, Figure 15 · p. 23
- An intermediate diffusion noise scale supports productive exploration, whereas very low scales encourage exploitation and very high scales produce random exploration.The ablation links noise scale to the variance of estimated entropy and the resulting exploration behavior.Source: §D.1 · p. 23
- Using one denoised sample is sufficient for informative rewards and high inference speed, while increasing samples from 1 to 4 reduces FPS from 87.7 to 45.9.Increasing samples from 1 to 3 has slight performance influence, whereas four samples can stall learning midway despite satisfactory asymptotic performance.Source: §D.1 · pp. 23–24
- A single coefficient balancing entropy-based and RND rewards has negligible effect, but the sparse task-reward coefficient must be appropriately scaled.The authors set β = 1 because a low value may fail to provide sufficient task-completion signals.Source: Table 3 · p. 24, Figure 16 · p. 24
D.2 More Comparison
- Visual pretraining does not remove the difficulty of vision-based manipulation, while the proposed reward-pretraining method remains superior in the Hammer comparison.The comparison contrasts the method with an alternative using VQ-GAN visual representation pretraining.Source: §D.2 · p. 24, Figure 18 · p. 24
- The comparisons include additional demonstration-count, generalization, and action-free-video settings, with Diffusion Reward outperforming VIPER on most unseen tasks.The supplied passages also emphasize that reward learning uses action-free expert-only videos rather than additional reward labels, action labels, or task-agnostic data.Source: Figure 17 · p. 24, Figure 19 · p. 25, §D.2 · p. 25
D.3 More Generalization Experiments
Additional generalization experiments test transfer across tasks and against AMP under a pseudo-generalization setup. Diffusion Reward outperforms competing methods in these comparisons while AMP shows fair transfer under a different data and interaction requirement.
D.3 More Generalization Experiments
- Diffusion Reward significantly outperforms VIPER on most of three additional unseen MetaWorld tasks.It also outperforms RND in Reach Wall during the initial training stage in terms of productive exploration.Source: §D.3 · p. 25
- AMP achieves fair transfer from Door Close to Drawer Open, but requires additional online interactions on the source task and remains below the other methods.This is termed pseudo-generalization because AMP does not use target-task expert videos but differs from the default setting through source-task interactions.Source: §D.3 · p. 25, Figure 20 · p. 25
E Visualization of Reward and Trajectory in Simulation and Real Robot
Visualization experiments show that the learned reward distinguishes expert-like from expert-unlike behavior in simulation and assigns higher rewards to successful real-robot demonstrations. A simulation failure mode suggests that difficult environmental dynamics remain a modeling boundary.
Simulation Results
- The learned reward greatly distinguishes expert-like and expert-unlike behaviors across 10 simulated tasks.The reward curve is visualized for tasks from MetaWorld and Adroit.Source: E Visualization of Reward and… · p. 25, Figure 22 · p. 30
- Two door-opening tasks show a return drop during final execution, attributed to difficulty modeling door dynamics.The authors identify explicit environmental-dynamics modeling as a future direction.Source: E Visualization of Reward and… · pp. 25–26
Real Robot Results
- On one real-robot bowl-picking task, Diffusion Reward assigns relatively higher rewards to expert demonstrations and lower rewards to random trajectories.The evaluation uses 20 trajectories collected with an Allegro hand, Franka arm, and RealSense D435i camera, split evenly between successful and random trajectories.Source: E Visualization of Reward and… · p. 26, Figure 23 · p. 31
Simulation Results
- The visualizations include ground-truth and predicted video frames distinguished by blue and orange borders.This figure provides a visual comparison of video prediction results rather than a reported reward outcome.Source: Figure 21 · p. 29