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

Authors: miles  hutson
isaac  kauvar
nick  haber
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 03:18:02

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.

Main paper: 
 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 a subset of the components and dynamics of the environment can be accurately modeled. In this setting, distracting stimuli can be particularly problematic, as they waste the capacity of the world model on useless details.To address the challenge of distractors, a number of 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-equivalent world modeling [Schrittwieser et al., 2020] , while environments such as Distracting Control Suite have been developed to assess distractor suppression [Stone et al., 2021] .In this paper we introduce a new method, Policy-Shaped 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-based loss weighting, use of a pre-trained segmentation model [Kirillov et al., 2023] and adversarial learning to create a distraction-suppressing agent that outperforms leading image-based MBRL agents. In addition to exhibiting similar performance in distraction-free settings and on a standard benchmark of robustness to distractions, our method markedly improves performance in the face of particularly challenging distractors that are intricate but entirely learnable. Because learnable distractors can be accurately modeled, they straightforwardly contribute to reducing the world model's error, but needlessly exhaust the capacity of the world model.In sum, we make the following key contributions:• We describe Policy-Shaped Prediction (PSP), a MBRL method that achieves strong distraction 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-inspired action prediction head that reduces sensitivity to self-linked distractions. • We introduce a challenging new benchmark for testing robustness to learnable distractions.• We demonstrate that PSP achieves 2x improvement in robustness against challenging distractions while maintaining similar performance in non-distracting settings. 
policy-shaped prediction: 
 We introduce PSP, a method to 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 to the input image to identify important pixels in the environment, and that we can aggregate these pixelwise 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-related notions of salience and aggregating an otherwise noisy gradient-based 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, we incorporated a way to explicitly mitigate distractions caused by actions of the agent itself.PSP employs (1) gradients of the policy with respect to image inputs to identify task-relevant elements of the image, (2) a 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.Figure 1 : Policy-Shaped Prediction in an environment with challenging distractions. (left) Training of an otherwise-unaltered DreamerV3 agent is modified in two ways: 1) A head is added to 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 2) the loss is scaled pixelwise by a 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.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-Gradient weighted Model loss (VaGraM) [Voelcker et al., 2022] , which weights the model loss according to the gradient of the value function 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-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 in order to run.To compute the policy-gradient weighting, we first sum across the dimensions of the action vector a = E(π(s)), where s is the latent state of the world model, to produce a scalar a = j a j , and then take the gradient with respect to the pixels of the input image x. 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 x.L image (ϕ) = xi ∂a ∂x i (x i -x i ) 2 (1)Gradient-based weighting of the world model's reconstruction ultimately is a 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 we introduce a second novel contribution: object-based aggregation of an explainability signal using a segmentation model (SEG). In principle many models should work, but to ensure we 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 a different model from being utilized, so long as it is of sufficient quality. This idea follows from the observation that the saliency map tends to 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 non-maximum suppression. We also include a mask to 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 a salience score that is aggregated within each segmentation mask. The weight of a pixel x i that has been assigned to segment SEG(x i ) is the mean absolute value weight of the pixels in SEG(x i ).W i = 1 ||SEG(x i )|| j∈SEG(xi) |∂a/∂x j | (2)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 i = 1 for all i, As a regularizer, we linearly interpolate between the salience weighting and a uniform weighting, with α = 0.9 for all our experiments, and using rescaledEQUATIONThis 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 to begin the train loop.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 we 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 own actions.The DreamerV3 world model consists of three main components: a convolutional neural network (CNN) image encoder z t ∼ q ϕ (z t |h t , e t ) with e t = CNN ρ (x t ), 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 h t = f ϕ (h t-1 , z t-1 , a t-1 ) and ẑt ∼ p ϕ (ẑ t |h t ) that is trained to simulate the progression of latent states given actions; and an image decoder, xt ∼ p ϕ (x t |h t , z t ) which reconstructs the image from the latent state. Problematically, 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, z t may source information about a t-1 directly through x t , despite a t-1 being an argument to f ϕ during the computation of h t . 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 actions, we add a small multilayer perceptron (MLP) head that is optimized to predict the previous action from the image embedding.EQUATIONWhen updating θ during world model training, we subtract the scaled gradient ϵ • ∇ θ L(â t-1 , a t-1 ) from the overall world model gradient, with ϵ = 1e3. This forces the latent state'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-based aggregation, and adversarial action prediction to world models other than our chosen DreamerV3 architecture.Question: Does the paper discuss the limitations of 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 do not believe are inherent to the method. Guidelines:• The answer NA means that the paper has no 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 a separate "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 on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 of 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
