Policy-shaped prediction: avoiding distractions in model-based reinforcement learning
Authors: Miles Richard Hutson
Isaac Kauvar
Nick Haber
Venue: NeurIPS 2024 poster (2024)
Keywords: machine learning
model based reinforcement learning
reinforcement learning
segment anything model
Primary area: reinforcement_learning
TL;DR: Reduce the impact of distractors on model-based RL using gradient-based interpretability with segmentation-based aggregation.

Abstract: 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 a 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.

Policy-shaped prediction: avoiding distractions in model-based reinforcement learning
Miles Hutson Stanford University hutson@stanford. edu
Isaac Kauvar Stanford University ikauvar@stanf ord. edu
Nick Haber Stanford University
nhaberOstanf oorttarrforr...... edu
Abstract
Model-based reinforcement learning (MBRL. is promising route 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 exhaust its capacity on meaningless content, the cost of ne- glecting important environment dynamics. While existing approaches attempt to solve this problem, we highlight its continuing impact on leading MBRL meth- ods includinn DreamerV3 and DreamerPro with novel environment where background distractions are intricate, predictable, and useless for planning future actions. To address this challenge we develop method for focusing the capacity of the world model through synergy of a pretrained segmentation model, task-aware reconstruction loss, and adversarial learning. Our method outperforms a variety of other approaches designed to reduce the impact of distractors, and an advance towards robust model-based reinforcement learning.
Introduction
Model-based reinforcement learning (MBRL) is a promising path to data-efficient policy learning, and recent advances show impressive performance with high dimensional sensory data [Hafner et al., 2023]. A central component of MBRL is a world model, which is trained to predict how an agent's actions impact future world states. However, the world is highly complex while the capacity of a world model is finite, and ultimately only subset of the components and dynamics of the environment be accurately modeled. this setting, distracting stimuli can particularly problematic, as they waste the capacity of the world model useless details.
To address the challenge of distractors, number MBRL. methods seek to isolate the most important components of an environment, including structural regularizations [Deng et al., 2022, Fu et al., 2021, Wang et al., 2022], pretraining the agent's visual encoder [Seo et al., 2022, Wu et al., 2023], and value-cquivalent world modeling [Schrittwieser et al., 2020]. while environments such as Distracting Control Suite have been developed to assess distractor suppression [Stone al., 2021].
In this paper we introduce a new method, Policy-Shaper Prediction (PSP), for identifying and focusing on the important parts of an image-based environment. Rather than relying on pre-imposed structural regularizations, PSP learns to prioritize information that is important to the policy. We synergize task-informed gradient-band
loss weighting, use of a pre-trained segmentation model [Kirillov et al., 2023] and adversarial learning to create iistractiionsuppressiin agent that outperforms leading image-based MBRL agents. addition to exhibiting similar performance in distraction-free settings and on standard benchmark of robustness to distractions, our method markedly improves performance the face of particularly challenging distractors that are intricate but entirely learnable. Because learnable distractors can be accurately modeled, they straightforwardly contribute reducing the world model's 's error, but needlessly exhaust the capacity the world model.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

In sum, we make the following key contributions:
We describe Policy-Shaped Prediction (PSP), MBRL method that achieves strong distrac- tion suppression by combining gradient-based loss weighting with a pretrained segmentation model to focus learning on important environment features.
We augment PSP with a biologically-inspree action prediction head that reduces sensitivity to self-linked distractions.
We introduce challenging new benchmark for testing robustness learnable distractions.
We demonstrate that PSP achieves 2x improvement in robustness against challenging dis- tractions while maintaining similar performance in non-distracting settings.
2 Policy-Shaped Prediction
We introduce PSP, a method reduce an agent's sensitivity to useless distractions by focusing on sensory stimuli that are most relevant to its policy, rather than seeking to model everything in the environment. Our guiding intuition was that we can use the gradient from the policy the input image to identify important pixels in the environment, and that salience signals to identify important objects by using image segmentation. Specifically, we extend the principles of VaGraM [Voelcker et al., 2022] to a high-dimensional vision model by using explainability-relane notions of salience and aggregating an otherwise noisy gradient-basse salience signal within objects. Additionally, inspired by the biological concept of efference copies [Crapse and Sommer, 2008], which are neural signals used to cancel out sensory consequences of an animal's actions, incorporated a way explicitly mitigate distractions caused by actions of the agent itself.
aggregate these pixelwise
PSP employs (1) gradients of the policy with respect to image inputs to identify task-relevant elements of the image, (2) segmentation model to aggregate gradients within each object in the image, and (3) an adversarial objective to the image encoder of the world model that discourages encoding of duplicate information about the previous action. Figure 1 illustrates the training modifications made by this method to the underlying DreamerV3 [Hafner et al., 2023] architecture. Notably, since these modifications only affect the training stage of the world model, the DreamerV3 agent remains unaltered during inference. Below, we describe each of the three key components in detail.
Policy-shaped world model learning
Policy to image gradient flow *...
Policy-shaped loss weight
Tt
Ut
at
1++1
Subtract gradient
à
Image Encoder
Image Decoder
Image Encoder
Segmentation
Action prediction
head
Policy-shaped loss weight
Aggregate by object
Weighted Loss
Policy-shaped loss weight
Action+Pixel
gradient
Figure 1: Policy-Shaped Prediction in an environment with challenging distractions. (left) Training of etherwise-unaltere DreamerV3 agent is modified in ways: A head is added predict the previous action based on the image encoding, and the gradient of the head is subtracted from the gradient of the image encoder, and the loss scaled pixelwise by policy-shaped loss weight. (right) The loss weight uses the gradient of the policy to the input pixels. The image is segmented, and the pixel weights are averaged within each segmented object. Dashed lines signify gradient flow.

---

2.1 Task-informed image reconstruction with policy-gradient weighting
Our approach builds upon the core idea that signals most important to the actor and/or critic should be given special importance in the world model. The concept of using the critic to inform model loss was applied in Value- GradieGr weighted Model loss (VaGraM) [Voelcker et al., 2022], which weights the model loss according to the gradient of the value funetion with respect to the state. We extend this concept to high-dimensional image inputs, which previous work did not demonstrate. This extension to the image domain is inspired by gradient- rradeent-based interpretability methods such as saliency maps [Simonyan et al., 2013, Shrikumar et al., 2017, Ancona et al., 2019].
By upweighting the reconstruction loss for parts of the image that inform value estimation, we might expect to improve the performance of a downstream policy that aims to maximize value. Going one step further, we propose using the gradient of the policy for weighting the model loss. While VaGraM focused solely on the value function, we hypothesize that the gradient of the policy may provide an even more informative signal because ultimately, the state representation must support effective action selection. We hypothesize that the set of signals informing action selection may be richer than those that inform value estimation, which might rely primarily on simple cues such as whether an agent has flipped over. In contrast, the signals needed to select actions can be more subtle, such as the distance of an agent's leg from the platform it pushes off of order to run.
To compute the policy-gradient weighting, we first sum across the dimensions of the action vector ({(())). where s the latent state of the world model, to produce scalar a Σ3 azj and then take the gradient with respect to the pixels of the input image To apply this weighting in the context of DreamerV3 [Hafner et al., 2023], we scale the image reconstruction loss term at each pixel i, for reconstructed image a.
Limage
Σ 22
da
(#; xz
(I)
Ox:
2.2 Object-based aggregation of gradient weights
Gradient-based weighting of the world model's reconstruction ultimately is form of applying model explainability methods, which attempt to highlight the most important elements of a model's input for its outputs. From model explainability literature, a known challenge with gradient- based weighting is its noisiness, which is likely caused by the presence of sharp but meaningless fluctuations in the derivative at small scales [Smilkov et al., 2017]. While this has been combated by more computationally demanding explainability approaches such as Integrated Gradients [Sundararajan et al., 2017] and SmoothGrad [Smilkov et al., 2017]. we found these to be infeasible to run within the train loop, since this would require taking the derivative of the function with respect to multiple varying inputs for every example in the original input batch. Instead, to combat this problem introduce a second novel contribution: object-based aggregation of an explainability signal using a segmentation model (SEG). In principle many models should work, but ensure had high quality segmentations, we used the Segment Anything Model (SAM) [Kirillov et al., 2023], a pre-trained and broadly applicable segmentation model. Nothing prevents different model from being utilized, so long as it is of sufficient quality. This idea follows from the observation that the saliency map tends show a broad, if noisy, correlation with relevant regions in the image (e.g. Figure 1). During data collection, we segment each image into object masks using the existing method laid out by the authors of SAM. We prompt the model with a grid of 256 points, filter the resulting masks with metrics for the SAM algorithm (Intersection- Over- Union and a "stability score"), "). followed by maximmm suppression. We also include a mask capture pixels that are not otherwise assigned to an object. Then, instead of weighting the image reconstruction loss with the raw gradient-based salience, we use salience score that is aggregated within each segmentation mask. The weight of a pixel x; that has been assigned to segment SEG(,;;) is the mean absolute value weight of the pixels in SEG(r;;)
W SEEG(x;))
10a/Onyl
(2)
(ES
3

---

To ignore any exploding gradients, we clip the raw salience map to the 99th percentile before aggregation. Gradients are sometimes near zero in the beginning of training, and in the rare case that all gradients are zero, we set W; = 1 for all i, As regularizer, we linearly interpolate between the salience weighting and a uniform weighting, with a = 0.9 for all our experiments, and using rescaled W = width height. / E W, to match the scale of the uniform background.
W = aW + (1-a)
This regularizer lets the world model maintain reasonable reconstruction of less-salient aspects of the environment, which the actor-critic function can use as it learns. Without some degree of reconstruction of less-salient regions, the model can become trapped in local minima with a bad policy and world model. Also, at the start of training, gradients from actor-critic functions are essentially random and often small. A uniform background offers a reasonable prior begin the train loop.
2.3 Adversarial action prediction head
We additionally sought to explicitly reduce the distracting sensory impact of an agent's own actions. As animals move, they experience sensory signals generated by their actions and the external environment, and they have evolved the ability to distinguish these signals using efference copies [Crapse and Sommer, 2008]. We hypothesized that could separate information about an agent's actions and its encoding of external stimuli through domain-adversarial training [Ganin et al., 2016]. To this end, we introduce an adversarial action prediction head that prevents the model from wasting capacity on irrelevant stimuli that are created by the agent's 's own actions.
The DreamerV3 world model consists of three main components: convolutional neural network (CNN) image encoder 2; ~ (2 l/, 6x) with e, = CNN, which processes the input image, serves as a prior during training, and encodes the environment state during inference; a recurrent state space machine (RSSM) consisting of h2 = fo (h,,1, 24-1, 4;-1) and & ~ p (hx) that is trained to simulate the progression of latent states given actions; and an image decoder, & ~ Po (āg h4. 2)) which reconstructs the image from the latent state. Problematicallyy. the encoder can capture information about previous actions from the image, despite this information already being provided directly to the RSSM through the action input. In other words, may source information about a directly through Tt, despite o4-1 being an argument to fa during the computation of hr- Unfortunately, our reconstruction loss weighting may not solve this problem, since during backpropagation from the actor-critic functions, we do not distinguish information about previous actions that comes from the image versus the action input to the RSSM.
To prevent the CNN encoder from wasting capacity on encoding duplicate information about an agent's 'ss i actions, we add a small multilayer perceptron (MLP) head that is optimized to predict the previous action from the image embedding.
â MLP (stop
(4)
LadvHedd (āg- (0-1);0-1) = (ā.-1 a=11²²
(5)
When updating 0 during world model training, subtract the scaled gradient VoL(ū, ar-1) from the overall world model gradient, with = 1e3. This forces the latent state's 's previous action information to come solely from the provided action vector.
Our training procedure for a DreamerV3 agent is shown in Algorithm 1. We note that it should be possible to apply these concepts of gradient-based weighting, segmentation semeentation-based aggregation, and adversarial action prediction to world models other than our chosen DreamerV3 architecture.
3 Experiments
To evaluate the model's performance we design our experiments around the following questions:
Q1. Is agent robust against distractors which are learnable by the world model, but of no utility for the actor-critic?
Q2. What aspects of the environment are assigned importance by our method?
Q3. Is our agent robust against distractors that are unrelated the agent's actions?
Q4. Does agent maintain performance in standard, lower-distracteo environments?
Q5. What are the contributions of each component of our method?

---

Algorithm Policy-Shaped Prediction training (for DreamerV3)
1: Input: World model parameterized by Φ, policy paramaterized by 0. image encoder parametrized by p. replay buffer with image transitions (#t-1, 0t-1, Lt, Tt. q)), SEG segmentation model (SAM, in our application), action prediction MLP parameterized by
2: for training iteration I, 2, do
3: Sample batch of transition sequences
4: G aa
# Gradient of policy with respect to input image pixels Segmentation of input image # Aggregate gradient using segmentation # Normalize weighting
5: S SEG(z)
6: W = (g
7: W width height- W E; W
8:
W aW' + (1 a) shape (W)
Ahape
# Linearly interpolate with uniform weighting 24, h4) In Pa(Ce 24, h4) # Weighted DreamerV3 prediction loss
9:
(o) (n 24, h4) @ W In (r
(0) = Eg E (predd Spedd (0) + Paya Ldya(() + Prep # DreamerV3 model loss
10:
â--1 MIPP (stop Sarp_grrdd()) (
Adversarial action prediction head
12:
① f Adam( VL (L 0--1 / /0p
13:
LadvHead (ā.- at-1) (dc-1 - at-
14:
f Adam (aamV/Ze LAdvHeadd W)
15:
end for
3.1
Experimental details
Baselines We test four Model-Based RL approaches as baselines: DreamerV3 [Hafner et al., 2023], and three methods specifically designed to handle distractions Task Informed Abstractions [Fu et al., 2021], Denoised MDP (method in their Figure 2b) [Wang et al., 2022], and DreamerPro [Deng et al., 2022]. Additionally, we choose DrQv2 [Yarats et al., 202 [a] as a representative baseline Model-Free Free approach. For all agents, we use random seeds per task, and default hyperparameters.
Environment details Visual observations are 64 x 64 x 3 pixel renderings. We test performance in three environments: DeepMind Control Suite (DMC) [Tassa et al., 2018J, Reafferent DMC (described below), and Distracting Control Suite [Stone al., 2021] (with background video initialized to random frame each episode, 2,000 grayscale frames from the "driving car" Kinetics dataset [Kay et al., 2017]). For each environment, we test two tasks: Cheetah Run and Hopper Stand. We selected these tasks because they present different levels of difficulty, allowing us to assess how instractionseessiiiiiy depends on task difficulty. For ablation experiments, test Cheetah Run.
3.2 Reafferent Deepmind Control Suite
In the natural world, distractions can be highly complex, but in many cases are also highly pre- dictable. For instance, the creaking sound rusty bicycle makes as you pedal, or the movement of your shadow you dance outside. We wanted an environment which would allow in- vestigation of how well existing methods would perform in scenarios where the representational complexity of the distractions are very high, but they cannot simply be ignored as "unlearnable" noise. Distinguishing this type of partly self- generated distraction requires identifying which parts of the world relevant taking action, not just those affected and unaffected by our ac- tion. To achieve this, we devised the Reafferent Deepmind Control environment, in which the distracting background images have substantial content, but they depend deterninistically the agent's previous action and the elapsed time (Figure 2). We build the Distracting Control Suite [Stone et al., 2021], using a background
Q--1 = action #1 Background #11
a; action #2 Background #2
Distracting background is complex, but entirely predictable based agent's previous action.
Figure 2: Schematic of the Reafferent Deepmind Control environment. The distracting background is entirely predictable based on the agent's 's previous action and the elapsed time in the episode.
the episode and are thus completely predictable
5

---

Cheetah Run (Reafferent DMC)
Hopper Stand (Reafferent DMC)
250
PSP (Ours)
400
DreamerV3
TIA
200
Denoised MDP
300
Dreamer Pro
150
Score
Score
yan W
200
100
100
50
0
01
0k
200k
400k
600k Steps
800k
1000k
200k
400k
600k Steps
800K
1000k
Figure 3: Training curve comparisons on Reafferent Deepmind Control. Mean + std. err.
consisting of 2,500 16 16 grids, with each grid cell filled by randomly chosen color. We then map a tuple of time (625 possible values) and a discretized version of the first action dimension (4 possible values) to assigned background. We devised this be analogous to the types of high complexity self-generated distractors found the natural world (e.g. one's ' shadow).
This environment allows to stress test the structural regularizations (and associated priors) that form the foundation of many existing distraction-avoidance methods. Many methods encode assumptions about the forms distractors will take (usually uncorrelated to agent actions, reward, or both), rather than a means of generally identifying and ignoring distractors. We hypothesize that a learning- based approach, in which we avoid distraction by learning what is actually important for the agent to get things done, has the potential to overcome even learnable-but
distractions.
We find that none of the baseline MBRL meth- ods perform well on the Reafferent Environment, relative to their published results on the Distract- ing Control Suite or their performance on the unmodified Deepmind Control Suite (Table 1, Figure 3). We find that the world models learn excellent reproductions of the distracting back- ground. However, the cost of this is that the reconstruction of the agent becomes less well- defined or even replaced by the background, es- pecially in positions where the outcome of movement is uncertain (see Figure and A1 for examples with DreamerV3 and A2 with De- noised MDPs). This matches our expectation that the model will waste capacity on trivially predictable dynamics, rather than on the much more important but uncertain agent dynamics. As expected, unlearnable distractions are less challenging (Figure A3).
True
Recon.
Diff.
DV3
PSP
Figure 4: Reconstructed image comparison, PSP vs. DreamerV3 on Reafferent Cheetah Run, same episode and time point. True, reconstructed, differ- ence (true recon.). DreamerV3 accurately repro- duces the background but not the cheetah.
Hopper
Cheetah
Notably, the model-free DrQv2 agent shows a reduction in performance from its previous performance the unmodified environment, but overall demonstrates quite robust performance (Table 1). This also matches our expectations, since the CNN encoder is learned as part of the policy in model-free learning, unlike with model-based, where the learning objective for the world model is separate from the learning objective for the policy.
Our method demonstrates substantial improvement over the existing baselines (Table 1, Figure 3). Although it shows a higher than desired level of variance between runs, especially the more challenging Hopper Stand task, it nevertheless achieves scores beyond the reach of any of the baselines. We believe this affirmatively answers Q1 by showing new
Figure 5: Example salience maps (policy- shaped loss weights) highlight the agent.

---

Cheetah Run (Distracting Control Suite)
Hopper Stand (Distracting Control Suite)
700
PSP (Ours)
600
DreamerV3
600
500
TIA
500
Denoised MDP
400
Dreamer Pro
Score
400
Score
300
300
200
200
100
100
0-
0-
Ok
200k
400k
600k
800k
1000k
0k
200k
400k
600k Steps
800k
1000k
Steps
Figure 6: Training curve comparison on Distracting Control. Mcan + std.
agent is in fact robust to challenging distractors. Addressing Q2, we also find that the salience maps, derived from the gradient the policy and used to weight the world model loss, highlight the regions of the image that we would expect (Figures 5 and A4). Interestingly,
that sometimes the cheetah's rear leg is highlighted when it is the only leg close to the ground, though in other instances the entire cheetah is highlighted (Figure A4).
3.3 Performance on unaltered DMC and Distracting Control Suite
On Distracting Control tasks, in which the background distractor is uncoupled from the agent's actions, PSP produced consistently improved performance relative to baseline DreamerV3, in contrast to the more variable performance of DreamerPro, TIA, and Denoised MDP (Table 1, Figure 6). This addresses Q3.
Importantly, PSP also shows comparable performance to other methods (including DreamerV3) on the unaltered Deepmind Control Suite, demonstrating that have not introduced tradeoff between performance on distracting and non-distracting environments (Table I, Figure A5), resolving Q4.
In sum, PSP exhibits similar performance to baseline methods in commonly used tests of distractor- suppression and in non- do-dstttacting environments, while also demonstrating unmatched performance on particularly challenging distractors that complex but learnable.
3.4 Ablation study
To understand the contributions of each sub-component of the method (i.e. address Q5), we conduct ablations on the reafferent and unaltered Cheetah Run (Table 2). We find that some ablations trade off performance between the environments, while our complete model has good performance on both.
For instance, the top-performing method the reafferent environment does incorporate the segmentation or adversarial components, and uses the value gradient rather than the policy gradient. However, the variance of its scores higher than any other approach, and more problematically, shows the worst performance of all experiments in the unaltered environment. We believe this occurs
Table 1: Performance comparison across environments. DrQv2 is model-free, all others are model- based. TIA is task-informed abstraction, dMDP is denoised MDP, mean + standard deviation.
Task
DrQv2
D3eamerV3
DreamerPro
TIA
dMDP
PSP
Reafferent Control
Cheetah Run
565.1 + 35.5 210.3 +: 353.8
158.4 + 45.7 4.6 3.9
39.7 + 9.0
200.4 + 203.9
6.7 4.3
383.1 + 23.8 128.5 215.7
Hopper Stund
1.0
0.9
1.7 2.5
Unmodified Deepmind Control
Cheetah Run
736.0 17.0
521.1 136.3 867 + 15.9
908.4 1.6 890 + 11.2
773.7 22.7
763.0 62.8
712.3 32.3
Hopper Stund
752.9 + 206.8
298.4 + 512
897.9 14.2
865.6 = 53.6
Distracting Docpmind Control
Cheetah Run
364.4 + 60.7 781.1 1:
243.8 81.2
179.1 + 24
548.5 +: 238.9
397.4 + 111.8
40K.6 = 125.1 118.9
Hopper Stand
173.7
561.8 1: 103.1
200.5 1 171.7
13.2 16.5
7

---

Table 2: Performance of ablated versions of PSP. Scores are shown for Cheetah Run in unaltered and reafferent Deepmind Control environments. The unablated PSP achieves good performance on both environments, and while some ablations achieve slightly better scores on either unaltered or reafferent, they trade off performance in the other environment.
Gradient
Gradient weighting with segmentation
Adversarial Action Head
Unaltered
Reafferent
weighting
Policy
712.3 + 32.3
383.1 +: 23.8
Policy
653,9 = 44.2
231.2 + 58.6
Policy
X
742.1 79.7
188.4 9.4
None
X
6
324.4 + 2.3
Policy
X
X
418.2 + 53.1
379.0 + 46.8
Value
X X
381.7 64.9
445.7 + 126.9
None
521.1 136.3
158.4 + 45.7
because of the flaws in using only the gradient as an explanation for pixels that explain the actor-critic output. These flaws are more evident in an environment where the background never changes, as the policy is not required to learn any robustness to shifts in the background. In other words, the gradient for changes in the static background may be quite large, since the model is immediately out of domain when the background changes. We find that iigmentationnbaseer aggregation is critical to improving our model's performance amid distractors, while also maintaining its performance in the non-adversarial baseline. Overall, the results of the ablations confirm that combining segmentation, policy gradient sensory weighting, and adversarial action prediction results in the best scores across the unaltered and reafferent environments.
We also investigated the importance of the weight interpolation (Equation 3). We find that interpo- lation produces the expected benefit of allowing the agent to construct a useful world model, even when the policy is not very good, and thus sidestepping the 'chicken-and-egg" problem where the agent has neither good policy nor good world model (Figure A6). Furthermore, we tested the ability of PSP adapt either task change (Figure A7) or change in the distractor (Figure A8), and we found that PSP was able to quickly adapt in both scenarios.
3.5 Additional segmentation models
We additionally tested the sensitivity of PSP to the segmentation model. Given that segmentation models are likely to continue improving over time, we won- dered 1) whether PSP could be compati- ble with other models besides SAM, and 2) how PSP performance might be mod- ulated by the performance the segmen- tation model. To investigate, we used the recently released SAM2 [Ravi et al., 2024], which has multiple model sizes that allow for trading off performance for segmentation speed, with as high 6x faster segmentation speeds than the orig- inal SAM. We updated PSP use the 'tiny' SAM2 model, the smallest and low- est accuracy of the provided model sizes. Our basic implementation with SAM2- tiny immediately improved segmentation speeds (and thus reduced the resources necessary for segmentation) by 2x. We found that PSP with SAM2-tiny yielded nearly identical performance as the original PSP with SAM on all three Cheetah environments (Table 3). On Hopper, substantially harder task, we observed increased sensitivity to the quality of the segmentation. PSP with SAM2-tiny performed the same as PSP with SAM in the unmodified environment. In the Reafferent environment, PSP with SAM2-tiny still outperformed all baselines, with one of three runs yielding successful policy (compared with zero out of twelve total runs across all basclines), though the successful run yielded lower score (81.7)
Table 3: Performance comparison of different segmen- tation models. SAM2 is the tiny model size of the new, faster Segment Anything model. SAM2-tiny performs similarly on Cheetah Run, but is worse on the more chal- lenging Hopper Stand. SAM2-large recovers some of the decreased performance on distracting variants of Hopper. The full training plots are included in A10 and All.
Task
SAM
SAM2 (tiny)
SAM2 (large)
Reafferent Control
Cheetuh Run Hopper Stand
383.1 23,8
398.7 # 70.5 29.4 1 45.3
130.3 214.1
41.4 = 74.4
Unmodified Deepmind Control
Cheetah Run Hopper Stand
712.3 :: 32.3 865.6 53.6
696.3 1 20.0 89 39.6
Distracting Deepmind Control
Cheetah Run Hopper Stand
+
417.7 1 118.9
187.4 1: 172.0
465.0 1 166.6

---

than the successful run with SAM (377.5), which was also one of three runs. Additionally, SAM2-tiny yielded significantly worse performance on Hopper in Distracting Control. We hypothesized that this environment is particularly challenging for segmentation because the distracting background is black and white video, which is difficult to discern from the platform. SAM2-tiny is less successful at segmenting the platform, which is problematic for learning good policy. We further hypothesized that this issue would be resolved by improving the performance of the segmentation model. We tested this by using the SAM2-large model, and found that indeed this recovered performance back to the level of the original SAM. yielding a score of 465 + 166.6. 6. We additionally tested SAM2-large on Reafferent Control Hopper and observed a similar pattern with, again, one out of three runs yielding a successful policy with a score (1 14.44) that improved on SAM2-tiny. Optimizations, including using SAM2's 's video segmentation capabilities and better utilization of the GPU, would likely further improve segmentation speed. These results also suggest that as long as the segmentation is 'good enough' to properly segment the environment, PSP is not very sensitive to the segmentation algorithm. We also note that objects can be over-segmented into multiple segments without causing problems (Figure A9), and thus adequate segmentation is not particularly stringent requirement.
4 Related Work
Distraction-sensiiivvty of model-based RL Recent advances in Model Based RL (MBRL) includ- ing World Models [Ha and Schmidhuber, 2018], SimPLe [Kaiser et al., 2019]. MuZero [Schrittwieser al., 2020], EfficientZero [Ye et al., 2021], DreamerV1 [Haffer et al., 2019], DreamerV2 [Hafner et al., 2020], and most recently DreamerV3 [Hafner et al., 2023] have surpassed model-free RL in settings such as Atari, Minecraft, and Deepmind Control Suite. One deficiency of current MBRL algorithms is susceptibility of the world model to become overwhelmed by easily predictable distractors, in part due to mismatch between the objectives of the policy (maximizing reward) and the world model (accurately predicting future states) [Lambert et al., 2020].
One line of work attempts to address the distractability of MBRL through structural regularizations. Deng et al. [2022] uses contrastive learning of prototypes instead of image reconstruction. Lamb al. [2022] introduces the Agent
State Discovery algorithm, which discards information not relevant to elements of the environment within the agent's 's control. Task Informed Abstractions (TIA) identifies task-relevant and task-irrelevant features via an adversarial loss on reward-relevant information [Fu et al., 2021]. Denoised MDPs extends TIA's factorization to include notions of controllability [Wang al., 2022]. Clavera et al. [2018] use meta-learning and an ensemble of dynamics models. These works form strong body of solutions, given prior knowledge of likely distractors, but they can struggle if distractor does not fall into the designed regularizations.
different approach instead learns what is important by using the actor-critic functions to scale the importance of various learned dynamics. VaGraM uses value gradients reweight state reconstruction loss [Voelcker et al., 2022], building on Lambert et al. [2020] and VAML [Farahmand, 2018], 88, but VaGram does not operate on visual tasks. Eysenbach et al. [2022] propose a single objective for jointly training the model and policy. Goal-Aware Prediction learns joint representation of the dynamics and a goal, by predicting a goal-state residual, although they describe this approach as likely still susceptible to distractions [Nair et al., 2020]. Seo et al. [2023] decouples visual representations and dynamics via an autoencoder, improving the performance of Dreamer on tasks involving small objects. Value-equivalent agents [Grimm et al., 2020], such as MuZero [Schrittwieser et al., 2020] or Value Prediction Networks [Oh et al., 2017], construct world model that only aims to represent dynamics relevant to predicting the value function, in contrast to methods such as Dreamer that aim to learn the broader dynamics of the environment. MuZero is very effective in settings with discrete actions such as Atari, Go, and chess. Adaptation to domains with complex action spaces such as Deepmind Control Suite [Hubert et al., 2021] have shown some success, however Dreamer- based agents that include image reconstruction for world model learning can still exhibit superior performance, and the image-related signals have been shown to be essential to their performance [Hafner et al., 2020]. Building on these methods, our work investigates how to combine the benefits of both image reconstruction and task-aware modeling, through policy-shape image- world modeling, by applying concepts from VaGraM the image-based MBRL setting.
Distraction-sensiiivity of model-free RL A parallel track of Model Free RL. (MFRL) has its own body of literature, with leading method DrQv2 [Yarats et al., 202 used in this paper

---

for comparison. DrQv2 is an off-policy actor-critic RL algorithm that operates directly on image observations, using DDPG as the base RL algorithm alongside random shift image augmentation for image generalization and sample efficiency. Although our work focuses on a solution to MBRL's iestraction-sensiivvvii it is worth noting analogous deficiencies can exist in MFRI, and there are a number of works addressing these. For instance, Mott et al. [2019] uses an attention mechanism to make the agent robust to environment changes, and Tomar et al. [2024] learn task-relevant inputs mask. Yarats et al. [2021b] [0], an inspiration for DreamerPro, creates prototypes that compress the sensory data, Grooten et al. [2023a] apply dynamic sparse training to the input layer of the actor and critic networks, and Grooten et al. [2023b] mask out non-salient pixels based on critic gradients.
5
Discussion
PSP combines three ideas to focus the capacity of an agent's world model on aspects of the environ- ment that are useful to its policy. First, the gradient of the policy with respect to the input image is used to identify pixels that influence the policy. Second, the importance of individual pixels are aggregated by object, using a segmentation model to identify objects. Third, wasteful encoding of the preceding action (which is known and does not need to be predicted) in the image embedding is removed using an adversarial prediction head. Together, these allow an agent to construct a world model that best informs its policy, and in doing so, the policy to shape what information is prioritized by its world model. The outcome of this process is an agent that is selective about what parts of the world it models, and that becomes resilient against enticingly learnable, but ultimately empty, distractions.
Our work draws a connection between the use of the value function gradient in VaGraM and related concepts from the vision model explainability literature. The value gradient can seen as analogous to saliency maps [Simonyan et al., 2013]. Other gradient-based attribution methods, such as those that multiply saliency maps by input intensities [Shrikumar et al., 2017] or Integrated Gradients (Sundararajan et al., 2017, Ancona et al., 2019] offer additional ways to perform attribution. Some gradient-based attribution methods, such as Integrated Gradients, can be computationally expensive due to the need to approximate incorporation of more advanced explainability methods such as these into PSP, and the concept of an agent 'interpreting itself" may exhibit broader utility. Finally, recent work that uses SAM combined with human supervision to improve the generalizability of model-free RL [Wang et al., 2023], together with our work, point towards the potential value of incorporating powerful object segmentation models into reinforcement learning systems.
integral over the input space. Future work may investigate
Limitations Limitations of PSP include its fundamentally object-centric view, which assumes that pixels belong to single objects, and that the objects can be ranked by their importance. Additionally, the SAM segmentation model requires significant compute, but these models will likely improve over time and can also be application-tailon inference of the world model and policy, only training. Finally, it is not yet clear how well PSP will adapt in environments where the reward structure or salient features change across time. PSP may make the world model more task-specific than other approaches. although it does keep some reconstruction weight on non-task-relevant features and we observed initial evidence of resiliency (Figures A7, A8).
Notably, however, segmentation is not necessary during
Outlook Our work finds headroom to improve the robustness of MBRL. to distractions by linking the actor-critic functions and the reconstruction loss and leveraging useful priors from pre-trained foundation models. The findings here open other lines of inquiry such as using better model explanation techniques or more explainable architectures, utilizing faster segmentation models, and utilizing segmentation models designed for videos, in order to do temporal aggregation. Substantial work likely remains to improve the speed of this technique and find extensions that allow it to reliably work for harder problems, such as applied robotics. The adversarial action prediction head's inspiration from the biological concept of efference copies also suggests there is still space MBRL to consider biological metaphors as helpful design principles for learning algorithms. In sum, we present PSP, a method for avoiding distractors by focusing the world model on the parts of the environment that are important for selecting actions.
10

---

References
M. Ancona, F.. Ceolini, Öztireli, and Gross. Gradient-based attribution methods. Explainable Al: Interpreting, explaining and visualizing deep learning, pages 169-191, 2019.
I. Clavera, J. Rothfuss, J. Schulman, Y. Fujita, T. Asfour, and P. Abbeel. Model-based reinforcement learning via meta-policy optimization, 2018.
T. B. Crapse and M. A. Sommer. Corollary discharge across the animal kingdom. Nature Reviews Neuroscience, 98):587-600, 2008.
Deng. I. Jang, and S. Ahn. Dreamerpro: Reconstruction-fees model-based reinforcement learning with prototypical representations. In International Conference on Machine Learning, pages PMLR, 2022.
B. Eysenbach, A. Khazatsky, S. Levine, and R. R. Salakhutdinov. Mismatched no more: Joint model- policy optimization for model-based rl. Advances in Neurul Information Processing Systems, 35: 2022.
A.-m Farahmand. Iterative value-aware model learning. Advances in Neural Information Processing Systems, 31, 2018.
X. Fu, G. Yang. P. Agrawal, and T. Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning, pages 3480-3491. PMLR, 2021.
Y. Ganin, E. Ustinova, H. Ajakan, P. Germain, H. Larochelle, F. Laviolette, M. March, and V. Lempit- sky. Domain-adversarial training of neural networks. Journal of machine learning research, 17 (59):1-35, 2016.
C. Grimm, A. Barreto, S. Singh, and D. Silver. The value equivalence principle for model-based reinforcement learning. Advances Neural Information Processing Systems, 33:5541-5553 2020.
B. Grooten, Sokar, S. Dohare, E. Mocanu, M. Taylor, M. Pechenizkiy, and D. C. Mocanu. Automatic noise filtering with dynamic sparse training in deep reinforcement learning. arXiv preprint wwXi:22302.06548 2023a.
B. Grooten, T. Tomilin, G. Vasan, M. Taylor, A. Mahmood, M. Fang, M. Pechenizkiy, and D. Mocanu. Madi: Learning to mask distractions for generalization in visual deep reinforcement learning. arXiv preprint arXiv:2312.15339, 2023b.
D. Ha and Schmidhuber. World models. arXiv preprint arXiv:1803.10122 2018.
D. Hafner, T. Lillicrap. J. Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603. 2019.
D. Hafner, T. Lillicrap. M. Norouzi, and J. Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02199 2020.
D. Hafner, J. Pasukonis, J. Ba, and T. Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301. 0.33000004044 0040 2023.
T. Hubert, Schritwwieser, I. Antonoglou, M. Barekatain, S. Schmitt, and D. Silver. Learning and planning in complex action spaces. In International Conference on Machine Learning, pages 4476 4486. PMLR, 2021.
L. Kaiser, M. Babaeizadeh, P. Milos, B. Osinski, R. H. Campbell, K. Czechowski, D. Erhan, C. Finn, P. Kozakowski, S. Levine, et al. Model-based reinforcement learning for atari. arXiv preprint 1Xi::0903300374, 2019.
W. Kay, J. Carreira, K. Simonyan, B. Zhang, C. Hillier, S. Vijayanarasimhan, F. Viola, T. Green, T. Back, P. Natsev, M. Suleyman, and A. Zisserman. The kinetics human action video dataset. CoRR,
2017. URL http://arxiv..rt atrg//bs/1705... 1701.06950.......
A. Kirillov, E. Mintun, N. Ravi, H. Mao, C. Rolland, L. Gustafson, T. Xiao, S. Whitehead, A. C. Berg, W.-Y. Lo, et al. Segment anything. Proceedings of the IEFF/CVF International Conference on Computer Vision, pages 4015-4026, 2023.

---

A. Lamb, R. Islam, Y. Efroni, A. Didolkar, D. Misra, D. Foster, L. Molu, R. Chari, A. Krishnamurthy, and J. Langford. Guaranteed discovery of models. arXiv preprint arXiv:2207.082229 2022.
latent states with multi-step inverse
N. Lambert, B. Amos, O. Yadan, and R. Calandra. Objective mismatch in model-based reinforcement learning, arXiv preprint arXiv:2002.04522 2020.
A. Mott, D. Zoran, M. Chrzanowski, D. Wierstra, and D. Jimenez Rezende. Towards interpretable re- inforcement learning using attention augmented agents. Advances in neural information processing systems, 32, 2019.
S. Nair, S. Savarese, and C. Finn. Goal-aware prediction: Learning to model what matters. In International Conference Machine Learning, pages 7207-7219. PMLR, 2020.
J. Oh, S. Singh, and H. Lee. Value prediction network. Advances in neural information processing systems, 30, 2017.
N. Ravi, V. Gabeur, Y.-T. Hu, R. Hu, C. Ryali, T. Ma, H. Khedr, R. Rädle, C. Rolland, L. Gustafson, et al. Sam 2: Segment anything in images and videos. arXiv preprint arXiv:2408.000 00714, 2024.
Schrittwieser, Antonoglou, T. Hubert, K. Simonyan, L. Sifre, S. Schmitt, A. Guez, E. Lockhart, D. Hassabis, T. Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-009 2020.
Y. Seo, K. Lee, James, and P. Abbeel. Reinforcement learning with action-free pre-training from videos, 2022.
Y. Seo, D. Hafner, H. Liu, F. Liu, S. James, K. Lee, and P. Abbeel. Masked world models for visual control. In Conference on Robot Learning, pages 132-1344. PMLR, 2023.
A. Shrikumar, P. Greenside, and A. Kundaje. Learning important features through propagating activation differences. In International conference on machine learning, pages 3145-3153. PMLR, 2017.
K. Simonyan, A. Vedaldi, and A. Zisserman. Deep inside convolutional networks: Visualising image classification models and saliency maps. arXiv preprini arXiv: arXiv: 131226004, 2013.
D. Smilkov, N. Thorat, B. Kim, F. B. Viégas, and M. Wattenberg. Smoothgrad: removing noise by adding noise. CoRR, ftbs/1006.0382.. 2017. URL http: http///axxvvv.e//ee//ee abtaaaas/////// 03825.
http:://aax....rrr/aa/////66
A. Stone, O. Ramirez, K. Konolige, and R. Jonschkowski. The distracting control suite-a challenging benchmark for reinforcement learning from pixels. arXiv preprint arXiv:2101.02722, 2021.
M. Sundararajan, A. Taly, and Q. Yan. Axiomatic attribution for deep networks. In International conference on machine learning, pages 3319-3328. PMLR, 2017.
Y. Tassa, Y. Doron, A. Muldal, T. Erez, Y. Li, D. d. L. Casas, D. Budden, A. Abdolmaleki, J. Merel,
A. Lefraneq, et al. Deepmind control suite. arXiv preprint arXiv: 2018.
M. Tomar, R. Islam, M. Taylor, S. Levine, and P. Bachman. Ignorance is bliss: Robust control via information gating. Advances in Neural Information Processing Systems, 36, 2024.
C. Voeleker, V. Liao, A. Garg, and A.-m. Farahmand. Value gradient weighted model-hased rein- forcement learning, arXiv preprint arXiv:2204. 00..004444 2022.
T. Wang, S. S. Du, A. Torralba, P. Isola, A. Zhang, and Y. Tian. Denoised mdps: Learning world models better than the world itself. arXiv preprint arXiv:2206./5477, 2022.
Z. Wang, Y. Ze. Y. Sun, Z. Yuan, and H. Xu. Generalizable visual reinforcement learning with segment anything model. arXiv preprint arXiv:2312.17116, 2023.
WI, H. Ma, C. Deng, and M. Long. Pre-training contextualized world models with in-the-wild videos for reinforcement learning, 2023.
12

---

D. Yarats, R. Fergus, A. Lazaric, and L. Pinto. Mastering visual continuous control: Improved data duta-augmented reinforcement learning. arXiv preprint arXiv:2107.09645, 2021a.
D. Yarats, R. Fergus, A. Lazaric, and L. Pinto. Reinforcement learning with prototypical rep- resentations. In International Conference on Machine Leurning, pages 11920-11931, PMLR, 2021b.
W. Ye, S. Liu, T. Kurutach, P. Abbeel, and Y. Gao. Mastering atari games with limited data, Advances in neural information processing systems, 34:25476-25444 2021.
13

---

True
Recon
Diff.
Loss Weight (our method)
DV3
PSP
Figure Al: Reconstructed image comparison, PSP vs. DreamerV3 on Reafferent Cheetah Run. This highlights how PSP focuses on modeling the hind leg (white arrow), while DreamerV3 focuses on modeling the background but fails to model the hind leg. From left: True, reconstructed, difference (true recon.), PSP salience loss weight.
True
Reconstructed
Figure A2: Denoised MDP reconstructs the background with high degree of fidelity, but does not clearly render the Cheetah agent.
DreamerV3: Cheetah Run
with white noise background distractor
500
DreamerV33
400
300
200
100
0
0k 200k 400K 600K B0OK
Steps
Figure A3: To more thoroughly show that the reafferent environment impacts DreamerV3 because of the learnable time & action mapping to backgrounds (and not purely because of the presence of the backgrounds themselves), we include this training curve from an environment which uses the same backgrounds, but with random choice of background at each timestep. This demonstrates effective policy learning in spite of the distracting (but unlearnable) background.
A Broader Impacts
At the current stage, this work remains reasonably far from any large socictal impacts, as it limited to agents interacting with small, simulated environments. Over the long term, however, if model-based RL algorithms are used control robots or internet-connected agents (such as large language model agents), the potential for both large positive and negative societal impacts becomes relevant. On the positive side, intelligent agents that are capable of modeling the world and avoiding distractors have the potential to aid humans in a wide variety of scenarios, from housework, to medical applications, the negative side, agents without proper safeguards have potential to inflict harm on humans and the environment, whether through negligence malfeasance. Ultimately, our work is targeted at producing the positive impacts, while still allowing for mitigation of the negative impacts.
to exploration, to internet research.
14

---

How salience maps change over training
Visual input
Sallence map
50K
100K
500K
950K
Training step
Figure A4: Improvement in specificity of the salience map across training.
Cheetah Run (Unmodified DMC)
Hopper Stand (Unmodified DMC)
800
800
600
600
Soore
Score
400
400
PSP (Ours)
DreamerV3
200
TIA
200
Denoised MDP
0
Dreamer Pro
0k
200k
400k
600k
800k
1000k
0k
200k
400k Steps
600k
800k
1000k
Steps
Figure A5: Training curve comparison on unmodified Deepmind Control. Mean + std. err. n=3 runs.
Cheetah Run (Reafferent Environment)
PSP
250
PSP without Interpolation
200
Score
150
100
50
01
100k
200k
3d0k 400k Steps
500k
Figure A6: Interpolation of the gradient based reconstruction loss and uniform reconstruction loss yields superior performance to gradient based reconstruction loss without interpolation, which manifests as faster initial rise in score with PSP. Interpolation allows the agent to construct a useful world model even when the policy is not yet very good (and hence minimizes the impact of the "chicken-and-egg" problem where the agent has neither a good policy nor a good world model). Mean + std. err., n=3 runs.
15

---

Adapt from Walker Run to Walker Stand (transition tasks at step 1M)
1000
PSP before Task Change
PSP after Task Change
800
Score
400
200
0
0k 200k 400k 600k 800K1000K
Steps
Figure A7: In order test the adaptability of model trained for task via PSP another task for the same agent, we switched from Walker Run to Walker Stand at step 1M, while leaving the Reafferent background unchanged. We that the learned model quickly adapt to the task, even with the Reafferent background. Mean + std. err., n=3 runs.
Adapt to dynamic distraction (distracting background changes 1M)
400
Score
Before Background Change
MSSP After Background Change
Ok
200k 400k 600k 800k1000k200k
Steps
Figure A8: order test the adaptability of PSP model to dynamic distraction, we switched the settings of the Reafferent distraction at step 1M. We see that the learned model can quickly adapt to the new distraction. Mean + std. err., n=2 runs.
Examples demonstrating that SAM segmentation is not perfect:
the cheetah is split into 2 (left), 3 (middle), or 4 (right) segments
Figure A9: How high quality must the segmentation be? We note while training that SAM did not have perfect segmentation performance our tasks. In fact, the Cheetah agent was usually segmented into several (and a varying number of) different regions. The algorithm proves robust. across these segmentation variations.
16

---

Cheetah Run (Reafferent Environment)
Hopper Stand (Reafferent Environment)
250
PSP
PSP with SAM 2 Tiny
400
200
PSP with SAM Large
300
150
Score
Score
200
100
100
50
PSP
0:
PSP with SAM 2 Tiny
01
0k
200k
400k
600K
800K 1000k
200k
400k
600k Steps
800K
1000k
Steps
Figure A10: Training curve comparisons on Reafferent Deepmind Control. Mean + std. of PSP with different segmentation algorithms: SAM, SAM Tiny, and SAM 2 Large. baselines are shown in Figure 3. Notably, all three PSP agents shown here achieve non-zero scores on Hopper Stand, in contrast to the baselines.
Cheetah Run (Distracting Control Suite)
Hopper Stand (Distracting Control Suite)
500
PSP
PSP
PSP with SAM 2 Tiny
500
PSP with SAM Tiny
400
PSP with SAM Large
400
300
Score
Score
300
200
200
100
100
0
0-
0k
200k
400k
600k
800k
1000k
200k
400k Steps
600k
800k
1000k
Steps
Figure A11: Training curve comparisons on Distracting Control Suite. Mean + std. eTr. of PSP with different segmentation algorithms: SAM, SAM 2 Tiny, and SAM Large. All baselines are shown in Figure 6.
B Computational Overhead
We characterize the computational cost of the PSP algorithm by ablating various components and measuring its training speed (Table A1). There four major components of PSP that affect performance. In order of decreasing computational cost: (1) Policy gradient- based weighting, (2) image segmentation (e.g. with SAM) (3) action adversarial head, and (4) segmentation-baset aggregation of the gradient weighting. These are all costs that apply only during training, not during
inference.
Policy gradient based weighting: For each latent state produced by the encoder RSSM in each step of a rollout, we take the gradient of the policy with regard to the image pixel inputs. In our implementation, this auto differentiation yields a complexity of O((E + R P) S2. -3.W- H), where E is the number of encoder parameters, R is the number of parameters of the RSSM, P is the number of parameters of the policy, and I/ is the image height.
the number of steps in a rollout, W is the width of the input image,
There are a couple of major opportunities for optimization in future implementations. First, for debugging and experimental reasons, we have been computing the gradient of the policy with respect to all rollout steps, instead of just the current step. This reveals an opportunity to reduce the computational burden by factor of S, where S is 64 in our implementation. Reducing this contribution could induce a significant speedup the performance of future implementations Second, we take the gradient of the policy with regard to the image during rollouts only to enable visualization
17

---

Table Al: Comparison of computational overhead. FPS stands for frames per second.
Adv Head??
SAM?
Policy VAML?
Interpolation?
FPS
Yes
Yes
Yes
Yes
5.0
Yes
No
Yes
Yes
5.9
Yes
No
No
N/A
17.1
No
No
N/A
19.0
for debugging and figures. This is strictly speaking unnecessary overhead and could be eliminated for a gain during training.
The details of the additional complexity depend on the specific algorithm used. With addition of results using the SAM2 tiny model, we demonstrate the viability of using different segmentation algorithms. Notably, the segmentation process is entirely parallelizable separately from the training process: images are segmented they are collected and stored in the replay buffer. We find that segmentation is not a bottleneck in training speed. Moreover we find that advances in segmentation algorithms (e.g. SAM2) reduce the computational resource requirements for this process.
Image segmentation:
For each training step, take the gradient of the loss of the action prediction head with regard to the parameters of the encoder. Therefore, the cost is O((E + R + A) S), where A is the number of parameters of the action prediction head. These gradients are added with the regular gradients during a train step, so the adversarial head does not introduce additional iterations during training.
Adversarial head:
Segmentation aggregation of gradients: We take the mean value of the gradient with respect to each mask by multiplying each mask by the gradient weighting and then taking the hadamard product of the resulting image (matrix) with the inverse of the sum of mask elements equal to This results in additional cost of ((M W H), where M is the number of masks.
C Experiments Compute Resources
Each trial of the PSP method used 4 Nvidia A40 GPUs to train the modified DreamerV3 model, and A40 GPUs to run the Segment Anything model in parallel. Given an estimated unique experiments for the final paper, 3 trials per experiment with our method, and about 1.5 days per training run, we used about 17 3 1.5 * GPUs 612 GPU days on A40 accelerators. Early experiments with this methodology likely used an additional 300. Baseline trials could be run on only single A40 GPU or a desktop NVIDIA 2070 SUPER, usually less than day, and accounted for a comparably negligible level of resources.
We believe this level of resource consumption could be easily reduced. The modifications the DreamerV3 model do not attempt to benchmark the most costly components. We suspect our method of parallelizing the new backpropagation from the policy the image could be optimized further from its naive Jax implementation. Additionally, SAM could be supplanted by a more efficient segmentation model. We focused on establishing the basic technique with SAM, and replacing it with more efficient methods should be the subject of future work.
D Code
The repository with code and instructions for reproducing these experiments available at this GitHub Repository.
18

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction clearly state our contributions and claims, and match the details in the results section.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Yes, the paper clearly points out what we believe to be core limitations and assumptions of our work, as well as present limitations that we not believe are inherent to the method.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
19

---

Answer:
[NA]
Justification:
The claims of this paper are tested empirically.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
The paper itself should include enough detail to reproduce our results with the open source SAM and Dreamerv3 models. Because the implementation is not trivial, we will also release the GitHub repository in the camera ready version, which includes an implementation of the core algorithm for DreamerV3 and a shared implementation of the test environments for DreamerV3 and every baseline. We have not included it in the review version the GitHub will identify the authors.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
20

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: An anonymized version of the code will be available the linked GitHub Repository for reviewers.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.e cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //n ctlpps..c//pbbic/guuu gwwdd/gu//////eeeeee /CodeSubmiee ssi onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We include all relevant details of our benchmark and (where diverge from the default) models.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification:
experiments are performed with three trials and std. dev is reported.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
21

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We have provided transparent and reasonable estimate of compute require- ments in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: cttt::/.e.....llleieccccccceeeee
Answer: [Yes]
Justification: The only new dataset involved in this work was generated by a small Python procedure and has privacy risks or ethical concerns. The implemented model modifies an open source repo and has had no interaction with human subjects.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Thstiiication:This is foundational work for improving MBRL and any societal impacts are at least one order removed, but we have outlined the possible societal impacts of improved MBRL in general.
22

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considcrations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning conrectly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper describes foundational change MBRL and introduces datasets that pose a risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All previous work is cited and proprietary code has been used beyond what is allowed by its license.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
23

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. cow comdaatasess has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Source code will be provided before publication if paper is accepted.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: No human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
25


Reviews:
[{"review_id": "GHf1knSDYf", "reviewer": "Reviewer_8n3B", "summary": "### Review Summary\n\nThis paper presents a novel approach to improving model-based reinforcement learning (MBRL) by identifying that detailed but irrelevant aspects of the world can exhaust the capacity of the world model, thus hindering the learning of important environment dynamics. The proposed method, Policy-Shaped Prediction (PSP), uses a pretrained segmentation model, a task-aware reconstruction loss, and adversarial learning to focus the world model's capacity on relevant environment features. The paper demonstrates that PSP outperforms other approaches, including DreamerV3 and DreamerPro, in environments with intricate but irrelevant background distractions.\n\n### Major Concerns\n\n1. **Effectiveness of Policy Gradient Information**: The utility of the policy gradient information when the policy is not sufficiently good is questionable. This raises a chicken-and-egg problem: if the policy is poor, the gradient information might not effectively guide the model's learning, but if the policy is already good, further gradient information might not be necessary. Clarifying this point is crucial to support the core contribution of the paper. I have the following guess given the positive results: either the gradient information provided early on is sufficiently informative or the policy and model improve together, even if initial gradients are inaccurate. The authors should provide further explanation on this aspect.\n\n2. **Motivation Comparison with Related Work**: The motivation emphasized in the paper—to distinguish between decision-relevant foreground and irrelevant background during model learning—has been similarly addressed in previous works on Block MDPs and invariant representations (e.g., [1], [2]). These works, although employing different technical approaches, share a similar motivation. The authors should compare their approach with these works, providing theoretical or experimental analysis to highlight the differences and advantages.\n\n3. **Dependence on Pretrained Segmentation Model**: Given that the segmentation model is pretrained, what is its range of applicability? How much do the final results depend on the effectiveness of the segmentation model? The authors only conducted experiments in two environments and did not explore this dependency, raising questions about the generalizability of the approach to new environments.\n\n4. **Experimental Section Clarity**: The experimental section starts with good questions and appears to be structured to answer them. However, the current presentation makes it difficult for readers to find the answers and understand which experiments address which questions. Revising the presentation for clarity would help make the experimental logic more apparent.\n\n5. **Applicability to Different Scenarios**: The three methods proposed (policy gradient with respect to state, segmentation, adversarial training loss) seem tailored to specific problems (e.g., decision depend on object segmentation, clear object boundaries). If the problem scenario changes slightly, will these methods still work? Even considering the current scenario, how many common RL problems can this approach cover? The experiments are primarily conducted on two tasks, which weakens the persuasiveness of the results.\n\n### References\n\n[1] Learning Invariant Representations for Reinforcement Learning without Reconstruction  \n[2] Invariant Causal Prediction for Block MDPs", "questions": "See above", "limitations": "See above", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "See above", "weaknesses": "See above", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments. \n\n**Effectiveness of Policy Gradient Information**: We do see some evidence of the chicken-and-egg problem. Importantly, though, our interpolated weighting term (see Eq. 3 in the manuscript) provides a route for escaping this problem by allowing the model and policy to always have access to cues from the entire state and to improve together. Thus, even if the initial policy-weighting is bad the model still has the chance to globally improve its predictions and recover from the initially bad weighting. As might be expected, following your great intuition, we find that the interpolation specified by Eq. 3 was essential to achieving consistently good performance. \nTo demonstrate the value of this interpolation step, we ran an experiment comparing agents with and without the interpolation, on Cheetah Run with Reafferent distraction. As shown in Fig. R1C, without interpolation (green line), the agent is stuck with poor performance for nearly 150K steps, until it eventually does manage to escape the chicken-and-egg problem. However, when including interpolation (blue line) we see that the agent is able to quickly achieve nonzero performance and begin steadily improving. \nAdditionally, in Fig. R1F, we demonstrate how the salience map improves throughout training, suggesting that the policy and world model can improve together with training.\n\n**Motivation Comparison with Related Work**: Thank you for bringing up the interesting connection to invariant representations. Indeed these methods are quite relevant, and it is valuable to consider methods for identifying state abstractions that leverage bisimulation [1] or the closely related concept of invariant causal feature prediction [2]. The DBC algorithm from [1] is most relevant to our context, as it is designed explicitly for high dimensional RL applications with distractions. On the other hand, the MISA algorithm from [2] requires multiple distinctly labeled experimental settings, and while they demonstrate imitation learning from high dimensional images, extension to high dimensional RL applications is not necessarily straightforward and could be a novel contribution in its own right. \n\nHowever, we do not think that it is apt to experimentally test DBC as a baseline in our setting for two key reasons: 1) The Denoised MDP work [Wang 2023] compares directly against DBC and demonstrates substantially better performance than DBC (see Table 1 in Wang 2023). In turn, our experiments compared directly against Denoised MDP and we observed that PSP exhibited substantially better performance than Denoised MDP. 2) Our focus is on model-based RL, but only a model-free version of DBC was published and it is not straightforward to consider how it would be applied in the context of Dreamer-type architectures.  For these reasons, we believe that our current experimental baselines are adequate to encompass these approaches. However, we think the connection is important and we will add discussion of invariant representations to the manuscript. Thank you!\n\n**Dependence on Pretrained Segmentation Model**:\nWe acknowledge that the use of a pretrained segmentation model has the potential to limit the range of applicability, however we also believe that foundational segmentation models such as SAM have grown robust enough to be considered for these applications. Moreover, future work will improve these models, with cheaper inference and higher quality and more generalizable segmentation – as evidenced by the recent release of SAM2.\nWe have now included experiments with an *additional* segmentation model, the 'tiny' variant of SAM2, and observe similarly good results as with SAM (Fig. R1A).\nAdditionally, we note the segmentation model does *not* have to be perfect. The SAM segmentation often splits apart what should be considered a single object. In Fig. R1B, for example, the cheetah is split into two, three, or four segments throughout a successful run. This suggests a degree of insensitivity to exact details of the segmentation. \nNotably, the Deepmind Control Suite setting is likely out of the domain of the data on which SAM or SAM2 were trained, and yet segmentations from these algorithms are still adequate to achieve good PSP performance. \n\n**Experimental Section Clarity**: Thank you. We will clarify the results section to explicitly indicate how we answer each of our listed experimental question (e.g. by adding references to Q1, Q2, etc. in the appropriate locations). Specifically: Q1 is answered in lines 199-228, Figure 3, and Table 1. Q2 is addressed by Figure 5 (and additional supplemental images of salience maps that we will include). Q3 is addressed by Figure 6 and Table 1 at line 233. Q4 is addressed by line 230, Figure A1, and Table 1. Q5 is addressed at lines 240-255, and Table 2. \n\n**Applicability to Different Scenarios**: We selected Deepmind Control Suite as the base environment to allow for direct comparison against a number of existing baselines, in the settings that they were demonstrated. Distracting Control Suite was the only environment shared across the DreamerPro, DenoisedMDPs, and Task Informed Abstractions baselines. Furthermore, we additionally incorporate the Reafferent environment to extend beyond Distracting Control Suite into more challenging scenarios. We note that the segmentation-based approach of PSP does assume object-based settings, however this assumption is relevant to a wide variety of problems. We acknowledge that there may be limitations of our approach in certain other scenarios, although it is not clear a priori what those limitations may be, and we look forward to pursuing future work to investigate this.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "b7PlF5B12L", "reviewer": "Reviewer_7m6w", "summary": "This paper presents a novel model-based reinforcement learning (MBRL) method that focuses on important parts of image-based environments with distractions, aiming to improve policy learning. The proposed method introduces gradient-based weighting, segmentation-based aggregation, and adversarial action prediction to world models. Experiments are conducted using the DeepMind Control Suite environment with varying levels of distractions.", "questions": "Please see the weaknesses section.", "limitations": "Limitations are discussed in the paper.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper introduces a new method that utilizes policy-gradient weighting and object-based weighting for image reconstruction.\n2. It devises a Reafferent DeepMind Control environment where the distracting background deterministically depends on the agent’s previous action and the elapsed time within an episode.", "weaknesses": "1. In Figure 1, is the arrow between the $z_i$ and the image decoder reversed? The notation of $i$ is confusing, as it is unclear whether it represents a time step or an image pixel (Line 86). The arrows between $z_i$ and $v_i$, $a_i$ seem to suggest that a value model and a policy model are learned in the world model learning phase. However, in my opinion, these two models should be learned in the downstream RL phase.\n2. I wonder if the policy used to learn policy-gradient weight is the downstream policy that aims to maximize value. If I understand correctly, the policy model is used, but its weights are not updated in the world model learning phase. Additionally, how the connection between the latent state $s$ and the image pixel in Eq (1) is established, considering that the input of the policy model is based on $s$.\n3. I have some doubts about the Segment Anything Model (SAM). Was it fine-tuned in the experiment? It would be helpful if the authors could provide more details on how SAM is used in this work and what the output of $SEG(x_i)$ is. SAM outputs multiple masks for different objects based on prompts, so how do you determine which mask to use?\n4. In Eq (4), inferring an action from the state of a single time step seems unreasonable. In general, actions are inferred from successive states. Moreover, I do not agree with the statement that the encoder extracting information about the agent's action is wasteful. A more reasonable explanation is needed.\n5. It is strongly encouraged for the authors to provide visual results of $W$ and segmentation masks in the experiments to facilitate clearer explanations and analysis.\n6. The experimental environment and tasks are too limited and simplistic to demonstrate the effectiveness of the proposed method. It would be beneficial to use a more realistic environment, such as CARLA, which also includes various distractions related and unrelated to the agent's actions.\n7. The performance of the proposed method appears relatively poor based on Figure 6, A1, and Tables 1 and 2. Specifically, it is challenging to conclude from Table 1 that the proposed components improve the model's performance.\n8. In Figure 5, the comparison images are not from the same point in time or the same episode, which does not support the stated conclusion. Additionally, for a more accurate performance comparison, it would be more appropriate to provide an MSE result that only computes the agent region.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed and helpful comments. \n\n**1.** We are grateful to the reviewer for catching this typo. In Figure 1, we should be using *t* as the subscript, not *i*. We will update the figure and text accordingly.  We will additionally make clear that while v_i and a_i are learned in a downstream RL phase, they depend on the latent state z_i.  \n\n**2.** Thank you for the opportunity to clarify. The policy used to weight reconstruction loss is the same policy trained downstream to maximize value. Model training in the Dreamer architecture consists of alternating phases of world model updating (based on action sequences in the replay buffer) and policy/value learning (using imagined action sequences that leverage the world model). The policy is dependent on the latent state, and the latent state is dependent on the image. The reviewer’s understanding is correct that neither the policy nor the value model weights are changed during the world model learning phase. However, the policy depends on the latent state, and the latent state depends on the image pixels. The gradient between the policy and the image pixel is established via standard auto-differentiation (with gradients that flow through the latent state). We will update the text to clarify, thank you. \n\n**3.** Thank you for the opportunity to clarify details about our use of SAM. In our experiments, SAM is not finetuned. For our implementation, in order to determine our segmentation masks, SAM is prompted with a grid of 256 points. The resulting masks are filtered via metrics computed by the SAM algorithm, specifically a prediction IOU (intersection over union) thresh and a “stability score\", followed by non-maximum suppression (NMS). Finally, pixels that have not been assigned to a segmented object are grouped into a single extra object, to ensure that they are not ignored entirely. The outcome of this procedure is a segment assignment for each pixel. We will update the text to clarify this procedure. \n\n**4.** Thank you for providing us the opportunity to further clarify the role of the action-prediction head specified in Eq. 4. The intuition is that the current state *may* contain information about the preceding action. You are completely correct that in general, a past action cannot be inferred from just a single timestep. However, in *some* cases they can. Our assertion is that maintaining this information in the image embedding is redundant and a waste of capacity, *since the the previous action is already provided as input to the world model*. This can become a particularly problematic source of wasted capacity when extensive or detailed aspects of the image are related to the previous action, and thus are redundant. The action prediction head allows us to remove this redundant information from the image embedding step. \n\nWe found that this intuition was supported by our empirical results. Notably, realize we failed to include results from a key ablation that we had run: policy gradient weighting + segmentation but without the action head. We have now added it (Fig. R1I, see the second row), and it makes it clear the performance improvements from including the action head (the first row), in both the unaltered and reafferent environments. \n\n**5.** Thank you for this great suggestion, we will provide additional examples of segmentations and weighting masks , such as Figs. R1B,F,G.\n\n**6.** We appreciate the reviewer's interest in application to more complex settings. We chose the current settings because they allowed for direct comparison across a wide variety of techniques that had been tuned for Deepmind Control Suite and related environments, and provided a good foundation for our new Reafferent environment. Extensions to environments such as CARLA are a great direction for future work that we are excited to pursue. \n\n**7.** Our claim based on the experimental results is that PSP yields: an unambiguously better performance on the most challenging distractors (e.g. Reafferent), a more consistently good result on the Distracting Control Suite (it is in second place on both tasks, while the first place method is different in each case), and adequate performance that is the same or better than baseline DreamerV3 on the unmodified Control Suite. The data in Figure 3, 6, A1 and Table 1 support these claims. A similar conclusion comes from Table 2, in which the combination of methods yields good performance on both the unaltered and reafferent environments. While we see that it is possible to achieve even better performance on Reafferent environment, this can come at the cost of dramatically reducing performance in non-distracting settings (as is the case with value-gradient or policy-gradient weighting alone, in which there is a stark tradeoff in performance between the distracting and non-distracting settings), and the full PSP model therefore performs better than the ablations across settings.\n\n**8.** Thank you for this suggestion. We have now included a comparison image at the time point in the same episode that clearly exhibits the superior model performance of the PSP agent on the Reafferent environment (Figure R1H).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WiV5plKm7z", "reviewer": "Reviewer_1AxW", "summary": "This paper introduces Policy-Shaped Prediction (PSP), a method in model-based reinforcement learning (MBRL) designed to focus on significant aspects of an environment by reducing the influence of distracting information. PSP incorporates a pre-trained segmentation model, a task-aware reconstruction loss, and adversarial learning to enhance the robustness and focus of MBRL agents. The method is evaluated using a benchmark tailored to assess its efficacy against intricate and predictable distractions that do not aid in decision-making.", "questions": "- Given the resource-intensive nature of PSP, how feasible is it to apply this method in real-world applications, such as robotics or autonomous driving, where computational resources and real-time performance are critical?\n- How does PSP perform in dynamic environments where the reward structure or salient features change over time? Is the method adaptable to such scenarios?\n- How sensitive is PSP to the quality of the segmentation model used? Would a less accurate segmentation model significantly impact the performance of PSP?", "limitations": "While the PSP shows impressive performance in the introduced benchmark, its dependency on a high-quality segmentation model could limit its deployment in varied or less-controlled environments. Furthermore, the paper lacks a detailed discussion on the potential computational overhead introduced by PSP components.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The introduction of a segmentation model to MBRL is innovative and appears effective based on the results presented.\n- The paper demonstrates improvements over existing approaches in terms of focusing learning on relevant environmental features and maintaining performance in distraction-filled settings.\n- Methodology extends well-known concepts with new mechanisms, potentially setting a foundation for further exploration in distraction-resistant MBRL.", "weaknesses": "- PSP exhibits high variance in performance, particularly noted in tasks like Hopper Stand, indicating that the method may not consistently achieve the best results across all scenarios.\n- The implementation of PSP, especially with the use of segmentation models like SAM and adversarial learning, is resource-intensive, which might limit its practical applicability in environments with limited computational resources.\n- The added complexity of PSP, including policy-gradient weighting and adversarial components, might pose challenges for implementation and stability during training.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable comments. \n\n**High variance in Hopper Stand:** Thank you for this observation. Hopper Stand is a challenging environment with fairly sparse reward. Indeed none of the other MBRL methods were ever able to achieve success on this task in the Reafferent environment (the maximum score of any method across this grand total of 12 runs was 8.7), whereas PSP achieved a score as high as 377.5. Thus, while all baseline methods consistently exhibited low performance on this task, across all runs in our experiments, PSP was uniquely able to reach a higher score. \n\nAdditionally, we note that we actually had a fourth run of PSP on Hopper Stand with Reafferent background that was quite promising but crashed (irrecoverably) before completing 1M steps, and was therefore not included in our tally, which achieved a score 389 by step ~900K. Thus, 2/4 runs with PSP yielded an agent that figured out how to successfully stand, whereas 0/12 MBRL baseline runs did (Fig. R1J).\n\n**Applicability in environments with limited computational resources:**\nThank you for this great point. We highlight that the added burden of PSP only occurs during training, and not policy inference, which is unchanged. Additionally, the adversarial learning, on its own, does not introduce significant additional computational burden, as the action head is much smaller than the image encoder. We see that the ablation in which only this head is added takes only an additional 11% training time (Fig. R1L). \n\nRegarding the expense of applying SAM, we acknowledge the burden this model introduces. We initially chose SAM as it was most likely have good segmentation quality across environments. However, this model has also spawned a number of cheaper segmentation approaches that could make this approach far less resource-intensive. In particular, the recently released SAM2 can achieve >6x inference speeds relative to SAM (Table 6 of Ravi et al 2024) and thus can reduce the required resources. We expect metrics like this to further improve as models are distilled, optimized, and improved. As highlighted before, we have now included results with the 'tiny' variant of SAM2, which exhibits similar PSP performance to SAM (Fig. R1A).\n\n**Complexity adding challenges for implementation and training stability**:\nThank you for raising these points. Empirically, we observe that PSP is stable on Cheetah Run. It is less stable on Hopper Stand, but this is a challenging environment in which the baselines are substantially less performant, and PSP does not appear to be less stable than the baselines. \n\nIn terms of implementation complexity, we introduce several important changes, which our ablations show to be critical. Our implementation is therefore admittedly more complex than the baseline. However, the changes primarily consist of adding a few extra terms to the loss function and incorporating a segmentation module, and the reference implementation we provide demonstrates that it is fairly straightforward to make these additions to an algorithm such as DreamerV3.\n\n**Feasibility for application to real-world**:\nWe appreciate the reviewer's interest in real-world applications, and we are excited about future work investigating such applications. We wish to highlight that the inference cost of a model trained with PSP is identical to its inference cost without PSP training; only the training cost is impacted. Therefore, any model-based RL method that is feasible for such applications already would still be feasible with the addition of PSP training techniques. We moreover point out that in general it is not trivial to apply Dreamer-type algorithms in real-world settings, but that there is encouraging precedent (e.g. DayDreamer, Wu, Escontrela, Hafner et al. 2022)\n\n**Dynamic environments with changing reward structure or salient features**:\nWe expect that PSP is likely capable of adapting to a changing reward structure, especially if the salient features remain the same. To test this, we trained a PSP agent on Walker Run for 1M steps, and then switched to a Walker Stand reward. We found that PSP was able to quickly adapt to the new reward function (Fig. R1D). Changes in salient features represent a more difficult challenge, but we think that PSP would still be able to adapt, in part because of the flexibility afforded by the mask interpolation of Eq. 3 (see Fig. R1C). It is likely that PSP will adapt more slowly then an algorithm that does not selectively prioritize parts of the scene (though the latter algorithms will not be robust to distractions). Methods such as Curious Replay (Kauvar et al 2023) may help mitigate these adaptation challenges. \n\n**Sensitivity to quality of segmentation model**:\nWe investigated this important question in two ways. First, we implemented PSP with the tiny variant of SAM2, and observe similar results as with SAM (Fig. R1A). Second, we investigated the segmentation masks directly, and we observe evidence that precise details of the segmentation may not matter much. In particular, as shown in Fig. R1B, the SAM segmentation does *not* segment the cheetah perfectly. In some cases, the cheetah is split into 2 segments, while in other cases it is split into 3 or 4 segments. Thus, PSP appears to be somewhat insensitive to the exact details of the segmentation. \nIn terms of deployment to varied or less-controlled environments, we again highlight that the segmentation is only used during the training step (which could be offline), and is not used during inference/deployment of the trained model/policy.\n\n**Computational overhead**:\nThank you for this suggestion. We have now detailed a break down of the computational overhead of the PSP components (Fig. R1L).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "X6Nd1OCHhq", "reviewer": "Reviewer_aBfS", "summary": "The paper addresses the problem of distractions in model-based RL by proposing a policy-shaped prediction (PSP) method, which combines segmentation and adversarial learning to accurately identify and prioritize policy learning on crucial parts of dynamics by incorporating saliency maps from image-based environments. Authors achieve this by finding relevant importance signals from the gradients of the policy w.r.t input images and aggregate importance signal into gradient-based weighting by SAM network. As the main motivation, along with the new distraction-suppression algorithm, the paper introduces a novel distraction benchmark. Overall, the authors show that the proposed PSP demonstrates up to 2x improvement in robustness from distractions, while attaining similar scores in distraction free settings as baselines.", "questions": "1)\tIt would be beneficial to include images of saliency maps during training and to discuss how they vary.\n2)\tCould the authors change SAM in PSP to another segmentation model? How would the final results change? \n3)\tHow well the algorithm performs on dynamically changing distractions? (e.g noisy TV)\n4)\tHow accurate are the SAM masks? How said inaccuracies affect the weight factor in Eq (2) on the overall objective in Eq (1)?\n5)\tCould the proposed method be applied in model-free RL setting?\n6)\tWhat are the benefits of using PSP against domain adaptation algorithms, which are also robust to perturbations as they learn general invariant representations of agent dynamics/environment (known to be robust to domain shifts)? (e.g., refer to Raychaudhuri et al, \"Cross-Domain Imitation Learning from Observations\") \n7)\tIn Table 2, the results for Reafferent environment with gradient weighting and with the adversarial action head seem to be very similar to the case when both of those are removed (383.1 +- 23.8 and 379.0 +- 46.8). Could the authors provide some reasoning for why is this the case?", "limitations": "Authors fully addressed possible limitations and potential societal impacts in Section 5 of the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The motivation behind proposed approach is sound and clearly explained. \n\nEmpirical results support the main claims of the paper, proving the efficiency of proposed method. \n\nMain novelty stems from enhancing the application of policy-gradient based weighting of the important features by reformulating weight factor in terms of mean of different segmentation masks, provided from powerful segmentation models. \n\nDetailed description of novel distraction benchmark is provided, showcasing the ability of PSP to be robust against static adversarial perturbations, surpassing previous methods. \n\nAuthors provide concise and dense experiments, ablation studies, along with all the necessary reproducibility details.", "weaknesses": "1)\tIt is not clear whether much simpler segmentation models would work, because the method seems to be agnostic to the choice of the segmentation model. Evaluation of additional segmentation models would support the paper and could potentially increase the overall speed of PSP training.\n2)\tIn the new benchmark description, Section 3.2, the authors claim that the background distractions are predictable and dependent on the agent’s actions and dynamics. However, the most natural distractions are independent from the agent and ignoring such dynamic case greatly limits the usefulness of the proposed benchmark. \n3)\tThe paper does not mention domain adaptation approaches in RL, which solve exactly same task of learning general policy, invariant to distractions or other kinds of domain shifts.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their excellent comments. \n\n**Evaluation of additional segmentation models**:\nThank you for this great suggestion. We now include evaluation the 'tiny' variant of the recently released SAM2 segmentation model. This increases the speed of the segmentation step by over 2x (in a very unoptimized implementation), and yields similarly performant results (Fig. R1A).\n\n**Dependent vs. independent distractions**:\nExisting benchmarks such as Distracting Control Suite target the described scenario -- with distractions that are independent from the agent, and natural videos as the distracting background. The goal of our Reafferent environment is to complement these existing benchmarks by focusing on scenarios where distractions are predictable and dependent on the agent. The Reafferent environment is not meant to replace the Distracting Control Suite, but rather to augment it in an important and challenging way. \n\n**Relation to domain-adaptation algorithms**:\nThank you for bringing up the interesting connection to domain adaptation algorithms. In general, the domain adaptation setting is somewhat different from our problem: in domain adaptation, (e.g. Raychaudhuri et al) there is an expert environment (which in our case might be an environment without distractions) and an agent environment (perhaps with distractions). The task is to adapt policies from the non-distracting environment to the distracting environment. This is an interesting and potentially fruitful approach, but it is also quite different from our setting in which the distractor is potentially always present. Additionally, even if the setting were directly applicable, Raychaudhuri et al. only uses a low dimensional joint-level state, and it is unclear how or if that would scale to high dimensional visual observations. Nevertheless, we will add discussion of the conceptual relevance of domain adaptation to the manuscript. Thank you. \n\n**Images of saliency maps during training**:\nThank you for the great suggestion, we now show example saliency maps from across training (Fig. R1F). Qualitatively, the saliency maps appear to improve their specificity across training. \n\n**Dynamically changing distractions (e.g. noisy TV)**:\nThank you for raising the interesting scenario of a noisy TV. Interestingly, existing algorithms such as baseline DreamerV3 actually have no problem dealing with such distractions: because the distraction is not predictable, the model learns to simply ignore it. Thus as the distractions become more dynamically changing and less predictable, they may tend to become less problematic. In Fig. R1K, we provide an example training curve from a baseline DreamerV3 agent in a cheetah environment with randomly-changing background distractions, showing that it is able to succeed in the face of this white noise distraction. \n\n**Accuracy of SAM masks**:\nWe investigated this important question in two ways. First, we implemented PSP with the tiny variant of SAM2, and observe similar results as with SAM. Second, we investigated the segmentation masks directly and observe evidence that precise details of the segmentation may not matter much.  In Fig. R1B for example, in a performant PSP agent, the cheetah is sometimes split in two, three, or four segments, which suggests a degree of insensitivity to exact details of the segmentation. Thus, PSP appears to be somewhat insensitive to the exact details of the segmentation. Indeed, because the segmentation only acts as a way to average together gradient signals (Eq. 2), so as long as it groups things of relatively similar importance, the precise details and exact accuracy do not appear to matter much. \n\n**Application to model-free RL**:\nThis is an interesting question, thank you. In model-free methods the latent state is already learned based on the gradient signals from the policy and value function, in contrast to the model-based settings (such as Dreamer) where the latent state is also informed by a reconstruction loss. Indeed, as we show, the model-free baseline DrQv2 performs quite well in the distracting settings. Our work is focused on bringing some of the advantageous features of the model-free approach into the model-based realm. \nHowever, it is quite possible that incorporating segmentation via an additional channel in the input image or adding an action-prediction head could clean up the image embedding in the model-free RL setting. This would be very interesting to investigate in future work. \n\n**Ablation results**:\nThis is a great observation that some of the ablation results appear to similar to each other. This observation touches on three key points.\nFirst, the latter score (379.0) *does* have policy gradient weighting, it just does not use the segmentation model to aggregate gradients. The comparison with no gradient weighting and no action head would be the unaltered DreamerV3 of Table 1 (which scores substantially worse, 158.4 ± 45.7), which we have now included in the ablation table.\nSecond, while there is a similar score in both cases on the Reafferent environment, there is a substantially better score with the full PSP approach on the Unaltered environment (712.3 vs 418.2). \nThird, while the gradient signal alone is adequate to boost performance on the reafferent environment, the gradient signal alone is noisy and this can disrupt the higher scores that are attainable on the unaltered environment. By incorporating the additional methods, especially the segmentation, we can reduce the noisiness of the gradient signal. We now include an additional ablation: policy gradient weighting + segmentation but without the action head (Fig. R1I), showing that the segmentation improves performance on the unaltered environment (compare the second and fifth row). Combining the gradient signal, segmentation, and adversarial action prediction as PSP then yields good performance on both environments.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
### Overview for all reviewers
We thank the reviewers for their excellent thoughtful and constructive comments. We agree with the reviewers' assessment that Policy-Shaped Prediction (PSP) is a novel and effective method to address a well-motivated and relevant problem: reducing the influence of distracting information in model-based RL. We will address each reviewer with individual responses.

We want to highlight that we have now extended PSP to use an **additional segmentation algorithm** and demonstrated ***similar performance but with substantially reduced resource consumption***. All four reviewers raised questions regarding the segmentation algorithm that we used (SAM), including the generalizability, required accuracy, and resource requirements of SAM. We believe we have addressed these questions with our new experiments. 

Our initial selection of SAM was guided by the belief that it represented just the first of many segmentation algorithms that are sure to follow with improved generalizability, performance, and resource consumption. This belief was corroborated by the recent release of SAM2 (on July 29), which according to the associated manuscript (Ravi, et al. 2024) exhibits both improved performance *and* as little as 17% the required resources (e.g. 6x the segmentation speed). We updated PSP to also be able to use SAM2, using the 'tiny' model, the smallest and lowest accuracy of the provided model sizes. Our basic implementation immediately reduced resource consumption by 2x, while yielding a score on Reafferent Cheetah of 360.2, within the range of our original results with SAM of 383.1 ± 23.8 (note: we only had time/resources available to run one seed, but we would include additional seeds in a final manuscript). Optimizations, including using SAM2's video segmentation capabilities and better utilization of the GPU, would likely further improve this resource consumption. These results also suggest that our implementation of PSP is *not* highly sensitive to the segmentation algorithm. 

In the attached pdf (Figure R1), we have also included additional ablations that highlight the importance of the mask interpolation step (Eq. 3) and the action prediction head (Eq. 4), and more visualizations of segmentation masks and salience maps. 

**Together, these results bolster the claim that PSP is an effective method for reducing the impact of challenging distractors**, and that it also has the potential to benefit from gains in related realms such as generalizable segmentation.

Meta-review: Scores 6 5 4 5

The paper proposes a new approach to avoid spending learning capacity on predictable features of the world that are irrelevant to the agent's decisions.  The initial concerns of the reviewers could be alleviated by the rebuttals also with new meaningful experiments.  A second concern is the added complexity for which the reviewers would like to have a more analytical analysis.  That would be definitely a useful addition and it would improve the paper.  Nonetheless, the reviews are quite positive about this work that has a good contribution and is well written and shows relevant empirical results with lots of details for reproducibility.  Thus I suggest to accept this paper.
Decision: Accept (poster)