{'fraction_non_alphanumeric': 0.02378469436868282, 'fraction_numerical': 0.00894432857372052, 'mean_word_length': 5.211509715994021, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Generalizing Reinforcement Learning through Fusing Self-Supervised Learning into Intrinsic Motivation

Despite the great potential of reinforcement learning (RL) in solving complex decision-making problems, generalization remains one of its key challenges, leading to difﬁculty in deploying learned RL policies to new environments. In this paper, we propose to improve the generalization of RL algorithms through fusing Self-supervised learning into Intrinsic Motivation (SIM). Speciﬁcally, SIM boosts represen- tation learning through driving the cross-correlation matrix between the embeddings of augmented and non-augmented samples close to the identity matrix. This aims to increase the similarity between the embedding vectors of a sample and its augmented version while minimizing the redundancy between the components of these vectors. Meanwhile, the re- dundancy reduction based self-supervised loss is converted to an intrinsic reward to further improve generalization in RL via an auxiliary objective. As a general paradigm, SIM can be implemented on top of any RL algorithm. Extensive evalua- tions have been performed on a diversity of tasks. Experimental results demonstrate that SIM consistently outperforms the state-of-the-art methods and exhibits superior generalization capability and sample efﬁciency.


Introduction
Reinforcement learning (RL) is capable of learning from previous experiences automatically without any labeled data, and has demonstrated great potential in solving various complex decision-making problems. By combining with convolutional neural networks (CNNs), vision-based RL has achieved prominent success in a broad range of areas, such as video games (Mnih et al. 2015), continuous control (Lillicrap et al. 2015), robot grasping (Kalashnikov et al. 2018), robot navigation (Wu et al. 2021) and so on. Despite its remarkable achievements, vision-based RL remains plagued by poor generalization capabilities, which significantly limits its deployment in real-world applications. Often, RL agents trained in one environment can hardly generalize well to unseen scenarios, even after they have been trained in a large number of diverse yet semantically similar environments (Gamrian and Goldberg 2019). The adaptation problems of vision-based RL agents can be further exacerbated due to the high-dimensional and partially-observable nature of pixel inputs.
As one of the well explored solutions, domain randomization typically creates a variety of environments with randomly sampled properties and finds the best policy which works across all of the environments (Tobin et al. 2017). Nevertheless, domain randomization can lead to large model complexity and is sensitive to sampling distribution and choice of randomized parameters. Another typical class of solutions is domain adaptation, which learns domaininvariant representations and mitigates the distribution discrepancy between the source and target domains (Bousmalis et al. 2018;Gamrian and Goldberg 2019). However, the target domain data are assumed to be accessible in domain adaptation settings while RL agents typically need to generalize to truly unknown environments.
Furthermore, the incorporation of data augmentations on input observations can also result in significantly improved generalization capability of RL agents Raileanu et al. 2020). Recently, studies have also demonstrated that various self-supervised learning methods can be applied to improve both sample efficiency and generalization of RL policies substantially via better representation learning Hansen et al. 2020;Hansen and Wang 2021). However, prior works only combine selfsupervised representation learning with reinforcement learning naively while it is promising to achieve greater improvement through linking self-supervised learning to reinforcement learning more ingeniously.
In this paper, we propose to improve the generalization of RL algorithms through fusing Self-supervised learning into Intrinsic Motivation (SIM). Similar to SODA (Hansen and Wang 2021), SIM decouples augmentation from policy learning to simplify RL optimization. It only uses augmented data to perform auxiliary representation learning while strictly using non-augmented samples for policy learning. In contrast to SODA, SIM does not require any prediction networks, momentum updates, or stop-gradients. It feeds the augmented and non-augmented data to two identical networks respectively, and makes the cross-correlation matrix between the two embeddings close to the identity matrix. This is to increase the similarity between the embeddings of a sample and its augmented version while mini-mizing the redundancy between the components of the vectors. In this way, the mutual information between the latent representations of augmented and non-augmented data is maximized so as to ignore features that are irrelevant to the RL tasks. Meanwhile, distinguished from all the existing methods, we also convert the redundancy reduction based self-supervised loss to an intrinsic reward to further improve RL generalization via an auxiliary objective. Since the selfsupervised loss can be used to measure the quality of the embeddings, this intrinsic motivation is capable of guiding RL agents to learn the latent representations more efficiently. Since SIM is generic and does not require any change to the underlying algorithms, it can be implemented on top of any RL method. Empirical evaluations have been performed on a diversity of tasks in the DeepMind Control suite benchmark (Tassa et al. 2018), and experimental results demonstrate the superiority of SIM which outperforms the baselines and improves RL generalization significantly.
The main contributions of our work are summarized as follows: • We propose a novel method, SIM, to improve the generalization of RL algorithms through fusing self-supervised learning into intrinsic motivation. • We introduce a redundancy reduction based selfsupervised learning method to enhance representation learning. Moreover, we convert this self-supervised loss to an intrinsic reward to further improve the generalization of RL agents via an auxiliary objective. • We show that SIM outperforms the state-of-the-art baselines by a large margin.

