How to Learn Domain-Invariant Representations for Visual Reinforcement Learning: An Information-Theoretical Perspective

1 Introduction

Visual Reinforcement Learning (VRL) has achieved widespread success in various domains such as robotics control [Huang et al., 2023], autonomous navigating [Wang et al., 2023; Hu et al., 2023], and video gaming [Zhou et al., 2023]. However, learning a well-generalized policy across different domains in VRL remains challenging due to overfitting to irrelevant and redundant information.

Previous VRL methods can roughly be categorized into two classes: 1) introducing image processing techniques to capture the foreground or learn a consistent representation [Wang et al., 2021; Hansen and Wang, 2021; Bertoin et al., 2022]; and 2) exploiting the structural properties of RL to learn a robust policy, such as investigating data augmentation in TD update [Yarats et al., 2021; Hansen et al., 2021b]. Although prior studies have demonstrated empirical improvements, theoretical motivation and support remain lacking.

∗Corresponding author: Kai Lv.

Different from previous works, we aim to introduce theoretical tools for VRL generalization. Information theory has garnered significant attention in deep learning, such as Information Bottleneck (IB) [Tishby and Zaslavsky, 2015]. IB believes that optimizing the Mutual Information (MI) between the representation and the input/output can lead to invariant representations with good generality. Then, [Xu and Raginsky, 2017] establishes a theoretical link between MI and the generalization capacities of deep learning. Nevertheless, utilizing information theory to explore generalization for VRL remains unexplored.

In this work, we introduce information theory to provide a novel perspective for VRL and identify three Mutual Information (MI) terms. Then, as depicted in Figure 1, we present three theoretical supports derived from the three MI terms for improving generalization. These theoretical supports are pivotal for acquiring domain-invariant representations, the sketch of the robotic arm, which is applicable across diverse domains. Specifically, these theoretical supports emphasize the necessity of preserving return and dynamic information under significant observation perturbation.

Based on the theoretical support, we propose a practical algorithm, Mutual Information-based Invariant Representation (MIIR). MIIR relaxes the MI terms and obtains three components: Q Predicting (QP), Dynamic Predicting (DP), and Deviation Perturbation Regularization (DPR). Specifically, QP maps various observations to the ground-truth Q values. DP is a dynamic transition model within the representation space independently of observations. DPR employs the pixel deviation within a batch to quantify the relevance of the action and perturbs the action-irrelevant pixels as much as possible.

We conduct extensive experiments across three benchmarks: DeepMind Control suite, Robotic Manipulation, and Carla (an autonomous driving environment). The agent is trained within a fixed background and tested in unseen environments with varying appearances. The results of numerous experiments demonstrate that MIIR has attained state-of-theart generalization performance and sample efficiency.

Our primary contributions are summarized as follows:

• We introduce information theory into VRL to understand the generalization. Specifically, we identify three MI terms for learning domain-invariant representations.
• In practice, we relax the MI terms and present a Mutual Information-based Invariant Representation (MIIR) algorithm comprising Q Predicting, Dynamic Predicting, and Deviation Perturbation Regularization.
• We conduct experiments across three benchmarks, demonstrating that MIIR achieves state-of-the-art generalization performance and sample efficiency.

2 RelatedWork

2.1 Data Augmentation in Visual RL

Data augmentation is widely adopted in visual reinforcement learning to improve sample efficiency and generalization. RAD [Laskin et al., 2020b] is the first study to investigate the influence of data augmentation in visual reinforcement learning. [Lee et al., 2020] introduce RandomConvolutions to alleviate overfitting to textures. SODA [Hansen and Wang, 2021] introduces a novel data augmentation, which utilizes additional images linearly added with observations. TLDA [Yuan et al., 2022a] suggests not to augment the taskrelevant pixels, which are measured by Lipschitz constants. SRM [Huang et al., 2022] proposes Frequency-based data augmentation as an effective method to improve generalization. CG2A [Liu et al., 2023] highlights the existence of gradient conflicts among various data augmentations and investigates how to integrate augmentation combinations better in visual RL. Unlike previous methods, we endeavor to perturb action-irrelevant pixels.

2.2 Generalization in Visual RL

VRL has become a research hot topic [Hafner et al., 2020; Zhang et al., 2021; Chebotar et al., 2023; Han et al., 2023] and numerous studies have investigated how to improve the generalization for VRL. Except for emphasizing data augmentation mentioned above, existing methods can be divided into two main paradigms: 1) methods based on visual processing and 2) methods focused on RL characteristics.

