Make the Pertinent Salient: Task-Relevant Reconstruction for Visual Control with Distractions

## Abstract

Abstract Recent advancements in Model-Based Reinforcement Learning (MBRL) have made it a powerful tool for visual control tasks.
Despite improved data efficiency, it remains challenging to train MBRL agents with generalizable perception.
Training in the presence of visual distractions is particularly difficult due to the high variation they introduce to representation learning.
Building on Dreamer , a popular MBRL method, we propose a simple yet effective auxiliary task to facilitate representation learning in distracting environments. Under the assumption that task-relevant components of image observations are straightforward to identify with prior knowledge in a given task, we use a segmentation mask on image observations to only reconstruct task-relevant components.
In doing so, we greatly reduce the complexity of representation learning by removing the need to encode task-irrelevant objects in the latent representation.
Our method, Segmentation Dreamer ( SD ), can be used either with ground-truth masks easily accessible in simulation or by leveraging potentially imperfect segmentation foundation models.
The latter is further improved by selectively applying the reconstruction loss to avoid providing misleading learning signals due to mask prediction errors.
In modified DeepMind Control suite (DMC) and Meta-World tasks with added visual distractions, SD achieves significantly better sample efficiency and greater final performance than prior work.
We find that SD is especially helpful in sparse reward tasks otherwise unsolvable by prior work, enabling the training of visually robust agents without the need for extensive reward engineering.

## 1 Introduction

Recently, model-based reinforcement learning (MBRL) (Sutton, ; Ha & Schmidhuber, ; Hafner et al., ; ; Hansen et al., ; ) has shown great promise in learning control policies, achieving high sample efficiency.
Among recent advances, the Dreamer family (Hafner et al., ; ; ) is considered a seminal work, showing strong performance in diverse visual control environments.
This is accomplished by a close cooperation between a world model and an actor–critic agent.
The world model learns to emulate the environment’s forward dynamics and reward function in a latent state space, and the agent is trained by interacting with this world model in place of the original environment.

Under this framework, accurate reward prediction is
all we should sufficiently require for agent training.
However, learning representations solely from reward signals is known to be challenging due to their lack of expressiveness and high variance (Hafner et al., ; Jaderberg et al., ).
To address this, Dreamer employs image reconstruction as an auxiliary task in world model training to facilitate representation learning.
In environments with little distraction, image reconstruction works effectively by delivering rich feature-learning signals derived from pixels.
However, in the presence of distractions, the image reconstruction task pushes the encoder to retain all image information, regardless of its task relevance. For instance, moving backgrounds in observations in Fig. 1 are considered distractions.
Including this information in the latent space complicates dynamics modeling and degrades sample efficiency by wasting model capacity and drowning the relevant signal in noise (Fu et al., ).

Figure: Figure 1: (Left) Providing mask example(s) and fine-tuning a mask model, or instrumenting a simulator, to obtain masks. (Right) An input observation in a distracting Meta-World with three alternative auxiliary task targets. Moving scenes in the background are considered distractions. (b) Observations including task-irrelevant information, disturbing world-model training. (c) and (d) Segmentation of task-relevant components using, respectively, a ground-truth mask and an approximate mask generated by segmentation models.
Refer to caption: x1.png

Distractions are prevalent in real-world visual control tasks.
A robot operating in a cluttered environment such as a warehouse may perceive much task-irrelevant information that it needs to ignore.
When training with domain randomization for added policy robustness, task-irrelevant information is actively added and must be denoised.
Prior approaches (Zhang et al., ; Nguyen et al., ; Deng et al., ; Fu et al., ; Bharadhwaj et al., ) address the noisy reconstruction problem by devising reconstruction-free auxiliary tasks, such as contrastive learning (Chen et al., ).
However, many of them suffer from sample inefficiency, requiring many trajectories to isolate the task-relevant information that needs to be encoded.
Training with such methods becomes particularly challenging in sparse reward environments, where the signal for task relevance is very weak.
Additionally, working with small objects, which is common in object manipulation tasks, poses difficulties for these methods because those objects contribute less to loss functions and are easily overlooked without special attention (Seo et al., ).

Inspired by these problems, we address the following question in this paper: How can we help world models learn task-relevant representations more efficiently? Our proposed solution takes advantage of the observation that identifying task-relevant components within images is often straightforward with some domain knowledge.
For instance, in a robotic manipulation task, the objects to manipulate and the robot arm are such task-relevant components, as shown in Fig. 1 (Left).
Given this assumption, we introduce a simple yet effective alternative auxiliary task to reconstruct only the task-related components of image observations.
We accomplish this by using segmentation masks of task-related objects which are easily accessible in simulations.
Specifically, we replace Dreamer’s auxiliary task to reconstruct raw RGB image observations with an alternative task to reconstruct images with a *task-relevance mask* applied to them. (Fig. 1c).
By doing this, the world model can learn features from a rich pixel-reconstruction loss signal without being hindered by the noise of visual distractions.

In contrast to previous work that incorporates segmentation masks in reinforcement learning (RL) as an input (James et al., ; So et al., ), we only use masks in an auxiliary task to improve representation learning. This brings about several advantages. First, we only need segmentation masks during training. The inputs for our method are still the original (potentially distracting) images, so masks are not needed at test time, making our method more computationally efficient. Moreover, the masks we use do not need to be perfect. Since they are used as a target for an auxiliary task, these masks can contain errors as long as they guide feature learning to be informative for the downstream task, which leaves room to replace a ground-truth (GT) mask with its approximation.

To this end, we present a way of training with our auxiliary task with segmentation estimates.
This can be useful in many practical cases where no GT mask is available during training.
This is made possible by recent advances in segmentation foundation models (Kirillov et al., ; Zhang et al., ; Xie et al., ). Specifically, we fine-tune segmentation models with annotated training data and use them to generate pseudo-labels for the auxiliary task.
Fig. 1d shows an example of an auxiliary target made from segmentation prediction.
Although the performance with segmentation estimates is impressive without further modification, we find that the training can sometimes be destabilized by incorrect learning signals induced by segmentation prediction errors.
Thus, we additionally provide a strategy to make training more robust to prediction errors and achieve higher performance.
Specifically, our strategy is to identify pixels where the foundation model mask prediction disagrees with a second mask prediction given by our world model. We ignore the RGB image reconstruction loss on these pixels to avoid training on potentially incorrect targets.

As previously mentioned, our method assumes that task-relevant parts are easy to identify in image observations with prior knowledge.
This is not a strong assumption in many object-centric and robotic domains, where image observations can often be decomposed into relevant and irrelevant regions.
However, there are scenarios beyond our scope, where this assumption may not hold because prior knowledge is unavailable or difficult to codify, such as Atari (Bellemare et al., ).

We evaluate our method on various robotics benchmarks, including DeepMind Control Suite (Tassa et al., ) and Meta-World (Yu et al., ), perturbing both with visual distractions.
We show that our method for reconstructing masked RGB targets using the ground-truth masks in the presence of distractions can reach the same level of performance as training in original environment with no distractions added.
Our method for training with approximate masks also shows impressive performance, often matching the performance of the ground-truth mask variant.
In both benchmarks, our approximate-mask method achieves higher sample efficiency and superior test returns compared to previous approaches.
Notably, this is accomplished with very few task-specific mask example data points (1, 5, or 10 used for fine-tuning), with much of its strength coming from the power of segmentation foundation models.
Our method effectively addresses the challenge of training agents in distracting environments by offloading the identification of task-relevant regions to out-of-the-box segmentation models, thereby achieving great sample efficiency and generalization ability.
Furthermore, our method proves particularly effective in sparse reward environments and those involving small objects, where prior approaches often struggle.

## 2 Related Work

