# Make the Pertinent Salient: Task-Relevant Reconstruction for Visual Control with Distractions

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in RL for robot locomotion and manipulation, one of 15 directions in Deep reinforcement learning methods.
- Builds on 48 of its 57 references that are in the map.
- First posted on arXiv 13 Oct 2024.

## Placed in

- Deep reinforcement learning methods / [RL for robot locomotion and manipulation](https://mlsanity.com/deeprl/rl-for-robot-locomotion-and-manipulation/)

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

- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [Mask R-CNN](https://mlsanity.com/paper/1703.06870/) (2017)
- [You Only Look Once: Unified, Real-Time Object Detection](https://mlsanity.com/paper/1506.02640/) (2015)
- [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://mlsanity.com/paper/2105.15203/) (2021)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [The Kinetics Human Action Video Dataset](https://mlsanity.com/paper/1705.06950/) (2017)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [On the Properties of Neural Machine Translation: Encoder-Decoder Approaches](https://mlsanity.com/paper/1409.1259/) (2014)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Co-teaching: Robust Training of Deep Neural Networks with Extremely Noisy Labels](https://mlsanity.com/paper/1804.06872/) (2018)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [R3M: A Universal Visual Representation for Robot Manipulation](https://mlsanity.com/paper/2203.12601/) (2022)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Reinforcement Learning with Unsupervised Auxiliary Tasks](https://mlsanity.com/paper/1611.05397/) (2016)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learning](https://mlsanity.com/paper/2107.09645/) (2021)
- [ManiSkill2: A Unified Benchmark for Generalizable Manipulation Skills](https://mlsanity.com/paper/2302.04659/) (2023)
- [Deep Visual Foresight for Planning Robot Motion](https://mlsanity.com/paper/1610.00696/) (2016)
- [Personalize Segment Anything Model with One Shot](https://mlsanity.com/paper/2305.03048/) (2023)
- [DeepMDP: Learning Continuous Latent Space Models for Representation Learning](https://mlsanity.com/paper/1906.02736/) (2019)
- [Sim-to-Real via Sim-to-Sim: Data-efficient Robotic Grasping via Randomized-to-Canonical Adaptation Networks](https://mlsanity.com/paper/1812.07252/) (2018)

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

- [Adapting World Models with Latent-State Dynamics Residuals](https://mlsanity.com/paper/2504.02252/) (2025)
- [Gaze on the Prize: Shaping Visual Attention with Return-Guided Contrastive Learning](https://mlsanity.com/paper/2510.08442/) (2025)
- [KAGE-Bench: Fast Known-Axis Visual Generalization Evaluation for Reinforcement Learning](https://mlsanity.com/paper/2601.14232/) (2026)
- [StarWM: Self-Supervised Trained Attention Routing for Robust World Models](https://mlsanity.com/paper/2609.30667/) (2026)
- [VIGOR: Zero-Shot Visual Generalization via Latent-Space Consistency in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2610.02801/) (2026)

## Abstract (arXiv)

Recent advancements in Model-Based Reinforcement Learning (MBRL) have made it a powerful tool for visual control tasks. Despite improved data efficiency, it remains challenging to train MBRL agents with generalizable perception. Training in the presence of visual distractions is particularly difficult due to the high variation they introduce to representation learning. Building on DREAMER, a popular MBRL method, we propose a simple yet effective auxiliary task to facilitate representation learning in distracting environments. Under the assumption that task-relevant components of image observations are straightforward to identify with prior knowledge in a given task, we use a segmentation mask on image observations to only reconstruct task-relevant components. In doing so, we greatly reduce the complexity of representation learning by removing the need to encode task-irrelevant objects in the latent representation. Our method, Segmentation Dreamer (SD), can be used either with ground-truth masks easily accessible in simulation or by leveraging potentially imperfect segmentation foundation models. The latter is further improved by selectively applying the reconstruction loss to avoid providing misleading learning signals due to mask prediction errors. In modified DeepMind Control suite (DMC) and Meta-World tasks with added visual distractions, SD achieves significantly better sample efficiency and greater final performance than prior work. We find that SD is especially helpful in sparse reward tasks otherwise unsolvable by prior work, enabling the training of visually robust agents without the need for extensive reward engineering.