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

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

## 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.

## 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 38)

- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [SAM 2: Segment Anything in Images and Videos](https://mlsanity.com/paper/2408.00714/) (2024)
- [Deep Inside Convolutional Networks: Visualising Image Classification Models and Saliency Maps](https://mlsanity.com/paper/1312.6034/) (2013)
- [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)
- [Learning Important Features Through Propagating Activation Differences](https://mlsanity.com/paper/1704.02685/) (2017)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [SmoothGrad: removing noise by adding noise](https://mlsanity.com/paper/1706.03825/) (2017)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learning](https://mlsanity.com/paper/2107.09645/) (2021)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [Reinforcement Learning with Prototypical Representations](https://mlsanity.com/paper/2102.11271/) (2021)
- [Masked World Models for Visual Control](https://mlsanity.com/paper/2206.14244/) (2022)
- [Value Prediction Network](https://mlsanity.com/paper/1707.03497/) (2017)
- [Model-Based Reinforcement Learning via Meta-Policy Optimization](https://mlsanity.com/paper/1809.05214/) (2018)
- [The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixels](https://mlsanity.com/paper/2101.02722/) (2021)
- [Objective Mismatch in Model-based Reinforcement Learning](https://mlsanity.com/paper/2002.04523/) (2020)
- [The Value Equivalence Principle for Model-Based Reinforcement Learning](https://mlsanity.com/paper/2011.03506/) (2020)
- [Towards Interpretable Reinforcement Learning Using Attention Augmented Agents](https://mlsanity.com/paper/1906.02500/) (2019)
- [Learning and Planning in Complex Action Spaces](https://mlsanity.com/paper/2104.06303/) (2021)
- [DreamerPro: Reconstruction-Free Model-Based Reinforcement Learning with Prototypical Representations](https://mlsanity.com/paper/2110.14565/) (2021)
- [Denoised MDPs: Learning World Models Better Than the World Itself](https://mlsanity.com/paper/2206.15477/) (2022)
- [Learning Task Informed Abstractions](https://mlsanity.com/paper/2106.15612/) (2021)
- [Goal-Aware Prediction: Learning to Model What Matters](https://mlsanity.com/paper/2007.07170/) (2020)

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

- [World Model Robustness via Surprise Recognition](https://mlsanity.com/paper/2512.01119/) (2025)
- [Temporally Centered SIGReg Improves LeWorldModel Representations for Robot Policy Learning](https://mlsanity.com/paper/2607.26924/) (2026)
- [Diagnosing JEPA World Models with Action-Conditioned Predictive Consistency](https://mlsanity.com/paper/2608.12939/) (2026)
- [HaM-World: Soft-Hamiltonian World Models with Selective Memory for Planning](https://mlsanity.com/paper/2605.05951/) (2026)
- [Dueling World Models: Advantage-Style Action Channels for Common-Mode Distractor Rejection](https://mlsanity.com/paper/2608.06706/) (2026)
- [Transferable Delay-Aware Reinforcement Learning via Implicit Causal Graph Modeling](https://mlsanity.com/paper/2605.12312/) (2026)

## 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.