##### Model-Based RL for Distracting Visual Environments.

Recent advances in MBRL have facilitated the learning of agents from image observations (Finn & Levine, ; Ha & Schmidhuber, ; Hafner et al., ; ; ; ; Schrittwieser et al., ; Hansen et al., ; ).
Nevertheless, learning perceptual representations in the presence of distractions remains challenging in these models (Zhu et al., ).
For effective representation learning, some works apply non-reconstructive representation learning methods (Nguyen et al., ; Deng et al., ), such as contrastive learning (Chen et al., ) and prototypical representation learning (Caron et al., ).
However, features learned with these methods do not necessarily involve task-related content since they do not explicitly consider task-relevance in feature learning.
Some other works design auxiliary objectives to explicitly use downstream task information (Zhang et al., ; Fu et al., ).
For example, DBC (Zhang et al., ) uses a bisimulation metric (Ferns et al., ) to encourage two trajectories of similar behaviors to become closer in a latent space.
Perhaps most relevant to our method is TIA (Fu et al., ) which explicitly separates task-relevant and irrelevant branches to distinguish reward-correlated visual features from distractions.
Features from each branch are combined later to reconstruct the original, distracting observation.
Recently, a few approaches proposed to leverage inductive biases such as predictability (Zhu et al., ) and controllability (Wang et al., ; Bharadhwaj et al., ) to learn useful features for visual control tasks.
These methods have shown to be more effective than using the reward signal alone, but many of them suffer from sample inefficiency, requiring many samples to implicitly identify what is task-relevant from data. In contrast, our work proposes to leverage domain knowledge in the form of image masks to provide an explicit signal for identifying task-relevant information. Notably, training in sparse reward environments with distraction has remained unsolved in the literature.
Several methods for robust representation learning have also been proposed for model-free RL (Laskin et al., ; Kostrikov et al., ; Yarats et al., ; Hansen et al., ; Hansen & Wang, ; Nair et al., ; Zhang et al., ).
However, the results suggest that MBRL is more powerful and sample efficient for visual control tasks, thus we focus on comparison with methods in the MBRL framework.

##### Segmentation for RL.

Segmentation models (He et al., ; Redmon et al., ) have been used in many downstream tasks, including RL, to assist in pre-processing inputs (Kirillov et al., ; Anantharaman et al., ; Yuan et al., ; James et al., ; So et al., ).
Recently, using segmentation masks in new domains has been made easier by the introduction of one/few-shot masks foundation models (Zhang et al., ; Xie et al., ) which can quickly adapt to new use-cases.
In the context of RL, the prevalent way to use segmentation models is to turn the input modality from RGB images to segmentation masks (James et al., ; So et al., ). By converting RGB images into semantic masks, agents can effectively handle complicated scenes and thus also be trained with domain randomization.
However, processing the input observation requires additional computation at execution time and an agent trained with predicted segmentation masks can easily break when the segmentation model malfunctions.
Our approach, on the other hand, leverages segmentation masks for an auxiliary task, removing the need for the segmentation model after deployment while also achieving higher test-time performance.
FOCUS (Ferraro et al., ) is another method that uses masked input as an auxiliary target.
However, this is primarily devised for disentangled representation, not for handling distractions. Moreover, it only provides preliminary results with segmentation models and lacks results and analysis on the effects on downstream tasks.

## 3 Preliminaries

We consider a partially observable Markov decision process (POMDP) formalized as a tuple $(\mathcal{S},\Omega,\mathcal{A,T,O},p_{0},\mathcal{R},\gamma)$, consisting of states $s\in\mathcal{S}$, observations $o\in\Omega$, actions $a\in\mathcal{A}$, state transition function $\mathcal{T}:\mathcal{S}\times\mathcal{A}\rightarrow\Delta(\mathcal{S})$, observation function $\mathcal{O}:\mathcal{S}\rightarrow\Omega$, initial state distribution $p_{0}$, reward function $\mathcal{R}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$, and discount factor $\gamma$.
At time $t$, the agent does not have access to actual world state $s_{t}$, but to the observation $o_{t}=\mathcal{O}(s_{t})$, which in this paper we consider to be a high-dimensional image.
Our objective is to learn a policy $\pi(a_{t}|o_{\leq t},a_{<t})$ that achieves high expected discounted cumulative rewards $\mathbb{E}[\sum_{t}\gamma^{t}r_{t}]$, with $r_{t}=\mathcal{R}(s_{t},a_{t})$ and the expectation over the joint stochastic process induced by the environment and the policy.

Dreamer (Hafner et al., ; ; ) is a broadly applicable MBRL method in which a world model is learned to represent environment dynamics in a latent state space $(h,z)\in\mathcal{H}\times\mathcal{Z}$, consisting of deterministic and stochastic components respectively, from which rewards, observations, and future latent states can be decoded.
The components of the world model are:

$$ Sequence model: $\displaystyle h_{t}$ $\displaystyle{}=f_{\phi}(h_{t-1},z_{t-1},a_{t-1})$ Observation encoder: $\displaystyle z_{t}$ $\displaystyle{}\sim q_{\phi}(z_{t}|h_{t},o_{t})$ Dynamics predictor: $\displaystyle\hat{z}_{t}$ $\displaystyle{}\sim p_{\phi}(\hat{z}_{t}|h_{t})$ (1) Reward predictor: $\displaystyle\hat{r}_{t}$ $\displaystyle{}\sim p_{\phi}(\hat{r}_{t}|h_{t},z_{t})$ Continuation predictor: $\displaystyle\hat{c}_{t}$ $\displaystyle{}\sim p_{\phi}(\hat{c}_{t}|h_{t},z_{t})$ Observation decoder: $\displaystyle\hat{o}_{t}$ $\displaystyle{}\sim p_{\phi}(\hat{o}_{t}|h_{t},z_{t}),$ $$

where the encoder maps observations $o_{t}$ into a latent representation, the dynamics model emulates the transition distribution in latent state space, the reward and continuation models respectively predict rewards and episode termination, and the observation decoder reconstructs the input. The concatenation of $h_{t}$ and $z_{t}$, *i.e*. $x_{t}=[h_{t};z_{t}]$, serves as the model state.
Given a starting state, an actor–critic agent is trained inside the world model by rolling out latent-state trajectories.
The world model itself is trained by optimizing a weighted combination of three losses:

$$ $\displaystyle\mathcal{L}(\phi)\mathrel{\dot{=}}{}$ $\displaystyle\mathbb{E}_{q_{\phi}}\left[\sum_{t=1}^{T}(\beta_{\text{pred}} \mathcal{L}_{\text{pred}}(\phi)+\beta_{\text{dyn}}\mathcal{L}_{\text{dyn}}( \phi)+\beta_{\text{rep}}\mathcal{L}_{\text{rep}}(\phi))\right]$ (2) $$

$$ $\displaystyle\mathcal{L}_{\text{pred}}(\phi)\mathrel{\dot{=}}{}$ $\displaystyle-\ln{p_{\phi}(o_{t}|z_{t},h_{t})}-\ln{p_{\phi}(r_{t}|z_{t},h_{t}) }-\ln{p_{\phi}(c_{t}|z_{t},h_{t}})$ (3) $\displaystyle\mathcal{L}_{\text{dyn}}(\phi)\mathrel{\dot{=}}{}$ $\displaystyle\max(1,\text{KL}[[\![q_{\phi}(z_{t}|h_{t},o_{t})]\!]\|\hskip 10.3 3327ptp_{\phi}(\hat{z}_{t}|h_{t}))])$ (4) $\displaystyle\mathcal{L}_{\text{rep}}(\phi)\mathrel{\dot{=}}{}$ $\displaystyle\max(1,\text{KL}[q_{\phi}(z_{t}|h_{t},o_{t})\hskip 7.74998pt\| \hskip 7.74998pt[\![p_{\phi}(\hat{z}_{t}|h_{t})]\!]]),$ (5) $$

