DiWA: Diffusion Policy Adaptation with World Models

TL;DR

DiWA introduces an offline reinforcement learning framework using a world model to efficiently fine-tune diffusion policies for robotic skills, significantly reducing real-world interactions needed compared to prior methods.

Contribution

DiWA is the first approach to fine-tune diffusion policies for robots entirely offline with a world model, improving sample efficiency and safety.

Findings

Achieves effective task performance on CALVIN benchmark with offline adaptation.

Requires orders of magnitude fewer physical interactions than model-free baselines.

Demonstrates practical and safe robot skill fine-tuning using offline data.

Abstract

Fine-tuning diffusion policies with reinforcement learning (RL) presents significant challenges. The long denoising sequence for each action prediction impedes effective reward propagation. Moreover, standard RL methods require millions of real-world interactions, posing a major bottleneck for practical fine-tuning. Although prior work frames the denoising process in diffusion policies as a Markov Decision Process to enable RL-based updates, its strong dependence on environment interaction remains highly inefficient. To bridge this gap, we introduce DiWA, a novel framework that leverages a world model for fine-tuning diffusion-based robotic skills entirely offline with reinforcement learning. Unlike model-free approaches that require millions of environment interactions to fine-tune a repertoire of robot skills, DiWA achieves effective adaptation using a world model trained once on a few hundred thousand offline play interactions. This results in dramatically improved sample efficiency, making the approach significantly more practical and safer for real-world robot learning. On the challenging CALVIN benchmark, DiWA improves performance across eight tasks using only offline adaptation, while requiring orders of magnitude fewer physical interactions than model-free baselines. To our knowledge, this is the first demonstration of fine-tuning diffusion policies for real-world robotic skills using an offline world model. We make the code publicly available at https://diwa.cs.uni-freiburg.de.