1 Introduction↩︎

Visual Reinforcement Learning (VRL) has achieved tremendous success in various fields, including video games [1], [2], robotic manipulation [3]–[5], and autonomous navigation [6], [7]. However, high-dimensional image observations inherently contain redundant, task-irrelevant details, which makes the agent prone to overfitting the training environment. This poses significant challenges in generalizing learned policies to new environments with visual changes. Therefore, it is crucial to improve the generalization ability of VRL.

To learn policies robust to visual changes, previous methods can be broadly categorized into two types. One is data augmentation  [8]–[12], aiming to cover sufficient environmental changes during training that also exist in the testing environment. The other is designing self-supervised auxiliary tasks [13]–[18] to assign certain properties to the representation. However, representations obtained from previous self-supervised auxiliary tasks don’t explicitly strip away task-irrelevant information. The policies based on these representations may also overfit to task-irrelevant features, thereby affecting policies’ generalization ability.

To address this problem, as shown in Figure  1, we propose the idea of decoupling the observations into task-relevant and task-irrelevant representations. Then, the task-relevant representations are utilized to support learning policies. This decoupling idea has two advantages: 1) The representations supporting policy learning are expected to exclude task-irrelevant information, making policies focus more on the task-relevant parts of the observations and thus promoting policy generalization; 2) By explicitly ignoring task-irrelevant features, the burden on the agent to process unnecessary information is reduced, thereby improving sample efficiency.

Based on the idea mentioned above, we introduce the Task-Relevant Representation Decoupling (T2RD) method for VRL, which is capable of zero-shot generalization to unseen visual scenarios. Specifically, to achieve the decoupling effect, T2RD consists of three components: task-relevant representation consistency, cross-reconstruction, and cross-dynamic prediction. The first two components are designed to separate content and style features visually. However, the resulting content representations may not necessarily be task-relevant. To further refine task-relevant features from the content representation, we designed the third component that integrates dynamic prediction, thereby achieving the decoupling of task-relevant and task-irrelevant features.

T2RD has two principal advancements: 1) Its representation decoupling module is self-supervised. Although there are similar decoupling ideas [19]–[22] in the field of Computer Vision (CV), most of them are supervised and require labeled datasets or even paired datasets. 2) It is more suitable for VRL tasks. Compared to methods that merely perform foreground-background segmentation [23] or content-style separation [24], our decoupling approach is more applicable to generalization tasks in VRL, achieving separation between task-relevant and task-irrelevant features.

Additionally, we propose a simple yet effective combined data augmentation technique for VRL. Recent methods mostly select a single data augmentation technique [11], [14], [15], [25]. However, these methods suffer from generalization bias, meaning that their generalization ability is heavily influenced by what data augmentation techniques they choose. By simply combining Random Convolution [14] and Random Overlay [25], we increase the diversity of augmentations, alleviating the generalization bias and thereby improving the algorithm’s generalization performance.

We conduct extensive experiments on two benchmarks: DeepMind Control Suite Generalization Benchmark (DMControl-GB)  [25], [26] and Robotic Manipulation  [27]. The agent is trained in a fixed environment and evaluated for its generalization ability to unseen environments during training. The experimental results demonstrate that our method exhibits state-of-the-art generalization performance and sample efficiency.

Our main contributions are emphasized as follows:

• We introduce the concept of decoupling observations into task-relevant features and task-irrelevant features in VRL, and then use only task-relevant features for policy learning.

• We propose a self-supervised representation decoupling framework suitable for VRL, which can effectively achieve the decoupling of task-relevant features and task-irrelevant features.

• We propose a simple yet effective combined data augmentation technique, which significantly improves the generalization performance of VRL, as evidenced by experiments.

• Extensive experiments prove that T2RD achieves state-of-the-art generalization performance and sample efficiency.

2 RELATED WORK↩︎

2.1 Data Augmentation in Visual RL↩︎

Data augmentation has been shown to improve the generalization ability of visual reinforcement learning algorithms. RAD  [8] pioneers the broad application of data augmentation in RL, improving sample efficiency and generalization. DrQ  [9] aggregates multiple data augmentations, regularizing learning directly from pixels by averaging across multiple image transformations for the Q-function and Q target function. SODA  [14] proposes a soft data augmentation method that decouples data augmentation from policy learning and introduces a novel data augmentation approach that linearly adds observations with additional images. TLDA  [11] introduces task-aware Lipschitz data augmentation, explicitly identifying task-relevant pixels with large Lipschitz constants and augmenting only task-irrelevant pixels. Compared with single data augmentations, our combined data augmentation method increases the diversity of augmentation, which can alleviate the issue of generalization bias. At the same time, using different augmentations improves the network’s robustness to visual changes.

2.2 Generalization in Visual RL↩︎

