Unsupervised Visual Attention and Invariance for Reinforcement Learning
arXiv:2104.02921
Abstract
Vision-based reinforcement learning (RL) is successful, but how to generalize it to unknown test environments remains challenging. Existing methods focus on training an RL policy that is universal to changing visual domains, whereas we focus on extracting visual foreground that is universal, feeding clean invariant vision to the RL policy learner. Our method is completely unsupervised, without manual annotations or access to environment internals. Given videos of actions in a training environment, we learn how to extract foregrounds with unsupervised keypoint detection, followed by unsupervised visual attention to automatically generate a foreground mask per video frame. We can then introduce artificial distractors and train a model to reconstruct the clean foreground mask from noisy observations. Only this learned model is needed during test to provide distraction-free visual input to the RL policy learner. Our Visual Attention and Invariance (VAI) method significantly outperforms the state-of-the-art on visual domain generalization, gaining 15 to 49% (61 to 229%) more cumulative rewards per episode on DeepMind Control (our DrawerWorld Manipulation) benchmarks. Our results demonstrate that it is not only possible to learn domain-invariant vision without any supervision, but freeing RL from visual distractions also makes the policy more focused and thus far better.
Accepted at CVPR 2021
References in corpus (20)
- A Simple Framework for Contrastive Learning of Visual Representations
- Playing Atari with Deep Reinforcement Learning
- Improved Baselines with Momentum Contrastive Learning
- Soft Actor-Critic Algorithms and Applications
- Towards Deep Learning Models Resistant to Adversarial Attacks
- DeepMind Control Suite
- Robust Adversarial Reinforcement Learning
- Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning
- Visual Foresight: Model-Based Deep Reinforcement Learning for Vision-Based Robotic Control
- Reinforcement Learning with Augmented Data
- Memory-based control with recurrent neural networks
- EPOpt: Learning Robust Neural Network Policies Using Model Ensembles
- Deep Reinforcement Learning from Self-Play in Imperfect-Information Games
- Unsupervised Learning of Object Landmarks through Conditional Image Generation
- Self-Supervised Policy Adaptation during Deployment
- Colorization as a Proxy Task for Visual Understanding
- Preparing for the Unknown: Learning a Universal Policy with Online System Identification
- Single Episode Policy Transfer in Reinforcement Learning
- Generalization in Reinforcement Learning by Soft Data Augmentation
- Unsupervised Feature Learning by Cross-Level Instance-Group Discrimination