Exploratory Diffusion Model for Unsupervised Reinforcement Learning

TL;DR

This paper introduces ExDM, a diffusion model-based approach for unsupervised reinforcement learning that enhances exploration and enables rapid adaptation to downstream tasks by modeling explored data distribution.

Contribution

The paper proposes the Exploratory Diffusion Model (ExDM), leveraging diffusion models for better data representation and exploration in unsupervised RL, with theoretical analysis and practical algorithms.

Findings

ExDM outperforms state-of-the-art methods in exploration efficiency.

ExDM enables faster adaptation to downstream tasks.

ExDM effectively handles complex environments.

Abstract

Unsupervised reinforcement learning (URL) aims to pre-train agents by exploring diverse states or skills in reward-free environments, facilitating efficient adaptation to downstream tasks. As the agent cannot access extrinsic rewards during unsupervised exploration, existing methods design intrinsic rewards to model the explored data and encourage further exploration. However, the explored data are always heterogeneous, posing the requirements of powerful representation abilities for both intrinsic reward models and pre-trained policies. In this work, we propose the Exploratory Diffusion Model (ExDM), which leverages the strong expressive ability of diffusion models to fit the explored data, simultaneously boosting exploration and providing an efficient initialization for downstream tasks. Specifically, ExDM can accurately estimate the distribution of collected data in the replay buffer with the diffusion model and introduces the score-based intrinsic reward, encouraging the agent to explore less-visited states. After obtaining the pre-trained policies, ExDM enables rapid adaptation to downstream tasks. In detail, we provide theoretical analyses and practical algorithms for fine-tuning diffusion policies, addressing key challenges such as training instability and computational complexity caused by multi-step sampling. Extensive experiments demonstrate that ExDM outperforms existing SOTA baselines in efficient unsupervised exploration and fast fine-tuning downstream tasks, especially in structurally complicated environments.