For visual processing, VAI [Wang et al., 2021] focuses on extracting foreground through unsupervised keypoint detection and visual attention, which is computational. CURL [Laskin et al., 2020a] introduces self-supervised learning to acquire consistent representations. SODA [Hansen and Wang, 2021] further incorporates a more advanced contrastive learning method, BYOL [Grill et al., 2020]. PIE-G [Yuan et al., 2022b] discuss that representations pre-trained on ImageNet can also contribute to the generalization. SGQN [Bertoin et al., 2022] utilizes saliency maps to focus on critical information for decision-making.

For RL characteristics, DrQ [Yarats et al., 2021] proposes augmenting both the target-Q network and Q network while optimizing TD loss. SECANT [Fan et al., 2021] introduces a distillation method, separating the processes of representation and policy learning. PAD [Hansen et al., 2021a] fine-tunes by leveraging the in-variance of inverse dynamics in the test environment. SVEA [Hansen et al., 2021b] only augments the Q network for updating stably. TRP [Wang et al., 2024] identifies policy divergence and Bellman error as critical factors impacting generalization. In this work, we focus on investigating VRL generalization from a theoretical perspective.

2.3 Information Theory in Generalization

Information theory is frequently employed to comprehend the representation process in deep learning [Tishby et al., 2000]. Consequently, Information Bottleneck [Tishby and Zaslavsky, 2015] has received widespread attention. [Xu and Raginsky, 2017] derives an upper bound for the generalization error and the input-output mutual information. [Shwartz-Ziv and Tishby, 2017] visualizes the deep learning training process and demonstrates that IB has the ability to compress input for deep learning to obtain effective representation. This theoretical conclusion is widely utilized in image processing [Li et al., 2022; Ahuja et al., 2021], time series forecasting [Jia et al., 2023], reinforcement learning [Eysenbach et al., 2019], and other fields. In Reinforcement Learning, [Laskin et al., 2022; Eysenbach et al., 2019] utilize MI to quantify the difference between trajectories and acquire diverse skills. [Eysenbach et al., 2021; Igl et al., 2019] demonstrate that IB theory can assist agents in generalizing across different environments whose dynamic transitions differ from the training environment. Inspired by previous works, we introduce information theory in VRL for appearance generalization.

3 Preliminaries

3.1 Visual Reinforcement Learning

Reinforcement Learning aims to learn a decision-making policy by interacting with the environment [Sutton and Barto, 2018]. Visual reinforcement learning can be formulated as a Partially Observable Markov Decision Process (POMDP) M=< O, S,A,P,R, γ >, where O is the observation space, S is the state space, A is the action space, P : S ×A ↦→S is the dynamic transition function, R : S ×A ↦→R is the reward function, and γ ∈ [0, 1) is the discount factor. The objective is to learn a policy πθ that maximize the expected cumulative return U πθ =

Eot,at∼τ [∑T-1 t=0 γ tR(ot, at)], where τ represents the trajectories sampled by the policy πθ.

In this work, we follow the previous works and utilize an off-policy paradigm that obtains optimal policy π∗ by estimating state-action value function Qϕ(ot, at) : S × A ↦→ R. Specifically, we minimize Temporal Difference (TD) error [Sutton, 1988] to update Qϕ(ot, at), formulated as:

where yt = R(ot, at)+γEπθ [Qtgt (ot+1, at+1)] and B is the replay buffer.

3.2 Generalization in Visual RL

In this work, we focus on generalization in Visual RL which is trained on a fixed environment and tested on unseen environments. Specifically, we consider a set of MDP M. Each Mi ∼ M shares the same underlying structure except its unique observation space. Our purpose is to find a policy π that can maximize the expected cumulative return on M ∼ M in a zero-shot manner. Therefore, the objective of visual RL generalization can be formulated as:

For convenience, we denote the training environment as Mtrain and the test environments asMtest. SinceMtest can not obtained during training, we often utilize empirical risk at Mtrain to learn the policy. Therefore, the generalization error can be formulated as:

3.3 Information Bottleneck

Information Bottleneck (IB) [Tishby and Zaslavsky, 2015] has been widely adopted in deep supervised learning. The key opinion of IB is that an invariant representation Z needs to compress the input information X as much as possible while mapping the label Y , which can be formulated as:

After that, [Xu and Raginsky, 2017] theoretically establishes the relationship between generalization error and mutual information:

where σ is the output variance and n is the sample size. This upper bound suggests that controlling the mutual information between the input and output of an algorithm can control its generalization error.

