Skip to content

qiwang067/GenReward

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

11 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

[CVPR 2026] Goal-Driven Reward by Video Diffusion Models for Reinforcement Learning

Qi Wang* · Mian Wu* · Yuyang Zhang* · Mingqi Yuan · Wenyao Zhang · Haoxiang You
Yunbo Wang · Xin Jin† · Xiaokang Yang · Wenjun Zeng

Teaser image

Reinforcement Learning (RL) has achieved remarkable success in various domains, yet it often relies on carefully designed programmatic reward functions to guide agent behavior. Designing such reward functions can be challenging and may not generalize well across different tasks. To address this limitation, we leverage the rich world knowledge contained in pretrained video diffusion models to provide goal-driven reward signals for RL agents without ad-hoc design of reward. Our key idea is to exploit off-the-shelf video diffusion models pretrained on large-scale video datasets as informative reward functions in terms of video-level and frame-level goals. For video-level rewards, we first finetune a pretrained video diffusion model on domain-specific datasets and then employ its video encoder to evaluate the alignment between the latent representations of agent's trajectories and the generated goal videos. To enable more fine-grained goal-achievement, we derive a frame-level goal by identifying the most relevant frame from the generated video using CLIP, which serves as the goal state. We then employ a learned forward–backward representation that represents the probability of visiting the goal state from a given state–action pair as frame-level reward, promoting more coherent and goal-driven trajectories. Experiments on various Meta-World tasks demonstrate the effectiveness of our approach.

Getting Started

1) Create an environment

conda create -n genreward python=3.11
conda activate genreward

2) Install dependencies

pip install -r requirements.txt

3) Environment-specific setup

For Meta-World environments:

pip install git+https://github.com/rlworkgroup/metaworld.git@a0009ed9a208ff9864a5c1368c04c273bb20dd06#egg=metaworld

Training

Meta-World Tasks

# Pick place task
bash run_gc_mt_pick.bash

# Pick out of hole task
bash run_gc_mt_pick_hole.bash

# Bin picking
bash run_gc_mt_pick_bin.bash

Citation

If you find this work useful in your research, please consider citing:

@inproceedings{wang2026goal,
  title={Goal-Driven Reward by Video Diffusion Models for Reinforcement Learning},
  author={Wang, Qi and Wu, Mian and Zhang, Yuyang and Yuan, Mingqi and Zhang, Wenyao and You, Haoxiang and Wang, Yunbo and Jin, Xin and Yang, Xiaokang and Zeng, Wenjun},
   booktitle={CVPR},
  year={2026}
}

Credits

This implementation builds upon:

Thanks to all the authors for their excellent work!

About

[CVPR 2026] PyTorch code for the paper "Goal-Driven Reward by Video Diffusion Models for Reinforcement Learning"

Topics

Resources

Stars

2 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors