Diffusion Reward: Learning Rewards via Conditional Video Diffusion

TL;DR

This work tackles the challenge of specifying dense rewards for vision-based reinforcement learning by learning from expert videos. It introduces Diffusion Reward, which trains a conditional video diffusion model on expert trajectories and uses the negative of the history-conditioned entropy as a reward, augmented with a novelty-based exploration term and the environment's sparse reward. Empirical results across ten robotic manipulation tasks in simulation and on real robots show strong gains over baselines and good zero-shot generalization to unseen tasks, with favorable offline RL performance. The approach leverages large-scale diffusion modeling to extract informative, generalizable reward signals from unlabeled videos, highlighting a promising direction for reward pretraining in robotics and beyond.

Abstract

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.

Figures (23)

• Figure 1: Overview of Diffusion Reward. (left) We present a reward learning framework in RL using video diffusion models pretrained with expert videos. We perform diffusion processes conditioned on historical frames to estimate conditional entropy as rewards to encourage RL exploration of expert-like behaviors. (right) The mean success rate of 10 visual robotic manipulation tasks demonstrates the effectiveness of our proposed Diffusion Reward over 5 runs. Shaded areas are standard errors.
• Figure 2: Analysis of video models and rewards.VIPER (CE) and Diffusion Reward (LL) replace original rewards with conditional entropy (CE) and log-likelihood (LL), respectively. CE-based rewards assign near-optimal rewards to unseen expert videos. Such a boost is enhanced by the strong modeling ability of diffusion models. Results are averaged over 7 MetaWorld tasks. Suboptimal represents videos with 25% randomly-taken actions.
• Figure 3: Reward analysis. (left) Our learned rewards are aggregately higher for expert behaviors over 7 tasks from MetaWorld. (right) Varied-quality trajectories can be distinguished from seen and unseen tasks. (top) Examples of trajectories.
• Figure 4: Diffusion Reward for online RL
• Figure 5: Task visualization. We evaluate methods on 10 challenging visual RL tasks from MetaWorld and Adroit with visual input and sparse rewards. Tasks are chosen to cover a wide range of manipulation skills.