Diffusion Dynamic Model for Unsupervised Reinforcement Learning.

Unsupervised reinforcement learning (URL) enables agents to adapt efficiently to novel tasks but relies on robust environmental modeling. Traditional world models enhance exploration but are prone to single-step errors, leading to distributional shifts. To address this challenge, we introduce the Diffusion Dynamic Model (DDM), incorporating diffusion models into the environmental modeling component of unsupervised reinforcement learning. During pre-training, DDM generates future observations conditioned on current observations and actions. These generated observations are used to design an intrinsic reward function that incentivizes the agent to explore diverse and high-uncertainty states. In the fine-tuning phase, DDM leverages its generative capabilities to perform data augmentation. By producing diverse, high-quality synthetic data, DDM expands the training dataset, enabling the agent to generalize better and adapt more rapidly to task-specific environments. Our approach is validated across three domains and twelve downstream tasks in the URLB benchmark, demonstrating superior exploration and adaptability compared to existing methods.