Learning Robust Representations for Visual Reinforcement Learning via Task-Relevant Mask Sampling
Authors: Vedant Dave
Ozan Özdenizci
Elmar Rueckert
Venue: Accepted by TMLR ()

Abstract: Humans excel at isolating relevant information from noisy data to predict the behavior of dynamic systems, effectively disregarding non-informative, temporally-correlated noise. In contrast, existing visual reinforcement learning algorithms face challenges in generating noise-free predictions within high-dimensional, noise-saturated environments, especially when trained on world models featuring realistic background noise extracted from natural video streams. We propose Task Relevant Mask Sampling (TRMS), a novel approach for identifying task-specific and reward-relevant masks. TRMS utilizes existing segmentation models as a masking prior, which is subsequently followed by a mask selector that dynamically identifies  subset of masks at each timestep, selecting those most probable to contribute to task-specific rewards. To mitigate the high computational cost associated with these masking priors, a lightweight student network is trained in parallel. This network learns to perform masking independently and replaces the Segment Anything Model~(SAM)-based teacher network after a brief initial phase (<10-25% of total training). TRMS  enhances the generalization capabilities of Soft Actor-Critic agents under distractions, achieves better performance on the RL-Vigen benchmark, which includes challenging variants of the DeepMind Control Suite, Dexterous Manipulation and Quadruped Locomotion tasks.

Published in Transactions on Machine Learning Research (09/2025)
Learning Robust Representations for Visual Reinforcement Learning via Task-Relevant Mask Sampling
Vedant Dave
vedant. Soaatt..doeeSnnilleeee..c.oo
Chaär of Cyber-Physioad Physical-Soi Syssical-Sysemms Technical University of Leoben, Austria
Ozan Özdenizci
oezdenizaci@tugraz.cc
Institute of Machine Learning and Newral Computation Graz University of Technology, Austria
Elmar Rueckert
rucckert@unileoben.ccd acc.al
Chair of Cyber- Phcete-Pyyilal--S Physical-Systemms Systems Technical University of Leoben, Austria
Reviewed on OpenReview: https: openreview. net/ forum? 2raNDaHutn
Abstract
Humans excel at isolating relevant information from noisy data to predict the behavior of dynamic systems, effectively disregarding non-informative, semporally-correlated noise. In contrast, cxisting visual reinforcement learning algorithms face challenges in generat- ing noise-free predictions within high-dimensional, noise-saturated environments, especially when trained on world models featuring realistic background noise extracted from natural video streams. We propose Task Relevant Mask Sampling (TRMS), a novel approach for identifying task-specific and reward-rolevant masks. TRMS utilizes existing segmentation models masking prior, which is subsequently followed by mask selector that dynami- cally identifies subset of masks at each timestep, selecting those most probable to contribute to task-specific rewards. To mitigate the high computational cost associated with these masking priors, a lightweight student network is trained in parallel. This network learns to perform masking independently and replaces the Segment. Anything Model (SAM)-based teacher network after brief initial phase (< 10 25% of total training). TRMS enhances the generalization capabilities of Soft. Actor-Critic agents under distractions, achieves bet- ter performance the RL-Vigen benchmark, which includes challenging variants of the Deep Control Suite, Dexterous Manipulation and Quadruped Locomotion tasks.
1 Introduction
Visual Reinforcement Learning (RL) has garnered considerable success in mastering complex behaviors derived directly from high-dimensional, imago-basod observations (Mnih et al., 2015; Levine et al., 2016; Lee al., 2020a; Laskin et al., 2020b). Conventionally, it is presumed that environmental observations, often obtained through hand-crafted features, contain only task-relevant information (Ha & Schmidhuber, 2018; Hafner et al., 2020; 2021; Hansen et al., 2024). This assumption enables RL algorithms to function in controlled settings with maximum efficiency, as it climinates exogenous noise (irrelevant or uncontrollable external factors), such as weather fluctuations or random background movements, that could disrupt or impede the learning process. In the real world, the landscape is vastly different, rich in complex visual information, much of which irrelevant to a specific task. The true challenge lies in accurately distinguishing task-relevant data while avoiding the unnecessary modeling of exogenous noise. Traditional RL approaches often fails to provide robust representations under noise, consequently failing to generalize. As a result.,

---

Published in Transactions on Machine Learning Research (09/2025)
they inadvertently incorporate irrelevant data into their representations, leading to the modeling of noise dynamics.
Recent approaches have sought to address this by selectively masking irrelevant parts of the input to focus learning on relevant information. Focus then Decide (FTD) by Chen et al. (2024) leverages the Segment Anything Model (SAM) (Kirillov et al., 2023) to select relevant segments via attention scores and attains high rewards. However, using SAM throughout the training process incurs high computational costs, making FTD impractical for complex tasks that require numerous masks. Contrary to this, methods like SGQN (Bertoin et al., 2022) use binarized attribution maps masks to enforce consistency between the Q-values of masked and original images, highlighting relevant arcas, but provides sub-optimal rewards and is extremely sensitivity to hyperparameter choices. InfoGating (Tomar et al., 2023) focuses only on oflline RL experiments, using a multi-step inverse dynamics model and U-Nets (Ronneberger et al., 2015) to mask irrelevant features. Finally, SAM-G (Wang et al., 2023) employs the SAM model but depends on human intervention for mask selection. However, it remains unclear how combining multiple encoders yields keypoints for task-relevant masks.
To achieve robust performance and yield high rewards in noise-prone environmental settings, we propose TRMS, a novel algorithm that leverages existing masking techniques to learn task-relevant masks and to filter out irrelevant. segments. This approach improves agent's 's generalization capabilities with respect to noise in multiple environments. The core idea of TRMS lies in employing pre-existing masking algorithms that extract meaningful substructures within an image by segmentation. Much like how the human brain processes visual information: decomposing scene into objects and selectively attending to task-relevant elements (Kaiser et al., 2016; Seidl et al., 2012a; Peters & Kriegeskorte, 2021), our method focuses on high- level abstractions, bypassing pixel-level relevance assessments. By segmenting the image into semantically meaningful subregions, we substantially reduce the complexity of the selector's task, requiring it to identify relevance from a more refined subset of the scene. This approach sharply contrasts with methods that attempt to infer such abstractions from raw pixel data (Bertoin et al., 2022; Hansen et al., 2021; Grooten et al., 2024), which inevitably sulfer from less efficient. learning (as shown in the evaluations through empirical results). Instead of evaluating every pixel individually, our method simplifies the process to determining whether a mask (representing a smaller subset of pixels) is correct not. A selector network, utilizing Convolutional Neural Network (CNN), provides binary output for each mask, classifying it as relevant irrelevant.
To improve the computational efficiency, we include the student network in our training procedure. The pre- trained segmentation model remains frozen throughout training and is used solely to generate output masks during inference. The approach is executed in two key phases: (i) Student Network used to learn the mask generation over a subset of the batch to mitigate the overhead of processing the entire batch. (ii) After Tirain steps, the student network, a lightweight CNN architecture, replaces the teacher network, enabling faster computations. An empirical evaluation of wall time is shown in Appendix D. Since the student network is co-optimized with the encoder, no auxiliary loss function is required beyond the initial masking loss during the first phase. The encoder's 's loss alonc is sufficient to guide the optimization process.
To evaluate the performance of TRMS, we conducted experiments across modified version of eleven envi- ronments from three benchmarks in RL-ViGen (Yuan et al., 2023): the Deepmind Control Generalization Benchmark (Hansen ct al., 2021), Quadruped Locomotion (Hansen et al., 2021) and Dexterous Manipula- tion (Rajeswaran et al., 2018). Moreover, TRMS outperforms eight well-established methods in various tasks in vision-based reinforcement learning (Hansen et al., 2021; Yuan et al., 2022b; Huang al., 2022; Bertoin et al., 2022; Yarats et al., 2021; 2022; Laskin et al., 2020a; Wang et al., 2022), achieving better performance across multiple environments.
Our main contributions are:
We propose TRMS, a novel algorithm with an actor-critic architecture that enhances task-relevant masking by leveraging pre-trained segmentation for providing semantically meaningful subregions, with selector that identifies task-relevant information from these subregions, improving both gen- eralization and robustness in visually complex environments.
2

---

