Self-Predictive Dynamics for Generalization of Vision-based Reinforcement Learning
This addresses generalization challenges in vision-based RL for robotics and autonomous driving, though it appears incremental as it builds on existing augmentation and prediction techniques.
The paper tackles the problem of vision-based reinforcement learning generalizing to unseen observations with distracting elements by proposing a Self-Predictive Dynamics method, which improves generalization performance in MuJoCo and CARLA tasks.
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.