experiments: 
 To evaluate the model's performance we design our experiments around the following questions:Q1. Is our 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 to the agent's actions? Q4. Does our agent maintain performance in standard, lower-distraction environments? Q5. What are the contributions of each component of our method? W ′ i = width • height • W i / i W i # Normalize weighting 8: W ′′ = αW ′ + (1 -α)1 shape (W ′ )# Linearly interpolate with uniform weighting 9:L pred (ϕ) = -ln p ϕ (x t | z t , h t ) ⊙ W ′′ -ln p ϕ (r t | z t , h t ) -ln p ϕ (c t | z t , h t ) # Weighted DreamerV3 prediction loss 10: L(ϕ) = E q ϕ T t=1 (β pred L pred (ϕ) + β dyn L dyn (ϕ) + β rep L rep (ϕ)) # DreamerV3 model loss 11: ât-1 = MLP ω (stop_grad(CNN ρ (x t )))# Adversarial action prediction head 12:ϕ ← Adam(∇L -ϵ * ∂L(â t-1 , a t-1 )/∂ρ, ϕ)13:L AdvHead (â t-1 , a t-1 ) = (â t-1 -a t-1 ) 2 14:ω ← Adam(∇L AdvHead , ω) 15: end forBaselines 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., 2021a] as a representative baseline Model-Free approach. For all agents, we use 3 random seeds per task, and default hyperparameters.Environment details Visual observations are 64 × 64 × 3 pixel renderings. We test performance in three environments: DeepMind Control Suite (DMC) [Tassa et al., 2018] , Reafferent DMC (described below), and Distracting Control Suite [Stone et al., 2021] (with background video initialized to a 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 distraction-sensitivity depends on task difficulty. For ablation experiments, we test on Cheetah Run. In the natural world, distractions can be highly complex, but in many cases are also highly predictable. For instance, the creaking sound a rusty bicycle makes as you pedal, or the movement of your own shadow as you dance outside. We wanted an environment which would allow investigation 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 selfgenerated distraction requires identifying which parts of the world are relevant to taking action, not just those affected and unaffected by our action. To achieve this, we devised the Reafferent Deepmind Control environment, in which the distracting background images have substantial content, but they depend deterministically on the agent's previous action and the elapsed time in the episode -and are thus completely predictable (Figure 2 ). We build on the Distracting Control Suite [Stone et al., 2021] , using a background consisting of 2,500 16 x 16 grids, with each grid cell filled by a 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 an assigned background. We devised this to be analogous to the types of high complexity self-generated distractors found in the natural world (e.g. one's shadow).This environment allows us 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-not-useful distractions. We find that none of the baseline MBRL methods perform well on the Reafferent Environment, relative to their published results on the Distracting 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 background. However, the cost of this is that the reconstruction of the agent becomes less welldefined or even replaced by the background, especially in positions where the outcome of a movement is uncertain (see Figure 4 and A1 for examples with DreamerV3 and A2 with Denoised 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 ). Notably, the model-free DrQv2 agent shows a reduction in performance from its previous performance on 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 new method demonstrates a substantial improvement over the existing baselines (Table 1 , Figure 3 ). Although it shows a higher than desired level of variance between runs, especially on 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 our new agent is in fact robust to challenging distractors. Addressing Q2, we also find that the salience maps, derived from the gradient of 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, we see 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 ).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 we have not introduced a tradeoff between performance on distracting and non-distracting environments (Table 1 , Figure A5 ), resolving Q4.In sum, PSP exhibits similar performance to baseline methods in commonly used tests of distractorsuppression and in non-distracting environments, while also demonstrating unmatched performance on particularly challenging distractors that are complex but learnable.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 on the reafferent environment does not incorporate the segmentation or adversarial components, and uses the value gradient rather than the policy gradient. However, the variance of its scores is higher than any other approach, and more problematically, it shows the worst performance of all experiments in the unaltered environment. We believe this occurs 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 segmentation-based 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 interpolation 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 a good policy nor a good world model (Figure A6 ). Furthermore, we tested the ability of PSP to adapt to either a task change (Figure A7 ) or a change in the distractor (Figure A8 ), and we found that PSP was able to quickly adapt in both scenarios. We additionally tested the sensitivity of PSP to the segmentation model. Given that segmentation models are likely to continue improving over time, we wondered 1) whether PSP could be compatible with other models besides SAM, and 2) how PSP performance might be modulated by the performance of the segmentation 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 as 6x faster segmentation speeds than the original SAM. We updated PSP to use the 'tiny' SAM2 model, the smallest and lowest accuracy of the provided model sizes.Our basic implementation with SAM2tiny 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, a 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 a successful policy (compared with zero out of twelve total runs across all baselines), though the successful run yielded a lower score (81.7) 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 a 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 a 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. 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 (114.4) that improved on SAM2-tiny. Optimizations, including using SAM2'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 a particularly stringent requirement.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