where $[\![\cdot]\!]$ denotes where gradients are stopped from backpropagating to the expression in brackets.

Critically, the first component of $\mathcal{L}_{\text{pred}}$ for reconstructing observations from world model states is leveraged as a powerful heuristic to shape the features in the latent space.
Under the assumption that observations primarily contain task-relevant information, this objective is likely to encourage the latent state to retain information critical for the RL agent.
However, the opposite can also be true.
If observations are dominated by task-irrelevant information, the latent dynamics may become more complex by incorporating features impertinent to decision-making. This can lead to wasted capacity in the latent state representation (Lambert et al., ), drown the supervision signal in noise, and reduce the sample efficiency.

##### Problem Setup.

We consider environments where the latter case is true and observations contain a large number of spurious variations (Zhu et al., ).
Concretely, we consider some features of states $s_{t}\in\mathcal{S}$ to be irrelevant for the control task. We assume that states $s_{t}$ can be decomposed into task-relevant components $s^{+}_{t}\in S^{+}$ and task-irrelevant components $s^{-}_{t}\in S^{-}$ such that $s_{t}=(s^{+}_{t},s^{-}_{t})\in\mathcal{S}=\mathcal{S}^{+}\times\mathcal{S}^{-}$.
We follow prior work in visual control under distraction and assume that
(1) the reward is a function only of the task-relevant component, *i.e*. $\mathcal{R}:\mathcal{S}^{+}\times\mathcal{A}\rightarrow\mathbb{R}$; and (2) the forward dynamics of the task-relevant part only depends on itself, $s_{t+1}^{+}\sim\mathcal{T}(s_{t+1}^{+}|s_{t}^{+},a_{t})$ (Zhu et al., ; Fu et al., ; Bharadhwaj et al., ).
Note that observations $o_{t}$ are a function of both $s_{t}^{+}$ and $s_{t}^{-}$, thus we have $\mathcal{O}:\mathcal{S}^{+}\times\mathcal{S}^{-}\rightarrow\Omega$.

Our goal is to learn effective latent representations $[h_{t};z_{t}]$ for task control. Ideally, this would mean that the world model will only encode and simulate task-relevant state components $s_{t}^{+}$ in its latent space without modeling unnecessary information in $s_{t}^{-}$. To learn features pertaining to $s_{t}^{+}$, image reconstruction can provide a rich and direct learning signal, but only when observation information about $s_{t}^{+}$ is not drowned out by other information from $s_{t}^{-}$. To overcome this pitfall, we propose to apply a heuristic filter to reconstruction targets $o_{t}$ with the criteria that it minimizes irrelevant information pertaining to $s_{t}^{-}$ while keeping task-relevant information about $s_{t}^{+}$.

## 4 Method

We build on Dreamer-V3 (Hafner et al., ) to explicitly model $s_{t}^{+}$ while attempting to avoid encoding information about $s_{t}^{-}$.
In Section 4.1, we describe how we accomplish this by using domain knowledge to apply a task-relevance mask to observation reconstruction targets.
In Section 4.2 we describe how we leverage segmentation mask foundation models to provide approximate masks over task-relevant observation components. Finally, in Section 4.3, we propose a modified decoder architecture and objective to mitigate noisy learning signals from incorrect mask predictions.

### 4.1 Using Segmentation Masks to Filter Image Targets

We first introduce our main assumption, that the task-relevant components of image observations are easily identifiable with domain knowledge.
In many real scenarios, it is often straightforward for a practitioner to know what the task-related parts of an image are, *e.g*. objects necessary for achieving a goal in object manipulation tasks.
With this assumption, we propose a new reconstruction-based auxiliary task that leverages domain knowledge of task-relevant regions. Instead of reconstructing the raw image observations (Fig. 1b) which may contain task-irrelevant distractions, we apply a heuristic task-relevance segmentation mask over the image observation (Fig. 1c) to exclusively reconstruct components of the image that are pertinent to control.

Since our new masked reconstruction target should contain only image regions that are relevant for achieving the downstream task, our world model should learn latent representations where a larger portion of the features are useful to the RL agent. By explicitly avoiding modeling task-irrelevant observation components, the latent dynamics should also become simpler and more sample-efficient to learn than the original (more complex, higher variance) dynamics on unfiltered observations.
In simulations, ground-truth masks of relevant observation components are often easily accessible, for example, in MuJoCo (Todorov et al., ), through added calls to the simulator API.
We term the method trained with our proposed replacement auxiliary task as Segmentation Dreamer (SD) and call the version trained with ground-truth masks SD${}^{\textrm{GT}}$.

Figure: Figure 2: Filtering $L_{2}$ loss to avoid training on false negatives in RGB labels. (Left) Estimated pixel locations (f) where the RGB target (c) is likely incorrectly masked out by the segmentation model (e). (Right) A world model equipped with two decoders, one for reconstructing task-relevant masked RGB images and the other for binary masks, the targets for which are generated by a segmentation model. RGB $L_{2}$ loss is selectively masked by the set difference between (d) and (e). Latent representations ($x_{t}$) in the world model are subjected to the training signal only from the RGB branch. The binary branch is only utilized for selective $L_{2}$ loss.
Refer to caption: x3.png

### 4.2 Leveraging Approximate Segmentation Masks

A simulator capable of providing ground-truth masks for task-relevant regions is not always available. For such cases where only RGB images are available from the environment, we propose to fine-tune a segmentation mask foundation model to our domain and integrate its predictions into the SD training pipeline.
Below, we describe our method for training with approximate task-relevance masks, termed SD${}^{\textrm{approx.}}$.

As an offline process before training the world model, we fine-tune a segmentation model with a small number of example RGB images and their segmentation masks annotations that indicate task-relevant image regions.
Thanks to recent advances in segmentation foundation models, we can obtain a new domain-specific mask model with a very small amount of training examples.
For our experiments, we use the Personalized SAM (PerSAM) (Zhang et al., ) using one-shot adaptation and SegFormer (Xie et al., ) fine-tuned with 5 and 10 examples. For the sake of controlled and reproducible evaluation, we extract these RGB and mask training pairs from simulators, however, this number of required samples is small enough that it can be collected with expert human annotation as well. Also, although we use these specific foundation models, our method should also be compatible with any semantic masking method.
Further details such as how we obtain the fine-tuning data can be found in Appendix I. Once fine-tuning is complete, we incorporate the segmentation model into the SD pipeline to create pseudo-labels for our proposed auxiliary task.

### 4.3 Learning in the Presence of Masking Errors

Although foundation segmentation models generalize well to new scenarios (e.g., different poses, occlusions), prediction errors are inevitable (Fig. 1d). Since each frame is processed independently, segmentation predictions can flicker along trajectories. False negatives in task relevance are particularly detrimental when using naive $L_{2}$ loss on image reconstruction. Missing relevant scene elements in reconstruction targets can lead the encoder to learn incomplete representations, dropping essential task-related information. This variability disrupts the learning of accurate representations and dynamics in the world model.

Despite noisy targets, neural networks can self-correct if most labels are accurate (Han et al., ). Additionally, Dreamer’s use of GRUs (Cho et al., ) provides temporal consistency even with flickering targets. However, as illustrated in Fig. 2 (b)&(c), it’s undesirable to propagate gradients from regions where the original image has been incorrectly masked out. Allowing gradients from these regions provides misleading signals. If we could identify the incorrect regions in the reconstruction target, we could nullify the decoder’s $L_{2}$ loss there—a technique we call selective $L_{2}$ loss.