Published in Transactions on Machine Learning Research (09/2025)
We further optimize TRMS through a two-phase training process, where lightweight student network incrementally replaces a teacher network, enabling faster mask learning and improving computational efficiency as compared to solely relying on heavy segmentation model (SAM).
We validate TRMS through comprehensive experiments across eleven modified environments from the RI-ViGen benchmarks, where it surpasses relevant existing methods in most of the environments in vision-based reinforcement learning tasks.
2 Related Work
Generalization in Visual RL. RL agents struggle with severe generalization limitations, where perfor- mance degrades sharply in unfamiliar environments due to overfitting and insufficient adaptability to unseen variations (Kirk et al., 2023; Jiang et al., 2023; Raileanu et al., 2021; Zhang et al., 2018; Cobbe et al., 2019). Numerous methods have been developed to improve generalization in reinforcement learning, including do- main adaptation (Xing et al., 2021b; Li et al., 2022; Sun et al., 2022), domain randomization (Mehta et al., 2020; Lee et al., 2020b; Tobin et al., 2017), and curriculum learning (Narvekar et al., 2020; Gupta et al., 2022). Contrastive learning (Laskin ct al., 2020a; Agarwal et al., 2021; Liu et al., 2023a: Yang et al., 2022), bisimulation metries (Ferns et al., 2011; Zhang al., 2021; Liu et al., 2023b; Sun et al., 2024; Zang et al., 2022), data augmentations (Hansen & Wang, 2021; Laskin et al., 2020b; Raileanu et al., 2021; Yarats et al., 2021; 2022; Mumuni & Mumuni, 2022), keypoints (Wang et al., 2021; 2023), and information-thooretie ap- proaches (Tommr et al., 2023; Fan & Li, 2022; Dave & Rueckert, 2024; You et al., 2022; Wang et al., 2024) improve state representations, whereas imitation learning builds policies robust to perturbations (Fan et al., 2021; Xing et al., 2021a; Wang & Hager, 2024).
Numerous works have proposed different ideas to mitigate the impact of task- insk-irreevvant distractors in rein- forcement. learning environments (Yarats al., 2022; Hansen et al., 2021; Huang al., 2022; Laskin et al., 2020a; Yang et al., 2023; Yuan et al., 2022a; Wang et al., 2024). SODA (Hansen & Wang, 2021) incorporates BYOL-like (Grill et al., 2020) architecture and augments data by lincarly combining supplementary images with observations. TLDA (Yuan et al., 2022a) takes a different approach, recommending the exclusion of task-critical pixels from augmentation, determined through the use of Lipschitz constants. PIEG (Yuan et al., 2022b) employs a pre-trained ResNet (He et al., 2016) as its backbone for generalization under distractors. SRM (Huang al., 2022) learns representations in frequency-domate and learns to discard certain frequency in the observation to address domain shifts. CG2A (Liu et al., 2023c) identifies potential conflicts among gradients generated by different augmentations and investigates how to better integrate these augmentations.
Masking Distractors in RL. Several approaches have been proposed to enhance the generaliza- tion of RL agents by selectively masking parts of the input. SGQN (Bertoin et al., 2022) proposes a saliency-guided method, where binarized altribution maps serve as input masks. It regularizes the value function by enforcing consistency between the Q-values of the masked and original state images, improving learning focus on relevant areas. MLR (Yu et al., 2022) introduces a self-supervised auxiliary objective that performs random masking and reconstructs masked information in the latent space, encouraging dynamic- relevant state ions. InfoGating (Tomar et al., 2023) utilizes multi-step inverse dynamics model as its primary objective and employs U-Nets to mask irrelevant information, with a focus on offline RL experiments. MaDi (Grooten et al., 2024) closely aligns to our approach in utilising a small CNN for masking (similar to our student network) and using reward-driven supervision to suppress irrelevant pixels, but learns masks end-to-end without leveraging structured segmentation priors. SAM-G (Wang et al., 2023) employs frozen Segment Anything Model (SAM) (Kirillov et al., 2023) model to generate observation masks. However, it fundamentally differs in its reliance foundation-model prompts, necessitating human intervention for mask selection. Moreover, the mechanism by which multiple encoders are combined to yield task-relevant masks remains unclear. FTD (Chen et al., 2024) also uses the SAM to select relevant segments via attention scores and regularizes the RL method with inverse dynamics and reward loss. However, the use of SAM throughout the training process results in high computational costs, making it impractical for complex tasks that require numerous masks. These methods rely on unstructured mask lcarning, costly vision models, or lack spatial selectivity. In contrast to these methods, TRMS combines segmentation priors from SAM with dynamic selector and lightweight student network.
3

---

Published in Transactions on Machine Learning Research (09/2025)
Relation to Cognitive Perception. Selecting visual information from cluttered real-world scenes requires aligning visual input with the observer's attentional set an internal representation of objects relovant to current behavioral goals- and as these goals shift, a new attentional set must be instantiated, necessitating the suppression of the previous set lo prevent distractions from irrelevant objects (Nasr et al., 2008; Scidl et al., 2012b). Segmentation allows for the tracking and prediction of objects, enabling cognitive functions like memory and action planning independent of sensory input (Scholl, 2001; Brady et al., 2011). Processes such as perceptual grouping, proto-objects, and object files underpin how humans segment and recognize relevant objects in complex scenes (Roelfsema & Ooyen, 2005; Gao et al., 2016). Inspired by human-like segmentation and attention mechanisms, we introduce a segment sampling strategy that leverages masking priors to learn robust and lask-relevant representations.
3 Preliminaries
3.1 Visual Reinforcement Learning
Reinforcement Learning (RL) aims to obtain optimal policies for sequential decision-making problems through iterative interactions with the environment (Sutlon, 2018). In contexts where agents receive high-dimensional sensory imputs, such as visual observations, these problems are aptly modeled as Par- tially Observable Markov Decision Processes (POMDPs). A POMDP is formally defined by the tuple M = (S, O R, where S is the set of latent states, A is the action space, O is the observa- tion space, P S x A + P(S) is a state transition function, Ω S + P(O) is the observation function, R : S x A + R is the reward function, and y E [0, 1) the discount factor, which attenuates future re- wards. To mitigate the challenges associated with partial observability (Kaclbling et al., 1998), we redefine the agent's state 8t as sequence of k consecutive observations, namely 8, = (ot, 0z-1..... 0c-k+1), where each E O. Although this window does not capture the full action-observation history, it often serves as a practical surrogate in visual domains, enriching the agent's 's perceptual input and partially recovering temporal dependencies lost in single-frame observations. The agent's objective is to obtain a policy Te.: parameterized by Par that maximizes the expected cumulative discounted return: ERNNT. E R(o,, at)]. where the trajectory induced by the undorlying dynamics of the POMDP.
3.2 Soft Actor-Critic (SAC)
Our methodology builds upon the Soft Actor-Critie (SAC) algorithm (Haarnoja et al., 2018), a model- free, off-policy RI. approach that integrates entropy maximization into the policy optimization framework to enhance exploration and stabilize learning. SAC employs a critic network Qo, to approximate the soft state-action value function, seeking to estimate the optimal action-value function Q" (*(ssa) in the context of stochastic policies. The actor is instantiated as a stochastic policy Te, that aims to maximize both the expected return and the entropy of the policy, thereby encouraging exploration of the action space. The shared encoder maps the high-dimenssioall observation space O into a lower-dimensional representations. To ensure the stability of the learning process, the critic and shared encoder have target networks start with the same parameters Ozgt = 0g. These target networks are updated via an exponential moving average (EMA), Oegs + - -1-c)) +cBg, where e E 10. The EMA update serves to temper abrupt fluctuations in parameter values, thereby contributing to the stability of the training process.
3.3 Generalization in Visual RL
Our work focuses on the challenge of generalization in visual reinforcement learning, where the agent is trained on the environment without distractors (including augmentations) and then evaluated on previously unseen environment with distractors. The goal is to obtain consistent behavious under domain (environment distribution) shift. Formally, we consider a family of POMDPs, denoted by M = M2.,, POMDP M; ~ M shares the same underlying dynamics and reward structure but differs in its observation space O,, typically due to variations in visual appearances. Our objective to learn a policy that maximizes the expected cumulative return across POMDPs sampled from M in zero-shot manner i.e.
Each
4

---

Published in Transactions on Machine Learning Research (09/2025)
The goal is to find policy that En 7 R(oo, We denote the training environment as Mirain and the set of test environments as Mestt The generalization performance of the policy can be quantified by the generalization gap (Kirk et al., 2023; Wang et al., 2024), defined Egen Manis (T)) MMMraam (T).
without additional training fine-tuning the test environments.
maximizes the expected discounted return: TM((T) = ((M,.))
4 Method
We propose Task-Relevan, Mask Sampling (TRMS), an algorithm designed to enhance generalization and robustness in noise-prone environments by masking out irrelevant segments. In this section, we provide a de- tailed overview of TRMS's architectural design, training procedure and its inclusion within the reinforcement learning framework.
1771
SAM
sampler
encoder
Ot
masker
encoder
TMS
m
encoder
Q
Figure Depiction of the Task-Relevant Mask Sampling (TRMS) architecture, detailing its components and sequential training phases. Initially, k-masks are generated by a frozen segmentation model, followed by sampling and logical OR operation to produce single mask m This mask is then applied the non- augmented image via a Hadamard product @, producing the masked image of which progresses through the encoder and Q-network to yield the corresponding g-value. Simultancously, m" is utilized to mask the augmented image of through the Hadamard operation. During the initial teaching phase, Treach, only the row in is activated. Beyond this phase, the student masking network replaces the segmentation and sampler, and is trained utilizing the row.
4.1 Masking Prior
TRMS explicitly leverages existing pre-trained segmentation models as a backbone to extract masks from non-augmented images. In this case, we employ the Fast Segment Anything Model (FastSAM) (Zhao et al., 2023)', known for its compact design that reduces memory usage and provides fast inference. As shown in Figure 1, this network remains frozen throughout the training phase, operating solely as an inference model.
Let or and 0:-1 denote the observation at time step and ++ respectively. Following the approach in existing methods (Hansen & Wang, 2021), we apply augmentation by overlaying random image onto these observations, resulting in augmented pairs (ōr, ō++1 (15(0)), (b)+1))) Since these augmented images
As MastSAM variant SAM, here use FastSAM and SAM interchangahly.

---