Researchers have studied how to improve the generalization ability of visual reinforcement learning from various aspects. Beyond data augmentation, existing works primarily pursue generalization by learning invariant representations, which can be broadly categorized into three main approaches. First, one category uses self-supervised objectives to learn representations. These methods impose auxiliary self-supervised tasks on visual data to encourage the model to learn general features beneficial for the task. Representative works include using contrastive learning (CURL [16]), consistency regularization (SODA [14]), domain-invariant representation learning [28], or foreground/saliency extraction (VAI [23], SGQN [29]). Second, another category leverages the intrinsic structure of the Markov Decision Process (MDP) and RL characteristics. This includes using bisimulation metrics [30] to group behaviorally similar states (DBC [17], Chen et al. [31]), achieving policy adaptation after deployment via inverse dynamics prediction (PAD [32]), stabilizing Q-learning through asymmetric data augmentation (SVEA [25]), bidirectionally predicting environment dynamics (BiT [15]), leveraging long-term information [33], or designing temporal contrastive losses [34]. Other strategies involve robustness-oriented learning and analysis [35]–[37] and designing intrinsic reward systems to achieve self-supervised domain adaptation [38]. While these methods effectively learn a single, robust representation by implicitly ignoring distractors, T2RD’s philosophy is different. We argue for explicitly decoupling an observation into task-relevant and task-irrelevant streams, then using dynamic prediction to ensure the task-relevant stream is behaviorally meaningful, which is a more direct path. Third, a powerful paradigm utilizes pre-trained Visual Foundation Models (VFMs). For instance, some works use SAM for object segmentation to remove background distractors [39], or directly employ general VFM features as input (PIE-G [40], Zhong et al. [41]). However, the features of the segmented objects themselves—their color, texture, or lighting—can still be task-irrelevant and cause overfitting. T2RD is positioned to address this problem by further decoupling an object’s task-relevant state from its superficial appearance.

2.3 Representation Decoupling↩︎

Human perceptual systems often distinguish between content and style to better understand observations [20], and in artificial intelligence, decomposed representations of content and style are also highly needed. In computer vision, for example, Lorenz et al. [42] separate appearance and shape by learning consistent parts across instances of the same category, while Wu et al. [43] introduce a two-branch autoencoder to disentangle objects into content and style representations. Ren et al. [44] propose a content-style disentanglement module by encoding dominant image factors. Yin et al. [24] model the content-style representation as a token-level bipartite graph.

This principle of learning informative representations also extends to visual recognition and spatio-temporal modeling [45], [46], as well as medical image analysis [47], and a related line of work leverages direct human perceptual cues to guide agents; for example, human gaze data can direct an agent’s attention for efficient imitation learning [48], and human-generated saliency has been used to guide value estimation in visual RL [49]. While these methods validate the benefit of isolating task-relevant information, they share a common limitation with most computer vision decoupling approaches: reliance on external supervision, whether paired datasets, labeled medical data, or direct human input. T2RD is distinct in that it learns this decoupling in a self-supervised manner tailored for VRL, using only the standard reward signal without external labels or human guidance.

3 Preliminaries↩︎

3.1 Visual Reinforcement Learning↩︎

Due to the partial observability of images in Visual Reinforcement Learning (VRL), the interaction between an agent and the environment is modeled as a Partially Observable Markov Decision Process (POMDP), represented by the tuple \(\mathcal{M} = \langle \mathcal{S}, \mathcal{O}, \mathcal{A}, \mathcal{P}, \mathcal{Z}, r, \mathcal{\gamma} \rangle\), where \(\mathcal{S}\) is the state space, \(\mathcal{O}\) is the high-dimensional observation space, \(\mathcal{A}\) is the action space, \(\mathcal{P}(\mathbf{s}_{t+1} \mid \mathbf{s}_{t}, a_{t})\) is the state transition function, \(\mathcal{Z}(\mathbf{o}_{t} \mid \mathbf{s}_{t})\) is the observation-emission distribution, \(r:\mathcal{S}\times \mathcal{A}\mapsto \mathbb{R}\) is the scalar reward function, and \(\gamma \in \left [ 0,1 \right )\) is the discount factor.

The objective is to learn the optimal policy \(\pi^*\) that maximizes the cumulative reward return: \[\pi^* = \arg \max_{\pi} \; \mathbb{E}_{\,\mathbf{o}_t \sim \mathcal{Z}(\cdot \mid \mathbf{s}_t),\; a_t \sim \pi(\cdot \mid \mathbf{o}_t),\; \mathbf{s}_{t+1} \sim \mathcal{P}(\cdot \mid \mathbf{s}_t, a_t)} \Big[ \sum_{t=1}^{T} \gamma^t \, r(\mathbf{s}_{t}, a_{t}) \Big]\] starting from the initial state \(\mathbf{s}_0 \in \mathcal{S}\) and taking actions \(a_t\) chosen by the observation-conditional policy \(\pi_{\theta}(\cdot \mid \mathbf{o}_t)\), while observations are generated via \(\mathcal{Z}(\mathbf{o}_t \mid \mathbf{s}_t)\) and the environment transitions to \(\mathbf{s}_{t+1}\) according to \(\mathcal{P}(\cdot \mid \mathbf{s}_t, a_t)\). The policy is parameterized by learnable parameters \(\theta\).

3.2 Generalization in VRL↩︎

In terms of generalization of visual reinforcement learning, we consider a set of POMDPs denoted by \(\mathbf{M} = \{ \mathcal{M}_1, \mathcal{M}_2, \ldots, \mathcal{M}_n \}\), where each POMDP \(\mathcal{M}_i\) has its own observation space \(O_i\) but shares a common underlying state space \(\mathcal{S}\) and state transition function \(\mathcal{P}\). During training, we only have access to a specific POMDP, denoted as \(\mathcal{M}_i\). Our goal is to train an agent on a specific scenario \(\mathcal{M}_i\) to learn a policy \(\pi^*_G\) that maximizes the expected cumulative reward across the entire set of POMDPs \(\mathbf{M}\), in a zero-shot generalization manner.

4 METHOD↩︎

4.1 Overview↩︎

As shown in Figure 2, we propose Task-Relevant Representation Decoupling (T2RD). Our architecture is inspired by successful self-supervised learning methods such as BYOL [13] and adopts a dual-branch design: an online branch and a target branch. The core idea is that the online branch is trained to learn from the representations produced by the target branch, which evolves slowly and thus provides a stable regression target. This setup mitigates training instability and helps prevent collapse to trivial solutions.

