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

We thank the reviewer for their excellent comments.

Evaluation of additional segmentation models: Thank 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).

Dependent vs. independent distractions: Existing 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.

Relation to domain-adaptation algorithms: Thank 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.

Images of saliency maps during training: Thank 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.

Dynamically changing distractions (e.g. noisy TV): Thank 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.

Accuracy of SAM masks: We 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.

Application to model-free RL: This 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. However, 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.

Ablation results: This is a great observation that some of the ablation results appear to similar to each other. This observation touches on three key points. First, 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. Second, 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). Third, 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.

We thank the reviewer for their valuable comments.

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.

Additionally, 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).

Applicability in environments with limited computational resources: Thank 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).

Regarding 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).

Complexity adding challenges for implementation and training stability: Thank 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.

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

Feasibility for application to real-world: We 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)

Dynamic environments with changing reward structure or salient features: We 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.

Sensitivity to quality of segmentation model: We 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. In 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.

Computational overhead: Thank you for this suggestion. We have now detailed a break down of the computational overhead of the PSP components (Fig. R1L).

We thank the reviewer for their detailed and helpful comments.

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.

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.

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.

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.

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

5. Thank you for this great suggestion, we will provide additional examples of segmentations and weighting masks , such as Figs. R1B,F,G.

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.

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.

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

We thank the reviewer for their insightful comments.

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. To 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. Additionally, in Fig. R1F, we demonstrate how the salience map improves throughout training, suggesting that the policy and world model can improve together with training.

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.

However, 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!

Dependence on Pretrained Segmentation Model: We 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. We have now included experiments with an additional segmentation model, the 'tiny' variant of SAM2, and observe similarly good results as with SAM (Fig. R1A). Additionally, 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. Notably, 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.

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.

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.