mlsanityData to 14 Oct 2026

Policy-shaped prediction: avoiding distractions in model-based reinforcement learning

arXiv 2412.05766 ↗First posted 8 Dec 2024cs.LG cs.AIcited by 6

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

Abstract (arXiv)

Model-based reinforcement learning (MBRL) is a promising route to sample-efficient policy optimization. However, a known vulnerability of reconstruction-based MBRL consists of scenarios in which detailed aspects of the world are highly predictable, but irrelevant to learning a good policy. Such scenarios can lead the model to exhaust its capacity on meaningless content, at the cost of neglecting important environment dynamics. While existing approaches attempt to solve this problem, we highlight its continuing impact on leading MBRL methods -- including DreamerV3 and DreamerPro -- with a novel environment where background distractions are intricate, predictable, and useless for planning future actions. To address this challenge we develop a method for focusing the capacity of the world model through synergy of a pretrained segmentation model, a task-aware reconstruction loss, and adversarial learning. Our method outperforms a variety of other approaches designed to reduce the impact of distractors, and is an advance towards robust model-based reinforcement learning.

Key facts. Cited by 6 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 38 of its 44 references that are in the map. First posted on arXiv 8 Dec 2024.

Builds on

38 of 44 references.

References · 38

  1. 2023Segment Anythingcited by 8,234
  2. 2017Axiomatic Attribution for Deep Networkscited by 3,384
  3. 2015Domain-Adversarial Training of Neural Networkscited by 3,044
  4. 2024SAM 2: Segment Anything in Images and Videoscited by 2,863
  5. 2013Deep Inside Convolutional Networks: Visualising Image Classification Models and Saliency Mapscited by 2,371
  6. 2017The Kinetics Human Action Video Datasetcited by 1,909
  7. 2019Dream to Control: Learning Behaviors by Latent Imaginationcited by 1,653
  8. 2018World Modelscited by 1,513
  9. 2017Learning Important Features Through Propagating Activation Differencescited by 1,458
  10. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,353
  11. 2023Mastering Diverse Domains through World Modelscited by 1,253
  12. 2017SmoothGrad: removing noise by adding noisecited by 1,109
26 more
  1. 2020Mastering Atari with Discrete World Modelscited by 1,007
  2. 2018DeepMind Control Suitecited by 956
  3. 2019Model-Based Reinforcement Learning for Ataricited by 502
  4. 2021Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learningcited by 372
  5. 2021Mastering Atari Games with Limited Datacited by 237
  6. 2021Reinforcement Learning with Prototypical Representationscited by 179
  7. 2022Masked World Models for Visual Controlcited by 170
  8. 2017Value Prediction Networkcited by 117
  9. 2018Model-Based Reinforcement Learning via Meta-Policy Optimizationcited by 101
  10. 2021The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixelscited by 100
  11. 2020Objective Mismatch in Model-based Reinforcement Learningcited by 96
  12. 2020The Value Equivalence Principle for Model-Based Reinforcement Learningcited by 94
  13. 2019Towards Interpretable Reinforcement Learning Using Attention Augmented Agentscited by 82
  14. 2021Learning and Planning in Complex Action Spacescited by 77
  15. 2021DreamerPro: Reconstruction-Free Model-Based Reinforcement Learning with Prototypical Representationscited by 74
  16. 2022Denoised MDPs: Learning World Models Better Than the World Itselfcited by 68
  17. 2021Learning Task Informed Abstractionscited by 58
  18. 2020Goal-Aware Prediction: Learning to Model What Matterscited by 45
  19. 2022Guaranteed Discovery of Control-Endogenous Latent States with Multi-Step Inverse Modelscited by 45
  20. 2023Pre-training Contextualized World Models with In-the-wild Videos for Reinforcement Learningcited by 45
  21. 2021Mismatched No More: Joint Model-Policy Optimization for Model-Based RLcited by 36
  22. 2022Value Gradient weighted Model-Based Reinforcement Learningcited by 32
  23. 2023Automatic Noise Filtering with Dynamic Sparse Training in Deep Reinforcement Learningcited by 13
  24. 2023Generalizable Visual Reinforcement Learning with Segment Anything Modelcited by 12
  25. 2023Ignorance is Bliss: Robust Control via Information Gatingcited by 10
  26. 2023MaDi: Learning to Mask Distractions for Generalization in Visual Deep Reinforcement Learningcited by 10

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 1; 2026: 5.

6 papers.

  1. 2025World Model Robustness via Surprise Recognitioncited by 4
  2. 2026Temporally Centered SIGReg Improves LeWorldModel Representations for Robot Policy Learningcited by 4
  3. 2026Diagnosing JEPA World Models with Action-Conditioned Predictive Consistencycited by 3
  4. 2026HaM-World: Soft-Hamiltonian World Models with Selective Memory for Planningcited by 2
  5. 2026Dueling World Models: Advantage-Style Action Channels for Common-Mode Distractor Rejectioncited by 2
  6. 2026Transferable Delay-Aware Reinforcement Learning via Implicit Causal Graph Modelingcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Policy-Driven World Model Adaptation for Robust Offline Model-based Reinforcement Learning
  2. 2022A Survey on Model-based Reinforcement Learning
  3. 2025Enter the Void - Planning to Seek Entropy When Reward is Scarce
  4. 2025Look Before Leap: Look-Ahead Planning with Uncertainty in Reinforcement Learning
  5. 2020A Game Theoretic Framework for Model Based Reinforcement Learning
  6. 2026All Models are Wrong, Knowing Where is Useful: On Model Uncertainty in Reinforcement Learning
  7. 2025Scalable Offline Model-Based RL with Action Chunks
  8. 2026Prioritized Rollouts for Efficient World Model-based Vision-Language-Action Policy Optimization