The online branch comprises three components, each serving a distinct purpose: an encoder \(f_{\theta}\) that extracts latent representations from pixels, a dual-head projection \(g_{\theta}\) that decouples the representation into task-relevant content and task-irrelevant style, and a predictor head \(h_{\theta}\) that learns to predict the target-branch representation from the online-branch representation. By requiring a non-trivial mapping between the online and target branches, this predictor helps avoid representation collapse. The target branch has a structurally identical encoder \(f_{\varphi}\) and dual-head projection \(g_{\varphi}\), but no predictor.

For policy learning, we adopt Soft Actor-Critic (SAC) [50]. As shown in Figure 3, the T2RD module serves as a representation learner: we first decouple visual inputs into task-relevant and task-irrelevant features, then feed only the task-relevant representations to SAC.

4.2 Self-Supervised Decoupling of Latent Representations↩︎

Our decoupling process begins with data preparation and augmentation. During training, we sample a batch of transitions \(\{o_t, a_t, r_t, o_{t+1}\} \sim \mathcal{B}\) from a replay buffer \(\mathcal{B}\). For each observation \(o_t\), we generate an augmented view \(o_t' = t(o_t)\) using our combined data augmentation, which samples a transformation \(t\) with equal probability from a set of augmentations \(\mathcal{T}\) (containing Random Overlay and Random Convolution). As depicted in Figure 2, the augmented observation \(o_t'\) is subsequently processed by the online branch, while the original observation \(o_t\) is fed into the target branch.

A crucial element of this dual-branch setup is preventing the model from collapsing. This is achieved by managing the target network’s updates. We apply a stop-gradient to the target branch’s output, preventing gradients from flowing into it. Instead of being trained via backpropagation, its parameters \(\varphi\) are updated as a slow-moving average of the online network’s parameters \(\theta\). This Exponential Moving Average (EMA) update ensures the target network provides a stable and consistent objective, which is essential for stabilizing the entire self-supervised training process: \[\varphi_{n+1} \leftarrow (1 - \rho) \varphi_n + \rho \theta_n,\] where \(\rho\) is the momentum coefficient for the EMA update, with its specific value listed in the hyperparameters of Table 4. Our self-supervised representation learning process is then driven by three key objectives, which we detail below.

4.2.1 Task-Relevant Representation Consistency.↩︎

The first objective ensures that the core “content” of an observation remains consistent even when its “style” is changed by data augmentation. The augmented observation \(o_{t}'\) and the original observation \(o_{t}\) share the same underlying state. Therefore, their task-relevant representations should be close. Thus, the content features \({\mu}_{{t}}\), \({\mu}_{{t}}'\) extracted from \({o}_{{t}}\) and \({o}_{{t}}'\) should be consistent. The objective function of task-relevant representation consistency can be expressed as: \[L_{\text{consis}} = \mathbb{E}_{\tau \sim D_{\text{env}}} \left[\| \mu_{t}' - \mu_{t} \|_2^2 \right], \label{Equ2}\tag{1}\] where \(D_{\text{env}}\) is the replay buffer. Equation 1 ensures the consistency of content features between the original and augmented images, thereby excluding style-related information from the obtained content representations \(\mu\). Such content representations possess domain-invariant properties, facilitating better generalization across different domains.

4.2.2 Cross-Reconstruction for Decoupling.↩︎

To decouple content and style, we introduce a reconstruction objective. The logic is as follows: while the consistency loss in Equation 1 drives the content representation \(\mu\) to become style-invariant, we must also ensure information sufficiency and separability. The reconstruction objective guarantees that the combination of \(\mu\) and \(\epsilon\) preserves the complete information of the original observation \(o\). Thus, as \(\mu\) is purified of style, the style information is naturally compressed into \(\epsilon\).

