Multi-task Self-Supervised Adaptation for Reinforcement Learning

Policy adaptation remains one of the key challenges for reinforcement learning (RL). Thus, RL agents often fail to generalize to unseen scenarios. In this paper, we propose to improve the generalization of RL algorithms through multi-task self-supervised adaptation (MSSA). The proposed method is a general paradigm that can be implemented on top of any RL algorithm. It better extracts high-level feature representations from augmented observations through incorporating multiple self-supervised learning tasks with complementary objectives. The selected self-supervision tasks include rotation prediction, inverse dynamics prediction and contrastive learning. It then performs control actions based on the extracted features. The proposed MSSA method consistently outperforms all the baseline methods on diverse complex tasks in the DeepMind Control suite benchmark and sets new state-of-the-art results without incurring longer inference time. It is demonstrated that MSSA has superior generalization capability and is robust to environmental changes.