related work: 
 Distraction-sensitivity of model-based RL Recent advances in Model Based RL (MBRL) including World Models [Ha and Schmidhuber, 2018] , SimPLe [Kaiser et al., 2019] , MuZero [Schrittwieser et al., 2020] , EfficientZero [Ye et al., 2021] , DreamerV1 [Hafner 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 a 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 et al. [2022] introduces the Agent Control-Endogenous State Discovery algorithm, which discards information not relevant to elements of the environment within the agent'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 et al., 2022] . Clavera et al. [2018] use meta-learning and an ensemble of dynamics models. These works form a strong body of solutions, given prior knowledge of likely distractors, but they can struggle if a distractor does not fall into the designed regularizations.A 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 to reweight state reconstruction loss [Voelcker et al., 2022] , building on Lambert et al. [2020] and IterVAML [Farahmand, 2018] , 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 a 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 a 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 Dreamerbased 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-shaped image-based world modeling, by applying concepts from VaGraM to the image-based MBRL setting.Distraction-sensitivity of model-free RL A parallel track of Model Free RL (MFRL) has its own body of literature, with a leading method DrQv2 [Yarats et al., 2021a] 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 distraction-sensitivity, it is worth noting analogous deficiencies can exist in MFRL 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] , 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. 
discussion: 
 PSP combines three ideas to focus the capacity of an agent's world model on aspects of the environment 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, use 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 be 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 an integral over the input space. Future work may investigate 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.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-tailored. Notably, however, segmentation is not necessary during 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 ).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 in 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. At the current stage, this work remains reasonably far from any large societal impacts, as it is limited to agents interacting with small, simulated environments. Over the long term, however, if model-based RL algorithms are used to 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, to exploration, to internet research. On the negative side, agents without proper safeguards have the potential to inflict harm on humans and the environment, whether through negligence or malfeasance. Ultimately, our work is targeted at producing the positive impacts, while still allowing for mitigation of the negative impacts. We characterize the computational cost of the PSP algorithm by ablating various components and measuring its training speed (Table A1 ). There are 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-based 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 ) • S 2 • 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, S is the number of steps in a rollout, W is the width of the input image, and H is the image height.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 a factor of S, where S is 64 in our implementation. Reducing this contribution could induce a significant speedup in the performance of future implementations. Second, we take the gradient of the policy with regard to the image during rollouts only to enable visualization for debugging and figures. This is strictly speaking unnecessary overhead and could be eliminated for a gain during training.Image segmentation: The details of the additional complexity depend on the specific algorithm used. With our addition of results using the SAM2 tiny model, we now demonstrate the viability of using different segmentation algorithms. Notably, the segmentation process is entirely parallelizable separately from the training process: images are segmented as 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.Adversarial head: For each training step, we 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.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 1. This results in an additional cost of O(M • W • H), where M is the number of masks.Each trial of the PSP method used 4 Nvidia A40 GPUs to train the modified DreamerV3 model, and 4 A40 GPUs to run the Segment Anything model in parallel. Given an estimated 17 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 * 8 GPUs = 612 GPU days on A40 accelerators. Early experiments with this methodology likely used an additional 300. Baseline trials could be run on only a single A40 GPU or a desktop NVIDIA 2070 SUPER, usually in less than a day, and accounted for a comparably negligible level of resources.We believe this level of resource consumption could be easily reduced. The modifications to the DreamerV3 model do not attempt to benchmark the most costly components. We suspect our method of parallelizing the new backpropagation from the policy to 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.The repository with code and instructions for reproducing these experiments is available at this GitHub Repository.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 in 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 it is clear that these goals are not attained by the paper.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: An anonymized version of the code will be available at the linked GitHub Repository for reviewers. Guidelines:• The answer NA means that paper does not include experiments requiring code. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or 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 considerations. • The conference expects that many papers will be foundational research and not 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 correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: This paper describes a foundational change to MBRL and introduces no new datasets that pose a risk for misuse. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 a best faith effort. 12. Licenses for existing assets Question: 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? Answer: [Yes] Justification: All previous work is cited and no proprietary code has been used beyond what is allowed by its license. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 not available online, the authors are encouraged to reach out to the asset's creators. 
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 of 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 is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
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 not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to 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. 
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 crowdsourcing 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. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_summary: ### Review Summary

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

### Major Concerns

1. **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.

2. **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.

3. **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.

4. **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.

5. **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.

### References

[1] Learning Invariant Representations for Reinforcement Learning without Reconstruction  
[2] Invariant Causal Prediction for Block MDPs
strength_weakness: strength_weakness: See above
questions: questions: See above
limitations: limitations: See above
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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$.
3. 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?
4. 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.
5. 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.
6. 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.
7. 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.
8. 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.
questions: questions: Please see the weaknesses section.
limitations: limitations: Limitations are discussed in the paper.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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.
- The added complexity of PSP, including policy-gradient weighting and adversarial components, might pose challenges for implementation and stability during training.
questions: 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?
- How does PSP perform in dynamic environments where the reward structure or salient features change over time? Is the method adaptable to such scenarios?
- 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: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1)	It 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.
2)	In 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. 
3)	The 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.
questions: questions: 1)	It would be beneficial to include images of saliency maps during training and to discuss how they vary.
2)	Could the authors change SAM in PSP to another segmentation model? How would the final results change? 
3)	How well the algorithm performs on dynamically changing distractions? (e.g noisy TV)
4)	How accurate are the SAM masks? How said inaccuracies affect the weight factor in Eq (2) on the overall objective in Eq (1)?
5)	Could the proposed method be applied in model-free RL setting?
6)	What 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") 
7)	In 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: limitations: Authors fully addressed possible limitations and potential societal impacts in Section 5 of the paper.