Published in Transactions on Machine Learning Research (09/2025)
hinder accurate mask extraction from SAM, rely non-angmenteed images for precise and consistent masks (as shown in in Fig 1). Provided an non-augmented image o, we utilize a frozen mask extractor (SAM), denoted by M, to obtain k-masks, {m", mm? m²)
4.2 Mask Sampling Network
These k masks are fed into CNN-based mask selection module, denoted as Gg(m), which generates the probability of selection for each mask w(m'). The goal of the mask selector is to identify task-relevant masks, specifically those with selection probabilities exceeding the threshold 1/k and to disregard all others. Formally, this selection represented as follows:
[1, ((mmm) 2 1/k, p(m') = (0, 0, otherwise.
The objective of the mask selector is to encourage the probabilities associated with task-relevant masks compule the Hadamard product across all selected masks, resulting in a
to exceed 1/k. Subsequently, final composite mask. This mask is then applied to both the original and augmented images, isolating the task-relevant regions of the observations,
m V [p(m') mmm]
where V denotes the Logical OR operation. Following this, the Hadamard product is calculated with the original and augmented observations, respectively, as follows,
o m @ Or.
(3)
o," m * 2 o..
(4)
where & represents the Hadamard product. of and of the images obtained by applying the sampled masks non-augmented and augmented observations. This process is applied to the observations at both time steps, t and ++1, 1, ensuring that the task-relevant features are preserved across temporal frames for both the original and augmented observations.
Multiple masks may relevant to the task; thus, the selector must be capable of selecting several masks simultaneously. Empirical observations suggest that using only a Softmax activation function the output layer often results in higher selection probabilities for only a few masks (typically one two). However, complex scenes froquently require the use of multiple masks. To address this limitation, employ the Gumbel-Soflmax distribution (Gumbel, 1954; Jang et al., 2017) at the output layer, which helps miligate this selection bias. By adjusting the temperature parameter we can encourage a more uniform probability distribution, making it easier for the selector by requiring only a slight increase in probability for the relovant masks to be chosen. The Gumbel-Softmax is defined as
exp ((log(z;) x4) + g:)/T)
J/1 Σー= exp (log((z) 9;)/T)
(5)
where Ør are i.i.d. samples from Gumbel distribution, typically computed as 9i = (--((-n((x)) with in ~ Uniform(O, 1). This formulation enables more robust selection in complex scenes with multiple relevant areas by reducing the chances of a single mask's probability dominating the distribution. More details are provided in the Appendix Since we perform a thresholding operation in Eq. (1), we apply Stroight-Through Estimators (STE) to address the backpropagation challenges associated with discrete operations (Bengio et al., 2013).
4.3 Q-relevant Mask Sampling
To effectively select task-relevant masks, we assume that the task already well-performed in the original, non-augmented environment. If the algorithm cannot solve the task without. augmentation, then solving it

---

Published in Transactions on Machine Learning Research (09/2025)
would be infeasible. Therefore, consider the Q-values obtained from the non-angmented environment expert Q-values. For each pair of masked augmented and non-augmented observations, denoted as 6 and omh, we obtain their corresponding state representations s and 3. The target function can be defined as Put = r(( ae) + + maxx; Quga (s+t1)) The Q-loss function for the non-augmented images is then given by
1 Lo B) (( 2 (goge (0, (+))
(6)
where B denotes the replay buffer. In a similar way, the q-target for the augmented images can be written as Pugt = =( ar) + maxa; Qra (8+41, af) and the corresponding Q-loss is
(goge - Q0, (s, (5)))
(0 3) E Bos
The Pask-Relevane Mask Selection (TRMS) loss is defined as the average these two losses:
1
( B) 2
Co(g
(8)
This TRMS loss simultaneously optimizes the critic's accuracy over both augmented and non-angmented observations, thereby promoting robustness to imput augmentations when identical masks are applied. If task-irrelevant segments are selected, the loss will increase. By discarding visually irrelevant segments, the mask selector reduces representationss ional variance from distractor pixels and encourages the policy to focus on task-relevant features. Consequently, this allows networks to consistently observe the same pixel value across different augmentations, thus maintaining a coherent state representation regardless of the augmentation applied. The Total Loss, a combination of TRMS and masking loss (Lmasker in Eq. (10)), of the overall architecture can be defined
Listal CTRMS + Lmasker.
4.4 Student Masking Networks
Although it is technically feasible to run the algorithm for 500K steps using the FastSAM models and a mask sampler, this approach is highly computationally intensive due the considerable time required by the masking models. These models, while accurate, not optimized for speed and may impose significant delays. Our empirical results shows that the model utilizing only SAM can take upto 3 days on the Walker Walk task (Appendix D). Moreover, FTD (Chen et al., 2024), which also relies on SAM, can require up to 3-6 days to process a single seed, depending on both the task complexity and the number of masks to be generated. To mitigate this computational burden, we introduce a Student Masking Network, a CNN-basod network that effectively mimics the behavior of the Prior Masking model but only for defined initial period, denoted as Teach. During this period, the student network learns directly from the teacher model, replicating its outputs.
The training objective for this student network capitalizes on the binary nature of the mask outputs. We employ Binary Gross-Entrops (BCE) loss to measure the discrepancy between the teacher network's output mask, m² (as defined in Eq. (2)), and the mask generated by the student network, (ot),, which is parameterized by a. Formally, this is expressed as
Lmasker
(10)
This BCE loss is then exclusively backpropagated through the student network, enabling it to gradually learn the teacher's masking stralegy during the initial training period Twach-- After completing these Treach steps, the teacher network is omitted, and only the student network utilised to generate task-relevant masks, maintaining operational efficicncy while significantly reducing computation time (Appendix D). An additional strategy to further optimize training time implemented after Treach; we increase the batch size. This adjustment. expedites the learning process by enabling the student network to process more data per
7

---

end for Published in Transactions on Machine Learning Research (09/2025)
Algorithm 1 Training Algorithm for SAC with TRMS
Require: Estep; 1b K Xno Xz, An
Variables Initialization
Require: Pa. By. M, Mc. Gg
> Networks Initialization
1: D 5-0
b Initialize replay buffer
2: for each initial collection step do
Oe ~ Frandom(--oo) (ae)
Sample random action
0141, + (ot+1, ae, rt+1)
Apply action
Append to buffer
end for
7:
for every training step do
8:
((o), or, re, ((+1) bLt D
D. Sample minibatch b> Sample action
9:
ae ~ (ailo)))))
10:
0++1, ~ (astop(at)
11:
+ (o O++1)
12:
oz, 01+1 (or),, v(0++1)
> Augmentation
13:
for each gradient step do
14:
if step I Traach then 0,"0" + o @ Go))Moo)) o. & Ga(M())+1)
15:
B SAM
16:
else
17:
0²", ,, oe @ ())))) θet1 & (1++))
D. CNN Masking
18:
end if
19:
Update Encoder and Mask Sampler (Eq. (8)) 0g s 0g X, LTRMS oc+1. 0T"h of 0q. @ + 0. - VLrLa
20:
21:
b. Update policy Updale Masker
22:
a a - Aa V mmaker (Ot, 0)
23: end for
24:
training iteration. Architectural details are provided in Appendix C. The details about the training of the entire architecture in Fig. is provided in Algorithm 12. We augmented the SAC algorithm with TRMS components, shown blue.
5 Experiments
In this section, we present our experimental evaluations conducted on generalization benchmarks from VL-ViGen (Yuan et al., 2023). These benchmarks were selected as they encompass a wide range of environ- ments: (1) DeepMind Control Generalization Benchmark (Hansen al., 2021) for evaluating continuous control agents across tasks with complex dynamics and diverse rewards, essential for testing generalization capabilities; (2) Quadruped Locomotion (Hansen et al., 2021), which includes Unitree quadruped robot, challenging agents with intricate balance and control requirements in dynamic environments; (3) Dexterous Manipulation (Rajeswaran et al., 2018), featuring multi-object interactions and sparse rewards that require precise control strategies to handle sophisticated manipulation tasks. We provide comprehensive description of our experimental configurations and compare the performance of TRMS against relovant existing approaches. This analysis demonstrates the effectiveness of TRMS in enhancing generalization across diverse tasks in vision-based reinforcement learning. An evaluation of wall time comparing TRMS with the Only-SAM method is presented in Appendix D. The analysis demonstrates that incorporating the student network substantially enhances computational efficiency compared to relying solely heavily parameterised segmentation model.
Temperature update SAC, double critics and target network updates omitted for clarity.
8

---

Published in Transactions on Machine Learning Research (09/2025)
Baselines. Our method is compared against several prominent Visual RL algorithms that specifically designed for generalization. DrQ (Yarats et al., 2021) improves SAC (Haarnoja et al., 2018) by augmenting visual inputs while updating the TD loss. DrQ-v2 (Yarats et al., 2022), a DDPG (Lillicrap et al., 2016) and DrQ-based model-free algorithm. CURL (Laskin et al., 2020a) enhances visual representations by using a contrastive learning approach similar to SimCLR (Chen et al., 2020) i.e. aligning augmented views of the same observation. SVEA (Hansen al., 2021) stabilizes learning by using un-augmented images for the target Q-value, while applying augmentation to reduce Q-value variance. SRM (Huang et al., 2022) learns representations in trequency-domain and learns to discard certain frequency in the observation to address domain shifts. PIE-G (Yuan et al., 2022b) incorporates ResNet (He et al., 2016) pre-trained models to enhance generalization, while SGQN (Bertoin et al., 2022) uses saliency maps to focus on key pixels crucial for decision-making. VRL3 (Wang el al., 2022) is SOTA algorithm for Adroit lasks, utilizing human demonstrations.
Augmentation. All of our baselines, except for SAC, leverage specific data augmentation techniques during training. TRMS uses an augmentation strategy inspired by SVEA (Hansen et al., 2021), which has proven effective for handling distracting video backgrounds through overlay augmentation. In this approach, we overlay a randomly selected image n from the Places dataset (Zhou et al., 2018) onto our observation frame as o, = 6 or + +1 (1-1 - 8) -n, where 0.5.
weighting factor that controls the degree image overlay.
all the experiments,
Zero-shot Evaluation. To assess generalization, perform zero-shot evaluations of the trained agents range of unseen environments with different distraction intensities. Specifically, evaluate performance on the rideo-easy and video-hard hard configurations across all environments, Each seed undergoes evaluation over 100 episodes, corresponding to the designated noise levels.
p 0.2873
p 0.1114
p 0.0604
p 0.0534
p = 0.0998
K
p = 0.0753
p 0.0730
p 0.0870
= 0.0592
= 0.0842
Figure 2: The figure on the left shows the probability of each mask out of the 10 masks provided. As there are 10 masks, only the masks above the probability of 0.1 (1/10) are selected, resulting into the image on the right.
5.1 Deepmind Control Suite
We evaluate our algorithm on the DMC-GB (Hansen Wang, 2021) benchmark, spanning six tasks: Walker Walk, Walker Stand, Ball in Cup Catch, Finger Spin, Cartpole Swingup, and Cheetah Run. RL-ViGen integrates DMC-GB with two difficulty levels: video-easy (10 background videos) and video-hard (100 back- ground videos without surface), where natural video backgrounds are used to rigorously test generalization under varying visual distractions.
Generalization Performance. We evaluate generalization by running five seeds per task and calculating the mean and standard deviation of returns. As shown in Table 1, TRMS surpasses all the basclines out of 12 environments. Notably, TRMS demonstrates an advantage in the video-hard setting, the most challenging environment due to its complex video perturbations, where it outperforms all the selected

---

Published in Transactions on Machine Learning Research (09/2025)
Table I: Performance on DMC Benchmark Environment in Video-Hard (VH) and Video- Easy (VE) settings.
S-up: Swingup.
Task (VE)
SAC
DrQ
DrQ-v2
CURL
SVEA
SRM
PIEG
SGQN
TRMS
A
Cartpole S-up
398+600
485+1005
267+41
404+677
782+27 27
724+75
482+51
717+335
787+33
+5 (0.63%)
Walker Walk
245+165
682+899
175+117
536+133
819+81
854+442
871+22
860+53
863+74
-8 (0.91%)
Walker Stand
389+131
873+-33
560++48
852+75
961+8
966+-42
957+12
955+9
967+5
(0.10%)
Ball in Cup
192+1577
318+157
171+-106
316+119
871+106
924+35
910+37
761+:171
938+110
+14 (1.15%)
Finger Spin
206+169
533+:119
456++15
502+:19
808+33
853+766
837+1077
609+-61
868+224
+-1 (3.70%)
Cheetah Run
87:21
102+:30
644:22
104+224
249+:20
257:221
287++20
269+:33
207
-80 (27.87%)
Average
253
499
457
456
757
763
724
697
772
+9 (1.18%)
Task (VH)
SAC
DrQ
DrQ-v2
CURL
SVEA
SRM
PIEG
SGQN
TRMS
A
Cartpole S-up
158:117
138:0
13013
114::15
393::45
475175
323:24
488118
5141102
+22 (5.33%)
Walker Walk
122:147
1041:22
34111
58118
377193
535::35
641
6551:45
747-63
+92 (14.05%)
Walker Stand
231::57
2891:49
151:133
45:5
8341:46
863-177
852-56
8511:24
906120
+43 (4.98%)
Ball in Cup
101::37
1001:40
971:27
1151:33
4031:74
56611335
773
7821:57
837120
+55 (7.05%)
Finger Spin
13110
91:113
21:14
271:11
3351588
419:322
762-599
55418
791154
+-29 (5.19%)
Cheetah Run
10+5
32+13
23+55
21+7
105+337
115+224
134+17
144+34
189+86
+-3 (22.72%)
Average
106
126
76
63
408
496
584
579
664
+88 (13.70%)
relevant baselines. In the video-easy setting, TRMS exhibits an improvement of 1.18% over the second-best performing method, SRM, and 5% over the average of the next four best methods SYYA,, and SGQN). Interestingly, in the video-hard settings, TRMS not only outperforms all baselines across every environment but also achieves SGQN, shown by the A in Table 1. Collectively, these results underscore TRMS's robust generalization capabilities, particularly under high-noise and complex video conditions, establishing it as a reliable solution across diverse lest environments.
impressive average improvement of 13.70% over the second-best method,
Mask Sampling Probabilities Visualization. The mask selection probability from the masker is illustrated in Figure 2. Given that there are 10 masks, the selector's objective is increase the probability of task-relevant segments above the threshold of 1/10 (i.e., 0.1) while reducing the probability of irrelevant segments below this threshold. this case, only the first two segments are selected (i.e., only their masks, without the overlayed image). These selected segments are then subjected to the Logical OR operation and the Hadamard product, as depicted in Eq. (2) and Eq. (3) respectively. The image on the right of the figure represents the resulting output, elcarly demonstrating the selective focus on relevant regions.
Representations under Distractors. To demonstrate that TRMS' capability of learning domain- invariant reprosentations, employ t-SNE (van der Maaten & Hinton, 2008) to visualize the features
SRM: 535
PIEG: 641
SGQN: 655
TRMS (ours): 719
Figure t-SNE visualization of clustering results for TRMS and three basclines. TRMS demonstrates a more distinct and well-separated clustering pattern, with each cluster representing identical agent poses with distinct backgrounds.
10

---

Published in Transactions on Machine Learning Research (09/2025)
Performance Comparison between Video Easy and Video Hard Mothods
Video Easy
Aw267
Video Hard
A=140
700
S=110
600
Average Performanee
400
300
200
100
DrQ-12
Mothod
Figure 4: For cach method, paired bars display the average performance under Video Easy and Video Hard conditions, and the annotated A indicates the absolute difference between the two settings. Standard deviations are omitted, as the underlying tasks differ significantly, making such variance comparisons unin- formative.
extracted by the encoder. We select 10 distinct observations from different states and replace their backgrounds with 40 unseen images. These images are then encoded using the learned representations from each algorithm. Observations corresponding to the same state are markod with identical col- ors. As illustrated in Figure 3, representations of images with varying backgrounds (distractors) but identical agent poses embedded into tight clusters, indicating that the learned embeddings capture pose-invariant features. Notably, TRMS exhibits the most compact clustering and achieves the highest re- wards under distracting backgrounds, suggesting a stronger ability to learn domain-invarian invariant representations.
Robustness to Distraction Levels. Figure 4 illustrates the absolute difference in average rewards as the distraction level increases from video-easy to video-hard. The figure reveals that methods experience considerable performance drops under higher distraction levels. Notably, TRMS exhibits the smallest fluctuation, closest to SGQN, while achieving considerable higher rewards than it, underscoring its resilience to elevated noise levels.
5.2 Locomotion
For Locomotion tasks, we utilise Unitree Series tasks (Hansen et al., 2021): Unitree Stand and Unitree Walk. The training and evaluation are performed a similar way described in Section for DMC settings.
Generalization Performance. We evaluate generalization by running three seeds per task and computing the mean and standard deviation of returns. As shown in Table 2, TRMS outperforms the baselines across out of 4 environments. In the video-easy setting, TRMS demonstrates a substantial 33.5% improvement in the Unitree Walk task, though it lags in Unitree Stand, resulting into lower average performance as compared
the baselines. However, as distractions intensify the vidco-hard sotting, TRMS consistently outperforms all baselines, showcasing remarkable resilience distractor noise. Overall, TRMS achieves 77.46% increase in performance in video-hard setting relative to the second-best. method.
11

---

Published in Transactions on Machine Learning Research (09/2025)
Table 2: Performance various methods on Unitree Walk and Unitree Stand tasks for Video- Easy (VE) and Vidco-Hare Hard (VH).
Task (VE)
DrQ
DrQ-v2
CURL
SVEA
SRM
PIEG
SGQN
TRMS (ours)
Walk
67.449.2
97.8+15.7
74.8+14.2
98.4+28.3
98.0=9.4 553.2127.99
140.2#63.9
151.7487.1
202.6+47.2 (+33.55%)
3414119.8
846677
431.4138.3
587.0139.6
379.6265.88
447.0150.3
1.5.41055.8 (-46.26%)
Average
204.4
236.3
253.1
342.7
325.6
259.9
332.0
259.0 (-24.42%)
Task (VH)
DrQ
DrQ-v2
CURL
SVEA
SRM
PIEG
SGQN
TRMS (ours)
Walk
39.6+22.3
83.0+24.2
25.2
73.8+52.2
72.4429.0
203.7#75.6
122846682
214.4=31.5 (+5.25%)
Stand
65.6112.7
95.8:3374
99.4:125.3
279.3110.77
30..0:34.5
202.0:43.1
139.8:44..0
305.4:1014 (+180%)
Average
52.6
89.4
80.3
176.6
186.2
202.8
131.3
259.9 (+77.46%)
5.3
Dexterous Manipulation
Adroit (Rajeswaran et al., 2018) is an euvironment specifically designed for complex dexterous hand ma- nipulation tasks, requiring substantial exploration and detailed feature extraction due to its sparse reward structure and the intricacy of its high-dimensional action space. In our experiments, we consider three of its tasks from a single view in RI-ViGen: Door, Hammer and Pen. This environment involves a magnitude of objects that needs to be masked, which makes the task extremely difficult.
TRMS achieved the highest average performance across tasks with scores of 59.5 in the video-easy setting and 54.9 in the video-hard setting. GQN matched TRMS at 59.5 in video-easy but scored lower in video-hard (30.3) due to heavy distractions. Notably, SGQN excelled specific tasks like Door and Hammer, where TRMS's scores were comparatively lower, suggesting that while TRMS provides robust overall performance. However, there is a room for improvement in these environments. See the future directions below.
Table 3: Performance comparison of various methods on Adroit tasks with Video Easy (VE) and Video Hard (VH) background.
Task (VE)
VRL3
SVEA
SGQN
PIE-G
TRMS
Pen
1.7+00.6
46.7#3.8
64.0#9.0
53.6=4.77
72.4+9.1 (+13.10%)
Door
0.0+0.00
44.8+8.55
58.2+12.3
56.6+11.1
50.7+12.3 5.. (-1.00%
(-12.90%)
Hammer
0.0.0..0
8.448.6
56.3-6.3
44.3+13.0
Average
0.6
33.3
59.5
51.5
59.5
Task (VH)
VRL3
SVEA
SGQN
PIE-G
TRMS
Pen
2.7+1.3
41.7+6.1
56.0+2.4
54.0=9.4
67.3+5.2
(+20.20%)
Door
0.0.0..0
7.6+1.8
20.31..1
52.6
43.715.4
(-17.0%)
Hammer
0.0.0..0
4.243.7
14.644.77
46.0
53.717.6
(+16.70%)
Average
0.9
17.8
30.3
50.8
54.9
(+8.10%)
6 Conclusion
To address the challenge of task-irrelevant distracting visual features in Visual Reinforcement Learning, we introduce TRMS, method that utilizes existing masking strategies extract masks from the visual scene. It segments and samples only the task-relevant masks. This approach eliminates the need for additional segmentation labels for individual tasks. We bypass the heavy computation time and resources by employing student network that learns these masks in few training steps. We evaluate TRMS on the RI-ViGen (Yuan et al., 2023) benchmark, covering lasks from the DeepMind Control Suite, Unitree locomotion, and Dexterous manipulation under varied distractions. TRMS achieves a 13.70% higher average reward in vidco-hard DeepMind tasks and surpassos basclines 10 out tasks. also yields a 77.46% improvement locomotion tasks and comparable performance in dexterous manipulation, all while demonstrating robust. resilience increasing noise levels compared to baselines.
12

---

Published in Transactions on Machine Learning Research (09/2025)
Limitations and Future Directions. Currently, the prior masking approach, FastSAM, generates masks independently for each state. Incorporating recent advancements like SAM2 (Ravi et al., 2024), which leverages temporal dependencics to refine mask extraction, could greatly enhance sampling efficiency. We need to select k-masks and initial teaching steps Treach for distinct environments, depending on the complexity of the environment. An extension for automated selection of this hyperparameter would be extremely useful. Utilizing segmentation models with human-in-the-loop guidance via natural language descriptions (Zhang et al., 2024) significantly enhance masking performance, especially in complex environments (Rajeswaran et al., 2018). Extending TRMS to non-agent-centric environments, such CARLA, presents unique challenges, as masking-based methods often face limitations in these domains, warranting further investigation. Additionally, future research will explore the applicability of TRMS in more complex, real-world robotic applications, an area that remains largely underexplored this domain.
Acknowledgments
This work was created as part of the research projects 'KIRAMET' (FO999899661), MUTAVIA (FO999922732) and NNATT (FO999907606)), which are funded by the Federal Ministry Republic of Austria Climate Action, Environment, Energy, Mobility, Innovation and Technology. It was also funded in part by the Austrian Science Fund (FWF) under grant 10.55776/COR12 (O. Oxdenizci).
References
Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Marc G Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In International Conference on Learning Representations, Virtual, 2021. PMLR. URL /ttpss//oppnreviee
http://ppaarrreeeeeeeeeaaaaaaaep net:/forum?2ddngal7-883333
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation, 2013.
David Bertoin, Adil Zouitine, Mehdi Zouitine, and Emmanuel Rachelson. Look where you look! saliency-uuided q-networks for generalization in visual reinforcement learning. In Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and Oh (eds.), Advances Neural In- formation Processing Systems, volume 35, pp. 30693 30706, New Orleans, Louisiana, 2022. Curran Associates, Inc. URL https://proceedings neurips neurips.cc/paper htcc//ppprfflll/sspeeeeeeeeeeeee wtw.:e//aaaae//222////./..//// Confeerreecceereeeeeoo pdf.
Timothy Brady, Talia Konkle, and George A Alvarez. A review of visual memory capacity: Beyond individual items and toward structured representations. Journal of vision, 11(5):4-4, 2011.
Chao Chen, Jiacheng Xu, Weijian Liao, Hao Ding, Zongzhang Zhang, Yang Yu, and Rui Zhao. Focus-then- decide: Segmenntation-asiisted reinforcement learning. Proceedings of the AAAI Conference Artificial Intelligence, 38(10):11240-1122) Mar. 2024. doi: 10.1609/ 10.1609/anai.v33t bnai.v38100220022 URL https://ojs. aaai......aaa............a.a.... aaai. org/index.u php/ http:///rticcce/ee///000000
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive lcarning of visual representations. In Hal Daumé III and Aarti Singh (eds.). Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1597-1607, Vienna, Austria, 13 18 Jul 2020. PMLR. URL sttps://proceedingg.oo htt.eeeddingn....ressssee//eeeee pwwpss/viisse// chen20j .ttml.
Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. In International conference on machine learning, pp. 1282-1289, California, 2019. PMLR.
Vedant Dave and Elmar Rucckert. Denoised predictive imagination: An information-thcorrtic approach for learning world models. In Seventeenth European Workshop on Reinforcement Learning, 2024. URL https://openrevier
13

---

Published in Transactions on Machine Learning Research (09/2025)
Jiameng Fan and Wenchao Li. DRIBO: Robust deep reinforcement learning via multi-view information bottleneck. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 6074-6102, Baltimore, USA, 17-23 Jul 2022. PMLR. URL https://proceeding mlr. wtw...preeevvve/fffa...
html.
Linxi Fan, Guanzhi Wang, De-An Huang, Zhiding Yu, Li Fei-Fei, Yuke Zhu, and Animashree Anandkumar.
Secant: Self-expert cloning for zero-shot generalization of visual policies. In Marina Meila and Tong Zhang (ods.), Proceedings of the 38th International Conference on Machine Learning, volume 139 Proceedings of Machine Leurning Research, pp. 3088 3099, Virtual, 18 24 Jul 2021. PMLR. URL mlr. pwwps/eessvveeeeeeeeee htt///v1///f/aa/ html.
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. STAM Journal on Computing, 40(6):1662-177 1714, 2011. doi: (0.1137/1008048/0 URL https:
//doi.org/10.. 1137/10080484X.
Zaifeng Gao, Shixian Yu, Chengfeng Zhu, Rende Shui, Xuchu Weng, Peng Li, and Mowei Shen. Object- based encoding in visual working memory: Evidence from memory-driven attentional capture. Scientific Reports, 6(1):22822) 2016.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, ko- ray kavukeuoglu, Remi Munos, and Michal Valko. Bootstrap your own latent a new approach to self-supervised learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.). Aduances in Neural Information Processing Systems, volume 33, pp. 227444,284 Virtual, 2020. Curran Associates, Inc. URL https: http:ss//ppoocednngs.... /ttp////a000E0000000000000000110
www.ffllleapaaeee palls//paper2020//
saaper.åff.
Bram Grooten, Tristan Tomilin, Gautham Vasan, Matthew E. Taylor, A. Rupam Mahmood, Meng Fang, Mykola Pechenizkiy, and Decebal Constantin Mocanu. Madi: Learning to mask distractions for gener- alization in visual deep reinforcement learning. In Proceedings of the 23rd International Conference on Autonomous Agents and Multiagent Systems, AAMAS '24, pp. 733-742, Richland, SC, 2024. International Foundation for Autonomous Agents and Multiagent Systems. ISBN 9798400704864.
Emil Julius Gumbel. Statistical theory of extreme values and some practical applications: series of lectures, volume 33. US Government Printing Office, USA, 1954.
Kashish Gupta, Debasmita Mukherjee, and Homayoun Najjaran. Extending the capabilities of reinforcement learning through curriculum: A review of methods and applications. SN Computer Science, 3:1-18, 2022.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In S. Ben- gio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems (NeurIPS). volume 31, Montréal, Canada, 2018. Curran Associates, Inc. URL http:://pprocedinng nwwsineuripp... neurips.ccppaper fwtes//paper/2011//// Recaa662223355577444998881222222 Papeee .appr.pdf..........pp......ppppp
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International conference machine learning, pp. 1861-881188 1870, Stockholm, 2018. PMLR.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, Addis Ababa, Ethiopia, 2020. OpenReview.nec, URL https://openrevier http:///pmmeeiieeeeeeeeoerrrraia net//orruuiitessS0r0A44S
Danijar Hafner, Timothy P Lillicrpp, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, Virtual, 2021. JMLR. URL
https://openreview httpa/r.....toro////d/aaaaaaa//a
14

---

Published in Transactions on Machine Learning Research (09/2025)
Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmentation. In 2021 TEEE International Conference on Robolics and Aulomalion (TCRA), pp.
Xi'an, China, 2021. JEEE Press. doi: 10.1100//CMA444502220....6666 201.44444000000002051100003 URL https://doi.org 100..111111///415066666
2021. 9561103.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision trans- formers under data augmentation. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wort- man Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 3680 3693, virtual, 2021. Curran Associates, Inc. URL https://proceedingg. wwwpreeeiig.....ii......... www.meee...pppppeeeeee faaaperrfiles/paperr .df.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpe2: Scalable, robust world models for continuous control. In International Conference on Learning Representations (ICLR), Vienma, 2024. PMLR.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual lcarning for image recognition. Proceedings of the IEEE Conference on Computer Vision and Pallern Recognition (CVPR), pp. 770 778, Nevada, United States, June 2016. IEEE.
Yangru Huang, Peixi Peng, Yifan Zhao, Guangyao Chen, and Yonghong Tian. Spectrum random masking for generalization in image-based reinforcement learning, In S. Koyejo, S. Mohamed, A. Agarwal, D. Bel- grave, K. Cho, and A. Oh (eds.), Advances in Neural Informalion Processing Systems, volume 35, pp. 20393 20406, New Orleans, USA, 2022. Curran Associates, Inc. URL https://proceedingg neurips. cc/paper fll_fffllspppppeeeeee http:/pppppp/2//////// /ttp://pr/..////a////////aaaaaa. or2222550..ttttll110000001111110 pdf.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gunbel-softnax. In Interna- tional Conference on Learning Representations, Toulon, France, 2017. JMLR. URL https://openrreiiee www.forrumni--rk38888568
Yiding Jiang, J. Zico Kolter, and Roberta Railcamu. On the importance of exploration for generaliza- tion in reinforcement learning. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 12951-12986, New Orleans, USA, 2023. Curran Associates, Inc. URL
reppppsccccppppppppp files/ pdf.
Leslie Pack Kaelbling, Michael L. Littman, and Authony R. Cassandra. Planning and acting partially observable stochastic domains. Artificial Intelligence, 101 134, 1998. ISSN 0004-3702. doi: https: URL https://ww..sceencddreeee http://ee...cccccceeee//rer.oiie
pis/800004370298000222.
Daniel Kaiser, Nikolaas N. Oosterhof, and Marius V. Peelen. The neural dynamics of attentional selection in natural scenes. Journal of Neuroscience, 36(41):10522-103333 2016. ISSN 0270-6474. doi: 10.1523/
URL https::/wwww.nnerroset http://o//ocottt/////////000
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint ar arXiv:1412.6980, 2014.
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollar, and Ross Girshick. Segment anything. In Proceedings of the IEEE/CV International Conference on Computer Vision (ICCV) pp. 4015-4026. Paris, France, October 2023. IEEE.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. Journal of Artificial Intelligence Research, 76:201-221 -214, 2023.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representations for reinforcement learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Reseorch, pp. 5639-5650, virtual, 13 18 Jul 2020a. PMLR. URL /tttps://proot /tppa//e//soccccos.....m. mlr. wttpss//119/lassi122000// html.
15

---

Published in Transactions on Machine Learning Research (09/2025)
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Picter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances in neural information processing systems, 33:19884-19895 2020b.
Alex X. Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. In Advances Neural Information Processing Systems (NeurIPS), virtual, 2020a. PMLR.
Kimin Lec, Kibok Lee, Jinwoo Shin, and Honglak Lec. Network randomization: A simple technique for generalization in deep reinforcement learning. In International Conference on Learning Representations, virtual, 2020b. PMLR. URL https: htttss//opeerewee //tp///epeeeeeeee/////ee nett/forum/.d-lllgeveeeet
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. Journal of Machine Learning Research, 17(39):1-40, 2016.
Dongfen Li, Lichao Meng, Jingjing Li, Ke Lu, and Yang Yang. Domain adaptive state representation aligument for reinforcement. learning. Information Sciences, 609:1353 1368, 2022.
Timothy Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Contimuous control with deep reinforcement learning. In Yoshua. Bengio and Yann LeCun (eds.), 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, pp. 14, San Juan, Puerto Rico, 2016. PMLR. URL http:/ attp://arxiv. arxiv. org/abs/// /aob///509/// 029711.
Minsong Liu, Luntong Li, Shuai Hao, Yuanheng Zhu, and Dongbin Zhao. Soft. contrastive learning with q-irrelevance abstraction for reinforcement learning. IEEE Transactions on Cognitive and Developmental Systems, 15(3):1463 144--1473 1473, 2023a. doi:
Qiyuan Liu, Qi Zhou, Rui Yang, and Jie Wang. Robust representation learning by clustering with bisimulation metrics for visual reinforcement learning with distractions. In Proceedings of the Thirty- Seventh AAAI Conference on Artificial Intelligence and Stérty-Fifth Conference on Innovative Applica- tions of Artificial Intelligence and Thirteenth Symposium on Educational Advances in Artificial Intelli- gence, 23AT23/1AAT33/ / EA Washington, USA, 2023b. AAAI Press. ISBN 978-1-57735-850-00 doi: http///tdoo..g.....// 16t.//...100////// /aaii. v3717 26063.
10.1609
10.1609/aaai.v337/ /at.aaaaai..777.. 26666. 26063. URL
Siao Liu, Zhaoyu Chen, Yang Liu, Yuzheng Wang, Dingkang Yang, Zhile Zhao, Ziqing Zhou, Xie Yi, Wei Li, Wenqiang Zhang, and Zhongxue Gan. Improving generalization in visual reinforcement learning via conflict-aware gradient agreement augmentation. Proceedings of the IEEE/CVI International Confer- ence on Computer Vision (ICCV) pp. 23436 23436-23446, Vancouver, Canada, October 2023c. IEEE.
Tly Toshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference Learning Representations, 2019. URL /ttps::/oppenrvviii
Bhairav Mehta, Manfred Diaz, Florian Golemo, Christopher J. Pal, and Liam Paull. Active domain random- ization. In Leslie Pack Kaclbling. Danica Kragic, and Komei Sugiura (ods.), Proceedings of the Conference on Robot Learving, volume 100 of Proceedings of Machine Learning Research, pp. 1162-1176, virtual, 30 Oct-O 01 Nov 2020. PMLR. URL bttps://proceeet
wttp...peeeeeeo//eeeeeeeeeee pptpss/r1100/eeeta200. html.
mlr,
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529 533, 2015.
Alhassan Mumuni and Fuseini Mumuni. Data augmentation: A comprehensive survey of modern approaches. Array, 16:1 16:100258, 2022. ISSN 2590-0056. doi: http:://liiooo////ooon///rr/rppp URL https: //www www..sciencedieee.... httpe/ttteectttccccieeeeeeeeeeee com http:///cccacceattte//////////oa
Sanmit Narvekar, Bei Peng, Mattoo Lconetti, Jivko Sinapov, Matthow E Taylor, and Peter Stone. Curriculum learning for reinforcement learning domains: A framework and survey. Journal of Machine Learning Research, 21 (181):1-50) 2020.
16

---

Published in Transactions on Machine Learning Research (09/2025)
Shahin Nasr, Ali Moceny, and Hossein Esteky. Neural correlate of filtering of irrelevant information from visual working memory. PLoS One, 3(9):e382, 2008.
Benjamin Peters and Nikolaus Kriegeskorte. Capturing the objects of vision with neural networks. Nature human behaviour, 5(9):1127-1144, 2021.
Roberta Raileanu, Maxwell Goldstein, Denis Yarats, Tlya Kostrikov, and Rob Fergus. Automatic data aug- mentation for generalization in reinforcement learning. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang. and Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, virtual, 2021. Curran Associates, Inc. URL https:// https::/prroeeddngsss www.:eiiiiinnnnnnn...... www.:eeuurrrcpppppeeeee Pawpe/dddd-apppppppppp..
pp. httpe//pprr0000///////////226610
Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. In Proceedings of Robotics: Science and Systems, Pittsburgh, Pennsylvania, June 2018.
Robotics: Science and Systems. doi:
Nikhila Ravi, Valentin Gabeur, Yuan- Ting Hu, Ronghang Hu, Chaitanya Ryali, Tengyu Ma, Haitham Khedr, Roman Rädle, Chloe Rolland, Laura Gustafson, Eric Mintun, Junting Pan, Kalyan Vasudev Alwala, Nicolas Carion, Chao- Yuan Wu, Ross Girshick, Piotr Dollár, and Christoph Feichtenhofer. Sam Segment anything in images and videos, 2024. URL
Starg/abs/44000 00714.
Picter RoeIfsema and Arjen van Ooyen. Attention-gated reinforcement learning of internal representations for classification. Newral compulation, 17(10):2176 2214, 2005.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-nel: Convolutional networks for biomedical image segmentation. In Nassir Navab, Joachim Hornegger, William M. Wells, and Alejandro F. Frangi (eds.), Medical Image Computing and Computer-Assiste Intervention - MICCAI 2015, pp. 234-241, Cham, 2015. Springer International Publishing. ISBN 978-3-919-2457444.
Brian Scholl. Objects and attention: The state of the art Cognition, 80(1-2):1 46, 2001.
Katharina N. Scidl, Marius V. Peelen, and Sabine Kastner. Neural evidence for distracter suppression during visual search in real-world scenes. Journal of Neuroscience, 32(34):11811-11189 2012a. doi: 10.1523/JNEERRRR /NE.//NNUUROT.....OOOO 699-112.12212 URL
visual search real-world scenes. Journal of Neuroscience, jpnearosc...rrg/ooteee//////eeee
Katharina Seidl, Marius V Peelen, and Sabine Kastner. Neural evidence for distracter suppression during 2012b.
Ruixiang Sun, Hongyu Zang, Xin Li, and Riashat Islam. Learning latent dynamic robust representations for world models. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 47234-47260, Vienna, 21-27 Jul 2024. PMLR. URL https://proceedeie
mlr. wal..ppess/vv22 20005//u22..... html.
Yanchao Sun, Ruijie Zheng, Xiyao Wang, Andrew E Cohen, and Furong Huang. Transfer RL across obser- vation feature spaces via model-based regularization. International Conference on Learning Represen- tations, pp. 23, virtual, 2022. PMLR. URL attps://openreviiet nett/Eoruu//dd777A000511133
Richard S Sutton. Reinforcement learning: An introduction. A Bradford Book, 2(2):548, 2018.
Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeicch Zaremba, and Picter Abbeel. Domain ran- domization for transferring deep neural networks from simulation the real world. In 2017 IEEE/RS) International Conference on Intelligent Robots and Systems (TROS), number in IEEE, pp. 23 300 Van- couver, British Columbia, Canada, 2017. IEEE. doi: 10. 000100//RRS..217.82222222
Manan Tomar, Riashat Islam, Matthew Taylor Sergey Levine, and Philip Bachman. Ignorance bliss: Robust control via information gating. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Newral Information Processing Systems, volume 36, pp. 38624-38641, New
17

---

Published in Transactions on Machine Learning Research (09/2025)
Orleans USA, 2023. Curran Associates, Inc. URL http:ppppp//pppoeeiggg..oo//a cwtprrips.ccpppepeeeeee wwwppper.ppppppcccccptcp files/
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of Machine Learning Research, 9(86):2579-280) 2605, 2008. URL http:://mlr..orgppppeeeeee
htmm..
Che Wang, Xufang Luo, Keith Ross, and Dongsheng Li. Vr13: A data-driven framework for visual deep reinforcement learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. USA, 2022. Curran Associates, Inc. URL https::/ neurips....... ctt.aps.ccpppppe htc/pappr_fileee wwwes//apepr.//// pdf.
New Orleans,
Shuo Wang, Zhihao Wu, Jinwen Wang, Xiaobo Hu, Youfang Lin, and Kai Lv. How to learn domain invariant representations for visual reinforcement learning: An information-thoortticl perspective. In Kate Larson (od.); Proceedings of the Thirty- Third International Joint Conference on Artificial Intelligence, IJCAI-24, pp. 1389 1397, Jeju, Korea, 8 2024. International Joint Conferences on Artificial Intelligence Organization. doi: 10.24963 21124463//jccii220 /jeai.2024/ 154. URL https://doi. org/ 10. 24963/ijcai 2024/154. Main Track.
Weiyao Wang and Gregory D. Hager. Domain adaptation of visual policies with single demonstration. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pp. 17208 17215, Yokohama, Tokyo, 2024. IEEE. doi: J0.1109/CRRA5714...0 27.2244...22.... 10010001160 00669.
Xudong Wang, Long Lian, and Stella X. Yu. Unsupervised visual attention and invariance for reinforcement learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognation (CVPR), pp. 6677 6687, Virtual, June 2021. IEEE.
Ziyu Wang, Yanjie Ze, Yifei Sun, Zhecheng Yuan, and Huazhe Xu. Generalizable visual reinforcement learning with segment anything model, 2023.
Eliot Xing, Abhinav Gupta, Sam Powers, and Victoria Dean. Kitchenshift: Evaluating zero-shot generaliza- tion of imitation-based policy learning under domain shifts. In NewrIPS 2021 Workshop Distribution
Shifts: Connecting Methods and Applications, Virtual, 2021a. NeurIPS 2021 Workshop on Distribution
Shifts: Connecting Methods and Applications. URL https: /opp///pppeeeeeeeee/eeee/eeeeeee pepperrreew..eet/oruutaatt. het/fforum:ddDDgg
Jinwei Xing, Takashi Nagata, Kexin Chen, Xinyun Zou, Emre Neftci, and Jeffrey L Krichmar. Domain adaptation in reinforcement learning via latent unified state representation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 10452- 00452-1049,, virtual, 2021b. AAAI Press.
Rui Yang, Jie Wang, Zijie Geng, Mingxuan Ye, Shuiwang Ji, Bin Li, and Feng Wu. Learning task-relevant representations for generalization via characteristic functions of reward sequence distributions. In Pro- ceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD '22, pp. 2242-2252. New York, NY, USA, 2022. Association for Computing Machinery. ISBN 9781450393850, doi: 3539999 URL https: / http:///ooooo///e///eeeee 10. 1e11111144448..... 3533333337783333391
Sizhe Yang, Yanjie Ze, and Huazhe Xu. Movie: Visual model-based policy adaptation for view gener- alization. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Ad- vances in Neural Information Processing Systems, volume 36, pp. 21507 21522,1523, New Orleans, USA, 2023. Curran Associates, Inc. URL https://prroceedings www.nnn......r....ee.. www.:ss......p..pppp.rrer. wwwess/paeee//aaeee pdf.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Represenlations, virtual, 2021. PMLR. URL http:p///pppprriiiieeeeeerorrres netf/forumPid--Y---6rrCC
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Im- proved data-augmented reinforcement learning. In International Conference on Learning Representations, virtual, 2022. pmlr. URL https: http////reeeeeeeee/eeeeeeeeeeeee Bicdd=_SJ-____ _SJ__yyes8..._--
18