Building on this, we introduce a novel twist to the standard reconstruction approach: cross-reconstruction. This technique forces a cleaner and more robust separation of content and style. Specifically, we reconstruct the full latent representation \(z_t\) of the original observation by combining the content representation \(\mu_t'\) from the augmented view with the style representation \(\epsilon_t\) from the original view: \[\hat{z}_t = d_\xi(\mu_{t}', \epsilon_t), \label{Equ3}\tag{2}\] The cross-reconstruction objective then minimizes the difference between the reconstructed \(\hat{z}_t\) and the original \(z_t\): \[L_{\text{recon}} = \mathbb{E}_{{\tau} \sim D_{\text{env}}} \left[\| z_t - \hat{z}_t \|_2^2 \right]. \label{Equ4}\tag{3}\] The “cross” design is critical: if the content representation \(\mu_t'\) still contains any style information from the augmented view, it will conflict with the style \(\epsilon_t\) from the original view, leading to a poor reconstruction and a high loss. This incentivizes the model to learn a pure, decoupled content representation.

4.2.3 Dynamic Prediction for Task-Relevance.↩︎

The first two components successfully decouple observations into a style-invariant “content” stream and a “style” stream. However, the obtained content representation is not guaranteed to be task-relevant. For example, the agent’s shadow might be considered content, but it is irrelevant for controlling the agent. To solve this, we introduce a dynamic prediction module. The intuition is that a truly task-relevant representation must contain information that is predictive of future states under the agent’s actions. By enforcing this predictive constraint, we refine the content representation to be specifically task-relevant. This crucial step elevates our framework from a simple content-style separator to a true decoupler of task-relevant and task-irrelevant features.

We also implement a cross-domain prediction technique here. Specifically, we use the task-relevant representation (\(\mu_t'\)) from the augmented view and the agent’s action (\(a_t\)) to predict the task-relevant representation (\(\mu_{t+1}\)) of the original next state: \[\hat{\mu}_{t+1} = P_v(\mu_{t}', a_{t}). \label{Equ5}\tag{4}\] This cross-domain objective is a powerful regularizer for generalization. It forces the representation \(\mu\) to be not only predictive but also domain-invariant, as it must bridge the visual gap between the augmented and original observation domains to succeed. The cross-dynamic prediction objective is: \[L_{\text{dynamic}} = \mathbb{E}_{{\tau} \sim D_{\text{env}}} \left[\| \hat{\mu}_{t+1} - \mu_{t+1} \|_2^2 \right]. \label{Equ6}\tag{5}\]

4.2.4 Overall Objective.↩︎

We jointly optimize the three objectives. The overall objective function for our self-supervised module is: \[L_{T2RD} = \lambda_1 \cdot L_{\text{consis}} + \lambda_2 \cdot L_{\text{recon}} + \lambda_3 \cdot L_{\text{dynamic}}.\] Here, \(\lambda_1\), \(\lambda_2\), and \(\lambda_3\) are scalar coefficients (with \(\lambda_i>0\)) that balance the contribution of each loss term. The specific values used in our experiments are summarized in Table 4.

4.3 Policy Learner↩︎

T2RD, as a representation decoupling module, is independent of the policy learner and applicable to any visual reinforcement learning (RL) algorithm. We use the Soft Actor-Critic (SAC) algorithm for policy learning. SAC is an actor-critic RL approach based on the maximum entropy framework, concurrently learning the Q function and the policy network \(\pi_{{\theta}}\), to maximize the expected cumulative reward. Specifically, as shown in Figure  3, the observation \({o}_{{t}}\) is decoupled into the task-relevant representation \(\mu_t\) and task-irrelevant representation \(\epsilon_t\) through the encoder \({f}_{{\theta}}\) and the dual-head mapping \({g}_{{\theta}}\). Then, the policy learner uses only the task-relevant representation \(\mu_t\) for policy learning. As shown in Algorithm 4, the optimization of representation learning \(L_{T2RD}\) alternates with the optimization of reinforcement learning \(L_{RL}\) with \({\omega}\) steps. The two optimization objectives jointly update the parameter \(\theta\).

4.4 Discussions↩︎

1) What sets T2RD apart from other self-supervised auxiliary methods? We explicitly aim to decouple task-relevant and task-irrelevant features in the observation, where the encoding processes of task-relevant and task-irrelevant features complement each other. This deepens the understanding of visual observations, thereby obtaining more precise task-relevant representations. For instance, knowing what is wrong can lead to a clearer understanding of what is right. Furthermore, during policy learning, by explicitly discarding task-irrelevant features, we can reduce the burden of information processing on the agent, thereby improving sample efficiency.

2) Why does T2RD greatly benefit from our combined data augmentation? As shown in the subsequent experimental section, while other algorithms also benefit from our combined data augmentation, the improvements are not as substantial as those seen with T2RD. We argue that while the combined data augmentation increases augmentation diversity, it may also lead to some interference and conflicts among them. Thanks to T2RD’s robust decoupling capabilities, it can fully exploit this increased diversity, maximizing the advantages of the combined data augmentation.

4.5 Theoretical Analysis↩︎

To ground our framework theoretically, we connect our approach to bisimulation metrics [17], [30], which quantify behavioral similarity between states in an MDP. The metric \(d_{\mathcal{M}}(s_i, s_j)\) is defined as the fixed point of the following Bellman-like operator: \[\label{eq:bisim} d_{\mathcal{M}}(s_i, s_j) = \max_{a \in \mathcal{A}} \left( |r(s_i, a) - r(s_j, a)| + \gamma W_1(\mathcal{P}(\cdot|s_i, a), \mathcal{P}(\cdot|s_j, a); d_{\mathcal{M}}) \right).\tag{6}\] where \(W_1\) is the 1-Wasserstein distance measuring the difference between next-state transition distributions. Crucially, this metric provides an upper bound on the optimal value function difference:

Theorem 1 (Ferns et al. [30]). For any two states \(s_i, s_j \in \mathcal{S}\), the difference in their optimal value functions is bounded by their bisimulation distance: \[\label{eq:ferns95bound} |V^*(s_i) - V^*(s_j)| \le d_{\mathcal{M}}(s_i, s_j).\tag{7}\] This provides a clear analytical path: if T2RD learns a representation \(\mu\) where the latent distance \(||\mu(o_i) - \mu(o_j)||_2\) approximates the bisimulation distance \(d_{\mathcal{M}}\), we can then bound the error of a value function learned upon this representation. Our complete objective is designed to provide a strong inductive bias for learning such a metric from high-dimensional observations.

4.5.1 Connecting T2RD Losses to the Bisimulation Metric↩︎

We first map high-dimensional observations to states and then establish the connection with the bisimulation metric.

State Invariance and Sufficiency (\(L_{\text{consis}}\) & \(L_{\text{recon}}\)). The bisimulation metric is defined over underlying states \(s\), while our agent sees observations \(o\). Therefore, the representation must be invariant to task-irrelevant visual features while remaining sufficient information for control.