Comments: title: Reply
Reply
TLDR: Reviewers please do acknowledge the rebuttals and react to them.
Clarification request
Official Comment by Reviewer 1AxW
Response to aBfS response
comment: Thank you for your detailed and thoughtful rebuttal. I want to start by acknowledging the significant contributions of this paper. The introduction of PSP is indeed an innovative approach to addressing distractions in model-based reinforcement learning, and the results you have demonstrated, particularly in challenging environments, are impressive. My current score reflects an appropriate evaluation of the paper's contributions and its potential impact on the field.

That said, there are a few remaining concerns, particularly regarding what I meant by an "analytical consideration of the empirical complexity." To clarify, my request pertains to a deeper exploration of how the computational complexity of the proposed PSP approach scales across different environments and hardware setups. While the empirical results you provided give insight into the performance under specific conditions, I believe a more detailed analysis could include:

1. **Scalability Analysis**: How does the complexity of each component of PSP (e.g., segmentation, adversarial learning) scale as the problem size increases, particularly in terms of state space, action space, and environment complexity? This could involve theoretical analysis or case studies showing performance across various levels of task complexity or hardware capabilities.

2. **Trade-off Analysis**: An exploration of the trade-offs between segmentation accuracy, computational cost, and overall performance. For instance, how does the choice of segmentation model (e.g., SAM vs. SAM2) affect performance when considering both accuracy and resource consumption? Understanding these trade-offs could help practitioners decide how best to apply PSP in environments with varying levels of available resources.