Since we cannot directly identify regions where the RGB target is incorrectly masked due to false negatives, we estimate them. Preliminary experiments show that a binary mask decoder from world model states (as an added auxiliary task) can be less prone to transient false negatives, unlike RGB prediction, which tends to memorize noisy labels. Therefore, we propose training a world model with two reconstruction tasks (Fig. 2, right): one decoding masked RGB images and the other predicting task-relevance binary masks. Both use the foundation model’s binary mask, $\textrm{mask}_{\mathrm{FM}}$, to construct targets. The RGB branch decodes masked RGB images, while the binary branch predicts $\textrm{mask}_{\mathrm{FM}}$. We denote the binary masks produced by the world model as $\textrm{mask}_{\mathrm{SD}}$, where pixels labeled *true* indicate task relevance.

To avoid training on incorrectly masked-out regions, we estimate where $\textrm{mask}_{\mathrm{FM}}$ may be falsely negative by finding disagreements with $\textrm{mask}_{\mathrm{SD}}$. Specifically, we selectively nullify RGB decoder $L_{2}$ loss for regions marked false in $\textrm{mask}_{\mathrm{FM}}$ but predicted true in $\textrm{mask}_{\mathrm{SD}}$. This prevents training on potentially falsely masked-out pixels still considered task-relevant by a second predictor. Formally, the mask for selective $L_{2}$ loss is the set difference between true pixel locations in $\textrm{mask}_{\mathrm{SD}}$ and $\textrm{mask}_{\mathrm{FM}}$:

$$ $\textrm{pixel}_{\mathrm{MaskOut}}=\textrm{pixel}_{\mathrm{SD}}\setminus\textrm {pixel}_{\mathrm{FM}}$ (6) $$

where $\textrm{pixel}_{\mathrm{MaskOut}}$ indicates pixels to nullify loss at, and $\textrm{pixel}_{\mathrm{SD}}$ and $\textrm{pixel}_{\mathrm{FM}}$ are pixels marked true in $\textrm{mask}_{\mathrm{SD}}$ and $\textrm{mask}_{\mathrm{FM}}$, respectively.

Fig. 2 (d–f) shows examples of $\textrm{mask}_{\mathrm{SD}}$, $\textrm{mask}_{\mathrm{FM}}$, and $\textrm{pixel}_{\mathrm{MaskOut}}$. See Appendix J for details on obtaining $\textrm{mask}_{\mathrm{SD}}$. Our experiments indicate that selective $L_{2}$ loss effectively overcomes noisy segmentation labels and improves downstream agent performance.

Lastly, we observe better performance when we prevent gradients from the binary mask decoding objective from propagating into the world model, so we apply a stop gradient to the inputs of the mask decoder head (see Appendix E for ablations).

## 5 Experiments

We evaluate our method on a variety of visual robotic control tasks from the DeepMind Control Suite (DMC) (Tassa et al., ) and Meta-World (Yu et al., ). Since the standard environments in these benchmarks have simple backgrounds with minimal distractions, we introduce visual distractions by replacing the backgrounds with random videos from the ‘driving car’ class in the Kinetics 400 dataset (Kay et al., ), following prior work (Zhang et al., ; Nguyen et al., ; Deng et al., ). Details about the environment setup and task visualizations are provided in Appendices F and B. In evaluation, we roll out policies over 10 episodes and compute the average episode return. Unless otherwise specified, we report the mean and standard error of the mean (SEM) of four independent runs with different random seeds.
We use default Dreamer-V3 hyperparameters in all experiments.

### 5.1 DMC Experiments

We evaluate SD on six tasks from DMC featuring different forms of contact dynamics, degrees of freedom, and reward sparsities.
For each task, models are trained for 1M environment steps generated by 500K policy decision steps with an action repeat of 2.

#### 5.1.1 Comparison with Dreamer

We compare our methods, SD${}^{\textrm{GT}}$ and SD${}^{\textrm{approx.}}$, to the base Dreamer (Hafner et al., ) method. Here, SD${}^{\textrm{approx.}}$ is denoted as SD${}^{\textrm{FM}}_{N}$, specifying the segmentation model used (FM) and the number of fine-tuning examples ($N$). All methods are trained in distracting environments, except for the Dreamer* baseline, which is trained in the original environment without visual distractions. In most cases, we consider Dreamer* as an upper bound for methods trained with distractions. Similarly, SD${}^{\textrm{GT}}$ serves as an upper bound for SD${}^{\textrm{approx.}}$, with the performance gap expected to decrease in the future as segmentation quality improves.

As shown in Fig. 3(a), Dreamer fails across all tasks due to task-irrelevant information in RGB reconstruction targets, which wastes latent capacity and complicates dynamics learning. In contrast, SD${}^{\textrm{GT}}$ achieves test returns comparable to Dreamer* by focusing on reconstructing essential features and ignoring irrelevant components. Interestingly, SD${}^{\textrm{GT}}$ outperforms Dreamer* in Cartpole Swingup, possibly because the original environment still contains small distractions (e.g., moving dots) that Dreamer* has to model.

A limitation of SD is its reliance on acccurate and correct prior knowledge to select task-relevant components. In Cheetah Run, SD${}^{\textrm{GT}}$ underperforms compared to Dreamer*, likely because we only include the cheetah’s body in the mask, excluding the ground plate, which may be important for contact dynamics. Visual examples and further experiments are in Appendices B and C.

For SD${}^{\textrm{approx.}}$, we test with two foundation models: PerSAM adapted with one RGB example and its GT mask, and SegFormer adapted with five such examples. Despite slower convergence due to noisier targets, both SD${}^{\textrm{PerSAM}}_{1}$ and SD${}^{\textrm{SegFormer}}_{5}$ achieve similar final performance to SD${}^{\textrm{GT}}$ in most tasks. A failure case for SD${}^{\textrm{PerSAM}}_{1}$ is Reacher Easy, where a single data point is insufficient to obtain a quality segmentation for the small task-relevant objects.

Figure: (a) Environment Steps vs. Expected Test Return
Refer to caption: x5.png

#### 5.1.2 Comparison with Baselines

We compare SD${}^{\textrm{approx.}}$ with state-of-the-art methods, including DreamerPro (Deng et al., ), RePo (Zhu et al., ), TIA (Fu et al., ), and TD-MPC2 (Hansen et al., ). DreamerPro incorporates prototypical representation learning in the Dreamer framework; RePo minimizes mutual information between observations and latent states while maximizing it between states and future rewards; TIA learns separate task-relevant and task-irrelevant representations which can be combined to decode observations; and TD-MPC2 decodes a terminal value function. Among these baselines, only TIA relies on observation reconstruction. Further details are in Appendix K.

Our results in Fig. 3(a) show that our method consistently outperforms the baselines in performance and sample efficiency. TIA underperforms in many tasks, requiring many samples to infer task-relevant observations from rewards and needing exhaustive hyperparameter tuning. Even with optimal settings, it may lead to degenerate solutions where a single branch captures all information. In contrast, our method focuses on task-relevant parts without additional tuning by effectively injecting prior knowledge. RePo performs comparably to ours in Cartpole Swingup but underperforms in other tasks and converges more slowly.

TD-MPC2 struggles significantly in distracting environments. We speculate that spurious correlations from distractions introduce noise to value-function credit assignment that hinders representation learning. Our method mitigates this by directly supervising task-relevant features using segmentation models, leading to more consistent and lower-variance targets.

Among these methods, DreamerPro is the most competitive, demonstrating the effectiveness of prototypical representation learning for control. However, it often requires more environment interactions and converges to lower performance.

In the Cartpole Swingup with sparse rewards, none of the prior works successfully solved the task, highlighting the challenge of inferring task relevance from weak signals. Our method achieves near-oracle performance, being the only one to train an agent with sparse rewards amidst distractions. This suggests the potential to train agents in real-world, distraction-rich environments without extensive reward engineering.

