mlsanityData to 14 Oct 2026

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

arXiv 2410.09972 ↗First posted 13 Oct 2024cs.LG cs.AI cs.CVcited by 5

On the map: Deep reinforcement learning methods → RL for robot locomotion and manipulation

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.

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.

Builds on

48 of 57 references.

References · 48

  1. 2020A Simple Framework for Contrastive Learning of Visual Representationscited by 12,049
  2. 2023Segment Anythingcited by 8,234
  3. 2017Mask R-CNNcited by 7,277
  4. 2015You Only Look Once: Unified, Real-Time Object Detectioncited by 4,852
  5. 2021SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformerscited by 2,940
  6. 2020Unsupervised Learning of Visual Features by Contrasting Cluster Assignmentscited by 2,857
  7. 2017The Kinetics Human Action Video Datasetcited by 1,909
  8. 2019Dream to Control: Learning Behaviors by Latent Imaginationcited by 1,653
  9. 2018World Modelscited by 1,513
  10. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,353
  11. 2014On the Properties of Neural Machine Translation: Encoder-Decoder Approachescited by 1,304
  12. 2018Learning Latent Dynamics for Planning from Pixelscited by 1,264
36 more
  1. 2023Mastering Diverse Domains through World Modelscited by 1,253
  2. 2019Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learningcited by 1,204
  3. 2018Co-teaching: Robust Training of Deep Neural Networks with Extremely Noisy Labelscited by 1,059
  4. 2020Mastering Atari with Discrete World Modelscited by 1,007
  5. 2018DeepMind Control Suitecited by 956
  6. 2020CURL: Contrastive Unsupervised Representations for Reinforcement Learningcited by 760
  7. 2022R3M: A Universal Visual Representation for Robot Manipulationcited by 694
  8. 2020Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixelscited by 585
  9. 2023TD-MPC2: Scalable, Robust World Models for Continuous Controlcited by 512
  10. 2020Learning Invariant Representations for Reinforcement Learning without Reconstructioncited by 436
  11. 2016Reinforcement Learning with Unsupervised Auxiliary Taskscited by 404
  12. 2022Temporal Difference Learning for Model Predictive Controlcited by 397
  13. 2021Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learningcited by 372
  14. 2023ManiSkill2: A Unified Benchmark for Generalizable Manipulation Skillscited by 327
  15. 2016Deep Visual Foresight for Planning Robot Motioncited by 285
  16. 2023Personalize Segment Anything Model with One Shotcited by 250
  17. 2019DeepMDP: Learning Continuous Latent Space Models for Representation Learningcited by 207
  18. 2018Sim-to-Real via Sim-to-Sim: Data-efficient Robotic Grasping via Randomized-to-Canonical Adaptation Networkscited by 174
  19. 2022Masked World Models for Visual Controlcited by 170
  20. 2015Variational Information Maximisation for Intrinsically Motivated Reinforcement Learningcited by 144
  21. 2018OCNet: Object Context Network for Scene Parsingcited by 140
  22. 2020Generalization in Reinforcement Learning by Soft Data Augmentationcited by 137
  23. 2021Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentationcited by 105
  24. 2021The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixelscited by 100
  25. 2020Objective Mismatch in Model-based Reinforcement Learningcited by 96
  26. 2021DreamerPro: Reconstruction-Free Model-Based Reinforcement Learning with Prototypical Representationscited by 74
  27. 2022Denoised MDPs: Learning World Models Better Than the World Itselfcited by 68
  28. 2023Learning to Model the World with Languagecited by 68
  29. 2021Learning Task Informed Abstractionscited by 58
  30. 2021Temporal Predictive Coding For Model-Based Planning In Latent Spacecited by 55
  31. 2020Goal-Aware Prediction: Learning to Model What Matterscited by 45
  32. 2022INFOrmation Prioritization through EmPOWERment in Visual Model-Based RLcited by 27
  33. 2023RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictabilitycited by 18
  34. 2023FOCUS: Object-Centric World Models for Robotics Manipulationcited by 17
  35. 2023Informed POMDP: Leveraging Additional Information in Model-Based RLcited by 11
  36. 2022Sim-to-Real via Sim-to-Seg: End-to-end Off-road Autonomous Driving Without Real Datacited by 9

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 2; 2026: 3.

5 papers.

  1. 2025Adapting World Models with Latent-State Dynamics Residualscited by 9
  2. 2025Gaze on the Prize: Shaping Visual Attention with Return-Guided Contrastive Learningcited by 1
  3. 2026KAGE-Bench: Fast Known-Axis Visual Generalization Evaluation for Reinforcement Learningcited by 0
  4. 2026StarWM: Self-Supervised Trained Attention Routing for Robust World Modelscited by 0
  5. 2026VIGOR: Zero-Shot Visual Generalization via Latent-Space Consistency in Model-Based Reinforcement Learningcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Dream to Generalize: Zero-Shot Model-Based Reinforcement Learning for Unseen Visual Distractions
  2. 2024Learning Latent Dynamic Robust Representations for World Models
  3. 2024DMC-VB: A Benchmark for Representation Learning for Control with Visual Distractors
  4. 2024ViSaRL: Visual Reinforcement Learning Guided by Human Saliency
  5. 2024Policy-shaped prediction: avoiding distractions in model-based reinforcement learning
  6. 2025MInCo: Mitigating Information Conflicts in Distracted Visual Model-based Reinforcement Learning
  7. 2026R2-Dreamer: Redundancy-Reduced World Models without Decoders or Augmentation
  8. 2026Task-Relevant Representation Decoupling for Visual Reinforcement Learning Generalization