3. **Resource-Constrained Scenarios**: While you've provided insights into the computational overhead during training, it would be helpful to understand more about how these factors might impact real-world applications, particularly in resource-constrained environments such as robotics or embedded systems. For instance, could there be strategies for optimizing or simplifying PSP to make it more feasible in such contexts?

These points are intended to provide additional clarity on the impact of your approach in a broader range of settings and could significantly enhance the practical utility of your work. Again, I appreciate the depth of your response and the additional experiments conducted. My evaluation is intended to reflect both the strengths of your contribution and the areas where further refinement could make the work even more impactful.
We are very appreciative of the reviewer's time and engagement with our responses. Thank you. 

While it is true that this may not apply broadly to decision-making scenarios that do not depend on object segmentation, there are many important scenarios that do (especially in the visual domain). Moreover, it is conceivable that our approach could still apply in those settings, with the segmentation simply serving as a reasonable method for reducing noise in the gradient-based weighting. We hope to investigate this in future work. 

In terms of the generalizability of the segmentation model, we think that our existing evidence is actually quite corroborative of good generalization to non-natural images. SAM and SAM2, it turns out, are quite effective (i.e. effective enough to work well with PSP) in the very non-natural Deepmind Control Suite and our even less-natural distracting variants. It is possible other settings may present more of a challenge, though we also expect these segmentation models to improve as the datasets used to train them continue to grow. Thank you again for your comments and insight, we really appreciate the discussion.
We are very grateful for the reviewer's time and engagement with our responses. Thank you.
Thank you to the author for the detailed response. Overall speaking, the author has addressed some of my concerns, so I have raised my score to 5. My remaining major concern is about the generalizability of the proposed method, which involves two aspects: 

1. If the key to the decision-making problem is not largely determined by object segmentation, is the proposed method still effective?
2. The generalizability of the Pretrained Segmentation Model—specifically, in what range can it still work if applied to non-natural images?
Dear reviewers,

thanks for your reviews.  Please do look at the rebuttals of the author and give the authors some feedback whether they could address your concerns.  This is really important for the authors.

Best regards Your AC.
Thank you for your detailed response. I have adjusted the score to 4 due to lingering doubts regarding the effectiveness of the proposed method.
We sincerely appreciate the reviewer's time and engagement with our responses. Thank you. We welcome further discussion and would like to seek clarification on a specific point. Could the reviewer please elaborate on what is meant by an "analytical consideration of the empirical complexity"?