#### 5.1.3 Ablation Study

We investigate the effects of the components in $\textsc{SD}^{\textrm{approx.}}$ by addressing: (1) the benefits of using segmentation models for targets vs. input preprocessing; (2) the effectiveness of the selective $L_{2}$ loss compared to the naive $L_{2}$ loss; and (3) the impact of the segmentation quality on RL performance. In these experiments, we fine-tune PerSAM with a single data point for segmentation mask prediction.

Table: Table 1: Final performance of SD variants. Mean over 4 runs with the standard error of the mean is reported. The highest means are highlighted.

##### Using segmentation masks for an auxiliary task vs. input preprocessing.

We create a variant of SD${}^{\textrm{PerSAM}}_{1}$ that uses masked observations for both inputs and targets, denoted in Tab. 1 by *As Input*. These results suggest that SD${}^{\textrm{PerSAM}}_{1}$, in addition to not requiring mask prediction at test-time, also achieves better test performance and lower variance. Using predicted masks as input is more prone to segmentation errors, restricting the agent’s perception when masks are incorrect and making training more challenging. In contrast, SD${}^{\textrm{approx.}}$ receives intact observations, with task-relevant filtering at the encoder level, leading to better state abstraction. Further analysis on test-time segmentation quality’s impact is in Appendix D.

Figure: Figure 4: (a)+(b) Qualitative comparison of SD trained with naive and selective $L_{2}$ loss. Trajectories are taken from each method’s train-time replay buffer, selected to have the same background. Frames with PerSAM error are highlighted. The model trained with the selective $L_{2}$ loss overcomes errors in the target, whereas the one trained with the naive $L_{2}$ loss memorizes target errors. (c)+(d) shows the precision and recall of PerSAM and the SD RGB decoder prediction. SD RGB predictions are binarized using a threshold to compute recall and precision w.r.t. the ground-truth mask. The data points used for plotting are from the same Cheetah Run training experiment as in (a)+(b). The selective $L_{2}$ loss significantly improves the recall with only a moderate impact on precision.
Refer to caption: x7.png

##### Selective L 2 subscript 𝐿 2 L_{2} italic_L start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT loss vs. naive L 2 subscript 𝐿 2 L_{2} italic_L start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT loss.

As shown in Tab. 1, SD${}^{\textrm{PerSAM}}_{1}$ consistently outperforms the *Naive $L_{2}$* variant, especially in complex tasks like Cheetah Run and Walker Run. Segmentation models often miss embodiment components (Fig. 4, third row). With the naive $L_{2}$ loss, the model replicates these errors, leading to incomplete latent representations and harming dynamics learning (Fig. 4a, fourth row). In contrast, SD${}^{\textrm{approx.}}$ self-corrects by skipping the $L_{2}$ computation where PerSAM targets are likely wrong (Fig. 4b, fourth row).
Fig. 4(c)&(d) show that the naive $L_{2}$ loss follows PerSAM’s trends, while the selective $L_{2}$ loss recovers from poor recall with only a moderate precision decrease.

##### Impact of segmentation quality on RL performance.

Fig. 3(b) plots the training-time segmentation quality against the RL agent’s test-time performance. Comparing three SD variants with different mask qualities (two estimated, one ground truth), we observe that better segmentation tends to lead to higher RL performance, as accurate targets better highlight task-relevant components. This suggests that improved segmentation models can enhance agent performance without ground-truth masks.
In Cartpole Swingup, one of two exceptions, the IoU difference between SD${}^{\textrm{PerSAM}}_{1}$ and SD${}^{\textrm{SegFormer}}_{5}$ is small, and the test returns may fall within the margin of error.
In Walker Run, the other exception, all variants show high segmentation quality and reach near-optimal performance. Here, we hypothesis that a small amount of noise in the target may act as a regularizer, contributing to marginally better downstream performance.

### 5.2 Meta-World Experiments

Figure: Figure 5: Learning curves on six visual robotic manipulation tasks from Meta-World. All curves show the mean over 4 seeds with the standard error of the mean shaded.
Refer to caption: x8.png

Object manipulation is a natural application for our method where prior knowledge can be applied straightforwardly by identifying and masking task-relevant objects and robot embodiments. We evaluate SD on six tasks from Meta-World (Yu et al., ), a popular benchmark for robotic manipulation.
Depending on the difficulty of each task, we conduct experiments for 30K, 100K, and 1M environment steps, with an action repeat of 2 (details in Appendix G). Preliminary tests showed that SegFormer performs well with few-shot learning on small objects. We fine-tune SegFormer with 10 data points to estimate masks in these experiments.

Fig. 5 suggests that our approach outperforms the baselines overall, with a more pronounced advantage in tasks involving small objects like Coffee-Button. Our method excels because it focuses on small, task-relevant objects, avoiding the reconstruction of unnecessary regions that occupy much of the input. In contrast, the baselines struggle as they often underestimate the significance of these small yet highly task-relevant objects. Among the baselines, RePo (Zhu et al., ) is the most competitive. However, RePo performs poorly in a sparse reward setup (see Appendix H).

## 6 Conclusion

In this paper, we propose SD, a simple yet effective method for learning task-relevant features in MBRL frameworks like Dreamer by using segmentation masks informed by domain knowledge. Using ground-truth masks, SD${}^{\textrm{GT}}$ achieves performance comparable with undistracted Dreamer with high sample efficiency in distracting environments when provided with accurate prior knowledge. Our main method, SD${}^{\textrm{approx.}}$, uses mask estimates from off-the-shelf one-shot or few-shot segmentation models and employs a selective $L_{2}$ loss. It learns effective world models that produce strong agents outperforming baselines.

To the best of our knowledge, our approach appears to be the first model-based approach to successfully train an agent in a sparse reward environment under visual distractions, enabling robust agent training without extensive reward engineering. This work also advances the integration of computer vision and RL by presenting a novel way to leverage recent advances in segmentation to address challenges in visual control tasks. The proposed method achieves strong performance on diverse tasks with distractions and effectively incorporates human input to indicate task relevance. This enables practitioners to readily train an agent for their own purposes without extensive reward engineering. However, SD has some limitations to consider in future work, which we further explore in Appendix M.

## Appendix A Code Release

We plan to make the code for Segmentation Dreamer publicly available upon acceptance.

## Appendix B Visualization of Tasks

### B.1 DeepMind Control suite (DMC)

Fig. 6 visualizes the six tasks in DMC (Tassa et al., ) used in our experiments. Each row presents the observation from the standard environment, the corresponding observation with added distractions, the ground-truth segmentation mask, and the RGB target with the ground-truth mask applied.
Cartpole Swingup Sparse and Cartpole Swingup share the same embodiment and dynamics.
Cartpole Swingup Sparse only provides a reward when the pole is upright, whereas Cartpole Swingup continuously provides dense rewards weighted by the proximity of the pole to the upright position.
Reacher Easy entails two objects marked with different colors in the segmentation mask, as shown in Fig. 6e 3rd column. Before passing the mask to SD, the mask is converted to a binary format where both objects are marked as *true* as task-relevant.

Figure: Figure 6: DMC tasks. Left to right: (1) standard environment observations, (2) distracting environment observations, (3) ground-truth segmentation masks, and (4) RGB observations with ground-truth masks applied. We use (4) as auxiliary reconstruction targets in SD${}^{\textrm{GT}}$.
Refer to caption: x9.png

### B.2 Meta-World

Fig. 7 shows the six tasks from Meta-World-V2 used in our experiments.
Meta-World is a realistic robotic manipulation benchmark with challenges such as multi-object interactions, small objects, and occlusions.

