Does Self-supervised Learning Really Improve Reinforcement Learning from Pixels?
Abstract
We investigate whether self-supervised learning (SSL) can improve online reinforcement learning (RL) from pixels.
We extend the contrastive reinforcement learning framework (e.g., CURL) that jointly optimizes SSL and RL losses and conduct an extensive amount of experiments with various self-supervised losses.
Our observations suggest that the existing SSL framework for RL fails to bring meaningful improvement over the baselines only taking advantage of image augmentation when the same amount of data and augmentation is used.
We further perform evolutionary searches to find the optimal combination of multiple self-supervised losses for RL, but find that even such a loss combination fails to meaningfully outperform the methods that only utilize carefully designed image augmentations.
After evaluating these approaches together in multiple different environments including a real-world robot environment, we confirm that no single self-supervised loss or image augmentation method can dominate all environments and that the current framework for joint optimization of SSL and RL is limited.
Finally, we conduct the ablation study on multiple factors and demonstrate the properties of representations learned with different approaches.
Introduction
Learning to act from image observations is crucial in many real-world applications.
One popular approach is online reinforcement learning (RL), which requires no human demonstration or expert trajectories.
Since all training samples are collected by the agent during policy learning in online RL, the collected data often has strong correlations and high variance, challenging the policy learning.
Meanwhile, the cost of interacting with environments requires the RL algorithms to have higher sample efficiency.
Compared to RL using state-based features, pixel-based RL continuously takes images as inputs, which usually come with a much higher dimensionality than numerical states.
Such properties pose serious challenges to image representation learning in RL.
Several recent works studied such challenges from various directions, including: (1) Inspired by the great success of self-supervised learning (SSL) with images and videos (e.g., [6, 7, 9, 11, 15, 16, 18, 22, 33, 34, 39, 42, 54, 56, 57, 63, 71]), some RL methods [1, 44, 48, 61, 65, 70, 82, 88] take advantage of self-supervised learning.
This is typically done by applying both self-supervised loss and reinforcement learning loss in one batch.
In this paper, we dub such joint optimization of the self-supervised loss and the RL loss as the joint learning framework .
(2) On the other hand, many papers [30, 45, 50, 60, 62, 76, 81, 83] investigate how online RL can take advantage of image augmentations.
Among them, RAD [45] and DrQ [81, 83] show significant improvements by applying relatively simple image augmentations to observations of RL agents.
Our objective is to study how well a single or combination of self-supervised losses and augmentations work under the current joint learning framework and to empirically identify their impact on RL systems.
In this paper, we extend such joint (SSL + RL) learning framework, conduct experiments comparing multiple self-supervised losses with augmentations, and empirically evaluate them in many environments from different benchmarks.
We confirm that a single self-supervised loss under such a joint learning framework typically fails to bring meaningful improvements to existing image augmentation-only methods.
We also computationally search for a better combination of losses and image augmentations for RL with the joint learning framework.
The experiments in different environments and tasks show inconsistency in self-supervised learning’s capability to improve reinforcement learning.
Given a sufficient amount of image augmentations, under the current framework, self-supervision failed to show benefits over augmentation-only methods regardless how many self-supervised losses are used.
With all our findings, we present this work as a thorough reference for investigating better frameworks and losses for SSL + RL and inspiring future research.
Our contributions can be summarized as follows:
We conduct an extensive comparison of various self-supervised losses under the existing joint learning framework for pixel-based reinforcement learning in many environments from different benchmarks, including one real-world environment.
We perform evolutionary searches for the optimal combination of multiple self-supervised losses and the magnitudes of image augmentation, and confirm its limitations.
We conduct the ablation study on multiple factors and demonstrate the properties of representations learned by different methods.
Preliminaries 2.1 Reinforcement Learning
In this paper, we extend the configurations of previous work [44, 82] and exploit SAC ( S oft A ctor
C ritic) [26, 27] and Rainbow DQN [35] for the environments with continuous action space and discrete action space respectively.
Soft Actor Critic [26, 27] is an off-policy actor-critic algorithm that takes advantage of the maximum entropy to encourage the agent to explore more states during the training.
It maintains a policy network π ψ and two critic networks Q ϕ 1 and Q ϕ 2 .
The goal of π ψ is to maximize the expected sum of rewards and a γ -discounted entropy simultaneously, where the entropy encourages the agent to explore during learning.
Rainbow DQN [35] is a variant of DQN [53] with a bag of improvements such as double Qlearning [31, 72], prioritized sampling [64], noisy net [20], distributional RL [4], dueling networks [78] and multi-step reward.
2.2 Pairwise Learning
We coin the term “pairwise” learning for the frameworks that learn visual representations based on semantic invariance between dual-stream encoder representations.
A general pairwise learning method first generates multiple augmented views by applying a series of random image augmentations to the input sample, then clusters views with the same semantics in the representation space.
Optionally in such frameworks, methods using contrastive losses repel samples with different semantics.
In this paper, we focus on four representative pairwise learning methods, MoCo [12, 13, 33], BYOL [23],
SimSiam [11] and DINO [7].
We have a detailed explanation and comparison of these methods in
Appendix A.1.
2.3 Representation Learning for Pixel-based RL
Previous works explore the possibility of learning better visual representation which may finally benefit policy learning.
One direction is using image augmentation for policy learning [30, 45, 50, 60, 62, 76, 81, 83], where RAD [45] and DrQ [81, 83] achieve significant performance using simple image augmentation.
Another direction is to combine SSL with RL [1, 44, 48, 61, 65, 70, 82, 88], in which there are two representative methods, SAC+AE [82] and CURL [44].
[Figure]
RAD ( R einforcement Learning with A ugmented D ata) [45] investigates the impact of different types of image augmentations for both image and state inputs.
By applying random translation or random crop to the input image, RAD significantly improves data efficiency solely through image augmentation without any auxiliary losses.
DrQ ( D ata- r egularized Q ) [81] further investigates the possibilities of utilizing image augmentation.
DrQ applies image augmentation twice on the input images and averages the Q value over two augmented images which is assigned as the Q value of the input images.
DrQ v2 [83], which is the successor of DrQ, switches to DDPG (Deep Deterministic Policy Gradient) [49] as the RL method, brings scheduled exploration noise to control the levels of exploration at different learning stages, and introduces faster implementations of the image augmentation and the replay buffer.
SAC+AE [82] takes advantage of a RAE (deterministic R egularized A uto E ncoder)
[21], in replacement of β − VAE [36] to improve learning stability.
The RAE is jointly trained with SAC by performing both SAC update and RAE update alternatingly in one batch.
CURL ( C ontrastive U nsupervised Representations for R einforcement
L earning)
[44] combines contrastive learning with an online RL algorithm by introducing an additional contrastive learning head at the end of the image encoder.
Similar to the aforementioned SAC+AE, here the contrastive loss and reinforcement learning loss are applied alternatively at training.
Self-supervision for Reinforcement Learning
To effectively evaluate different self-supervised losses, we extend the well-known joint learning framework widely used in previous papers [1, 44, 48, 65, 82] by adding a general self-supervised learning head to the RL framework.
We keep the same RL method in CURL [44]: we use SAC [27] in tasks with continuous action space and use Rainbow DQN [35] in tasks with discrete action space.
3.1 General Joint Learning Framework
With SAC Fig.
1a shows a general joint learning framework, using SAC as the RL method.
The unmodified SAC contains an online encoder f q , a target (or momentum) encoder f k , and an actor head A p .
Each encoder is also followed by two critic heads.
Besides that, we attach an additional self-supervised head g q after the online encoder.
For pairwise learning losses, we concatenate a momentum SSL head g k after the target encoder when needed.
For every sampled batch of transitions, we first apply image augmentation to both the current state s and the next state s ′ and update the SAC model ( f q , Q i =1 , 2 q , A p ) using the augmented images.
Note that for stability concerns, we do not update the parameters of the image encoder when updating the actor head A p .
Then, the target networks are updated by Exponential Moving Average (EMA).
This is followed by also performing an EMA update of the SSL head if required.
Finally, the online encoder f q and the self-supervised head g q are updated by the self-supervised loss.
By alternatingly performing RL and SSL in every batch, we jointly train all the components in the framework.
The pseudo-code of SAC update alternating RL and SSL is provided in Algorithm 1.
Green: additional operations for SSL; Orange: only for BYOL and DINO.
[Table]
With Rainbow DQN Fig.
1b demonstrates how to jointly apply SSL to Rainbow DQN.
The unmodified Rainbow DQN maintains an online encoder f q and a target encoder f ′ , followed by two state value heads Q q and Q ′ .
We introduce an additional momentum encoder f k and self-supervised heads g q
and g k as suggested in CURL.
For each batch, the self-supervised losses are computed using augmented images, while the RL loss is computed using the original data.
Finally, the online encoder f q and the self-supervised head g q are updated by the self-supervised loss.
The pseudo-code of Rainbow DQN update can be found at Algorithm 2.
Green: additional operations for SSL; Orange: only for BYOL and DINO.
[Table]
3.2 Losses for Self-supervised Learning
The self-supervised losses we investigated can be categorized into four classes: pairwise learning, transformation awareness, reconstruction, and reinforcement learning context prediction.
Pairwise Learning We investigate three representative pairwise learning methods: BYOL [23],
DINO [7] and SimSiam [11], along with existing CURL whose framework is similar to MoCo [33].
BYOL, DINO, and SimSiam only explicitly pull positive samples closer without the need for a large number of negative samples.
CURL uses a contrastive loss taking both positive and negative samples into consideration.
Given the general joint learning framework described in Sec. 3.1, by substituting the self-supervised head and loss, we can easily formulate different agents w.r.t.
self-supervised losses.
For BYOL, as shown in Fig.
10a, a projector and a predictor are appended to the online encoder sequentially while a momentum projector is attached on top of the target/momentum encoder.
DINO (Fig. 10c) maintains only projector in both online and target branches.
Similar to BYOL, the momentum projector in
DINO is also updated by EMA.
The two encoders in BYOL and DINO operate on two augmented views of the data respectively whereas SimSiam (see Fig. 10b), uses only the online network and a projector for processing both the augmented views.
We also test two methods that introduce RL-specific variables to this pairwise learning framework,
CURL-w-Action and CURL-w-Critic .
CURL-w-Action is based on CURL while the contrastive loss is applied to the concatenation of image representation and output of the actor network, instead of the image representation only.
Similarly, CURL-w-Critic concatenates the critic network output with the existing image representation for contrastive loss.
Transformation Awareness
Recent works (e.g., [14, 22, 38, 41, 47, 54]) have shown that the awareness of transformations (like rotation, Jigsaw puzzle, and temporal ordering) improves many downstream tasks in computer vision like image classification and action recognition.
Typically such awareness can be acquired by explicitly asking a classifier to identify the applied transformation from the pixel representation.
Therefore, we investigate two simple classification losses, rotation classification ( RotationCLS ) and shuffle classification ( ShuffleCLS ), and set a two-layer MLP classifier as the self-supervised head in the joint learning framework.
RotationCLS represents the methods that encourage spatial transformation awareness.
Inspired by
RotNet [22] and E-SSL [14], we rotate the input image after augmentation by 0°, 90°, 180°and 270°.
The classifier predicts the rotation angle from the visual representation and it is trained by cross-entropy loss.
Shuffle Tuple [52] encourages the encoder to develop an awareness of action causality by predicting if two frames appear in order.
We adapt Shuffle Tuple by randomly shuffling the current state image and next state image in a state transition tuple and predicting whether it is shuffled or not.
The classifier also takes action into consideration because some of the transitions are reversible.
The overall architecture of ShuffleCLS is shown in Fig.
Reconstruction Reconstructing the input image with an hourglass architecture has been shown to be an effective way to learn image representation [21, 36, 42].
We simply extend SAC+AE by changing the input and reconstruction target to be augmented images.
The reconstruction loss and regularization from RAE [21] are left untouched.
Recent study on M asked A uto E ncoder [34] (MAE) adapts the reconstruction task for patch-based
Vision Transformers [19].
The objective in MAE includes reconstructing the entire image from input masked image patches.
Inspired by this, we adapt SAC+AE into SAC+MAE by replacing the augmented input image with its masked version and only penalizing the reconstruction error for the masked patches.
RL Context Prediction Besides the self-supervised learning methods that are specifically designed for pixels, we investigate the losses using attributes naturally collected during the RL process.
For any state transition that is not the end of a trajectory, it contains four components: current state s , next state s ′ , action a and reward r , with the trajectory termination signal omitted.
Inspired by Shelhamer et al.
[69], we concatenate the visual representation of the current state s and another representation h as the input.
Without loss of generality, the second input representation h can be any of these three representations of s ′ , a , and r .
Then, we predict the remaining components using a two-layer
MLP.
For continuous outputs, mean-squared error (MSE) loss is applied, while for the discrete target (e.g., action in discrete action space), we use cross-entropy loss.
The architecture of this group of self-supervised losses is shown in Fig.
From the combination of inputs and outputs, we define nine losses whose I/
O specifications are provided in Table 1.
For those losses whose outputs include two components, two target prediction networks share the same SSL head except the last task-specific layer.
3.3 Evolving Multiple Self-supervised Losses
Besides a single self-supervised loss or handcrafted combination of two losses, we further investigate how multiple self-supervised losses affect the policy learning together with the joint learning
[Table]
Rep. of s ′
Input
InputOutput OutputInput
Output
Output Action a
OutputOutput -
Input Input
Output Output
Input Reward rOutput
Input
InputOutput
Output
Input
Output framework.
In such a configuration, the agent maintains multiple SSL heads at the same time and we apply losses to their corresponding head individually.
We formulate the combination of multiple losses as a weighted sum L Combo =
(cid:80) N l i =1 w i · L i where w i is the weight of a specific loss L i and
N l is the total number of losses in the search space.
In the joint learning framework, we apply both self-supervised L Combo and RL losses jointly to the networks for every mini-batch.
Considering that the policy learning is quite sensitive to hyper-parameters, it is non-trivial to find each weight for every
SSL loss.
ELo ( E volving Lo sses) [59] shows promising results in unsupervised video representation learning [57, 71], by using evolutionary search to automatically find optimal combination of many self-supervised losses.
In the spirit of ELo, we turn to evolutionary search to automatically find the optimal solution.
Assume an unknown objective function whose inputs are weights of multiple losses w i and the magnitudes of image augmentation m j =1 , 2 for the online encoder and momentum encoder.
The function output is the score achieved by the trained agent in its environment with a certain random seed: R seedenv ( m j =1 , 2 , w i =1 , 2 ,...N l ) .
Essentially, the objective function maps a set of w i and m j to the reward achieved by a corresponding agent, and w i and m j stay unchanged during the agent learning process.
The optimization algorithm approaches the maximum value of the objective function by repeatedly updating w i and m j and testing the value of the objective function, which in our case is the training and evaluation of an agent with the given parameters (i.e., the input of the objective function).
We choose an off-the-shelf optimization algorithm PSO ( P article S
warm
O ptimization)
[40] for its simplicity.
For each set of inputs, we find it critical to run with multiple random seeds and report IQM (interquartile mean) 1 for a stable and robust search.
The optimization process is presented as:
[Equation]
Note that we are also implicitly searching for the balance between the self-supervised loss and the
RL loss by performing this search, as it has the capability to adjust the absolute weights of the selfsupervised losses overall.
We search on DMControl with SAC using three different configurations named ELo-SAC, ELo-SACv2 and ELo-SACv3 respectively.
ELo-Rainbow performs a search on
Atari with Efficient Rainbow.
Please refer to Sec.
A.2.4 for our detailed configurations and search results.
Experiments
We conduct experiments in three major directions, in order to better understand how we should integrate SSL with RL.
First, we demonstrate how different self-supervised losses affect the RL process, by trying them on multiple challenging tasks.
Then, we dive into detailed ablations on multiple factors, and finally, we perform empirical analysis on the visual representations learned with the joint learning framework (Sec. A.6).
In addition, we benchmark a pretraining framework as an alternative to the joint learning framework (Sec. A.7).
Evaluation Scheme Thorough evaluation of reinforcement learning algorithms is challenging due to the high variances between each run and the extensive requirement of computation.
Consequently, we run all experiments with multiple different random seeds and report the interquartile mean and the standard deviation of the scores as suggested by Agarwal et al.
For a quantitative comparison of the different methods mentioned in Section 3.2, in addition to the absolute scores, we assign a
Relative Score to each method.
We denote the interquartile mean of scores achieved by agent A in environment e ∈ E as IQM A,e and denote the collection of all interquartile mean scores achieved in environment e by different agents as IQM e .
The Relative Score of agent
A is computed as
S A Relative = (cid:80) e ∈ E ( IQM A,e − mean ( IQM e )) / std ( IQM e ) .
DMControl Experiments DMControl
( D eep M ind C ontrol suite) contains many challenging visual continuous control tasks, which are widely utilized by recent papers.
We evaluate all the methods introduced in Sec. 3, along with two important baselines, SAC-NoAug and SAC-Aug(100), in six environments of DMControl that are commonly used in previous papers [44, 45, 81, 82].
Other methods that only take advantage of image augmentation, like RAD [45] and DrQ [81] are also benchmarked for comparison.
Note that we feed the augmented images to SAC+AE [82] for a fair comparison, which makes a different configuration from the original paper.
Please refer to
Appendix A.2.5 for a detailed comparison of method variants and the exact data augmentation they applied.
We mainly follow the hyper-parameters and the test environments reported in CURL, except that we use the same learning rate 10 − 3 in all environments for simplicity.
All the methods are benchmarked at 100k environment steps, with training batch size 512 under 10 random seeds, and they share the same capacity of policy network.
The relative score of each tested algorithm on DMControl is reported as Fig.
We also strongly encourage readers to check full results at Table 11 and results in two additional harder environments at Table 12 for a full picture.
From the first glance at Fig.
4, no tested SSL-based method within the joint learning framework achieves better performance than DrQ and RAD which are carefully designed to take the best advantage of specific image augmentations.
Compared to the baseline SAC-Aug(100), approaches with a self-supervised loss frequently (11 out of 19) fail to improve reinforcement learning.
Some SSL methods (like SimSiam, ShuffleCLS) ruin the policy learning resulting in performance even worse than SAC-NoAug, which suggests that improper use of self-supervised loss can damage the benefits brought by image augmentation.
Then, regarding combining losses, Guess-AF and Predict-FR, which are manually designed to combine two individual losses, are not better than the single self-supervised loss in their combinations (check Guess-Action and Predict-Reward in Fig. 4).
ELo-SAC and ELo-SACv2 find the desired combination by searching in one task.
Such combination generalizes to other environments on DMControl with better overall performance than any approach in the search space.
In the ‘cheetah run’ where the search was performed, they obtained the best result among the approaches with SSL (Table 12).
This demonstrates the feasibility of ELo-SAC and implies that the obtained combination through evolutionary search has the potential to generalize to other environments in DMControl.
However, weaker performance in ‘finger, spin’ and ‘reacher, easy’ made ELo-SAC relatively worse than DrQ (which does not use any self-supervision) on average.
Interestingly, there is a similar performance pattern between ELo-SAC and ELo-SACv2 though they have different search spaces.
By contrast, ELo-SACv3 finds an overall better combination by searching in six environments simultaneously.
Though it achieves highest score in ‘walker, walk’ and ‘reacher, easy’, it performs worse in ‘cartpole, swingup’ and ‘cheetah, run’ than ELo-SAC and
ELo-SACv2.
Such observations could be a clue to the properties of different tasks and self-supervised methods.
Ablations Our observations with SAC-Aug(88), SAC-Aug(100), and RAD suggest the importance of augmentation hyper-parmeters, given the only difference between these three methods is the augmentation applied.
We conduct an ablation study on the image augmentation random crop [44] in cheetah run, DMControl.
All the hyper-parameters are as noted in Table 2 except that the environment step is set to 400 k and the batch size is reduced to 128 .
Fig.
6 shows how the magnitudes of random crop and translate contribute to the score that the agent achieved.
The image size before the random crop is linear to the magnitude of the random crop when using a fixed crop size: the larger the image size, the stronger the augmentation.
There is a trend that the score first increases and then decreases as the image augmentation gets stronger.
In summary, it is critical to engineering image augmentation carefully when designing an RL system with or without SSL.
Then we investigate a different visual encoder backbone ResNet [32] by replacing the last two convolutional layers with a residual block that has the same number of layers and channels as the
CNN baseline.
The ResNet backbone slightly improves all these methods (see Fig. 5).
We also encourage the readers to check more ablations regarding image augmentation (e.g. random translate), learning rate, encoder layers, and activation function in Appendix A.3.
Atari Game Experiments Atari 2600 Games are also challenging benchmarks but with discrete action space.
We choose seven games in this benchmark for selected methods.
All the methods use
Efficient Rainbow [73] as the RL method, which is a Rainbow [35] variant with modifications for better data efficiency.
Note that Efficient Rainbow, as a baseline, does not take advantage of image augmentation.
Therefore, we also benchmark Rainbow-Aug which is essentially Efficient Rainbow taking the augmented images for policy learning instead.
We use the same image augmentation and hyper-parameters reported by CURL for all applicable methods.
For a fair comparison, the augmentation for DrQ* is also adopted from CURL, which is different from what the original DrQ paper suggested.
We denote our setting as DrQ* to distinguish it from the original DrQ. Similarly,
Rainbow+AE takes augmented images.
For each game, we run 20 random seeds and benchmark the agent at 400K environment steps (100K model steps with a frame skip of 4).
We report interquartile mean, standard deviation, and Relative Scores same as DMControl (See Table 13).
Figure 7 shows a summary of the seven different tasks in Relative Score.
Firstly, compared to vanilla baseline Efficient Rainbow which does not have any image augmentation or self-supervised learning,
Rainbow-Aug performs worse overall with additional image augmentation for RL.
This suggests that the image augmentation used for self-supervised learning in CURL does not easily transfer.
Similarly, DrQ* achieves compromised performance than Efficient Rainbow, showing that using image augmentation for Rainbow on Atari does not benefit policy learning unlike SAC on DMControl.
Based on the inconsistent impacts of image augmentation, further investigation is required when applying image augmentation to RL on Atari.
As for the self-supervised losses, BYOL, Rainbow+AE, Extract-Reward, and Predict-Reward gain better performance than CURL.
However, only Rainbow+AE shows significant improvement on
Efficient Rainbow and outperforms all the other tested methods, which interestingly is inconsistent with SAC+AE on DMControl.
Predict-FR-Balanced, which shows considerable improvements on DMControl by manually combining two self-supervised losses, even fails to surpass PredictReward on Atari.
ELo-Rainbow, which searches in Frostbite, improves the baseline only in demon attack and frostbite.
The high variance on this benchmark made the evolutionary search extremely difficult.
Further, there are huge image domain gaps between games, which makes it even harder for
ELo-Rainbow to work across multiple games on Atari.
Real Robot Experiments We further conduct experiments in a real-world robot environment, uArm reacher.
Similar to Burgert et al.
[5], the goal is to move the actuator close to a target object as fast as possible.
Our autonomous training environment and results are shown in Figs.
8 and(Please check Appendix A.5 for environment setup details).
We benchmark all methods with five different random seeds, using the same hyper-parameters as DMControl experiments unless reported in Table 3.
Results are shown as Fig.
Surprisingly, in this real-world environment, the agent fails to learn an effective policy without any image augmentation.
ELo-SAC with the optimal combination found in cheetah run, achieves significantly better performances than DrQ and CURL, which is quite different from our previous observations on DMControl.
The image augmentation alone (i.e., SAC-Aug(100)) was sufficient to outperform CURL using self-supervision.
But it needs some engineering to design the image augmentation that helps (see DrQ and DrQ(100) in Fig. 9).
Related Works
Self-supervised learning can fit in robot policy learning in multiple fashions and at different stages.
Some works [25, 66–68, 70, 74, 80, 86] use SSL for representation learning in a pre-training stage before policy learning.
Others [24, 28, 37, 44, 46, 48, 51, 56, 65, 82, 84, 85, 87, 88] jointly optimize the self-supervised loss with policy learning.
Specifically, Transporter [43] and VAI [77] train an unsupervised keypoint detector to discover critical objects in the image for control.
RRL [67] and
VRL3 [74] also benefit from pre-training a deeper visual encoder on large datasets like ImageNet [17].
TCN [66] and CURL [44] take advantage of contrastive learning.
After the agent is deployed, SSL can be used to continuously improve the policy [29].
Shelhamer et al.
[69] study several self-supervised losses within both the pretraining framework and the joint learning framework, while their selection of losses, the number of runs, and test environments are limited from a current point of view.
Chen et al.
[10] focus on imitation learning and test multiple SSL objectives for representation learning in various environments.
They confirmed the critical role of image augmentation in imitation learning and showed inconsistencies in performance across environments.
Our investigation supports some of their observations, beyond that, our evolving loss, real robot environment, and representation analysis provide unique perspectives for online reinforcement learning.
Discussion
From DMControl and the real robot experiments, we empirically show that compared to the image augmentation, the role of existing self-supervised losses with the joint learning framework is usually limited, even with the help of evolutionary search.
While results on Atari show a different trend from DMControl, once again we confirm that there is no golden self-supervised loss or image augmentation that generalizes across environments.
At the same time, it is usually challenging to conclude a consistent trend that one method is meaningfully better than others across multiple tasks.
One should cautiously decide the design choice of image augmentation or self-supervised loss for a specific RL task.
We are excited to see future works that introduce more self-supervised losses designed specifically for RL, as well as novel training frameworks that can benefit policy learning.
Acknowledgments
We thank Kumara Kahatapitiya, Ryan Burgert and other lab members of Robotics Lab for valuable discussion.
We thank Hanyi Yu and Rui Miao for their helpful feedback.
This work is supported by Institute of Information & communications Technology Planning &
Evaluation (IITP) grant funded by the Ministry of Science and ICT (No.2018-0-00205, Development of Core Technology of Robot Task-Intelligence for Improvement of Labor Condition.
This work is also supported by the National Science Foundation (IIS-2104404 and CNS-2104416).
A Appendix A.1 Background on Pairwise Learning
We coin the term “pairwise” learning for the frameworks that learn visual representations based on semantic invariance between dual-stream encoder representations.
A general pairwise learning method first generates multiple augmented views by applying a series of random image augmentations to the input sample, then clusters views with the same semantics in the representation space.
Optionally in such frameworks, methods using contrastive losses repel samples with different semantics.
Previous works not only have built various self-supervised tasks that benefit representation learning but also show that learned representations can benefit different downstream tasks.
In this paper, we focus on four representative pairwise learning methods, MoCo [12, 13, 33],
BYOL [23], SimSiam [11] and DINO [7].
Specifically, MoCo takes advantage of the contrastive loss and negative samples in the mini-batch, while BYOL, SimSiam, and DINO focus on the similarity of the same image across diverse augmentations.
MoCo Mo mentum Co ntrast (MoCo) takes advantage of a contrastive loss function InfoNCE [56] with dot product similarity.
It starts from two identical encoder networks, an online encoder f q and a momentum encoder f k .
At each training step, a mini-batch of N images x are uniformly sampled from a training set D .
Given two distributions of image augmentations T and T ′ , two image augmentations t ∼ T and t ′ ∼ T ′ are sampled respectively and applied to x , resulting in 2 N samples.
Augmented images, v = t ( x ) and v ′ =
t ′ ( x ) , are called views .
Then, v and v ′ are fed to two encoders to generate queries q = f q ( v ) and keys k = f k ( v ′ ) .
For each view v i in v and its corresponding query q i =
f q ( v i ) , the contrastive loss is formulated as:
[Equation]
where sim( q i , k i ) =
exp( q i · sg [ k i ] /τ ) , sg [ · ] implies stop gradients and τ is a temperature hyperparameter.
This loss encourages q i to be similar to its corresponding key k i (called positive ), but dissimilar to other keys (called negatives ) in the mini-batch.
The online encoder f q with parameters θ q is updated by the above contrastive loss.
The momentum encoder f k with parameters θ k , an
Exponential Moving Average (EMA) of f q , is updated by
[Equation]
where m ∈ [0 , 1) is a momentum coefficient that controls how fast θ k updates towards the online network θ q .
Finally, f k will be discarded once the training completes.
BYOL Similar to MoCo, in addition to f q and f k , BYOL maintains two identical projection networks g q , g k and one prediction networks p q (See Fig. 10a).
BYOL also starts from inputs v and v ′ but calculates the projection z 1 = g q ( f q ( v )) and
z 2 = g k ( f k ( v ′ )) , and tries to regress z 2 from z 1 using the prediction network p q .
After applying l 2 -normalization to the prediction p q ( z 1 ) and the target projection z 2 , a mean squared error is measured as:
[Equation]
whose value is low when p q ( z 1 ) is close to z 2 .
Similarly, by swapping v and v ′ , another symmetric loss can be applied on top of z ′ 1 =
g q ( f q ( v ′ )) and z ′ 2 =
g k ( f k ( v )) , as L BYOL , 2 = ∥
p q ( z ′ 1 ) − sg [ z ′ 2 ] ∥ 22 .
The total loss is L BYOL =
( L BYOL , 1 +
L BYOL , 2 ) / 2 .
The parameters of f k and g k are also the EMA of f q and g q respectively.
Finally, f k , g q , g k and p q will be discarded once the training completes.
SimSiam SimSiam ( Sim ple Siam ese) shares the same architecture as BYOL, while the parameters of the ‘momentum’ branch of SimSiam are always tied to the ‘online’ branch (See Fig. 10b).
Therefore,
SimSiam only maintains one branch, including an encoder f , a projector g , and a predictor p .
[Figure]
SimSiam uses negative cosine similarity to encourage the predicted representation h
= p ( g ( f ( v )) to be similar to the projected representation of another view g ( f ( v ′ )) , as follows:
[Equation]
Another symmetric loss term can also be derived as L SimSiam , 2 = − h
[Equation] [Equation]
sg [ g ( f ( v ))]
∥ sg [ g ( f ( v ))] ∥ 2 , where h ′ =
p ( g ( f ( v ′ ))) .
The total loss is L SimSiam =
( L SimSiam , 1 + L SimSiam , 2 ) / 2 .
And g q will be discarded once the model is trained.
DINO DINO shares a similar overall structure as MoCo which contains two encoders f q and f k , and f k is the EMA of f q (See Fig. 10c).
The outputs of both encoder networks are normalized as probability distributions over K dimensions by applying softmax with a temperature parameter τ t , and K is the dimension of f q ( v ) .
DINO also maintains a centering vector C with dimension K .
Following the formulation of knowledge distillation, a cross-entropy loss is applied to encourage the output distribution of f q to become similar to a centered distribution from f k , as follows:
[Equation]
(6) where P ( x ) = softmax ( x/τ t ) .
By swapping v and v ′ in Eq. 6, another loss L DINO , 2 which is symmetric to L DINO , 1 can be derived.
And the total loss is the mean value of L DINO , 1 and L DINO , 2 .
After each step of optimization, f k is updated by Eq.
C also gets updated in a similar manner:
[Equation]
Here m c ∈ [0 , 1) is another momentum coefficient.
A.2 Implementation Details
Here we present the implementation details in all settings.
A.2.1 General Joint Learning Framework
The general joint learning framework starts from the official implementation of CURL [44] for
DMControl and Atari.
For different self-supervised learning losses, we only replace the contrastive learning head of CURL with a different SSL head and update the loss calculation.
All the hyperparameters are left untouched, except that we use the learning rate 10 − 3 for all DMControl environments.
The detailed hyper-parameters can be found at Table 2 (DMControl) and Table 4 (Atari).
We keep the most of hyper-parameters from DMControl for real-world robot experiments.
The modified configuration is listed as Table 3.
We use the official implementations of DrQ [81] and RAD [45] for DMControl benchmark.
The re-implementation of DrQ (denoting as DrQ*) has the same joint learning framework and image augmentation from CURL.
[Table]
[Table]
A.2.2 Losses for Self-supervised Learning
Pairwise Learning In this section, we replace the contrastive learning head with projectors and encoders depending on the exact loss.
All the projectors and predictors are two-layer MLPs with
ReLU in the middle.
The input dimension which is the output dimension of the encoder, ( 50 on
DMControl and uArm Reacher, and 576 on Atari).
The hidden dimension of the MLP is 256 and the output dimension is 128 .
We use the same encoder EMA update rate ( τ = 0 .
[Table]
Transformation Awareness We use a two-layer MLP with ReLU as the classifier for both rotation classification and shuffle classification.
The hidden dimension of the MLP is 1024 and the classifier is supervised by a cross-entropy loss.
The output dimension is 4 for four-fold rotation classification and 1 for binary shuffle classification.
Reconstruction In this section, we follow the official implementation of SAC+AE [82] and apply the same image augmentation from CURL.
The decoder has one fully connected layer and the same number of transposed convolutional layers as the convolutional layers in the encoder.
When the output image from the decoder is smaller than the ground truth, we crop the ground truth to the size of the decoder output from the upper left corner.
For MAE we start from augmented SAC+AE and first divide the augmented image into nonoverlapping patches in the spatial domain with a size of 4 × 4 .
Then we randomly mask 50% of the patches by setting the pixel value of the masked patches to zero.
Finally, the reconstruction loss is modified to calculate MSE only over the masked patches.
Other regularization losses are left untouched.
RL Context Prediction For all kinds of losses
, the dimensions of all the fully connected layers and hidden layers in MLPs are 1024 .
A.2.3 Manually Balance Two Self-supervised Losses
In this section, we further explore the ways to manually combine two self-supervised losses.
ExtractAR, Guess-AF, and Predict-FR are methods manually designed to combine two individual losses
However, Guess-AF and Predict-FR are not better than the single self-supervised loss in their combinations (see Guess-Action and Predict-Reward in Table 11 and Fig. 4).
Considering that
Extract-AR, Guess-AF, and Predict-FR concatenate both the outputs and apply supervision by averaging loss per element of the output, the target with a higher dimension will naturally get more penalty due to the larger number of elements in the output.
We further test the ‘Balanced’ configuration, where we only modify how the supervision is applied.
Take Extract-AR as an example, in the ‘Balanced’ setting, we first calculate loss regarding action prediction and reward prediction separately, then the total self-supervised loss is the average of both the action prediction loss and the reward prediction loss.
By adjusting the combination weights, the ‘Balanced’ trick brings overall improvements on top of all three methods as shown in Table 5.
Such observation suggests that we need to carefully design how the two losses are combined, which is getting trickier as the number of combined losses increases.
[Table]
A.2.4 Evolving Multiple Self-supervised Losses
We choose PSO ( P article S warm O ptimization) [40] for the optimal combination of hyper-parameters including N w weights of losses w i =1 , 2 ,...,N w and two magnitudes of augmentation m j =1 , 2 for the online networks and the target networks respectively.
Each m j varies from [84 , 116] .
We limit the range of each w i to [0 , 10] for ELo-SAC and ELo-Rainbow while a range of [10 − 4 , 10 4 ] for
ELo-SACv2 and ELo-SACv3.
During the evolutionary search, we use a batch size of 128 for ELo-SAC and ELo-SACv2, each combination is trained with 5 different random seeds.
As for ELo-SACv3, the batch size is set to 64 and the number of random seeds is set to 3 to save computation.
ELo-Rainbow also train with 5 random seeds during the evolutionary search.
Other hyperparameters used in the search and all hyperparameters for evaluation are identical to Table 2 and Table 4.
ELo-SAC ELo-SAC maintains a population of 50 for DMControl and each particle evolvesgenerations in “cheetah, run”.
Before the search, the first i th particles are initialized with m j =1 , 2 = 88 , and each particle only has one weight set to 1 and other weights set to 0 .
In another word, these first i th particles start with the existing single self-supervised loss method in the search space.
Other particles are randomly initialized.
Table 6 shows the combination Elo-SAC found in cheetah run.
The columns in Table 6 show the search space.
The first six columns denote the optimal weight w i of its corresponding loss obtained with the evolutionary search, while the last two columns denote the original image size before random crop (image augmentation magnitude m j =1 , 2 ).
[Table]
ELo-SACv
Compared with ELo-SAC, ELo-SACv2 has the following major improvements:
Initialization Define the set of image augmentation magnitude M = { ( m 1 = t, m 2 = t ) | t = 86 , 88 , 92 , 100 , 116 } , and the set of SSL weights W = { ( w i =
t = 1 , w i ̸
= t = 0) | t
= 1 , 2 , ..., 6 } .
The first |M| × |W| particles are initialized from the Cartesian product of M and W .
Other particles are randomly initialized.
Search space
The search space of self-supervised losses is updated based on the loss performance at Table 11.
We empirically choose the losses from different categories that have a relatively strong performance when it is applied solely to RL. 3.
Weight range
The weight of each self-supervised loss is presented on a log scale so that the search can cover a larger range.
Besides the improvements above, ELo-SACv2 evolves 45 generations and the optimal combination is chosen from the top 10 combinations regarding the overall performance.
The optimal combination found by ELo-SACv2 is shown in Table 7.
ELo-SACv2 slightly improves the results of ELo-SAC with all the modifications (see Figure 4 and Table 11).
[Table]
ELo-SACv
Since ELo-SAC and ELo-SACv2 only search in one DMControl environment, ‘cheetah run’, and both the found solutions perform weaker on ‘finger, spin’ and ‘reacher, easy’, we further extend ELo-SACv2 to search in multiple environments at the same time, named ELo-SACv3.
The optimization process of ELo-SACv3 is presented as:
[Equation]
where ˆ R = R/R DrQ is the original agent reward R normalized by the score of DrQ R DrQ reported in [81], and envs is the set of six DMControl environments listed in Table
[Table]
We let ELo-SACv3 evolve for 25 generations and chose the loss combination with best performance among the top 10 records.
The found parameters are listed in Table
[Table]
[Table]
ELo-Rainbow ELo-Rainbow has a population of 30 and the initialization is similar to ELo-SAC.
The search is performed on Frostbite only for 10 generations and the found combination is shown in
Table 9.
[Table]
Interestingly, we find that the optimal combination found by ELo-SAC is relatively sparse, where
BYOL and Predict FR are the only two major losses.
Similarly, ELo-SACv2 relies more on PredictFR-Balanced and RotationCLS, while ELo-SACv3 relies on Extract-AR-Balanced mostly.
However, for ELo-Rainbow, the magnitudes of all the weights are relatively similar.
The difference between the found results reflects the different properties of different environments.
Our further experiments in
DMControl confirm the generalization ability to evolve losses; i.e., the obtained solution of weights in one environment achieves relatively good performance on other environments in the same benchmark.
However, results on Atari are much inconsistent with DMControl.
We cover detailed observations and discussions in Section 4 and Appendix A.4.
The code for ELo-SACv3 is available at https://github.com/LostXine/elo-sac , and the code for ELo-Rainbow is available at https://github.com/LostXine/elo-rainbow .
A.2.5 Comparison of method variants
In Section 4, several variants of the existing methods are introduced.
The difference between these methods, especially on image augmentation, can be summarized as follows:
SAC-NoAug is the original pixel-based SAC [26, 27].
SAC-Aug(88) and SAC-Aug(100) use the random crop as the only image augmentation, where (88) means the original image has a size of 88 × 88 before randomly cropping to 84 × 84 and (100) means the original image has a size of 100 × 100 .
These two methods should be regarded as variants of RAD with different augmentation choices.
The random crop from 100 × 100 to 84 × 84 is the default image augmentation method for all the methods introduced in Sec. 3.2, including SAC+AE, which makes it different from the configuration in the original paper.
Essentially, if we remove their self-supervised loss, they will fall back to SAC-Aug(100).
Similarly, we test DrQ variants by replacing its default random shift augmentation with the random crop, reported as DrQ(88) and DrQ(100).
Meanwhile, RAD uses random translate by default except on walker walk; ELo-SAC and ELo-SACv2 first crop the center of the input image to the found optimal sizes.
Then two central patches are randomly cropped to 84 × 84 as the inputs for the online networks and the target networks respectively.
For the policy learning part, all the methods share the same model.
However, DrQ, DrQ(88),
DrQ(100), and SAC+AE apply an additional tanh activation after the visual encoder.
We also study the effect of the activation function in the coming Section A.3.4.
A.3 Ablations
Besides random crop and encoder backbone investigated in Section 4, we further perform detailed ablations on more image augmentation, learning rate, encoder architecture, and activation function in this section.
The default test environment is identical to ablations in Section 4.
A.3.1 Image Augmentation
We study the effect of random translate, an image augmentation method which is widely used in
RAD [45].
Similar to random crop, the image size for translate is linear to the magnitude of the translate, when using a fixed crop size: the larger the image size, the stronger the augmentation.
As shown in Fig.
11, image size for translate has a similar pattern for most tested methods (with an exception of RotationCLS).
In summary, it is critical to engineering image augmentation carefully when designing an RL system with or without SSL.
A.3.2 Learning Rate
Fig.
12 shows how the learning rate of self-supervised loss contributes to the performance.
In this group of ablations, we only change the learning rate for SSL and leave the RL part untouched.
SAC-NoAug and SAC-Aug(100) are both baselines for reference without any self-supervised losses.
The results suggest that a smaller learning rate for SSL may improve performance.
Therefore, it is necessary to search for the absolute weights of losses like ELo [59], which is equivalent to searching for the learning rate.
A.3.3 Encoder Architecture
We further investigate the effect of additional linear layers in the visual encoder.
Additional linear layers with ReLU activation are appended to the end of the visual encoder.
All additional layers have a latent dimension of 128 .
Fig.
13 shows that additional layers usually bring downgraded performance, which could be caused by limited data.
Another important aspect of the encoder architecture design is how to merge or separate two branches.
As Visionary [3] suggests, where and how to merge visual representation with action representation is critical when designing an efficient value network.
Similarly, we hypothesize that the point where split the representation for the SSL branch and the RL branch is also important.
Figure 14 demonstrates two separation point configurations, named A and B. Figure 15 shows how the performance of different approaches changes in such two configurations.
A.3.4 Activation Function
Though all the methods we tested in Table 11 share the same visual encoder architecture, DrQ and
SAC+AE apply an additional activation function tanh to the visual representation.
To make a fair comparison and study the effect of such an activation function, we conducted detailed ablations on six DMControl tasks using the hyper-parameters identical to Table 2.
Results in Table 10 confirm that the default design choice of all three methods is better than their alternatives
[Table] [Table]
4 Detailed Results on DMControl and Atari
DMControl Table 11 notes the Interquartile mean, standard deviation and Relative Scores of tested algorithms in six DMControl environments.
The score distribution of tested algorithms over six environments is summarized as Figure 20.
Table 12 and Figure 21 include results of two additional harder environments in DMControl.
The figures of reward curve v.s. environment step are grouped as
Figure 23-32 and Figure 33-37 by the learning method.
Atari Similar to Table 11 and Figure 20, Table 13 Figure 22 are results in seven Atari environments.
A.5 Real-world Robot Experiments
To further evaluate methods in the real world applications, we set up a continuous robot arm control environment, uArm reacher.
With the help of some simple techniques in computer vision and robotics, our environment can autonomously randomly reset and keep the agent training without any human input.
The environment requires a robotic arm with a suction cup actuator, two fixed RGB cameras, and a cube that can be picked up by the suction cup as the target, as shown in Fig.
The goal is to move the actuator close to the target as fast as possible.
The observation comes from two cameras with a native resolution of 640 × 480 .
The images are then resized to 100 × 100 , stacked along channel axis, and finally randomly cropped into 84 × 84 , resulting in an 84 × 84 × (3 + 3) image observation before fed to the network.
The action space is a 3D vector ranging from -1 to 1, and it will be mapped to the actuator position movement in a 3D robot Cartesian coordinates whose original point is the center of the robot base.
The robot’s motion range is manually limited for safety concerns while avoiding the actuator moving the target in one episode.
Following reacher in DMControl, we define a very simple reward function.
The reward function returns 1 when the 3D Euclidean distance between the actuator and the target is lower than a threshold, otherwise, it returns − 1 e − 3 .
The length of each episode is set to 200 steps, which limits the range of the episode accumulated reward to [ − 0 . 2 , 200] .
To enable automatic reward generation, we make an automatic calibration framework to get the target location in 3D, and calibrate the top-down camera before any experiments.
We use AprilTag [55] to locate the robot position in the image plane, and read the 3D robot coordinates directly from the robot.
By doing so, we can build a map between 2D image coordinates and 3D robot coordinates.
The 2D coordinates of the target is first extracted by a simple color threshold.
Then, given the constant
[Table] [Table]
193.1±186.
DrQ 0.424±1.
258.95±205.
S e l fs up .
CURL 0.076±0.
115.5±116.
BYOL 0.025±0.
49.725±126.
DINO 0.25±0.
200.333±178.
RotationCLS 0.031±0.
210.05±117.
SAC+AE 0.061±0.
140.567±185.
ELo-SAC 0.116±0.
81.592±53.
ELo-SACv 0.147±0.
152.858±70.
ELo-SACv 0.177±0.
98.208±78.
height of the target, we can obtain 3D target location from its 2D image coordinates according to the 2D ↔ 3D map.
The environmental reset process is also automatic.
At the beginning of each episode, the robot arm will pick up the target cube, and randomly release the cube at a certain height like throwing dice, in order to randomly initialize the cube location.
The new location of the target cube is saved for generating rewards.
After the robot arm automatically moves to a fixed pre-assigned starting point, the environmental reset is done and then the RL agent takes over the control.
The RL agent can perform regular online training until the episode ends.
Finally, after each episode ends, the environment repeats the reset process to initialize the next episode.
[Table]
-19.2±2.
A.6 Empirical Analysis on the Learned Representations
To further understand the role of self-supervised loss and image augmentation in an online reinforcement learning system with the joint learning framework, we empirically show the properties of representations learned by different losses.
We first follow Wang et al.
[75] and measure the three metrics Dynamic Awareness, Diversity, and
Orthogonality, extending them from discrete action space to continuous action space.
Dynamics Awareness means two states that are adjacent in time should have similar representations and states further apart should have a low similarity.
Diversity measures a ratio between state and state-value differences.
High diversity means two states have two different representations to be distinguished even when they have similar state values.
Orthogonality reflects the linear independence of the representation, in another word, the higher the orthogonality, the lower the redundancy in the representations.
Assume an image observation x i is taken when the intrinsic system state is s i .
Denoting the visual representation of x i generated by the encoder from the critic networks as ϕ i , and Critic ( ϕ i , · ) is the learned critic network output.
Eq.
9 shows how to compute the three representation metrics.
[Equation]
where N is the total number of samples, U (1 , N ) means uniformly sample
from [
1 , N
, d s,i,j
=
∥ ϕ i − ϕ
j
∥ and d v,i,j = | max a Critic ( ϕ i , a ) − max a Critic
ϕ j a ) | . Predict State from Visual Representation Besides the three metrics on visual representations and state-values, we further measure the quality of visual representation ϕ
i by predicting the system state s i only using ϕ i .
The intuition is that a better visual representation should be able to capture the intrinsic system state more precisely.
We utilize a two-layer MLP to regress the system state s i on its corresponding visual representation
ϕ
i .
Mean squared error is applied to supervise the network as well as to evaluate the network on the test set.
To properly measure all these metrics, We first collect a dataset in cartpole swingup from DMControl using state-based SAC, which is different from any methods we’ll benchmark to avoid bias.
We run state-based SAC with five random seeds, and take the replay buffer of each run to form a dataset.
The whole dataset has 12500 × 5 = 62500 state transitions.
We measure Dynamics Awareness and
Orthogonality on the full dataset, while Diversity is calculated for one run due to computational cost.
For state prediction, we use the first four runs as the training set and the last run is held for the evaluation.
Finally, we benchmark selected methods with five different random seeds on cartpole swingup, and reports the metrics above every 100 model update step.
We demonstrate how the four metrics correlate to the environment step and agent performance as Fig.
16 and Fig.
Fig.
16 shows how metrics change as training.
Most of the methods converge on a similar Orthogonality, Dynamic Awareness, and Diversity value.
SAC-NoAug has a low Dynamic Awareness measure which could be used to explain its low performance.
While a higher Dynamic Awareness measure does not bring extra scores for BYOL and SAC+AE.
Similarly, the lower Diversity value of DrQ and
ELo-SAC do not hurt their performance either.
Meanwhile, most of the metrics become relatively stable after the first 4000 steps.
Therefore, we confirm that the shallower layers of the neural networks in visual reinforcement learning converge faster as observed by Chen et al.
Fig.
17 shows the correlation between metrics and the agent performance.
We report the Pearson correlation coefficient as Table 14.
As Wang et al.
[75] suggested, these metrics only measure certain properties of the visual representation, and they do not suggest that a property is necessary for better policy learning.
However, we find that the state prediction error is correlated to the agent performance to some extent, which may be valuable in some cases.
[Table]
A.7 Observation on Pretraining Framework
Besides the joint learning framework used in CURL and SAC+AE, Shelhamer et al.
[69] investigate a pretraining framework to combine SSL with RL and use the self-supervised loss as an intrinsic reward to further boost performance during online learning.
Recent works on policy learning (e.g., [58, 67, 74, 80, 86]) also take advantage of the self-supervised learning in a multi-step framework and show its great potential in solving challenging visual-based problems.
This pretraining framework is similar to how self-supervision has been benefiting supervised Computer Vision tasks ([7, 9, 11, 16, 33, 39, 59]): pretrain with self-supervised losses, and then finetune with the downstream task loss.
Motivated by them, in this section, we design and benchmark the two-stage pretraining framework, replacing the joint learning framework used in CURL and SAC+AE.
In the first stage, we use data collected by training a SAC-Aug(100) agent on the same task and update the visual encoder only using self-supervised loss.
We name this stage pretraining which means to use self-supervised losses to update the model and to be downstream task agnostic.
Then in the second stage i.e., the online training stage, we only keep the trained encoder from the first stage and train an agent using SAC-Aug(100).
The only difference between this stage and training an agent from scratch is that here the visual encoder has been “initialized” with the pre-trained weights while it is randomly initialized in SAC-Aug(100).
This also means that the image encoder can be tuned by
RL loss in the online training stage to match the online sample distribution.
Fig.
18 compares two training frameworks, in which the rounded rectangle means to update the model with the labeled loss for one step.
The methods using the pretraining framework have the prefix ‘Pretrain’.
‘Pretrain-Random’ means the data used for pretraining is collected by a random policy.
In both cases, the pretraining framework has the same model update steps as the joint learning framework baseline.
But note that the pretraining model has access to extra data collected by other policies, which makes it an unfair comparison.
To this end, we test another joint learning configuration named with the prefix ‘Longer’.
Here we match the total number of environment steps (or collected data) to its pretraining variants.
Similarly, three methods without any self-supervised learning are benchmarked with ’Longer’ configuration.
We compare two frameworks in six DMControl environments, Relative Scores are reported as Fig.and the full results are shown as Table 15.
In general, given the same amount of model updates, the pretraining framework performs better than the joint learning framework except ELo-SAC (we believe this is because ELo-SAC search was done only under the joint learning framework).
But such an advantage of the pretraining framework may come from the extra data used in the pretraining stage.
When the same amount of data is given, the longer joint-learning configuration usually performs better than the pretraining methods except when
AutoEncoder is the self-supervised loss.
Such observations imply that the learning framework has different impacts on policy learning even if the same self-supervised loss is applied.
It might not be the best practice to directly use the existing self-supervised losses designed for joint learning framework with the pretraining framework.
However, only Longer-ELo-SAC achieves comparable results compared to image augmentation based methods with ‘Longer’ configuration.
On the other hand, we argue that on DMControl, the advantages of the pretraining framework come from the access to extra data instead of the framework itself.
When the total environment step is limited and no previous data has been collected, the joint learning framework can better solve DMControl problems.
A.8 Limitations
In this paper, we focus on SAC for environments with continuous action space and Rainbow for environments with discrete action space.
Though both methods are generic, it will be interesting to see how self-supervised losses work with other RL methods and image augmentations in more challenging environments.
Meanwhile, both RAD [45] and DrQ [81] investigate many image augmentation approaches for their learning methods.
We only focus on random crop and translate because of their positive effects, and more combinations of image augmentation and self-supervised learning methods worth further investigation.
In addition, the search space of ELo-based methods are relatively limited.
They may achieve better scores with a larger search space (more losses) and a more representative searching environment.
A.9 Computation Information
Training one DMControl agent for 50k model update steps usually takes 3 hours on one NVIDIA
A5000 GPU.
It takes around 1.5 hours to train five Atari agents in parallel using an Apple M1 Max
CPU.
[Table]