Related Works
Self-Supervised Learning Self-supervised learning empowers the learning of effective feature representations from unlabeled data which can benefit a wide variety of downstream tasks. So far, many auxiliary data prediction tasks have been proposed for selfsupervised learning, such as relative position prediction (Doersch, Gupta, and Efros 2015), jigsaw puzzle (Noroozi and Favaro 2016), colorization (Zhang, Isola, and Efros 2016), context encoder (Pathak et al. 2016), rotation prediction (Gidaris, Singh, and Komodakis 2018), and so on. As an important subclass of self-supervised learning, contrastive learning has gained growing popularity due to its encouraging progress. The rationale of contrastive learning is to minimize the distances between positive pairs while enlarging distances between negative pairs. For example, the CPC approach translates a generative modeling problem to a classification problem which distinguishes feature representation from a set of unrelated samples (Oord, Li, and Vinyals 2018). SimCLR learns visual representations via maximizing agreement between different augmented views of the same sample in the latent space (Chen et al. 2020). By introducing a queue-based dynamic dictionary, MoCo trains the representation encoder through matching the encoded query to a dictionary of encoded keys using a contrastive loss (He et al. 2020). Different from the above contrastive learning approaches, BYOL (Grill et al. 2020) learns invariant feature representations without using negative samples, and instead adopts a bootstrapping mechanism. It relies on two sets of networks, where the target network has the same architecture as the online one but is updated through polyak averaging. It then trains the online network on one augmented view to predict the representations of a different augmented view generated by the target network. Through solving the problem from a different perspective, Barlow Twins (Zbontar et al. 2021) feeds two augmented versions of the training samples into the same network and tries to make the cross-correlation matrix between these two group of embeddings close to the identity. The objective is to keep the embeddings of different augmented versions of the same samples similar while minimizing the redundancy between these representation vectors. In this way, self-supervised representation learning can be achieved without negative samples, prediction networks, momentum updates, or stop-gradients.