(1) The consistency loss (\(L_{\text{consis}}\)) enforces state invariance. Since an observation \(o_t\) and its augmentation \(o_t'\) correspond to the same underlying state \(s_t\), their task-relevant representations should be identical. Minimizing \(||\mu_{t}^{\prime}-\mu_{t}||_{2}^{2}\) achieves precisely this.

(2) The cross-reconstruction loss (\(L_{\text{recon}}\)) ensures information sufficiency and separation. First, reconstruction guarantees that the learned representations (\(\mu\), \(\epsilon\)) retain sufficient information from the high-dimensional observations \(o\). Second, by forcing the decoder to reconstruct the original latent state \(z_t\) from the augmented task-relevant content \(\mu_t'\) and the original task-irrelevant style \(\epsilon_t\), it ensures that the task-relevant content representation \(\mu\) is purged of style information. This mechanism acts as a practical information bottleneck, ensuring that the learned distance \(||\mu_i - \mu_j||_2\) reflects differences between underlying states (\(s_i, s_j\)), not superficial visual distractors.

Dynamics and Rewards (\(L_{\text{dynamic}}\) & \(L_{\text{RL}}\)). The bisimulation metric in Equation 6 has two components: reward difference and transition dynamics difference. Our framework addresses both.

(1) The cross-dynamic prediction loss (\(L_{\text{dynamic}}\)) encourages representations \(\mu\) to respect the MDP’s transition dynamics, acting as a practical, sample-based proxy for minimizing the Wasserstein distance term in the bisimulation metric.

(2) The reward term is implicitly modeled by the downstream policy learner (\(L_{\text{RL}}\)). As described in Algorithm 4, the encoder’s parameters \(\theta\) are updated by gradients from both \(L_{\text{T2RD}}\) and the reward-driven SAC loss \(L_{\text{RL}}\) [50]. The gradients from the SAC critic, which minimizes Bellman error, force the representation \(\mu\) to distinguish between states that lead to different rewards.

This joint optimization ensures our learned representation \(\mu\) is sensitive to both rewards and dynamics, thereby learning a functional approximation of the full bisimulation metric.

4.5.2 Bound on the Action-Value Function↩︎

Based on the argument that our learning objective guides the representation to satisfy the properties of a bisimulation metric, we can establish a bound on the action-value function’s error.

Theorem 2. Under the assumption that the learned representation \(\mu\) yields a distance metric \(d_{\mu}(s_i, s_j) = ||\mu(o_i) - \mu(o_j)||_2\) that approximates the true bisimulation metric up to a bounded error \(\epsilon_{\text{T2RD}}+ \epsilon_{\text{RL}}\) (i.e., \(|d_{\mu}(s_i,s_j) - d_{\mathcal{M}}(s_i, s_j)| \le \epsilon_{\text{T2RD}}+ \epsilon_{\text{RL}}\)), let \(\hat{Q}(\mu(o), a)\) be the action-value function learned on the representations produced by an encoder trained with the combined T2RD and SAC objectives, and let \(Q^*(s, a)\) be the true optimal action-value function for the underlying state \(s\). The suboptimality is bounded as follows: \[\label{eq:our95bound} \max_{s, a} |Q^*(s, a) - \hat{Q}(\mu(o), a)| \le \frac{\epsilon_{\text{T2RD}} + \epsilon_{\text{RL}}}{1-\gamma}\tag{8}\] where \(\epsilon_{\text{T2RD}}\) is the representation error, which our framework is designed to minimize via the \(L_{\text{T2RD}}\) objective, and \(\epsilon_{\text{RL}}\) is the Bellman error from the policy learner, controlled by the converged \(L_{\text{RL}}\) loss.

5 Experiments↩︎

In this section, we conduct extensive experiments on DeepMind Control Suite Generalization Benchmark (DMControl-GB)  [26] and Robotic Manipulation tasks  [27] to validate the generalization performance and sample efficiency of our method. The agent is trained in a fixed environment, and we evaluate the algorithm’s generalization ability to unseen environments with more complex backgrounds or color variations.

Baselines. We compare our approach with other SOTA algorithms: 1) DrQ  [9] applies data augmentation to both Q-values and target Q-values during TD updates. 2) SODA  [14] proposes soft data augmentation and employs a BYOL-like architecture to learn consistent representations. 3) SVEA  [25] only applies data augmentation to the Q values without augmenting target Q values. 4) TLDA  [11] introduces task-aware Lipschitz data augmentation. 5) PIE-G  [40] utilizes ImageNet  [51] pre-trained encoders to provide a fairly general representation for RL. 6) SGQN  [29] learns masks through saliency maps and consistency regularization to exclude irrelevant pixels. 7) BiT  [15] introduces a bidirectional transition model. 8) TRP  [52] introduces the truncated return prediction task, promoting cross-domain policy consistency. 9) Visarl [49] leverages human-generated saliency maps to guide visual RL, using human perceptual cues to focus learning on task-relevant regions.

Data Augmentations. Random Overlay  [14] is an augmentation that performs linear interpolation between the original observation and another image randomly sampled from Places dataset  [53], a dataset containing ten million scene photographs; Random Convolution  [54] is an augmentation method that integrates a convolutional layer with random parameters into the observation, preserving the overall shape while adding a blurring effect to the texture features. The effects of the two data augmentations are shown in Figure 5. We adopt a combination of the above two data augmentations, that is, each batch randomly selects Random Overlay and Random Convolution with equal probability.

5.1 DeepMind Control Suite↩︎