---

Published in Transactions on Machine Learning Research (09/2025)
Bang You, Jingming Xie, Youping Chen, Jan Peters, and Oleg Arenz. Self-supervised sequential information bottleneck for robust exploration in deep reinforcement learning. arxiv preprint or (r.Xi:2209.05555 00339, 1(1): 1, 2022.
Tao Yu, Zhizheng Zhang, Cuiling Lan, Yan Lu, and Zhibo Chen. Mask-based latent reconstruction for reinforcement learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Newral Informalion Processing Systems, volume 35, pp. 25117-2511,, New Orleans, USA, 2022. Curran Associates, Inc. URL https://proceeddngg neurips wourips.cc/pappe.. wiles/paper/ PaperrConfereecee .df.
Zhecheng Yuan, Guozheng Ma, Yao Mu, Bo Xia, Bo Yuan, Xueqian Wang, Ping Luo, and Huazhe Xu. Don't touch what matters: Task-aware lipschitz data augmentation for visual reinforcement learning, In Lud De Raedt (ed.). Proceedings of the Thirty-First International Joint Conference on Artificial Intelli- gence, IJCAT-22, pp. 3702-3708, Vienna, 7 2022a. International Joint Conferences on Artificial Intelligence Organization. doi: 10.24963/ijoee /ttp//aaa.222/////2 URL attps://dof org/10. 24963/ijcai. htt.iiaaiiii/........ 2022/514. Main Track.
Zhecheng Yuan, Zhengrong Xue, Bo Yuan, Xueqian Wang, YI WU, Yang Gao, and Huazhe Xu. Pre-trained image encoder for generalizable visual reinforcement learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural In- formation Processing Systems, volume 35, pp. 13022 13037, New Orleans, Louisiana, 2022b. Curran Associates, Inc. URL https://proceedings mew.rddiinnnnnurip.....e.... wwwrips.cc/papee.eeeeeeeeoopeeee cc/paper ceappr_ffles/peeeeeeeeee filess/apprr/0022/52 1aae244496c100ddoeeeeeeee/ddaap
Zhecheng Yuan, Sizhe Yang, Pu Hua, Can Chang, Kaizhe Hu, and Huazhe Xu. Rl-vigen: reinforcement learning benchmark for visual generalization. A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 6720-6747, New Orleans, USA, 2023. Curran Associates, Inc. URL https: httpp::/prroceeiings.. nww.wps.cccrcee.eeeeee. wwwpcceppppppeeeeeeeeeeeeeeeeeee files/ //tpr////0////////////eeeee00000 wwwpe//eeeee000000000000000ee0e. hatprerreaaaaaaaeeaaaaaaaaaaaaaa Bww.Beanhhhhhheeeeee pdf.
Hongyu Zang, Xin Li, and Mingzhong Wang. Simsr: Simple distance-based state representations for deep re- inforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence, 36(8))8977-9005, Jun. 2022. doi: 111609/aaa.v3668... 20883.
CRL https:/ //jjs.//......//// aaai oerili.rg/ninddeeeeeeee http:pppp//Atttaaieeeeeeeeee
Amy Zhang, Nicolas Ballas, and Joelle Pineau. A dissection of overfitting and generalization continuous reinforcement learning. arXio preprint Xrxiv:180 (00:180007997, 1(1), 2018.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine, Learning invariant reprosentations for reinforcement learning without reconstruction. In International Conference on Learning Representations, pp. 17, Virtual, 2021. PMLR. URL https: htttp::///penreviii
Yuxuan Zhang, Tianheng Cheng, Lianghui Zhu, Rui Hu, Lei Liu, Heng Liu, Longjin Ran, Xiaoxin Chen, Wenyu Liu, and Xinggang Wang. Evf-sam: Early dision-languar fusion for text-prompted segment anything model. XrXiu preprint Stt::2406200077 2024.
Xu Zhao, Wenchao Ding, Yongqi An, Yinglong Du, Tao Yu, Min Li, Ming Tang, and Jinqiao Wang. Fast segment anything, 2023.
Bolei Zhou, Agata Lapedriza, Aditya Khosla, Aude Oliva, and Antonio Torralba. Places: A 10 million image database for scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, 40(6): 1452 1464, 2018. doi: 10.100//PAAM12211.22222
19

---

Published in Transactions on Machine Learning Research (09/2025)
A Implementation Details
We provide additional details on our implementation of TRMS. Table 4 summarizes the hyperparameters utilized in our study. For the DMC Generalization Benchmark (Hansen & Wang, 2021) and Locomotion Task (Hansen et al., 2021), we used the default hyperparametens specified RI-ViGen (Yuan et al., 2023) for the respective baselines.
Table 4: Hyperparameters for TRMS in DMC-GB.
Hyperparameter
Values
Input size
84 x 84
Optimizer
Adam (Kingma, 2014)
Learning Rate (Actor,Crite Masker)
10-4
Adam B1, 82 (All Networks)
0.9, 0.999
Discount (7)
0.99
Frame Stack
3
Action Repeat
2
Initial Batch Size
Batch Size
256
Feature Dimension
256
Initial Sampling Steps
2000
Replay Buffer Size
150K
Environment Steps
500K
Gradient Steps Per Training Step
Target Update Interval
2
Target Smoothing Cocfficient Critic
0.01
Target Smoothing Coefficient Encoder Initial Temperature (a in SAC)
0.05
0.1
Temperature Learning Rate
10
Mask Sampler
AdamW (Loshchilov Hutter, 2019)
Learning Rate
1
Number of Masks
4
Gumbel Softmax temperature (r)
5.0
Segmentation Model
FastSAM (Zhao el al., 2023) (default)
Image Size
640
IoU Threshold
0.75
Confidence Threshold
0.40
Overlap Mask
False
Initial Teaching Steps (Treach)
25K (DMC-GB) and 50K (Otherwise)
20

---

Published in Transactions on Machine Learning Research (09/2025)
The common Reinforcement Learning Hyperparameters were kept consistent with those used for TRMS, while method-specifi parameters followed the configurations provided in their respective papers. Details on the specific hyperparameters are listed in Table 5.
Table 5: Hyperparameters for basclines in DMC-GB.
Hyper-parameters
Value
Feature dim
DrQ-v2, CURL: 50; otherwise: 256
N-step return
DrQ: 1; otherwise: 3
Optimizer
Adam
Hidden dim
1024
Frame stack
SGQN Quantile Threshold
0.95 0.98
Critic Weight Decay
10-5
SGQN Auxiliary Learning Rate
8e-5
For the Dexterous Manipulation tasks, the hyperparameterss utilised for the respective baselines mentioned in their papers and for environments mentioned in RL-ViGen (Yuan et al., 2023). They are described in the Table 6.
Table 6: Adroit Hyperparameters.
Hyper-parameter
Task
Value
Training Frames
Hammer
106
Door
106
Pen
x 106
Learning Rate
Hammer
10-4
Door
10-4
Pen
10
k Masks
Hammer
15
Door
15
Pen
10
SGQN Quantile
Hammer
0.9
Door
0.9
Pen
0.9
SGQN Critic Weight
Hammer
Door
0.5
Pen
0.9
SGQN Auxiliary Learning Rate
Hammer
8 x 10-5
Door
8 x 10
Pen
8 x 10-5
Soft Actor-Critic Algorithm
The Soft Actor-Critic (SAC) (Haarnoja et al., 2018) algorithm an off-policy reinforcement learning method that maximizes cumulative rewards while promoting entropy to encourage exploration. The objective for SAC is given by = arg max E (e Te.
Σ r((t, az) + aHl(TS, (:1st)
(11)
21

---

Published in Transactions on Machine Learning Research (09/2025)
where is temperature parameter that balances reward and entropy terms. SAC employs two critic networks, Qu and Qo,.: trained to minimize the soft. Bellman residual
(()g.) = ((s) ((00, (s, (
min Qony (5'.a') - ((o))=(0o7((d())) 27
(12)
The actor nelwork ne. is updated to maximize the expected Q-value regularized by an entropy term, defined
(0,) = Bs-Daaamse [ox log To. (a]s)= - Q0, (0, a)],
(13)
where a is the entropy temperature coefficient. The corresponding policy gradient is given by
(o)) Es,a &V log To. (als) (a - Qo, (s, (s,a))].
(14)
To balance exploration and exploitation, SAC adapts a by minimizing the temperature objective:
[(o) EaNto. a
log Te. (a|s) - -aH]
(15)
where denotes the target entropy, encouraging diverse action sampling high-dimensional action spaces.
C Architecture Details
Encoder. The encoder network consists of two main components: a shared convolutional module and a subsequent linear projection. The shared convolutional module is an 11-layer network designed to process input observations composed of 3 stacked RGB frames, with dimensions [9, 84, 84], ultimately generating spatial feature maps. The first layer employs 3 x 3 convolutional kernel with stride of 2 and 32 output channels, allowing for an early reduction in spatial resolution. The remaining layers are structured as sequential ReL.U-convoltion blocks, each composed of ReLU activation followed by a 3 x 3 convolution with a stride of 1 and maintaining 32 channels across all layers. This uniform channel depth preserves consistency in feature representation throughout the network.
The final convolutional output is then flattened into feature vector of size 32 x 21 21. This vector is subsequently passed through linear projection layer, which reduces the dimensionality to 512, thus producing condensed latent representation suitable for further processing. This final representation serves as the input to the policy and value networks, allowing for efficient and effective state encoding. Furthermore, input frames are normalized by scaling to the range [-0.5, 0.5]
Actor. The actor network comprises feature extractor and a policy head. The feature extractor maps the 512-dimensional imput reprosentation 256-dimensional latent space via a fully connected layer, followed by layer normalization and Tanh activations for normalized, non-linear transformations. The resulting features are then passed through two hidden layers with 1024 units and ReLU activations, capturing complex action- value mappings.
The final layer outputs action means, 14, which are scaled by a Tanh activation to enforce bounded action outputs. The standard deviation, a, constant and scaled by an input parameter, std. Together, l and o parameterize Truncated Normal distribution for continuous action sampling.
Critic. The critic network comprises two parallel Q-networks, Q1 and Q2, which are employed to estimate state-action values and mitigale overestimation bias. Initially, observations are passed through a shared feature extraction module that projects the input representation of dimension 512 into a 256-dimensional feature voctor. This transformation is accomplished through a fully connected layer, followed by layer normalization and Tanh activation, which ensures normalized outputs and reduces the likelihood of activation saturation.
22

---

Published in Transactions on Machine Learning Research (09/2025)
The resulting 256-dimensional feature vector is concatenated with the action input, forming joint representation that is processed by each Q-network independently. Both Q1 and Q2 are structured with two hidden layers, each containing 1024 units, and use ReLU activations to introduce non-linearity, enabling the networks to model complex value functions effectively. The final layer each Q-network outputs a single scalar, representing the Q-value for the given stale-action pair. By using two independent. Q-networks, the critic can take the minimum of both Q-value estimates, which reduces overestimation an issue commonly encountered in value- reinforcement learning. This design contributes to the stability and robustness of the learned policy.
Mask Sampler. The Mask Sampler network is designed to process multiple input masks and output selection probabilities for each. A single input has the shape (k, 84, 84), where the number of masks. The network comprises three convolutional layers, each with spatial dimensions. The first layer maps the input channels, corresponding to the number of masks, to 32 feature maps, followed by batch normalization and ReLU activation. The feature depth is then sequentially increased to 64 and 128 channels by the subsequent convolutional layers, each followed by batch normalization and ReLU to enhance spatial feature extraction.
kernels and padding of to maintain
The output of the final convolutional layer undergoes global average pooling to reduce the spatial dimensions to x 1, yielding a 128-dimensional vector. This vector is then passed through a fully connected layer with units and a ReLU activation, followed by a final linear layer that outputs logits corresponding to the number of masks. A Snumee-Softtax activation with subsequent temperature applied these logits.
Masker. The Masker network is a convolutional architecture designed to produce single-channel mask from RGB imputs. The network consists of five convolutional layers, each with a 3 3 kernel and padding of to maintain spatial dimensions. The first two layers map the input image, with three color channels, to 64 feature channels through successive applications convolution, ReLU activation, and Batch Normalization, promoting feature extraction while stabilizing training.
The third and fourth layers reduce the feature depth to 32 channels, employing similar ReLU and batch normalization operations to preserve spatial information while refining feature representations. The final convolutional layer outpuis a single-channel feature map, which represents the generated mask.
The output feature map is then passed through sigmoid activation function, designed to constrain output values between 0 and 1. All convolutional layers are initialized with Xavier uniform initialization, ensur- ing balanced weight distributions. This architecture allows Masker effectively learn spatial patterns for producing accurate binary masks from RGB input data.
23

---

Published in Transactions on Machine Learning Research (09/2025)
D
Wall Time
As TRMS incorporates both SAM and a student network to enhance computational efficiency, here we compare the wall time required for convergence between TRMS and the SAM-only variant. Evaluation is performed across three seeds of the Walker Walk task, where the teacher network is employed for the first 25k steps.
As shown, TRMS achieves an average wall time of approximately 15 hours, whereas the SAM-only variant reaches up to 3 days and 1 hour (~73 hours). This substantial reduction demonstrates that TRMS not only accelerates learning early but also sustains faster convergence through its student network (due to its small architecture), which absorbs knowledge from the heavy parameterised teacher network.
SAM Only
2h
Method
TRMS
16h
10
20
30
50
60
Wall Time (Hours)
Figure 5: Comparison of TRMS and SAM-Only Wall time for Walker Walk on three seeds.
24

---

Published in Transactions on Machine Learning Research (09/2025)
E
Additional Visual Results
The following visual demonstrations showcase selected frames from the video-hard environment. On the left, each image displays the actual observation, while the corresponding masked observation is presented on the right. These comparisons highlight the complexity of the environment, where the masking process isolates relevant features, facilitating the agent's focus on relevant task elements amidst challenging distractors.
Finger Spin
Cartpole Swingup
Walker Stand
Walker Walk
(e) Walker Walk
(f) Door Open
Figure 6: Visual examples from the video-hard environment, showing actual observations (left) and corre- sponding masked observations (right), highlighting relevant feature isolation.
F Baseline Results:
Baseline results presented in this paper obtained as follows. For the DMC-GB environments, implemented most baselines ourselves, with the exceptions of DrQ, DrQ-v2, and SVEA, whose performance numbers were directly cited from the PTE-G paper (Yuan el al., 2022b). For the PIE-G baseline, we con- ducted experiments using the authors' original implementation. SGQN (Bertoin et al., 2022) results were reproduced using the official implementation provided with the RL-ViGen (Yuan et al., 2023) benchmark. For Locomotion and Manipulation environments, all baseline performances are directly cited from the RL- ViGen benchmark 3,
25


Reviews:
[{"review_id": "zzj5iuke3z", "reviewer": "Reviewer_qsNd", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths:\n- This paper targets at an essential problem - improving the robustness and generalization of visual reinforcement learning.\n- The empirical results are sufficient by involving multiple benchmarks and baselines. The proposed algorithm demonstrates state-of-the-art performance.\n\n## Weaknesses:\n- The novelty is limited, which combines SAC and SAM through a mask selector.\n- There should be a comparsion with related works for generalization in visual RL in Section 2, to showcase the novelty and research gap.\n- It would be important to explicitly testifiy the following statement through illustrative exmaples: \"This occurs because selecting an irrelevant segment means that, during each state sampling from the replay buffer, an insignificant pixel in a given state might assume varying values, despite the pixel contributing nothing meaningful to the task value. As a result, the mask selection process is encouraged to discard these irrelevant pixels. \", as this is the core motivation.\n- The baselines are relatively old - from 2022 and before.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Please check the strengths.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please check the weaknesses.", "broader_impact_concerns": "N\\A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "lDNHDEfxI7", "reviewer": "Reviewer_MqPz", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nThe paper is well written and it proposes TRMS, a computationally efficient method for segmentation-driven masking in RL, leveraging FastSAM and a student network. Furthermore, the  authors provide a  substantial comparison with baselines to showcase the performance of TRMS. \n\nWeaknesses:\nHow is hyperparameter tuning done in the paper? The method requires selecting the number of masks and the duration of the teaching phase, which are environment-dependent. Is there systematic way to choose those hyperparameters\n\nIt is not clear where the loss $\\mathcal{L}_{\\pi}$ in algorithm 1 is defined\n\nWhy is there no comparison with FTD and MaDi? They seem to be closely related methods. If there is a particular reason, then it should be mentioned otherwise they should be part of the baselines.\n\nThe experiments use low-resolution inputs. It is unclear how TRMS would scale to high-resolution observations, where segmentation and masking could become computationally prohibitive.\n\nHow is diversification ensured when training the student network and how is over-fitting avoided? This part of the paper needs some more clarity.\n\n$\\mathcal{L}_{masker}$ appears before it is defined.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The TRMS method introduces a masking technique that identifies and preserves task-relevant image regions, leading to enhanced generalization and robustness in noise-prone environments. Furthermore, the authors train a student network to reduce computational complexity of FastSAM. Finally,the authors showcase the performance of their method in a number of diverse simulation environments and they compare it with a number of baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Kindly address the weaknesses listed above.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "s4WucxNgLx", "reviewer": "Reviewer_n92v", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\nClear Motivation and Human-Inspired Perspective:\nThe paper addresses an important challenge in visual reinforcement learning—robustness and generalization in noisy, high-dimensional environments—by leveraging human-like segmentation and attention mechanisms. This conceptual approach is well-motivated and grounded in both cognitive science and machine learning literature.\n\nComprehensive Related Work and Experimental Scope:\nThe literature review is thorough, and the comparisons cover a wide range of recent and relevant baseline methods. The experimental evaluation spans multiple well-known benchmarks and difficult tasks, making the results broadly credible and informative.\n\n\n\n---\n\nWeaknesses\n\nUnclear Baseline Numbers and Potential Reproducibility Issue:\nThe most significant concern is the lack of clarity regarding the source of the baseline results reported in the experimental tables. It is not explicitly stated whether these baseline numbers are reproduced by the authors under consistent experimental settings, or if they are taken directly from the original papers.\nIn several cases, the baseline performances (e.g., for DrQ, SVEA, SRM, PIEG, SGQN) in challenging settings appear substantially lower than those reported in the respective original publications. This discrepancy raises concerns about the experimental fairness and reproducibility of the results.\nWithout a clear explanation, it is difficult to judge whether the observed improvements stem from genuine methodological advances, or from potential inconsistencies in baseline reproduction (e.g., differences in hyperparameters, training steps, seeds, or evaluation protocols).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a novel method, Task-Relevant Mask Sampling (TRMS), to improve the generalization and robustness of visual reinforcement learning (RL) agents in environments with high-dimensional observations and severe background distractions. The main contributions of this work are as follows:\n\n1. Task-Relevant Masking Framework:\nThe authors introduce TRMS, which leverages pre-trained image segmentation models (such as FastSAM) to generate multiple candidate object masks for each observation. A mask selection network then dynamically selects the subset of masks most relevant to the task at each timestep, allowing the RL agent to focus on task-relevant visual regions and filter out distractors.\n\n\n2. Efficient Student Network Distillation:\nTo address the computational overhead of running segmentation models throughout training, the authors develop a lightweight student network. This network is trained to imitate the masking process of the pre-trained segmentation model during an initial phase and subsequently replaces the heavy segmentation module for the remainder of training, greatly improving computational efficiency.\n\n\n3. Comprehensive Empirical Evaluation:\nThe method is evaluated on a suite of challenging vision-based RL benchmarks (RL-ViGen), including variants of DeepMind Control Suite, Quadruped Locomotion, and Dexterous Manipulation with video-based background distractions. TRMS demonstrates superior generalization and robustness compared to a broad set of prior methods, achieving state-of-the-art performance in most settings, especially under high-noise conditions.\n\n\n4. Ablation and Analysis:\nThe paper includes visualizations and ablation studies to illustrate the effectiveness of mask selection and the learned representations' invariance to visual distractors.\n\n\n\nOverall, the work presents a practical and effective framework for isolating task-relevant information in visual RL, and sets a new benchmark for robust generalization in distraction-rich environments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Recommendation for Improvement:\nThe paper would benefit significantly from a detailed clarification of how the baseline results were obtained. If the baselines were re-implemented, please specify the exact experimental settings and confirm that they match those of the original papers. If the numbers are cited, please clearly indicate the source and discuss any known discrepancies. Providing open-source code and/or additional ablation studies to support reproducibility would also greatly strengthen the contribution.", "broader_impact_concerns": "no", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)