Research map / Reinforcement learning / Visual and representation RL
Policy-shaped prediction: avoiding distractions in model-based reinforcement learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Policy-Shaped Prediction (PSP), which reweights DreamerV3's image reconstruction loss using policy-gradient saliency aggregated over SAM segments plus an adversarial action-prediction head, and introduces a Reafferent DMC benchmark with learnable, action-dependent distractors.
Problem
Reconstruction-based model-based RL wastes world-model capacity on predictable, task-irrelevant distractors, including distractors that are learnable and correlated with the agent's own actions.
Approach
Weight the world model's pixel reconstruction loss by policy-gradient saliency aggregated within SAM-segmented objects (extending VaGraM to images), and subtract the gradient of an adversarial previous-action predictor from the encoder so action information is not duplicated in the image embedding.
Key results
- On Reafferent Cheetah Run, PSP scores 383.1 ± 23.8 versus 158.4 ± 45.7 for DreamerV3 and 565.1 ± 35.5 for model-free DrQv2.
- On Reafferent Hopper Stand, PSP scores 128.5 ± 215.7 while all model-based baselines score below 5 except DreamerPro at 3.8 and TIA at 0.9 (DrQv2 scores 210.3).
- On unmodified DMC Cheetah Run, PSP scores 712.3 ± 32.3, comparable to DreamerV3 (521.1) and other methods, showing no trade-off on non-distracting tasks.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use PSP instead of DreamerV3 | DreamerV3 · Reafferent Control Cheetah Run | score: 158.4 $\pm$ 45.7 -> 383.1 $\pm$ 23.8 | Helped |
| Use PSP instead of DreamerV3 | DreamerV3 · Reafferent Control Hopper Stand | score: 4.6 $\pm$ 3.9 -> 128.5 $\pm$ 215.7 | Helped |
| Use PSP instead of DreamerV3 | DreamerV3 · Distracting DeepMind Control Cheetah Run | score: 243.8 $\pm$ 81.2 -> 408.6 $\pm$ 125.1 | Helped |
| Use PSP instead of DreamerV3 | DreamerV3 · Distracting DeepMind Control Hopper Stand | score: 173.7 $\pm$ 160.9 -> 417.7 $\pm$ 118.9 | Helped |
| Add segmentation-based gradient aggregation (ablation) | Policy-gradient weighted DreamerV3 with adversarial action head · Cheetah Run, unaltered and Reafferent Control | score: 742.1 $\pm$ 79.7 -> 712.3 $\pm$ 32.3; 188.4 $\pm$ 9.4 -> 383.1 $\pm$ 23.8 | Mixed |
| Add the adversarial action prediction head (ablation) | Policy-gradient weighted DreamerV3 with segmentation · Cheetah Run, unaltered and Reafferent Control | score: 653.9 $\pm$ 44.2 -> 712.3 $\pm$ 32.3; 231.2 $\pm$ 58.6 -> 383.1 $\pm$ 23.8 | Helped |
| Add policy-gradient weighting (ablation) | DreamerV3 with segmentation and adversarial action head · Cheetah Run, unaltered and Reafferent Control | score: 674.2 $\pm$ 50.7 -> 712.3 $\pm$ 32.3; 324.4 $\pm$ 2.3 -> 383.1 $\pm$ 23.8 | Mixed |
| Add segmentation and adversarial action prediction (ablation) | Policy-gradient weighted DreamerV3 · Cheetah Run, unaltered and Reafferent Control | score: 418.2 $\pm$ 53.1 -> 712.3 $\pm$ 32.3; 379.0 $\pm$ 46.8 -> 383.1 $\pm$ 23.8 | Mixed |
| Replace policy-gradient weighting with value-gradient weighting (ablation) | DreamerV3 without segmentation or adversarial action head · Cheetah Run, unaltered and Reafferent Control | score: 521.1 $\pm$ 136.3 -> 381.7 $\pm$ 64.9; 158.4 $\pm$ 45.7 -> 445.7 $\pm$ 126.9 | Mixed |
| Interpolate gradient-based reconstruction weighting with uniform weighting (ablation) | PSP · Cheetah Run | score | Helped |
| Replace SAM with SAM2-tiny (ablation) | PSP · Cheetah Run environments | score: 383.1 $\pm$ 23.8 -> 398.7 $\pm$ 70.5; 712.3 $\pm$ 32.3 -> 696.3 $\pm$ 20.0; 364.4 $\pm$ 60.7 -> 352.0 $\pm$ 60.4 | No clear effect |
| Replace SAM with SAM2-tiny (ablation) | PSP · Hopper Stand, unmodified and distracting DeepMind Control | score: 865.6 $\pm$ 53.6 -> 891.2 $\pm$ 39.6; 417.7 $\pm$ 118.9 -> 187.4 $\pm$ 172.0 | Mixed |
Surprising findings
- Model-free DrQv2 outperforms all model-based baselines on the Reafferent environment, including DreamerV3, DreamerPro, TIA and Denoised MDP.
- Removing segmentation or the adversarial head while keeping policy-gradient weighting can improve Reafferent scores in some ablations but degrades unmodified performance.
Limitations
- Object-centric assumption that pixels belong to single objects ranked by importance.
- SAM segmentation is computationally expensive during training; results on Hopper Stand show high variance across seeds.
- Unclear how well PSP adapts when reward structure or salient features change over time.
Experimental setup
Datasets: DeepMind Control Suite Distracting Control Suite (Kinetics driving car video backgrounds)
Benchmarks: Reafferent DeepMind Control Suite DeepMind Control Suite Distracting Control Suite
Metrics: episode return
Baselines: DreamerV3 DreamerPro Task Informed Abstractions Denoised MDP DrQv2
Builds on: DreamerV3 VaGraM Segment Anything Model Task Informed Abstractions Denoised MDP
Closest papers by meaning
- VIGOR: Zero-Shot Visual Generalization via Latent-Space Consistency in Model-Based Reinforcement Learning (NeurIPS 2026 · similarity 0.88)
- On Rollouts in Model-Based Reinforcement Learning (ICLR 2025 · similarity 0.87)
- Learning Latent Dynamic Robust Representations for World Models (ICML 2024 · similarity 0.87)
- Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning (ICML 2025 · similarity 0.87)
- DreamSmooth: Improving Model-based Reinforcement Learning via Reward Smoothing (ICLR 2024 · similarity 0.87)
- Towards General-Purpose Model-Free Reinforcement Learning (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.