We evaluate the generalization performance and sample efficiency of T2RD across five tasks in DMControl-GB: “Walker Walk”, “Walker Stand”, “Ball in Cup Catch”, “Finger Spin”, and “Cartpole Swingup”. We conduct these assessments under three environmental settings: “Color Hard”, “Video Easy”, and “Video Hard”. In the “Color Hard” setting, the evaluation is carried out in environments where the colors of task-relevant objects (e.g., the agent’s body) and the background and ground are randomized. This setup directly tests the policy’s ability to generalize to changes in the appearance of the agent itself, not just background variations. In the “Video” setting, the evaluation is conducted in environments with various natural video backgrounds. Figure 6 provides examples of observations in both the training and three test environments. We train the agent for a total of 500,000 steps, conducting evaluations every 10,000 steps by calculating the average return over 30 episodes. We stack three frames to form an observation with dimensions of 9 \(\times\) 100 \(\times\) 100. We then apply the random crop to the observation, resulting in dimensions of 9 \(\times\) 84 \(\times\) 84.

T2RD exhibits the best generalization performance. Table 1 shows the generalization results of T2RD and other SOTA algorithms in DMControl-GB tasks. Our algorithm performs exceptionally well, achieving the best results in 9 out of 15 tasks. Particularly in the “Cartpole Swingup” task of the “Video Hard” setting, compared to the suboptimal algorithm, our method achieves a 33% improvement. Overall, T2RD demonstrates outstanding generalization performance across all tasks in various testing environments.

T2RD has the highest sample efficiency. In Figure 7, we present the training curves of our method compared to SGQN, BiT, and TRP across five different tasks, with each algorithm running 5 seeds per task. Our method exhibits superior sample efficiency and asymptotic performance. Moreover, as illustrated in Figure  8, our method maintains excellent sample efficiency and asymptotic performance when generalized to previously unseen environments.

T2RD successfully learns task-relevant representations. In Figure  9, we use t-SNE to visualize the features learned by our method and by SGQN, SVEA, and PIE-G. We first randomly select 10 observations originating from different states in the “Cartpole Swingup” task and then replace their backgrounds with 40 unseen images. The 400 processed images are then embedded by the learned encoder of each algorithm. Observations from the same state are represented by the same color. From Figure  9, we can observe that T2RD achieves the most effective aggregation for observations from the same state and effective separation between different states.

T2RD can effectively achieve decoupling. To validate that T2RD can effectively decouple task-relevant and task-irrelevant features, we conduct both quantitative and qualitative analyses.

First, we present t-SNE visualizations in Figure 10 to quantitatively analyze the learned representations over time. For this experiment, we take a single observation, create 80 variants by replacing its background with unseen images, and feed them into the T2RD encoder at different training checkpoints. As the figure illustrates, the task-relevant representations (orange dots) become progressively more concentrated into a cluster. This indicates that the model learns to map the same underlying state to a consistent representation, regardless of the visual background. Conversely, the task-irrelevant representations (blue dots) become more dispersed and gradually separable from the task-relevant features. This indicates that T2RD has learned the ability to decouple task-relevant and task-irrelevant features within visual observations.

Second, to provide more direct, qualitative proof, we conduct a cross-reconstruction experiment, visualized in Figure 11. This involves reconstructing an original observation \(o_t\) by combining the task-relevant representation \(\mu_t'\) from a visually distinct augmented view \(o_t'\) with the task-irrelevant representation \(\epsilon_t\) from the original view \(o_t\). The result is a powerful demonstration of decoupling: the reconstructed image successfully captures the agent’s posture (task-relevant) from the original observation, while completely ignoring the distracting background or color perturbation (task-irrelevant) introduced in the augmented view. Together, these visualizations provide strong evidence that T2RD effectively decouples task-relevant content from task-irrelevant style.

T2RD accurately identifies task-relevant regions. In Figure 12, we visualize the attention regions of T2RD under different backgrounds in DMControl-GB: “Training”, “Color Hard”, “Video Easy”, and “Video Hard”. T2RD can accurately perceive task-relevant regions in visual observations, regardless of the presence of background interference. This provides the agent with reliable task-relevant representations for learning generalizable policies.

5.2 Robotic Manipulation↩︎

In addition to DeepMind control tasks, we also conduct experiments on Robotic Manipulation tasks. We consider two goal-reaching tasks: 1) Reach, where the robotic arm attempts to reach a red mask on the tabletop; 2) Peg in box, where the objective is to control the robotic arm to insert a peg into a box. The dimensions of the observation are 3 \(\times\) 100 \(\times\) 100. We perform a random crop on the observation to obtain a cropped observation with dimensions of 3 \(\times\) 84 \(\times\) 84. The agent is trained in a fixed background and tested in environments where key information is retained while the texture or color of the background is disturbed. We train the agent for 500,000 steps, conducting tests every 10,000 steps by calculating the average episode return over 30 episodes. For each task, we employ four test environments: Test1, Test2, Test3, and Test4. Figure 13 provides examples of observations in both the training and four test environments. We compare our algorithm against SAC, SODA, SVEA, SGQN, BiT, and TRP.

The results are shown in Table ¿tbl:table:robot?, where T2RD achieves state-of-the-art performance in 7 out of 8 test environments. In the Reach task, T2RD shows the best performance in Test1, Test3, and Test4 environments, and even in Test2, we are the second-best algorithm. Compared to the second-best algorithm, our average performance improves by 40.1%. In the Peg in box task, T2RD outperforms all other algorithms in all test environments. Compared to the second-best algorithm, our average performance improves by 57.1%. The above experimental results demonstrate the superiority of our algorithm. Additionally, T2RD proves applicable to various VRL tasks.

5.3 Ablation Study↩︎