4 Methodology

In this section, we first introduce information theory in visual RL. We identify three Mutual Information (MI) terms that contribute to the generalization. In practice, we relax the MI terms and propose a Mutual Information-based Invariant Representation (MIIR) with three components.

4.1 Information Theory in Visual RL

Equation (4) and Equation (5) show that well-generalized representations require controlling the input-output mutual information. Therefore, we can introduce the mutual information into the visual RL for learning domain-invariant representations with a certain policy πθ:

where z is the representation.

The first term implies that the representation and cumulative return sequences should exhibit a robust correlation when given an action sequence. To derive this term, we utilize variational inference techniques [Alemi et al., 2017]:

Therefore, we obtain a lower bound of term1 by adding a non-negatived KL divergence. Note that the marginal distribution lnp(U πθ 1:T |a1:T ) is constant since it is independent of representation.

The second term indicates that the representation sequence should exhibit maximal independence from the observation sequence when given the action sequence. We also utilize variational inference to derive an upper bound for term2:

Note that Equation (8) introduces a variational approximation q(zt|zt-1, at-1) to construct a solvable variational upper bound. This upper bound suggests that the representation should be able to characterize the dynamic transition function independently of observations.

The third term implies the necessity to decrease the dependence of actions on observations. We can express the term3

using the following formula:

H(ot|at) reveals that learning domain-invariant representations requires increasing the entropy of observations without affecting the action.

In summary, we identify three theoretical supports from MI terms about how to learn domain-invariant representations. These objectives entail 1) maximizing the correlation between the representation and return, 2) preserving dynamic transition without depending on observations, and 3) increasing the entropy of the action-irrelevant pixels.

4.2 Mutual Information-based Invariant Representation

In this section, we propose Mutual Information-based Invariant Representation (MIIR) to learn domain-invariant representations. MIIR comprises three components: Q Predicting (QP), Dynamic Predicting (DP), and Deviation Perturbation Regularization (DPR). These components are derived by relaxing the above MI terms under certain mild assumptions. Figure 2 depicts the framework, and Algorithm 1 outlines the steps of our method.

Q Predicting Based on term1

Equation (7) shows that the representation zt should retain sufficient information to predict the cumulative return U πθ based on specific action sequences a1:T . Since the Q-function represents the expected cumulative return for taking a particular action in a given state, the Qfunction and the expected cumulative return satisfy: U πθ t = Eat∼π[Qπθ ϕ (st, at)]. Therefore, Equation (7) can be refor-

Algorithm 1 Mutual Information-based Invariant Representation

ω, ϕ, ξ, θ: randomly initialized network parameters

α, β: learning rates

B: empty replay buffer

• 1: for t = 1 · · · T do
• 2: Sample action at ∼ πθ(ot)
• 3: Interact with the environment ot+1 ∼ P(ot, at)
• 4: Update transitions B ←B ∪ (ot, at,Rt, ot+1)
• 5: Sample a batch transitions (ot, at,Rt, ot+1)
• 6: Perturb ot as oaug t with DPR
• 7: if t mod update frequency = 0: then
• 8: Update policy network θ ← θ -α∇θU π (θ)
• 9: Update Q Predicting ϕ←ϕ-α∇ϕLq(ϕ)
• 10: Update Dynamic Predicting ξ ← ξ -β∇ξLdy(ξ)
• 11: end if
• 12: end for

mulated as:

Furthermore, we assume that the output of q(·) follows a Gaussian distribution N(µq, σ2 q ), where µq is the predicted mean of Q value and σ2 q is a pre-given variance. Therefore, the maximum likelihood objective in Equation (10) can be reformulated as a Mean Squared Error (MSE) loss, which is:

where representations z is embedded by a encoder fω(·) from observations o or augmented observations aug(o).

Dynamic Predicting Based on term2

Equation (8) implies that a well-generalized representation can elucidate the dynamic transition function independently of observations. Specifically, predicting representation zt only depended on last action at-1 and representation zt-1. In practice, we employ a multivariate Gaussian distribution N(µdy, σ2dy ) to characterize the output of the learned dynamic transition function. µdy represents the predicted mean of zt and σ2dy is a pre-given variance. Additionally, we assume that the dimensions of the multivariate Gaussian distribution are independent. Then, we can transform the KL divergence into MSE objective as:

where zt = fξ(zt-1, at-1) is the learned dynamic transition function. Note that zt-1 is embedded from observations ot-1 or augmented observations aug(ot-1). We keep ot unaugmented for stabilizing the learning process of fξ(zt-1, at-1).

