Diffusion Spectral Representation for Reinforcement Learning

TL;DR

This paper introduces Diffusion Spectral Representation (Diff-SR), a novel framework that leverages diffusion models for reinforcement learning by focusing on representation learning, enabling efficient policy optimization without costly sampling.

Contribution

The paper proposes Diff-SR, a new method that extracts value function representations from diffusion models, bypassing sampling costs and improving efficiency in RL tasks.

Findings

Diff-SR achieves robust performance across various RL benchmarks.

It enables efficient policy optimization without extensive diffusion sampling.

The approach works well in both fully and partially observable environments.

Abstract

Diffusion-based models have achieved notable empirical successes in reinforcement learning (RL) due to their expressiveness in modeling complex distributions. Despite existing methods being promising, the key challenge of extending existing methods for broader real-world applications lies in the computational cost at inference time, i.e., sampling from a diffusion model is considerably slow as it often requires tens to hundreds of iterations to generate even one sample. To circumvent this issue, we propose to leverage the flexibility of diffusion models for RL from a representation learning perspective. In particular, by exploiting the connection between diffusion models and energy-based models, we develop Diffusion Spectral Representation (Diff-SR), a coherent algorithm framework that enables extracting sufficient representations for value functions in Markov decision processes (MDP) and partially observable Markov decision processes (POMDP). We further demonstrate how Diff-SR facilitates efficient policy optimization and practical algorithms while explicitly bypassing the difficulty and inference cost of sampling from the diffusion model. Finally, we provide comprehensive empirical studies to verify the benefits of Diff-SR in delivering robust and advantageous performance across various benchmarks with both fully and partially observable settings.