Goal-Conditioned Data Augmentation for Offline Reinforcement Learning

TL;DR

This paper introduces GODA, a goal-conditioned diffusion-based data augmentation method that improves offline reinforcement learning by generating higher-quality, goal-oriented samples to overcome suboptimal datasets.

Contribution

GODA is a novel goal-conditioned diffusion approach with return-based guidance and adaptive gating, enhancing data quality for offline RL from limited demonstrations.

Findings

GODA outperforms state-of-the-art augmentation methods on D4RL benchmarks.

GODA improves policy performance in traffic signal control tasks.

The method effectively generates higher-return samples from suboptimal datasets.

Abstract

Offline reinforcement learning (RL) enables policy learning from pre-collected offline datasets, relaxing the need to interact directly with the environment. However, limited by the quality of offline datasets, it generally fails to learn well-qualified policies in suboptimal datasets. To address datasets with insufficient optimal demonstrations, we introduce Goal-cOnditioned Data Augmentation (GODA), a novel goal-conditioned diffusion-based method for augmenting samples with higher quality. Leveraging recent advancements in generative modelling, GODA incorporates a novel return-oriented goal condition with various selection mechanisms. Specifically, we introduce a controllable scaling technique to provide enhanced return-based guidance during data sampling. GODA learns a comprehensive distribution representation of the original offline datasets while generating new data with selectively higher-return goals, thereby maximizing the utility of limited optimal demonstrations. Furthermore, we propose a novel adaptive gated conditioning method for processing noisy inputs and conditions, enhancing the capture of goal-oriented guidance. We conduct experiments on the D4RL benchmark and real-world challenges, specifically traffic signal control (TSC) tasks, to demonstrate GODA's effectiveness in enhancing data quality and superior performance compared to state-of-the-art data augmentation methods across various offline RL algorithms.