Deviation Perturbation Regularization Based on term3

As shown in Equation (9), increasing the entropy of actionirrelevant pixels in observations is essential to acquire domain-invariant representations. We employ two mild assumptions [Misra et al., 2024]: 1) action-relevant pixels exhibit a different frequency of change compared to others, and

2) the domain style remains i.i.d. throughout training. Therefore, pixels with high-frequency changing across transitions are salient and more likely to be action-relevant.

To measure the frequency of pixels, we introduce a batch kernel k ∈ RN×N×B inspired by the Mean Filter. N denotes the kernel size, and B represents the batch size of observations. First, we utilize the batch kernel to traverse all pixels and channels (with padding set to 1) to obtain a smooth feature map ¯o ∈ RH×W×C . Notably, the smooth feature map has the same dimensions as the observation.

Then, we subtract the smooth feature map ¯o from the observations o to compute the deviation of each pixel. Specifically, a larger deviation indicates the pixel contains more high-frequency information which is more significant. To select more important pixels within each observation, we set a ρ threshold and obtain a binarized feature map Mρ. Mρ(i, j) = 1 if the pixel o(i, j) belongs to the ρ-quantile of highest deviation for o, and 0 otherwise.

Finally, we perturb the pixels with lower deviation values by replacing observations to images ˜o sampled from Places [Zhou et al., 2017]. Note that Places is a widely used dataset in Random Overlay (a data augmentation technique in visual RL). Therefore, DPR can be expressed as:

where ⊙ denotes the Hadamard product.

5 Experiments

In this section, we perform extensive experiments on DeepMind Control suite (DMControl) tasks [Tassa et al., 2018], Robotic Manipulation tasks [Jangir et al., 2022], and Carla [Dosovitskiy et al., 2017] to evaluate generalization. We train the policy on a fixed environment and evaluate its generalization in environments with more complex backgrounds. The code has been released at: https://github.com/Rebel-Uranus/MIIR.

Baselines. We compare our method with current state-ofart approaches for generalization in visual reinforcement learning: 1) DrQ [Yarats et al., 2021], a method adopts data augmentation for both Q-values and target Q-values when updating the TD loss; 2) SODA [Hansen and Wang, 2021], introducing BYOL [Grill et al., 2020] and Random Overlay (a new data augmentation technique) to learn consistent representations; 3) TLDA [Yuan et al., 2022a], an approach to avoid touching the action-relevant pixels when using data augmentation; 4) SVEA [Hansen et al., 2021b], stabilizing Q learning by leaving target Q-values unaugmented during TD update; 5) PIE-G [Yuan et al., 2022b], an algorithm that utilizes pre-trained ResNet [He et al., 2016] representation into visual RL. 6) SGQN [Bertoin et al., 2022], focusing on crucial pixels by introducing the saliency map; 7) CG2A [Liu et al., 2023], alleviating gradient conflict when utilizes various data augmentations together.

5.1 DeepMind Control suite

We evaluate our algorithm on DMControl-GB [Hansen and Wang, 2021] across 6 tasks, i.e., Walker Walk, Walker Stand, Ball in cupCatch, Finger Spin, Cartpole Swingup, and Cheetah Run. DMControl-GB provides video-easy and video-hard settings, where the backgrounds are replaced by various natural videos.

Generalization Performance. For each task, we run 5 seeds to compute the mean and standard deviation of return for evaluating generalization. As illustrated in Table 1, MIIR outperforms the baselines in 11 out of 12 cases regarding the return. Remarkably, our method attains the highest performance in the video-hard setting (the hardest test setting with complex video perturbation). In particular, MIIR obtains a 4% improvement compared with the second-best method in video-easy setting. Moreover, in video-hard setting, our algorithm improved the average return across all tasks by 22%. In conclusion, MIIR achieves remarkable generalization performance across all tasks in diverse test environments.

Sample Efficiency. To evaluate sample efficiency, we compare MIIR with SVEA and SGQN in the training domain across six tasks (each algorithm runs 5 seeds on each task).

In Figure 3, the solid lines depict the average return, and the shaded regions represent one standard deviation. MIIR achieves the best sample efficiency and asymptotic performance across all tasks. Despite MIIR employing DPR to introduce noise into the observations, sample efficiency has not been adversely impacted. In particular, MIIR significantly improves in Cheetah Run.

