Uncertainty-Aware Data Augmentation for Offline Reinforcement Learning

One of the key challenges in Offline Reinforcement Learning is that it cannot conduct further environment exploration and performs poorly in terms of out-of-distribution generalizations. Data augmentation is commonly used to solve the issue of limited coverage of the full state-action space in static offline dataset. However, the existing data augmentation methods for proprioceptive observation suffer from the dilemma where the data coverage is often limited by tight constraints, while aggressive methods may exacerbate the performance. At the heart of this phenomenon are the diverged action distribution and the high uncertainty of the value function. In this paper, we propose to extend the static offline datasets during training by adding gradient-based perturbation to the state and utilizing the estimated uncertainty of the value function to constrain the range of the gradient. The estimated uncertainty of the value function works as a guidance to adjust the range of augmentation automatically, ensuring the adaptability and reliability of the state perturbation. The proposed algorithm Uncertainty-Aware Data Augmentation(UADA), is plugged into various standard offline RL algorithms and evaluated on several offline reinforcement learning tasks. The empirical results confirm that UADA substantially improves the performance and achieves better model stability compared with the original algorithms.