Figure: Figure 7: Meta-World tasks. Left to right: (1) standard environment observations, (2) distracting environment observations, (3) ground-truth segmentation masks, and (4) RGB observations with ground-truth masks applied. We use (4) as auxiliary reconstruction targets in SD${}^{\textrm{GT}}$. Masks with multiple classes for different objects are converted to binary masks (all non-background regions are *true* and task-relevant) before use with SD.
Refer to caption: x10.png

## Appendix C The Impact of Prior Knowledge

We investigate the impact of accurate prior knowledge of task-relevant objects. Specifically, we conduct additional experiments on Cheetah Run—the task showing the largest disparity between Dreamer* and SD${}^{\textrm{GT}}$ in Fig. 3(a). In our primary experiment, we designated only the cheetah’s body as the task-relevant object. However, since the cheetah’s dynamics are influenced by ground contact, the ground plate should have also been considered task-relevant.

Fig. 8 (a–c) illustrates the observation with distractions, the auxiliary target without the ground plate, and with the ground plate included, respectively. Fig. 8d compares SD${}^{\textrm{GT}}$ trained with different selections of task-relevant objects included in the masked RGB reconstruction targets. We show that including the ground plate leads to faster learning and performance closer to that of the oracle. This highlights the significant influence of prior knowledge on downstream tasks, suggesting that comprehensively including task-relevant objects yields greater benefits.

Figure: Figure 8: The impact of prior knowledge on Cheetah Run. (d) The mean over 4 seeds with the standard error of the mean (SEM) is shaded.
Refer to caption: x11.png

## Appendix D The Impact of Test-Time Segmentation Quality on Performance

We investigate how test-time segmentation quality affects SD${}^{\textrm{approx.}}$ as well as the *As Input* variation that applies mask predictions to RGB inputs in addition to reconstruction targets.
For this analysis, we use PerSAM fine-tuned with a single data point for segmentation prediction.
To measure segmentation quality, we compute episodic segmentation quality by averaging over frame-level IoU.
In Fig. 9 we plot episode segmentation quality versus test-time reward on the evaluation episodes during the last 10% of training time.

Fig. 9 illustrates that SD${}^{\textrm{approx.}}$ exhibits greater robustness to test-time segmentation quality compared to the *As Input* variation, with the discrepancy increasing as the IoU decreases. This disparity primarily arises because *As Input* relies on observations restricted by segmentation predictions, and thus its performance deteriorates quickly as the segmentation quality decreases. In contrast, SD${}^{\textrm{approx.}}$ takes the original observation as input and all feature extraction is handled by the observation encoder, informed by our masked RGB reconstruction objective.
Consequently, SD${}^{\textrm{approx.}}$ maintains resilience to test-time segmentation quality.

An intriguing observation is that a poorly trained agent can lead to poor test-time segmentation quality. For instance, Cartpole Swingup (Sparse) exhibits different segmentation quality distributions between SD${}^{\textrm{approx.}}$ and *As Input*. This discrepancy occurs because the sub-optimal agent often positions the pole at the cart track edge, causing occlusion and hindering accurate segmentation prediction by PerSAM.

Figure: Figure 9: Test-time episodic reward vs PerSAM episodic IoU for SD${}^{\textrm{PerSAM}}_{1}$ and *As Input* (SD${}^{\textrm{PerSAM}}_{1}$ with masked RGB observations as input). SD${}^{\textrm{PerSAM}}_{1}$ is more robust to test-time segmentation prediction errors.
Refer to caption: x13.png

## Appendix E Ablation without Stop Gradient

Should the $\textsc{SD}^{\textrm{approx.}}$ world model be shielded from gradients of the binary mask decoder head?

To estimate potential regions on RGB targets where task-relevant regions are incorrectly masked out, we train a binary mask prediction head on the world model to help detect false negatives in masks provided by the foundation model.
We see better performance when gradients from this binary mask decoder objective are not propagated to the rest of the world model. Thus, the default $\textsc{SD}^{\textrm{approx.}}$ architecture is trained with the gradients of the binary mask branch stopped at its $[h_{t};z_{t}]$ inputs, and the latent representations in the world model are trained only by the task-relevant RGB branch in addition to the standard Dreamer reward/continue prediction and KL-divergence between the dynamics prior and observation encoder posterior.
Tab. 2 shows that the performance drops significantly when training without stopping these gradients.

We also examine masks predicted by the binary mask decoder head in Fig. 10. Predictions are coarser grained than their RGB counterparts, lacking details important for predicting intricate forward dynamics. Overall, reconstructing RGB observations with task-relevance masks applied demonstrates itself as a superior inductive bias to learn useful features for downstream tasks compared to binary masks or raw unfiltered RGB observations.

Table: Table 2: Final performance of SD and SD without stop gradient.

Figure: Figure 10: From the top row to the bottom row: (1) ground-truth segmentation masks, (2) SD${}^{\textrm{approx.}}$ binary mask predictions, and (3) SD${}^{\textrm{approx.}}$ RGB predictions.
Refer to caption: x14.png

## Appendix F Distracting DMC Setup

We follow the DBC (Zhang et al., ) implementation to replace the background with color videos.
The ground plate is also presented in the distracting environment.
We used hold-out videos as background for testing. We sampled 100 videos for training from the Kinetics 400 training set of the ’driving car’ class, and test-time videos were sampled from the validation set of the same class.

## Appendix G Distracting Meta-World Setup

We test on six tasks from Meta-World-V2.
For all tasks, we use the corner3 camera viewpoint.
The maximum episode length for Meta-World tasks is 500 environment steps, with the action repeat of 2 (making 250 policy decision steps).
We classify these tasks into easy, medium, and difficult categories based on the training curve of Dreamer* (Dreamer trained in the standard environments).
Coffee Button, Drawer Close, and Handle Press are classified as easy, and we train baselines on these for 30K environment steps. Button Press Topdown (medium) is trained for 100K steps, and Door Open and Drawer Open (difficult) are trained for 1M environment steps.

## Appendix H Results on Meta-World with Sparse Rewards

We also evaluate on sparse reward variations of the distracting Meta-World environments where a reward of 1 is only provided on timesteps when a *success* signal is given by the environment (e.g. objects are at their goal configuration). Rewards are 0 in all other timesteps.
The maximum attainable episode reward is 250.

The sparse reward setting is more challenging because the less informative reward signal makes credit assignment more difficult for the RL agent.
Fig. 11 shows that our method consistently achieves higher sample efficiency and better performance, showing promise for training agents robust to visual distractions without extensive reward engineering.
In Meta-World experiments, TIA (Fu et al., ) is not included as it requires exhaustive hyperparameter tuning for new domains and is the lowest-performing method in DMC in general.

Figure: Figure 11: Learning curves on six visual robotic manipulation tasks from Meta-World with sparse rewards.
Refer to caption: x15.png

## Appendix I Fine-tuning PerSAM and SegFormer

In this section, we describe how we fine-tune segmentation models and collect RGB and segmentation mask examples to adapt them.

PerSAM.
Personalized SAM (PerSAM) (Zhang et al., ) is a segmentation model designed for personalized object segmentation building upon the Segment Anything Model (SAM) (Kirillov et al., ).
This model is particularly a good fit for our SD use case since it can obtain a personalized segmentation model without additional training by one-shot adapting to a single in-domain image. In our experiments, we use the model with ViT-T as a backbone.

SegFormer.
We use 5 or 10 pairs of examples to fine-tune SegFormer (Xie et al., ) MiT-b0.

To collect a one-shot in-domain RGB image and mask example for DMC and MetaWorld experiments, we sample a state from the initial distribution $p_{0}$ and render the RGB observation.
In a few-shot scenario, we deploy a random agent in to collect more diverse observations from more diverse states.

