Diffusion Self-Weighted Guidance for Offline Reinforcement Learning

TL;DR

This paper introduces Self-Weighted Guidance (SWG), a diffusion-based method for offline reinforcement learning that simplifies score computation by integrating weights directly into the diffusion model, achieving competitive results.

Contribution

The paper proposes a novel diffusion model framework that directly incorporates weight functions for offline RL, eliminating the need for additional network training.

Findings

SWG generates desired action distributions in toy examples.

SWG performs comparably to state-of-the-art on D4RL benchmarks.

Ablation studies validate the scalability and effectiveness of SWG.

Abstract

Offline reinforcement learning (RL) recovers the optimal policy $\pi$ given historical observations of an agent. In practice, $\pi$ is modeled as a weighted version of the agent's behavior policy $\mu$, using a weight function $w$ working as a critic of the agent's behavior. Though recent approaches to offline RL based on diffusion models have exhibited promising results, the computation of the required scores is challenging due to their dependence on the unknown $w$. In this work, we alleviate this issue by constructing a diffusion over both the actions and the weights. With the proposed setting, the required scores are directly obtained from the diffusion model without learning extra networks. Our main conceptual contribution is a novel guidance method, where guidance (which is a function of $w$) comes from the same diffusion model, therefore, our proposal is termed Self-Weighted Guidance (SWG). We show that SWG generates samples from the desired distribution on toy examples and performs on par with state-of-the-art methods on D4RL's challenging environments, while maintaining a streamlined training pipeline. We further validate SWG through ablation studies on weight formulations and scalability.