Representation Visualization. To validate that MIIR can learn domain-invariant representations, we utilize t-SNE [Van der Maaten and Hinton, 2008] to visualize the learned features. First, we randomly select 10 different observations originating from different states and replace the background with 40 unseen images. Then, the 400 pre-processed images are embedded by the learned encoder of each algorithm. Observations from the same state are represented using the same color. Figure 4 demonstrates that algorithms with better invariant representations can achieve better generalization. Specifically, MIIR achieves the most effective aggregation for observations from the same state and optimal separation between different states.

5.2 Robotic Manipulation

To emphasize the generality of our method, we conduct experiments on Robotic Manipulation tasks introduced by [Jangir et al., 2022]. The benchmark involves two goal-reaching tasks: 1) Reach, controlling a robotic arm to reach a red mask on the table; 2) PegBox, attempting to insert the peg into a box. Each task can be evaluated in 5 test environments. Only the texture or color of the background varies during testing. Since CG2A has not published the code, we utilize the results reported in their work. We compare our method with SAC, SODA, SVEA, SGQN, and CG2A.

As illustrated in Table 2, MIIR achieves the SOTA in 8 out of 10 test environments. In the Reach task, MIIR achieves the best performance in Test1, Test4, and Test5 environments. Specifically, MIIR demonstrates improvements of 6× and 18× in Test4 and Test5, respectively. In the PegBox task, MIIR outperforms other methods in all test environments. In particular, MIIR obtains 3.9× and 3.4× improvements in Test4 and Test5, respectively. Notably, our method achieves the highest return value with the smallest performance gap across various test environments on all tasks.

5.3 Autonomous Driving in Carla

We employ a more complex and realistic autonomous driving environment (Carla) to demonstrate the effectiveness of MIIR. The agent is trained on a "train" weather condition and tested in 8 unseen weather conditions. Each weather condition is characterized by a distinct combination of rain, clouds, wind, and sunshine. The goal is to steer the vehicle as far as possible within an episode. We run 5 seeds for each algorithm to calculate the average return and the standard deviation. Since the appearance change is mainly caused by different weathers, we adopt Random Convolution for SVEA and SGQN. In addition, the perturbation images ˜o of MIIR are the observations o augmented by Random Convolution.

As shown in Table 4, MIIR achieves SOTA in 7 out of 8 test environments. Notably, MIIR exhibits stable performance across all test environments. On the contrary, SVEA and SGQN show more fluctuations in performance. SGQN exhibits weaker performance than SVEA and fails to maintain the advantages observed in the other benchmarks. The primary reason may be that SGQN needs to capture key pixels, which is challenging due to the complexity of Carla.

In conclusion, MIIR demonstrates the advancement across various tasks. Furthermore, MIIR exhibits high stability when generalized across different visual appearances.

5.4 Ablation Study

To illustrate the effectiveness of each component, we conduct experiments on DMC-GB, comparing different combinations of components. All combinations are trained for 500, 000 steps, evaluated across 6 tasks, and tested in "video easy" and "video hard" settings. Table 3 presents the results with one stand derivation.

We remove each component from MIIR to investigate its role, respectively. "QP & DPR" and "DP & DPR" is competitive with MIIR in both settings. However, "QP & DPR" and "DP & DPR" show obvious differences with MIIR on Cheetah Run. MIIR consistently outperforms, indicating that the return and dynamic transition information are both important. Notably, "QP & DPR" and "DP & DPR" demonstrate different advantages across different tasks.

Furthermore, we replace DPR with Random Overlay, denoted as "QP & DP", to investigate the importance of perturbing observations as much as possible. In "video easy", "QP & DP" performs comparably with MIIR on most tasks. On the contrary, "QP & DP" drops dramatically in "video hard". The disparity under different settings highlights the importance of increasing the entropy of observations. In other words, perturbing action-irrelevant pixels is crucial for generalizing to environments with significant interference.

Although the roles of components may overlap, experiments have shown that each component contributes differently. In brief, the three components work synergistically to promote the generalization of VRL.

6 Conclusion

We introduce information theory into Visual Reinforcement Learning (VRL) to explore how to learn domain-invariant representations. Specifically, we identify three MI terms and proposed a practical algorithm named Mutual Informationbased Invariant Representation (MIIR). MIIR contains three components obtained by relaxing each corresponding MI term. We conduct extensive experiments to validate the generalization capabilities across the DeepMind Control suite, Robotic Manipulation, and autonomous driving Carla. Meanwhile, MIIR achieves the best sample efficiency.

Acknowledgments

This work was supported by the National Natural Science Foundation of China (Grant No. 62206013).