Score Regularized Policy Optimization through Diffusion Behavior

TL;DR

This paper introduces a method that combines diffusion models with policy optimization in offline reinforcement learning, significantly speeding up action sampling without sacrificing performance.

Contribution

It proposes a novel approach to extract deterministic policies from diffusion models, avoiding slow sampling while leveraging diffusion's generative power.

Findings

Speeds up action sampling by over 25 times in locomotion tasks

Maintains state-of-the-art performance in D4RL benchmarks

Effectively regularizes policy gradients using diffusion behavior models

Abstract

Recent developments in offline reinforcement learning have uncovered the immense potential of diffusion modeling, which excels at representing heterogeneous behavior policies. However, sampling from diffusion policies is considerably slow because it necessitates tens to hundreds of iterative inference steps for one action. To address this issue, we propose to extract an efficient deterministic inference policy from critic models and pretrained diffusion behavior models, leveraging the latter to directly regularize the policy gradient with the behavior distribution's score function during optimization. Our method enjoys powerful generative capabilities of diffusion modeling while completely circumventing the computationally intensive and time-consuming diffusion sampling scheme, both during training and evaluation. Extensive results on D4RL tasks show that our method boosts action sampling speed by more than 25 times compared with various leading diffusion-based methods in locomotion tasks, while still maintaining state-of-the-art performance.