Equivariant Data Augmentation for Generalization in Offline
  Reinforcement Learning
Cristina Pinneri Sarah Bechtle Markus Wulfmeier Arunkumar Byravan Jingwei Zhang William F. Whitney Martin Riedmiller
We present a novel approach to address the challenge of generalization in
offline reinforcement learning (RL), where the agent learns from a fixed
dataset without any additional interaction with the environment. Specifically,
we aim to improve the agent's ability to generalize to out-of-distribution
goals. To achieve this, we propose to learn a dynamics model and check if it is
equivariant with respect to a fixed type of transformation, namely translations
in the state space. We then use an entropy regularizer to increase the
equivariant set and augment the dataset with the resulting transformed samples.
Finally, we learn a new policy offline based on the augmented dataset, with an
off-the-shelf offline RL algorithm. Our experimental results demonstrate that
our approach can greatly improve the test performance of the policy on the
considered environments.