We conduct ablation experiments on the DMControl-GB to assess the impact of each objective function proposed by T2RD. The experimental results are shown in Figure 14, where “T2RD (w/o consistency loss)” marks the absence of the task-relevant representation consistency module; “T2RD (w/o reconstruction)” indicates the exclusion of the cross-reconstruction module; “T2RD (w/o dynamic prediction)” denotes the omission of the cross-dynamic prediction module; “Baseline” represents the fundamental version of the T2RD algorithm, which omits these three objective functions but retains the combined data augmentation.

From Figure  14, it’s clear that T2RD achieves superior sample efficiency and asymptotic performance. The comparison between “T2RD (w/o dynamic prediction)” and T2RD underscores the critical role of dynamic prediction in learning task-relevant representations within our decoupling framework. Nonetheless, even without dynamic prediction, methods that separate content and style at the visual level remain effective, continuing to surpass other SOTA algorithms. Additionally, the “T2RD (w/o reconstruction)” curve shows that the decoupling of content and style facilitates the learning of representations that are robust to visual changes. Finally, “T2RD (w/o consistency loss)” also achieves excellent experimental results, but there is still a certain gap compared to T2RD. This suggests that the task-relevant representation consistency module contributes to learning domain-invariant representations, which promotes generalization.

Additionally, we conduct ablation experiments on three modules in Robotic Manipulation tasks, with the results shown in Table ¿tbl:table:Robotic32ablation?. The experimental results show that removing any one of the three modules leads to a decrease in performance. In summary, these three modules are complementary and work together to achieve the decoupling of task-relevant and task-irrelevant features, thereby promoting the generalization of VRL.

5.4 Variant Experiments↩︎

To provide deeper insights into our key design choices beyond the module-level ablations in Section 5.3, we conduct a series of more fine-grained variant experiments. The results are shown in Table ¿tbl:table:variation? and Table [tbl:tab32varient32augmentation32results]. We analyze structural variants, key mechanisms, and data augmentation methods.

5.4.1 Analysis of Structural Variants↩︎

We first explore several structural variants of T2RD. “T2RD-Alt. Branch” employs an alternative reconstruction branch using \(\mu_{t}\) and \(\epsilon_{t}'\) to reconstruct \(z_{t}'\). In Table ¿tbl:table:variation?, “T2RD-Dual” indicates the simultaneous use of both the primary and alternative reconstruction branches. “T2RD-Pixel Recon.” denotes pixel-based reconstruction, meaning that instead of reconstructing back to the latent representation \(z_{t}\), we reconstruct back to the observation \(o_{t}\).

