Unsupervised Visual Attention and Invariance for Reinforcement Learning

2021/04/07 by Xudong Wang, Long Lian, Wang, Xudong +3 · 3 citations
Computer Science · #Computer Vision and Pattern Recognition (cs.CV) #Domain Adaptation and Few-Shot Learning #FOS: Computer and information sciences #Multimodal Machine Learning Applications #Reinforcement Learning in Robotics #Robotics (cs.RO) #cs.CV #cs.RO

paper · pdf · doi:10.48550/arxiv.2104.02921

Accepted at CVPR 2021

openalex publication_date 2021/04/07 · arxiv created 2021/04/16 · arxiv updated 2021/04/20 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

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.

Citations

• Playing Atari with Deep Reinforcement Learning
• Visual Domain Adaptation: A survey of recent advances
• Memory-based control with recurrent neural networks
• EPOpt: Learning Robust Neural Network Policies Using Model Ensembles
• Preparing for the Unknown: Learning a Universal Policy with Online System Identification
• Robust Adversarial Reinforcement Learning
• Towards Deep Learning Models Resistant to Adversarial Attacks
• DeepMind Control Suite
• Soft Actor-Critic Algorithms and Applications
• Colorization as a Proxy Task for Visual Understanding
• Single Episode Policy Transfer in Reinforcement Learning
• Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning
• Learning dexterous in-hand manipulation
• A Simple Framework for Contrastive Learning of Visual Representations
• Improved Baselines with Momentum Contrastive Learning
• Reinforcement Learning with Augmented Data
• Momentum Contrast for Unsupervised Visual Representation Learning
• Unsupervised Feature Learning by Cross-Level Instance-Group Discrimination
• Generalization in Reinforcement Learning by Soft Data Augmentation
• CURL: Contrastive Unsupervised Representations for Reinforcement Learning

Cited by

• Merging and Disentangling Views in Visual Reinforcement Learning for Robotic Manipulation
• How Ensembles of Distilled Policies Improve Generalisation in Reinforcement Learning
• Maximum Likelihood Reinforcement Learning

Related

• Dream to Generalize: Zero-Shot Model-Based Reinforcement Learning for Unseen Visual Distractions
• Attention-Privileged Reinforcement Learning
• Pre-Trained Image Encoder for Generalizable Visual Reinforcement Learning
• Light-weight probing of unsupervised representations for Reinforcement Learning
• Learning to Perceive in Deep Model-Free Reinforcement Learning