# Unsupervised Visual Attention and Invariance for Reinforcement Learning

arXiv: https://arxiv.org/abs/2104.02921 · HTML: https://mlsanity.com/paper/2104.02921/

## Key facts

- Cited by 30 papers in the map (data through 14 Oct 2026).
- Builds on 47 of its 70 references that are in the map.
- First posted on arXiv 7 Apr 2021.

## Builds on (most cited first, 30 of 47)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Perceptual Losses for Real-Time Style Transfer and Super-Resolution](https://mlsanity.com/paper/1603.08155/) (2016)
- [Improved Baselines with Momentum Contrastive Learning](https://mlsanity.com/paper/2003.04297/) (2020)
- [Unsupervised Feature Learning via Non-Parametric Instance-level Discrimination](https://mlsanity.com/paper/1805.01978/) (2018)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness](https://mlsanity.com/paper/1811.12231/) (2018)
- [Unsupervised Learning of Visual Representations by Solving Jigsaw Puzzles](https://mlsanity.com/paper/1603.09246/) (2016)
- [Learning Transferable Features with Deep Adaptation Networks](https://mlsanity.com/paper/1502.02791/) (2015)
- [Adversarial Discriminative Domain Adaptation](https://mlsanity.com/paper/1702.05464/) (2017)
- [Soft Actor-Critic Algorithms and Applications](https://mlsanity.com/paper/1812.05905/) (2018)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Unsupervised Visual Representation Learning by Context Prediction](https://mlsanity.com/paper/1505.05192/) (2015)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [End-to-End Training of Deep Visuomotor Policies](https://mlsanity.com/paper/1504.00702/) (2015)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Learning Dexterous In-Hand Manipulation](https://mlsanity.com/paper/1808.00177/) (2018)
- [Sim-to-Real Transfer of Robotic Control with Dynamics Randomization](https://mlsanity.com/paper/1710.06537/) (2017)
- [Learning multiple visual domains with residual adapters](https://mlsanity.com/paper/1705.08045/) (2017)
- [Using Self-Supervised Learning Can Improve Model Robustness and Uncertainty](https://mlsanity.com/paper/1906.12340/) (2019)
- [Robust Adversarial Reinforcement Learning](https://mlsanity.com/paper/1703.02702/) (2017)
- [Reinforcement Learning with Unsupervised Auxiliary Tasks](https://mlsanity.com/paper/1611.05397/) (2016)
- [Deep Visual Foresight for Planning Robot Motion](https://mlsanity.com/paper/1610.00696/) (2016)
- [Direct and Indirect Effects](https://mlsanity.com/paper/1301.2300/) (2013)
- [Long-Tailed Classification by Keeping the Good and Removing the Bad Momentum Causal Effect](https://mlsanity.com/paper/2009.12991/) (2020)
- [Visual Reinforcement Learning with Imagined Goals](https://mlsanity.com/paper/1807.04742/) (2018)

## Built on by (most cited first, 30 of 30)

- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation](https://mlsanity.com/paper/2107.00644/) (2021)
- [Pre-Trained Image Encoder for Generalizable Visual Reinforcement Learning](https://mlsanity.com/paper/2212.08860/) (2022)
- [Crossway Diffusion: Improving Diffusion-based Visuomotor Policy via Self-supervised Learning](https://mlsanity.com/paper/2307.01849/) (2023)
- [Maximum Likelihood Reinforcement Learning](https://mlsanity.com/paper/2602.02710/) (2026)
- [Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning](https://mlsanity.com/paper/2209.09203/) (2022)
- [A Comprehensive Survey of Data Augmentation in Visual Reinforcement Learning](https://mlsanity.com/paper/2210.04561/) (2022)
- [Does Self-supervised Learning Really Improve Reinforcement Learning from Pixels?](https://mlsanity.com/paper/2206.05266/) (2022)
- [3D Implicit Transporter for Temporally Consistent Keypoint Discovery](https://mlsanity.com/paper/2309.05098/) (2023)
- [Learning Deep Sensorimotor Policies for Vision-based Autonomous Drone Racing](https://mlsanity.com/paper/2210.14985/) (2022)
- [A Recipe for Unbounded Data Augmentation in Visual Reinforcement Learning](https://mlsanity.com/paper/2405.17416/) (2024)
- [Generalizable Visual Reinforcement Learning with Segment Anything Model](https://mlsanity.com/paper/2312.17116/) (2023)
- [R2-Dreamer: Redundancy-Reduced World Models without Decoders or Augmentation](https://mlsanity.com/paper/2603.18202/) (2026)
- [Local Feature Swapping for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2204.06355/) (2022)
- [Conceptual Reinforcement Learning for Language-Conditioned Tasks](https://mlsanity.com/paper/2303.05069/) (2023)
- [DEAR: Disentangled Environment and Agent Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2407.00633/) (2024)
- [Disentangled World Models: Learning to Transfer Semantic Knowledge from Distracting Videos for Reinforcement Learning](https://mlsanity.com/paper/2503.08751/) (2025)
- [Understanding What Affects the Generalization Gap in Visual Reinforcement Learning: Theory and Empirical Evidence](https://mlsanity.com/paper/2402.02701/) (2024)
- [Merging and Disentangling Views in Visual Reinforcement Learning for Robotic Manipulation](https://mlsanity.com/paper/2505.04619/) (2025)
- [AACC: Asymmetric Actor-Critic in Contextual Reinforcement Learning](https://mlsanity.com/paper/2208.02376/) (2022)
- [The Role of Diverse Replay for Generalisation in Reinforcement Learning](https://mlsanity.com/paper/2306.05727/) (2023)
- [Self-supervised Scene Text Segmentation with Object-centric Layered Representations Augmented by Text Regions](https://mlsanity.com/paper/2308.13178/) (2023)
- [OrthoNets: Orthogonal Channel Attention Networks](https://mlsanity.com/paper/2311.03071/) (2023)
- [Efficient Training of Generalizable Visuomotor Policies via Control-Aware Augmentation](https://mlsanity.com/paper/2401.09258/) (2024)
- [Training on Irrelevant States Implies Data Augmentation: Generalization in Contextual MDPs](https://mlsanity.com/paper/2410.03565/) (2024)
- [How Ensembles of Distilled Policies Improve Generalisation in Reinforcement Learning](https://mlsanity.com/paper/2505.16581/) (2025)
- [Explore-Go: Leveraging Exploration for Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2406.08069/) (2024)
- [Domain Adaptation with Adaptive Imagination for Visual Reinforcement Learning under Limited Target Data](https://mlsanity.com/paper/2606.30192/) (2026)
- [Task-Relevant Representation Decoupling for Visual Reinforcement Learning Generalization](https://mlsanity.com/paper/2607.00796/) (2026)
- [StarWM: Self-Supervised Trained Attention Routing for Robust World Models](https://mlsanity.com/paper/2609.30667/) (2026)

## Abstract (arXiv)

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.