🤖 AI Summary
Vision-based reinforcement learning exhibits poor generalization under unseen image corruptions (e.g., shadows, cloud cover, illumination shifts). To address this, we propose Self-Predictive Dynamics (SPD), the first method to jointly integrate bidirectional (forward and inverse) dynamics prediction with weak/strong dual-path data augmentation for task-agnostic, robust representation learning. SPD employs contrastive self-supervised modeling to explicitly disentangle task-relevant dynamics from corruption-invariant features. Evaluated on MuJoCo vision-based control and CARLA autonomous driving benchmarks, SPD significantly improves policy generalization across corrupted environments—achieving a 23.6% performance gain over state-of-the-art methods under unseen corruptions. The implementation is publicly available.
📝 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.