Generalization in RL
Generalization remains one of the key challenges for RL. So far, different strategies have been proposed to address this issue. One of the leading solutions is domain randomization, which creates a variety of environments with randomly sampled properties and figures out the best policy that works across all the environments. For example, the simulation environment is rendered with random textures to train the adaptive policies (Tobin et al. 2017). Despite its effectiveness, domain randomization is sensitive and can lead to large model complexity. Another widely explored solution is domain adaptation, which reduces the distribution divergence between the source and target domains. By adopting domain adaptation in RL, adaptive robot grasping policy is learned through Sim2Real transfer (Bousmalis et al. 2018) and observation adaptation for video games is achieved through image-to-image translation (Gamrian and Goldberg 2019). Nevertheless, domain adaptation assumes the target domain data are accessible while RL agents often need to generalize to unknown scenarios. Moreover, a variety of works on data augmentation, such as RAD , DrQ (Kostrikov, Yarats, and Fergus 2020) and DrAC (Raileanu et al. 2020), prove that data augmentation can be very effective as well as efficient to improve the generalization of RL agents as well.
In recent years, studies have demonstrated that it is promising to improve RL generalization more significantly through leveraging the advances in self-supervised learning. For instance, PAD explores the use of self-supervision tasks, such as inverse dynamics prediction, to enable continued training during deployment and thereby improve the generalization of vision-based RL (Hansen et al. 2020). CURL adopts the MoCo mechanism in RL and enforces consistencies in observations via matching embeddings of two augmented versions of the raw observation . More recently, by using the concepts of BYOL in RL, SODA learns invariant representations through bootstrapping (Hansen and Wang 2021 Figure 1: SIM architecture. SIM trains the RL network jointly with an auxiliary self-supervised network. The two networks share a common encoder and the self-supervised loss is converted as an intrinsic reward for RL agent's progress toward better representation learning. Left: representation learning via redundancy reduction based self-supervised learning with augmented observations. Right: policy learning with non-augmented observations. augmented images to their non-augmented counterparts in latent space to ensure that data augmentation is only used for representation learning in order to reduce the difficulty of RL optimization. In this paper, we employ the mechanism of Barlow Twins to improve the representation learning in RL without requiring any negative samples, prediction networks, momentum updates, or stop-gradients. Moreover, we propose to fuse self-supervised learning into intrinsic motivation through converting the self-supervised loss to an intrinsic reward to further improve the generalization of RL via an auxiliary objective.

Method
We propose to improve the generalization of RL through fusing Self-supervised learning into Intrinsic Motivation (SIM). Since SIM does not require any modification to the underlying base RL method, it is a general paradigm that can be implemented on top of any RL algorithm. In this work, we train SIM alongside the Soft Actor-Critic (SAC) algorithm (Haarnoja et al. 2018) to demonstrate its effectiveness. Generally, the objective of SIM is to achieve effective representation learning so that the important information shared between an observation and its augmented counterpart is well encoded.

Architecture Overview
The network architecture of SIM is depicted in Figure 1. Generally, SIM trains the reinforcement learning network jointly with an auxiliary self-supervised learning network. The RL policy network with parameters θ can be split into two parts, i.e., an encoder f θ and a policy π θ , so that given an input observation o, the action can be derived as a = π θ (f θ (o)). The encoder f θ contains 11 convolutional layers with 32 filters and the policy π θ contains 4 fully-connected layers. Meanwhile, the proposed auxiliary self-supervised learning task shares a common encoder f θ with the policy network and is composed of the shared encoder f θ and a projector h θ . The projector h θ has 3 fully-connected layers where the first two layers are followed by a batch normalization layer and rectified linear units. Particularly, the outputs of the encoder and the projector are referred as representations and embeddings respectively. The representations are used for the downstream RL tasks and mapped to actions through the policy network, while the embeddings are normalized along the batch dimension and fed to the loss function of the self-supervised learning task.
During training, observations are sampled from the replay buffer. Similar to SODA, SIM decouples the training data flow so that augmented data are only employed for representation learning to reduce the difficulty of RL optimization. In addition, instead of contrasting two augmented instances of the same image to a batch of negative samples, SIM learns invariance through maximizing similarity between the embeddings of the augmented observations and their non-augmented counterparts without the demand for negative samples. Therefore, both the non-augmented and augmented observations are fed to the self-supervised learning network and mapped to two sets of embeddings via a pair of identical networks. As shown in Figure 2, we augment the observations using the random overlay method (Hansen and Wang 2021) which produces an augmented view through linearly interpolating between an observation and an image. With this strong data augmentation, f θ is encouraged to learn features that are shared across different views while discarding factors of variation that are irrelevant to the RL tasks. Since it is common in the RL setting to stack consecutive frames as observations to infer temporal information, we apply augmentations randomly across the batch while the same image is applied to all frames of a given observation to retain the temporal information.

Self-Supervised Representation Learning
Inspired by the self-supervised learning method proposed in (Zbontar et al. 2021), we also apply the redundancy reduction based method to learn invariant feature representations. However, our goal is to align the embeddings of augmented and non-augmented observations rather than those of two different augmented views. Without negative samples, we achieve self-supervised representation learning through maximizing similarity between the embedding vectors while reducing redundancy between their components.
First, given an original observation o, a data augmentation can be sampled so that t ∼ T . Then, by applying data augmentation t, an augmented observation o = t(o) can be generated so that o and o can be regarded as two different views of the same underlying state. Considering a batch of observations O, a batch of augmented observation O can be produced by applying augmentations randomly across O. As shown in Figure 1, these two batches of different views are fed to the encoder f θ followed by the projector h θ to generate two batches of embedding vectors Z and Z respectively. It is worth mentioning that the two embeddings are normalized along the batch direction so that each unit has zero mean over the batch.
Let C be a cross-correlation matrix computed between the two sets of embedding vectors, the loss function of the selfsupervised task can be expressed as: In these equations, λ is a positive constant used to trade off the two terms of the loss, b is the batch index, and i, j are indexes of the network output vector dimension. Therefore, C is a square matrix with the size same as the dimensionality of the network's output and each of its entry calculates the cosine similarity between z b,i and z b,j . In this way, the first term can drive the embedding invariant to the applied data end if 13: end for augmentation while the second term can decorrelate the different vector components to avoid the inclusion of redundant information.
Without modifying either the architecture or inputs, the RL objective is optimized directly on non-augmented original observations o. Therefore, the policy network and selfsupervised network are continuously optimized in an alternative manner during training while only the encoder f θ and the policy π θ are employed during evaluation.

Fusion of Self-Supervised Learning into Intrinsic Motivation
In RL, intrinsic motivation is originally proposed to facilitate the exploration and its rationale is to predict the consequences caused by taking an action (Pathak et al. 2017;Burda et al. 2018). However, in this work, our goal is to utilize intrinsic motivation to further facilitate the representation learning via better leveraging the useful information provided by the self-supervised network. Since the selfsupervised auxiliary task contains adequate information to measure the quality of the embeddings, an intrinsic reward converted from the self-supervised loss can be promising to assist the RL agents to learn the feature representations more efficiently as well as effectively. As a result, the supervision signals during RL training are contributed from both the extrinsic and intrinsic rewards so that the total reward function R is defined as the sum of these two types of rewards expressed as: where R e t , R i t and μ represent the extrinsic reward, intrinsic reward, and weighting coefficient, respectively. To better stimulate the representation learning, we define the intrinsic reward as: where θ t and θ t−1 denote the current network parameters and those at the previous step, respectively. Since L SSL is computed based on the whole batch, the same intrinsic reward is assigned to all the samples in the batch. By introducing this auxiliary objective, the RL agent is motivated to optimize the encoder more efficiently to achieve improved representation learning. We summarize the training procedure of SIM in Algorithm 1.

Experiments Setup
DeepMind Control Suite (DMControl) (Tassa et al. 2018) is a widely used benchmark dataset for vision-based RL algorithm comparison. It contains a variety of challenging and diverse continuous control tasks. We evaluate the RL agents on 5 tasks from DMControl in which the agents observe only raw pixels without access to state information.
For benchmarking, we compare the SIM with five different baselines, i.e., CURL , RAD , DrQ (Kostrikov, Yarats, and Fergus 2020), PAD (Hansen et al. 2020) and SODA (Hansen and Wang 2021). All the methods are implemented on top of the SAC algorithm. In particular, RAD and DrQ achieve policy adaptation through data augmentations while CURL, PAD and SODA generalize the policy via adding self-supervised objectives.
All the RL agents are trained in a fixed training environment and evaluated in unseen environments as shown in Figure 3. Particularly, the color hard environment randomizes the colors of background, foreground and the agent itself, while the video easy and video hard replace the background with natural videos. We evaluate the performance of different methods at 500k environment steps because all methods have achieved their optimal performance on most environments at this point. For each test environment, the methods are evaluated across 100 randomly initialized episodes per seed.
During training, we use a batch size of 128 for both RL and SSL tasks. The three linear layers of the projector h θ contain 2048, 4096, 4096 output units, respectively. Both L RL and L SSL are optimized using Adam (Kingma and Ba 2014). The SSL update frequency is set to 2 so that the SSL objective is only optimized after every second RL update. The other hyperparameters are listed in Table 1.

Generalization to Unseen Environments
After training, we evaluate the generalization capability of SIM in the unseen color hard, video easy and video hard environments and compare it to a number of recently proposed state-of-the-art approaches. Generally, SIM is much more superior compared to the baseline methods. The experimental results in the color hard environment are shown in Table 2, where the best result for each task is highlighted in bold. It is worth noting that SIM outperforms the baseline methods in all tasks. Particularly, it improves the performance by as much as 16% on the walker walk task and 21% on the finger spin task. Similarly, SIM also significantly outperforms the baseline methods in the video easy and video hard as demonstrated in Table 3 and Table 4, respectively. In the video easy environment, SIM outperforms the state-ofthe-art methods in all tasks except ball in cup catch. And its maximum performance improvement is 17% on the finger spin task. In the extremely challenging video hard environment, an improvement of 20% is achieved on both the walker walk and finger spin tasks.
However, although SIM outperforms all previous methods except SODA dramatically in the video hard environment, there is still large room for improvement on the majority of tasks and it performs worse on the cartpole swingup and ball in cup catch tasks compared to SODA. We conjecture that this performance gap can be addressed by further improving the removal of task-irrelevant features so as to reduce observational overfitting. This can be achieved through improving the encoder architecture, introducing better self-supervised tasks, modifying the RL update rules and enhancing the interaction between RL and SSL objectives.
Although sample efficiency is not the key consideration of our method, the training and evaluation curves are depicted in Figure 4 for completeness. It is observed that while SAC is able to converge in training, its generalization ability is poor during evaluation. Compared to SAC, both SIM and SODA substantially improve RL generalization. With the best performance, SIM exhibits a similar sample efficiency compared to SODA. Therefore, although strong augmentations can improve generalization at the cost of sample     efficiency (Tian et al. 2020), SIM achieves better generalization capability without compromising the sample efficiency.

Ablation Study
Furthermore, we conduct ablation study to verify the effectiveness of fusing self-supervised learning into intrinsic motivation, and denote the SIM method without intrinsic reward as SIM (w/o IM). As shown in Table 2, Table 3 and  Table 4, SIM outperforms SIM (w/o IM) consistently in all tasks except on the cartpole swingup task in the video hard environment. To illustrate the competency of intrinsic motivation more intuitively, we calculate the mean episode rewards averaged across all 5 tasks from DMControl during training and evaluation. As shown in Figure 5, although SIM and SIM (w/o IM) converge to similar performance during training, the superiority of SIM becomes notable in unseen test environments. Moreover, SIM also exhibits slightly better sample efficiency compared to SIM (w/o IM). Therefore, it demonstrates the importance of fusing self-supervised learning into intrinsic motivation in improving the generalization capability of RL agents. Our code and more experimental results are available at https://github.com/KerryWu16/SIM.

Conclusion
In this paper, we propose to address the generalization challenge for RL agents via fusing Self-supervised learning into Intrinsic Motivation (SIM). We employ a redundancy reduction based self-supervised learning method to learn invariant representations through increasing the similarity between the embeddings of a sample and its augmented counterpart while minimizing the redundancy between the components of the vectors. Meanwhile, we also convert the selfsupervised loss to an intrinsic reward to further improve RL generalization via an auxiliary objective. SIM can be implemented on top of any RL method without incurring any change to the underlying algorithm and experimental results have demonstrated its remarkable superiority.