Goal-Driven Reward by Video Diffusion Models for Reinforcement Learning

TL;DR

This paper introduces a novel reward mechanism for reinforcement learning that leverages pretrained video diffusion models to provide goal-driven signals without manual reward design, improving task generalization.

Contribution

It proposes using pretrained video diffusion models as reward functions, employing their latent representations and CLIP-based frame relevance for more effective goal guidance in RL.

Findings

Effective in Meta-World and Distracting Control Suite environments.

Improves goal achievement without manual reward engineering.

Utilizes domain-specific fine-tuning of video diffusion models.

Abstract

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 Meta-World and Distracting Control Suite demonstrate the effectiveness of our approach.