To generate the associated masks for these states, we make additional queries to the simulation rendering API. We represent the pixel values for background and irrelevant objects as *false* and task-relevant objects as *true*. In multi-object cases, we may perform a separate adaptation operation for each task-relevant object, resulting in more than 2 mask classes. In such cases, before integrating masks with SD${}^{\textrm{approx.}}$, we will combine the union of the mask classes for all pertinent objects as a single *true* task-relevant class, creating a binary segmentation mask compatible with our method.

In cases where example masks cannot be programmatically extracted, because such a small number of examples are required (1-10), it should also be very feasible for a human to use software to manually annotate the needed mask examples from collected RGB images.

## Appendix J Details on Selective L 2 subscript 𝐿 2 L_{2} italic_L start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT Loss

The binary mask prediction branch in SD${}^{\textrm{approx.}}$ is equipped with the sigmoid layer at its output.
In order to obtain binary $\textrm{mask}_{\mathrm{SD}}$, we binarize the SD binary mask prediction with a threshold of 0.9.

## Appendix K Details on Baselines

It is known that RePo (Zhu et al., ) outperforms many earlier works (Fu et al., ; Hansen et al., ; Zhang et al., ; Wang et al., ; Gelada et al., ) and that DreamerPro (Deng et al., ) surpasses TPC (Nguyen et al., ).
However, theses two groups of works have been using slightly different environment setups and have not been compared with each other despite addressing the same high-level problem on the same DMC environments.
In our experiments, we evaluate the representatives in each cluster on a common ground (See Appendix F) and compare them with our method.

In our experiments, we use hyperparameters used in the original papers for all the baselines, except RePo (Zhu et al., ) in Meta-World.
RePo does not have experiments on Meta-World in which case we use hyperparameters used for Maniskill2 (Gu et al., ) which is another robot manipulation benchmark.

## Appendix L Extended Related work

There are several model-based RL approaches which introduce new auxiliary tasks.
Dynalang (Lin et al., ) integrates language modeling as a self-supervised learning objective in world-model training.
It shows impressive performance on benchmarks where the dynamics can be effectively described in natural language.
However, it is not trivial to apply this method in low-level control scenarios such as locomotion control in DMC. Informed POMDP (Lambrechts et al., ) introduces an information decoder which uses priviledged simulator information to decode a sufficient statistic for optimal control.
This shares an idea of using additional information available at training time with our method SD${}^{\textrm{GT}}$.
Although this can be effective on training in simulation where well-shaped proprioceptive states exist, it cannot be applied to cases where such information is hard to obtain. In goal-conditioned RL, GAP (Nair et al., ) proposed to decode the difference between the future state and the goal state to help learn goal-relevant features in the state space.

## Appendix M Limitations

Segmentation Dreamer achieves excellent performance across diverse tasks in the presence of distractions and provides a human interface to indicate task relevance. This capability enables practitioners to readily train an agent for their specific purposes without suffering from poor learning performance due to visual distractions. However, there are several limitations to consider.

First, since SD${}^{\textrm{approx.}}$ harnesses a segmentation model, it can become confused when a scene contains distractor objects that resemble task-relevant objects. This challenge can be mitigated by combining our method with approaches such as InfoPower (Bharadhwaj et al., ), which learns controllable representations through empowerment (Mohamed & Jimenez Rezende, ). This integration would help distinguish controllable task-relevant objects from those with similar appearances but move without agent interaction.

Second, our method does not explicitly address randomization in the visual appearance of task-relevant objects, such as variations in brightness, illumination, or color. Two observations of the same internal state but with differently colored task-relevant objects may be guided toward different latent representations because our task-relevant ”pixel-value” reconstruction loss forces them to be differentiated. Ideally, these observations should map to the same state abstraction since they exhibit similar behaviors in terms of the downstream task.
Given that training with pixel-value perturbations on task-relevant objects is easier compared to dealing with dominating background distractors (Stone et al., ), our method is expected to manage such perturbations effectively without modifications. However, augmenting our approach with additional auxiliary tasks based on behavior similarity (Zhang et al., ) would further enhance representation learning and directly address this issue.

Finally, our approximation model faces scalability challenges when task-relevant objects constitute an open set. For instance, in autonomous driving scenarios, obstacles are task-relevant but cannot be explicitly specified. While our method serves as an effective solution when task-relevant objects are easily identifiable, complementary approaches should be considered when this assumption does not hold true.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: How to ensure that the segmentation masks correctly identify task-relevant parts of the image, especially when the foundation model (used for segmentation) does not inherently understand what is important for the task. In some cases, the background details may be crucial for the task, and it is not clear how the model can distinguish between what is task-relevant and what is not. How can you ensure that the masking process does not remove essential information, especially in environments where the task might require nuanced understanding of the background or context? In other words, how can the sparse rewards guide good segmentation masks prediction, when they are not enough at guiding good policies?

The variance in the rewards observed in the Meta-World experiments could be indicative of instability in the model. Can the authors provide more insights into why this happens and how this issue might be mitigated? For example, are there any hyperparameters or training strategies that could be tuned to improve stability across environments?
soundness: 2
strengths: 1. Clear Presentation: The paper is well-written, with clear and easy-to-understand figures, text, and mathematical formulations. The explanation of the model and its components is thorough, making it accessible to both newcomers and experienced researchers in the field.

2. Intuitive Idea: The core idea of the paper—introducing task-relevant segmentation as an auxiliary task—is intuitively appealing. It draws inspiration from human perception, where individuals tend to focus on the most relevant parts of their environment when performing a specific task. This concept feels natural and aligns with cognitive science.

3. Solid Experimental Comparison: The authors conduct experiments for comparing with some strong baselines, and provide detailed ablations for SD's designs.
summary: This paper proposes Segmentation Dreamer, a model-based reinforcement learning method that extends the Dreamer framework by introducing an auxiliary task focused on predicting task-related segmentation masks. Specifically, the reconstruction task in Dreamer is modified to work with images that contain task-relevant parts, rather than reconstructing the entire image. The authors evaluate the method on two RL benchmarks: DeepMind Control and Meta-World, demonstrating its effectiveness in these environments.
weaknesses: 1. Limited Scope of Contribution: One of the main weaknesses of the paper is that the contribution feels somewhat narrow. While the introduction of task-relevant segmentation masks is an interesting idea, it could potentially have a broader application beyond Dreamer or visual distraction tasks. Other model-based RL frameworks, such as transformer-based models (e.g., Trajectory Transformer), also use reconstruction as an auxiliary task, meaning the innovation may not be entirely novel. The applicability of this approach to other RL models and tasks remains underexplored, which limits its perceived impact.

2. Stability Concerns in Meta-World: The experimental results on Meta-World show large reward variance, which raises concerns about the stability of the model. This suggests that the proposed approach may struggle with robustness, particularly in more complex or dynamic environments. A deeper analysis of the reasons for this instability and potential solutions would strengthen the paper.

3. Limited Scope of Experiments: The experiments are conducted on relatively simple and controlled environments (DMC and Meta-World), which might not fully capture the complexity and variability of real-world tasks. The lack of experiments in more complex or high-dimensional settings limits the generalizability of the findings. It would be helpful to see how the method performs on more challenging tasks or in real-world applications where the environment is less predictable.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. How does the quality of the mask affect the results of the proposed method? The authors propose a selective L2 loss during training. Can the authors provide results about using native L2 loss and selective L2 loss for training respectively?
soundness: 3
strengths: * This paper is well-written and easy to follow. The idea of using the segmentation to capture the task-relavent information is reasonable.

* The literature review is sufficient and the authors have adequately discussed the major difference between the proposed method and existing works.

