MGDA: Model-based Goal Data Augmentation for Offline Goal-conditioned Weighted Supervised Learning

TL;DR

This paper introduces MGDA, a model-based goal data augmentation method that improves goal data sampling for offline goal-conditioned reinforcement learning, enhancing trajectory stitching and overall performance.

Contribution

The paper proposes a novel MGDA approach that uses a learned dynamics model and Lipschitz continuity to improve goal augmentation in GCWSL methods.

Findings

MGDA outperforms previous augmentation techniques in maze datasets.

MGDA enhances trajectory stitching capabilities.

Empirical results confirm improved goal-reaching performance.

Abstract

Recently, a state-of-the-art family of algorithms, known as Goal-Conditioned Weighted Supervised Learning (GCWSL) methods, has been introduced to tackle challenges in offline goal-conditioned reinforcement learning (RL). GCWSL optimizes a lower bound of the goal-conditioned RL objective and has demonstrated outstanding performance across diverse goal-reaching tasks, providing a simple, effective, and stable solution. However, prior research has identified a critical limitation of GCWSL: the lack of trajectory stitching capabilities. To address this, goal data augmentation strategies have been proposed to enhance these methods. Nevertheless, existing techniques often struggle to sample suitable augmented goals for GCWSL effectively. In this paper, we establish unified principles for goal data augmentation, focusing on goal diversity, action optimality, and goal reachability. Based on these principles, we propose a Model-based Goal Data Augmentation (MGDA) approach, which leverages a learned dynamics model to sample more suitable augmented goals. MGDA uniquely incorporates the local Lipschitz continuity assumption within the learned model to mitigate the impact of compounding errors. Empirical results show that MGDA significantly enhances the performance of GCWSL methods on both state-based and vision-based maze datasets, surpassing previous goal data augmentation techniques in improving stitching capabilities.