Our experiments demonstrated significantly improved performance over the baseline DreamerV3 in distracting environments while maintaining performance in less-distracting settings. This considerable increase in performance only reduced the unoptimized wall clock speed by less than a factor of four. While optimization efforts could likely improve this speed, we believe the novel ideas and performance improvements introduced in our work are the primary focus.

Furthermore, we showed that ongoing innovations in segmentation algorithms yield equivalent performance on PSP while significantly reducing computational resource requirements.
We would greatly appreciate if the reviewer could specify what additional information they are seeking beyond what we provided, including in our response to aBfS. Your insights are very valuable to us, and we look forward to addressing any remaining concerns or questions.
Thank you for your detailed and thoughtful rebuttal. I have carefully reviewed your responses, including the additional experiments and clarifications provided in response to Reviewer aBfS's comments. I appreciate the analytical breakdown of PSP's complexity and the steps taken to mitigate computational overhead, especially regarding the segmentation process with SAM2. However, I concur with Reviewer aBfS that further analytical consideration of the empirical complexity is needed to better understand its impact on broader adoption, particularly in resource-limited environments. While your response and contributions are valuable, I will keep my original score, as addressing these complexity concerns more thoroughly would strengthen the paper’s contribution.
### Overview
We deeply appreciate the reviewer for taking the time and effort to review our additional results, and we thank this reviewer for noting the meaningful addition these results contribute to the work. **We believe we have addressed each of this reviewer's stated concerns**, including with three new experiments (SAM2, noisy TV, and additional ablations) and additional data (example saliency maps, example segmentations). Furthermore, we believe these *new and expanded result address most of the concerns of all reviewers* and should be accounted for if considering others' scores when determining whether to raise this reviewer's score. Thank you again for your effort and for your thoughtful and constructive comments.

### Analytical complexity
In terms of the response to 1AxW, we are happy to discuss the analytical properties of PSP. There are four major components of PSP that affect performance. In order of additional compute: (1) Policy gradient based weighting, (2) image segmentation (e.g. with SAM)  (3) action adversarial head, and (4) segmentation-based aggregation of the gradient weighting. 

*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. This auto differentiation yields a complexity of O((E + R + P)\*S^2\*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, S is the number of steps in a rollout, W is the width of the input image, and H is the image height.

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 a factor of S, where S is 64 in our implementation. We note reducing this contribution could likely induce a radical speedup in the performance of future implementations. Second, we take the gradient of the policy with regard to the image during rollouts only to enable visualization for debugging & figures. This is strictly speaking unnecessary overhead and could be eliminated for a gain during training.

*Image segmentation*: The details of the additional complexity depend on the specific algorithm used. With our addition of results using SAM2 tiny, we now demonstrate the viability of using different segmentation algorithms. Notably, the segmentation process is entirely parallelizable separately from the training process: images are segmented as 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. 

*Adversarial head*: For each training step, we 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.

*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 & then taking the hadamard product of the resulting image (matrix) with the inverse of the sum of mask elements that are 1. This results in an additional cost of O(M\*W\*H), where M is the number of masks.

While these analytical properties can guide further improvements to the algorithm, we ultimately note that the set of hardware and software used is a very important component of actual wall clock performance. Many elements of complexity are parallelizable on a GPU, and thus may not actually contribute greatly to wall clock time. Additionally, Jax & CUDNN make choices about the exact algorithms that implement the high level graph, which in our experience can impact not only wall clock time but even numerical stability. Therefore, we do believe that any assessment of performance must include concrete code and empirical assessments, and this is why we focused on presenting empirical comparisons. Thank you for the opportunity to clarify.
I appreciate all the clarifications. As a reviewer who gave the highest score to this work, I do *not* think I should be increasing the score any higher. For the discussion period, however, I do want to underscore that the authors conducted a lot of new meaningful experiments for their rebuttal. On the other hand, the response to 1AxW on the added complexity appears to be too empirical, without any attempts of assessing the complexity analytically (even if taking into account a sound ablation effort). Hence, I keep the original score.