# MaDi: Learning to Mask Distractions for Generalization in Visual Deep Reinforcement Learning

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 53 references that are in the map.
- First posted on arXiv 23 Dec 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Solving Rubik's Cube with a Robot Hand](https://mlsanity.com/paper/1910.07113/) (2019)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Learning to reinforcement learn](https://mlsanity.com/paper/1611.05763/) (2016)
- [Target-driven Visual Navigation in Indoor Scenes using Deep Reinforcement Learning](https://mlsanity.com/paper/1609.05143/) (2016)
- [Quantifying Generalization in Reinforcement Learning](https://mlsanity.com/paper/1812.02341/) (2018)
- [Robust Deep Reinforcement Learning against Adversarial Perturbations on State Observations](https://mlsanity.com/paper/2003.08938/) (2020)
- [Learning to Navigate in Complex Environments](https://mlsanity.com/paper/1611.03673/) (2016)
- [A Study on Overfitting in Deep Reinforcement Learning](https://mlsanity.com/paper/1804.06893/) (2018)
- [Deep Spatial Autoencoders for Visuomotor Learning](https://mlsanity.com/paper/1509.06113/) (2015)
- [Generalization in Reinforcement Learning by Soft Data Augmentation](https://mlsanity.com/paper/2011.13389/) (2020)
- [Self-Supervised Policy Adaptation during Deployment](https://mlsanity.com/paper/2007.04309/) (2020)
- [Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2101.05265/) (2021)
- [Generalization and Regularization in DQN](https://mlsanity.com/paper/1810.00123/) (2018)
- [Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation](https://mlsanity.com/paper/2107.00644/) (2021)
- [The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixels](https://mlsanity.com/paper/2101.02722/) (2021)
- [Improving Generalization in Reinforcement Learning with Mixture Regularization](https://mlsanity.com/paper/2010.10814/) (2020)

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

- [Run-time Observation Interventions Make Vision-Language-Action Models More Visually Robust](https://mlsanity.com/paper/2410.01971/) (2024)
- [Revisiting Sparse Rewards for Goal-Reaching Reinforcement Learning](https://mlsanity.com/paper/2407.00324/) (2024)
- [Policy-shaped prediction: avoiding distractions in model-based reinforcement learning](https://mlsanity.com/paper/2412.05766/) (2024)
- [HASARD: A Benchmark for Vision-Based Safe Reinforcement Learning in Embodied Agents](https://mlsanity.com/paper/2503.08241/) (2025)
- [World Model Robustness via Surprise Recognition](https://mlsanity.com/paper/2512.01119/) (2025)
- [Merging and Disentangling Views in Visual Reinforcement Learning for Robotic Manipulation](https://mlsanity.com/paper/2505.04619/) (2025)
- [What to Ignore, What to React: Visually Robust RL Fine-Tuning of VLA Models](https://mlsanity.com/paper/2605.13105/) (2026)
- [Common Benchmarks Undervalue the Generalization Power of Programmatic Policies](https://mlsanity.com/paper/2506.14162/) (2025)
- [DeGuV: Depth-Guided Visual Reinforcement Learning for Generalization and Interpretability in Manipulation](https://mlsanity.com/paper/2509.04970/) (2025)
- [How (and How Not) to Use Data Augmentation in VLA Post-Training](https://mlsanity.com/paper/2610.05994/) (2026)

## Abstract (arXiv)

The visual world provides an abundance of information, but many input pixels received by agents often contain distracting stimuli. Autonomous agents need the ability to distinguish useful information from task-irrelevant perceptions, enabling them to generalize to unseen environments with new distractions. Existing works approach this problem using data augmentation or large auxiliary networks with additional loss functions. We introduce MaDi, a novel algorithm that learns to mask distractions by the reward signal only. In MaDi, the conventional actor-critic structure of deep reinforcement learning agents is complemented by a small third sibling, the Masker. This lightweight neural network generates a mask to determine what the actor and critic will receive, such that they can focus on learning the task. The masks are created dynamically, depending on the current input. We run experiments on the DeepMind Control Generalization Benchmark, the Distracting Control Suite, and a real UR5 Robotic Arm. Our algorithm improves the agent's focus with useful masks, while its efficient Masker network only adds 0.2% more parameters to the original structure, in contrast to previous work. MaDi consistently achieves generalization results better than or competitive to state-of-the-art methods.