* The evaluations of the proposed method are comprehensive. The authors compare the proposed methods with several state-of-the-art algorithms and evaluations are conducted on two benchmarks.
summary: This paper addresses the challenge of training Model-Based Reinforcement Learning (MBRL) agents in visually distracting environments, where standard approaches struggle with representation learning due to high visual variability. Building on the DREAMER framework, the authors propose Segmentation Dreamer (SD), an auxiliary task that uses segmentation masks to focus representation learning on task-relevant elements, thereby simplifying latent encoding by excluding irrelevant objects. In experiments with visually complex tasks, SD significantly improves sample efficiency, performance, and robustness, particularly in sparse reward scenarios, enabling effective agent training without extensive reward engineering.
weaknesses: * Although the authors conduct experiments on several benchmarks, I still have concerns about the generalization ability. In DMControl tasks, the agent’s visual state consistently appears at the center of the image, and the fixed camera makes segmenting the agent's body relatively straightforward. Could the authors provide additional visualized results for Meta-World and other more complex tasks? For instance, how would segmentation handle the robotic arm and objects in these more challenging scenarios?

* The segmentation process will inevitably introduce additional computational overhead during the model training. For example, the proposed method need masks of the agent body, which is obtained from a fine-tuned SAM model.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Since segmentation quality is important to task performance (Figure 3b), what will happen if more samples, than 1 or 5-10, are used in fine-tuning pre-trained segmentation models to downstream domains? Will the proposed selective l2 loss become less important if segmentation model improves with more training samples?
2. How is binary mask predictor head trained, i.e. what is loss function objective here? What will happen if binary mask decoder is trained but selective l2 loss is not used in learning world model, i.e. is the mask prediction task or addressing segmentation model errors more important in this framework?
3. In terms of $mask_\mathbf{SD}$ and $mask_\mathbf{FM}$, could authors visualize how these segmentations look like in order to understand what are similarity and differences in predictions from adapting pre-trained models and those from the process of learning world models?
soundness: 3
strengths: 1. Paper is well-written, offering clarity and easy to understand
2. The motivation of only reconstructing task-relevant parts in image is reasonable. The proposed approach can handle potential error of adapting pre-trained segmentation model to visual RL domains
3. Extensive experiments including baseline comparisons and ablation analysis demonstrate benefits of proposed approach in improving sample-efficiency
summary: This paper proposes Segmentation Dreamer (SD), an approach to learn task relevant features in model-based RL framework by using segmentation mask from prior domain knowledge as an auxiliary objective. The proposed approach only reconstruct task relevant components instead of entire original image to avoid noisy learning signals from distractions. Additionally a selective l2 loss is proposed to handle error predictions from segmentation masks. Experiments are performed on visual control benchmarks including DMC and Meta-World with visual distractions by replacing them with diverse background scenes. Experiments demonstrate that SD outperforms Dreamer and other baselines in better sample efficiency and convergence performance. Ablations show effectiveness of each component, specifically focusing on segmentation quality and usefulness of proposed selective l2 loss.
weaknesses: 1. This framework requires significant prior knowledge about which part is task-relevant in the image. This may limit the scope of potential usage of the proposed approach, where domain knowledge may be ambiguous or hard to obtain.
2. Since only parts of image is reconstructed specified by mask, this problem setting is simpler than the original, which may be hard to directly compare to other baselines
3. Experiments are limited to performance of visual control tasks; we do not know how good pre-trained segmentation models are adapted to downstream domains.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Could the authors elaborate on how their approach of integrating segmentation masks into MBRL differs from and improves upon existing methods that combine foundation models with RL?
2. What are the authors' thoughts on the generalizability of their method to real-world tasks where distractions are inherently complex and closely intertwined with task-relevant features? 
3. Can the authors discuss potential biases in their model that may arise from errors in segmentation masks? What strategies could mitigate the impact of segmentation errors in untested environments?
4. In environments where distractions and task-relevant features are dynamic, how does the method ensure consistent performance?
5. Is there a mechanism for the model to adapt to new environments or changing task dynamics without complete retraining?
soundness: 2
strengths: To address the limitations of predicted masks, the authors propose a selective L2 loss function.
This is a thoughtful addition to improve the robustness of the model to segmentation errors.
summary: The paper introduces an approach to enhance MBRL in environments with visual distractions. Building upon the DREAMER framework, the authors propose a method called Segmentation Dreamer (SD). This method utilizes segmentation masks to focus on task-relevant components of the visual input, reducing the complexity of representation learning. To address the limitations of predicted masks, the authors introduce a selective L2 loss function.
weaknesses: 1. While the paper presents an integration of segmentation masks into MBRL, it lacks a thorough comparison with existing works that also combine foundation models with RL. A more detailed analysis against these related works would strengthen the paper's claims and provide a clearer picture of the method's advantages and potential limitations[1, 2].
2. The generalizability of the approach across a wider range of real-world scenarios beyond the controlled environments should be considered. In real-world scenarios, distractions may closely resemble task-relevant features, potentially occluding critical information. The paper could benefit from evaluating the method in more realistic settings with complex distractions[3, 4].
3. The paper addresses some limitations, particularly regarding the dependence on accurate segmentation masks. However, a more detailed discussion on potential model biases, especially how segmentation errors might affect performance in untested environments, would be beneficial.



Reference:
[1] "Generalizable Visual Reinforcement Learning with Segment Anything Model." arXiv, 2024
[2] "Empowering Embodied Visual Tracking with Visual Foundation Models and Offline RL." European Conference on Computer Vision, 2024
[3] "Towards distraction-robust active visual tracking." International Conference on Machine Learning. PMLR, 2021.
[4] "Anti-distractor active object tracking in 3D environments." IEEE Transactions on Circuits and Systems for Video Technology 32.6 (2021): 3697-3707.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Can you also compare against other methods also utilizing  segmentation model? e.g., using SAM to get relevant part and feed that to some RL algorithm?  

Can you discuss how would some extension of the proposed method handle cases such as change of viewing angle, color, lighting, texture etc.?
soundness: 3
strengths: Originality: Propose selective L2 loss to handle not perfect segmentation input
Quality: A good approach to introduce segmentation model to help MBRL 
Clarity: Visualizations easy to read
Significance: Deal with distraction in RL policy deployment, which is important
summary: The paper presents Segmentation Dreamer (SD), which improves Model-Based Reinforcement Learning (MBRL) in visually distracting environments. It focuses on task-relevant components using segmentation masks, and achieves better sample efficiency and performance, especially in sparse reward settings and when exposed to significant visual distractions.
weaknesses: The main weakness is that proposed algorithm mainly tested with visual distraction with swapped 2D background shift. This set of benchmarks are surrogate to real domain shifts but the segmentation model directly aims to remove that. Thereby, it would be hard to generalize to other cases such as change of viewing angle, color, lighting, texture etc.

Meta-review: This paper introduces Segmentation Dreamer (SD), an extension of the DREAMER framework designed to improve Model-Based Reinforcement Learning (MBRL) in visually distracting environments. SD focuses representation learning on task-relevant components of the visual input by using segmentation masks, simplifying latent encoding and avoiding noisy signals from irrelevant objects. The method is evaluated on tasks from the DeepMind Control and Meta-World benchmarks, demonstrating improved sample efficiency and performance.

The strengths of this paper include a clear presentation with easy-to-understand explanations and visualizations, as well as an intuitive idea of using segmentation masks to focus on task-relevant components. The main weaknesses include limited generalizability, as the experiments are primarily conducted in controlled environments, the method's dependence on prior knowledge to identify task-relevant components, and its computational cost.

This paper is decided to be rejected, with an average score of 5.4, including two borderline positive scores and three borderline negative scores.

Decision: Reject