Self-Predictive Dynamics for Generalization of Vision-based Reinforcement Learning

TL;DR

This paper introduces Self-Predictive Dynamics (SPD), a novel method that enhances the generalization of vision-based reinforcement learning by extracting task-relevant features through predictive representations, especially under unseen and distracting visual conditions.

Contribution

SPD employs parallel augmentations and transition prediction to improve feature extraction and generalization in vision-based reinforcement learning tasks.

Findings

SPD outperforms previous methods on MuJoCo visual control tasks.

SPD significantly improves generalization in unseen observations.

SPD effectively handles distracting elements like shadows and clouds.

Abstract

Vision-based reinforcement learning requires efficient and robust representations of image-based observations, especially when the images contain distracting (task-irrelevant) elements such as shadows, clouds, and light. It becomes more important if those distractions are not exposed during training. We design a Self-Predictive Dynamics (SPD) method to extract task-relevant features efficiently, even in unseen observations after training. SPD uses weak and strong augmentations in parallel, and learns representations by predicting inverse and forward transitions across the two-way augmented versions. In a set of MuJoCo visual control tasks and an autonomous driving task (CARLA), SPD outperforms previous studies in complex observations, and significantly improves the generalization performance for unseen observations. Our code is available at https://github.com/unigary/SPD.