The experimental results in Table ¿tbl:table:variation? reveal that both the “T2RD-Alt. Branch” and “T2RD-Dual” are less effective than our primary T2RD design. We argue that the task-irrelevant representations of the augmented observations (\(\epsilon_t'\)) contain too many distracting elements, leading to algorithm divergence. Furthermore, we find that reconstruction based on latent representations is more effective than pixel-based reconstruction (“T2RD-Pixel Recon.” in Table ¿tbl:table:variation?), a finding also supported by previous research [55].

5.4.2 Analysis of Key Design Mechanisms↩︎

To provide a more fine-grained analysis, we evaluate the specific impact of our cross-domain strategy and the EMA update mechanism.

Cross-Domain vs. Same-Domain. To isolate the effect of our cross-domain strategy, we test two variants. “T2RD-NonCross-Recon.” performs same-domain reconstruction (using \(\mu_{t}\) and \(\epsilon_{t}\) to reconstruct \(z_{t}\)). “T2RD-NonCross-Pred.” performs same-domain prediction (using the augmented representation \(\mu_t'\) to predict the next augmented representation \(\mu_{t+1}'\)). As shown in Table ¿tbl:table:variation?, T2RD with its cross-domain objectives significantly outperforms both of these variants. This demonstrates that enforcing consistency across different visual domains (original vs. augmented) is a powerful constraint that promotes the learning of more robust and generalizable representations.

Importance of EMA. We introduce “T2RD (w/o EMA)”, where the target network parameters are updated via a direct copy from the online network instead of the EMA update. The results in Table ¿tbl:table:variation? show a noticeable drop in performance compared to the standard T2RD. This empirically confirms that the EMA is a crucial component for stabilizing the self-supervised training process by providing a consistent and slowly evolving target, which helps prevent representational collapse.

5.4.3 Analysis of Data Augmentation Methods↩︎

In addition to structural variant experiments, we conduct variant experiments on data augmentation, comparing the experimental results of the T2RD algorithm using Random Overlay, Random Convolution, and combined data augmentation. The experimental results are shown in Table [tbl:tab32varient32augmentation32results], where T2RD (conv) indicates the use of Random Convolution as a replacement for the combined data augmentation technique in the T2RD algorithm. Similarly, T2RD (overlay) refers to the use of Random Overlay data augmentation. The experimental results demonstrate that combined data augmentation significantly improves the generalization performance of the algorithm compared to single data augmentation methods. As discussed in the methods section, T2RD greatly benefits from combined data augmentation due to its strong decoupling ability. T2RD can fully exploit the advantages brought by increased augmentation diversity, alleviating the issue of generalization bias. Additionally, thanks to its strong decoupling capabilities, T2RD is less affected by conflicts between different augmentations, thus maximizing the benefits of combined data augmentation.

5.5 Analysis of Combined Data Augmentation and Generalization Bias↩︎

Many VRL algorithms exhibit generalization bias by relying on a single data augmentation, where performance becomes overly dependent on the specific transformations seen during training. For instance, an agent trained only on “color” augmentation (like Random Convolution) is good at handling color changes but poor at handling background changes, and vice versa.

To empirically demonstrate this, we conduct an experiment where we train SODA [14] and SVEA [25] with each of these single augmentations separately. The results, presented in Table 2, clearly confirm the existence of this bias. In the “Color-Hard” setting, agents trained with Random Convolution generally outperform those trained with Random Overlay. Conversely, in the “Video-Easy” setting, agents trained with Random Overlay are significantly superior. This confirms that relying on a single augmentation leads to policies that are brittle to other forms of visual change.

Our proposed combined data augmentation method effectively mitigates this bias. In Table ¿tbl:tab:combined95data95augmentation?, we apply the combined augmentation strategy to several baseline algorithms. As shown, the performance of baselines such as SODA [14] and SVEA [25] not only improves significantly but also becomes more balanced across the three distinct test environments (“Color Hard”, “Video Easy”, and “Video Hard”). This provides strong evidence that our combined strategy alleviates generalization bias by breaking the agent’s dependence on a single type of visual perturbation.

However, this analysis also highlights the unique synergy between our architecture and the augmentation strategy. While the baselines benefit from the combined augmentation, T2RD’s performance (as shown in Table ¿tbl:tab:combined95data95augmentation?) remains markedly superior. We argue that although increasing augmentation diversity is beneficial, it can also introduce interference or conflicts among different transformations. Thanks to T2RD’s robust decoupling capability, it can fully exploit this increased diversity—maximizing the advantages of combined data augmentation where other methods may struggle. For example, the slight performance drop observed for SGQN [29] on one task suggests that, without a strong decoupling mechanism, increased augmentation diversity may sometimes interfere with methods that rely heavily on fine-grained visual cues such as saliency maps. This demonstrates that the observed gains arise not from augmentation alone, but from the powerful combination of our novel architecture and the augmentation strategy.

Finally, we explore the impact of the sampling probability \(p\) for Random Overlay within our combined approach. The specific experimental results are shown in Table 3. The results indicate that when \(p = 0.5\), i.e., when Random Convolution and Random Overlay are chosen with equal probability, we observe the best experimental outcomes, reinforcing the idea that balanced diversity is key to effective data augmentation.

6 Reproducibility↩︎

Implementation Details. To ensure a fair comparison, we use the same hyperparameter settings as most methods, such as SODA and SVEA, in the DeepMind Control Suite. In Robotic Manipulation tasks, our hyperparameters follow those used in SGQN. The specific hyperparameter settings are detailed in Table 4.

Network Architecture Details. We adopt a BYOL-inspired architecture for representation learning. The observation encoder \(f_{\theta}\) is an 11-layer CNN that maps stacked observations (9 \(\times\) 84 \(\times\) 84) to 14112-dimensional latent vectors, which are projected to 200 dimensions and then split into task-relevant and task-irrelevant features (100 each) by a dual-head MLP \(g_{\theta}\). A predictor network \(h_{\theta}\), consisting of a linear layer with batch normalization and ReLU, encourages consistency between augmented and original task-relevant representations. To achieve decoupling, we use two auxiliary modules: (i) a cross-reconstruction network with two linear layers and one ReLU, which combines \(\mu_t'\) and \(\epsilon_t\) to reconstruct \(z_t\); and (ii) a cross-dynamic prediction network with three linear layers and two ReLUs, which predicts the next-step task-relevant features \(\mu_{t+1}\) from \(\mu_t'\) and action \(a_t\). For policy learning, we adopt SAC, where the actor and critic are standard three-layer MLPs with ReLU activations. The policy learner shares the encoder and projection with T2RD, ensuring that only task-relevant features are used for control. The SAC objective \(L_{\mathrm{RL}}\) and the self-supervised objective \(L_{\mathrm{T2RD}}\) are jointly optimized using Adam, with alternating updates every \(\omega\) steps.

7 Conclusion↩︎

We introduce the concept of decoupling observations into task-relevant and task-irrelevant features in Visual Reinforcement Learning (VRL), making the learned policy more robust to task-irrelevant visual changes. Based on this idea, we propose Task-Relevant Representation Decoupling (T2RD), a self-supervised algorithm for VRL. T2RD consists of three components: task-relevant representation consistency, cross-reconstruction, and cross-dynamic prediction. The first two components achieve decoupling of content and style features at the visual level, while the third part learns task-relevant representations through dynamic prediction. In addition, we propose a simple yet effective combined data augmentation approach. We conduct extensive experiments to verify the generalization ability and sample efficiency of T2RD across the DeepMind Control Suite and Robotic Manipulation tasks.

8 Limitations and Future Work↩︎

Despite the strong generalization performance of T2RD, we acknowledge several limitations that also point to future directions.

First, T2RD assumes visual observations can be decoupled into task-relevant “content” and task-irrelevant “style” streams. While this assumption holds in many benchmarks, it may fail when task-relevant cues are intertwined with style, such as navigation relying on specific terrain textures. In such cases, the model might treat crucial information as irrelevant. Future work could explore causal inference or adaptive mechanisms to dynamically identify task-relevant features.

Second, although T2RD is trained from scratch, its principle of decoupling remains highly relevant in the era of Visual Foundation Models (VFMs). Object-centric features extracted by VFMs still mix relevant states (e.g., pose) with irrelevant appearance (e.g., texture, lighting). Applying T2RD as a lightweight, self-supervised fine-tuning module to these features is a promising direction.

Finally, extending T2RD to real-world robotics is essential. The simulation-to-reality gap in lighting, textures, and object appearance often causes policy failures. Our decoupling approach is well-suited to mitigate this challenge, and validating T2RD on physical robots is a key future step.

References↩︎