Efficient and Uncertainty-Aware Diffusion Framework for Offline-to-Online Reinforcement Learning

Offline-to-Online Reinforcement Learning (O2O-RL) leverages an offline, pre-trained policy to minimize costly online interactions. Although data-efficient, O2O-RL is susceptible to shifts between offline and online distributions. Existing work aims to mitigate the harm of this shift by finetuning the policy on trajectory data sampled from a diffusion model. Inspired by this line of work, we propose DUAL: an efficient Diffusion Uncertainty-Aware framework for offline-to-online reinforcement Learning. DUAL utilizes the prior knowledge of the diffusion model to distill a fast-sampling diffusion actor policy and transition model in the offline phase. DUAL also employs a Laplace approximation and distance transition-state-shift detection, thereby using uncertainty quantification to improve exploration versus exploitation in the online phase. We formally show that our actor loss with the Laplace approximation provides a proxy for a principled estimate of epistemic uncertainty. Empirically, DUAL improves the online expected return over O2O-RL baselines across multiple settings and environments.

Training AI agents to make good decisions, e.g., controlling a robot, typically requires enormous amounts of costly trial-and-error in the real world. A promising shortcut is to first train on pre-recorded data, then refine through live experience. However, the live world often differs from the recorded data, causing the agent to make overconfident mistakes and get stuck in poor habits. We introduce DUAL, a diffusion-based framework that gives the agent two key abilities: a fast, planning-aware decision policy, and a principled sense of "how uncertain am I?" — knowing when to explore new behaviors versus exploit what it already knows. Across a range of robotic control tasks, DUAL outperforms existing approaches, especially when the world's dynamics shift between the pre-training and live phases.