Aligning Diffusion Behaviors with Q-functions for Efficient Continuous Control

TL;DR

This paper introduces a novel offline reinforcement learning approach that leverages diffusion models and Q-function alignment to improve continuous control, achieving superior performance with minimal labeled data.

Contribution

It proposes Efficient Diffusion Alignment (EDA), a new method combining diffusion models with Q-function alignment for improved offline continuous control.

Findings

EDA outperforms all baselines on D4RL benchmark.

Maintains 95% performance with only 1% Q-labeled data.

Diffusion policies enable effective behavior modeling and adaptation.

Abstract

Drawing upon recent advances in language model alignment, we formulate offline Reinforcement Learning as a two-stage optimization problem: First pretraining expressive generative policies on reward-free behavior datasets, then fine-tuning these policies to align with task-specific annotations like Q-values. This strategy allows us to leverage abundant and diverse behavior data to enhance generalization and enable rapid adaptation to downstream tasks using minimal annotations. In particular, we introduce Efficient Diffusion Alignment (EDA) for solving continuous control problems. EDA utilizes diffusion models for behavior modeling. However, unlike previous approaches, we represent diffusion policies as the derivative of a scalar neural network with respect to action inputs. This representation is critical because it enables direct density calculation for diffusion models, making them compatible with existing LLM alignment theories. During policy fine-tuning, we extend preference-based alignment methods like Direct Preference Optimization (DPO) to align diffusion behaviors with continuous Q-functions. Our evaluation on the D4RL benchmark shows that EDA exceeds all baseline methods in overall performance. Notably, EDA maintains about 95\% of performance and still outperforms several baselines given only 1\% of Q-labelled data during fine-tuning.