ROMA: Reverse Model-Based Data Augmentation for Offline Reinforcement Learning

One of the main challenges of offline Reinforcement Learning is that the difference between learning policy and behavior policy leads to the possibility that the agent may need to estimate the value of unseen action, which usually results in value overestimation and training instability. The Imitation-learning-based method, which is easy to implement and scale up, bypasses this problem by performing some kind of imitation learning on the dataset and enjoying a more stable training process. However, it’s hard for them to splice sub-optimal trajectories, which limits its performance to the quality of the optimal trajectories in the dataset. In order to improve the performance ceiling of the IL-based method while maintaining its advantages, we propose the Reverse Offline Model-based Augmentation (ROMA) framework to optimize offline trajectory data. ROMA learns a reverse model of the environment in a self-supervised way and can splice shortcuts inside each offline trajectory to augment their return value. Experiments on two navigation tasks demonstrate the effectiveness of ROMA, it can effectively optimize trajectory quality and significantly improve the performance of multiple IL-based methods by providing better imitation objects and reducing reward condition mismatch.