S ELF -S UPERVISED S PATIAL R EPRESENTATIONS FOR D EEP R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Recent reinforcement learning (RL) methods have found extracting high-level features from raw pixels with self-supervised learning to be effective in learning policies.
However, these methods focus on learning global representations of images, and disregard local spatial structures present in the consecutively stacked frames.
In this paper, we propose a novel approach that learns self-supervised spatial representations ( S 3 R) for effectively encoding such spatial structures in an unsupervised manner.
Given the input frames, the spatial latent volumes are ﬁrst generated individually using an encoder, and they are used to capture the change in terms of spatial structures, i.e. , ﬂow maps among multiple frames.
To be speciﬁc, the proposed method establishes ﬂow vectors between two latent volumes via a supervision by the image reconstruction loss.
This enables for providing plenty of local samples for training the encoder of deep RL.
We further attempt to leverage the spatial representations in the self-predictive representations (SPR) method that predicts future representations using the action-conditioned transition model.
The proposed method imposes similarity constraints on the three latent volumes; warped query representations by estimated ﬂows, predicted target representations from the transition model, and target representations of future state.
Experimental results on complex tasks in Atari Games and DeepMind Control Suite demonstrate that the RL methods are signiﬁcantly boosted by the proposed self-supervised learning of spatial representations.
The code is available at https://sites. google.com/view/iclr2022-s3r .
I NTRODUCTION
Deep reinforcement learning (RL) has been an appealing tool for training agents to solve various tasks including complex control and video games (François-Lavet et al., 2018).
While most approaches have focused on training deep RL agent under the assumption that compact state representations are readily available, this assumption does not hold in the cases where raw visual observations ( e.g. images) are used as inputs for training the deep RL agent.
Without designing an effective algorithm that uses model-based policy and value improvement operators (Schrittwieser et al., 2021), or without attempting to use additional image augmentation (Laskin et al., 2020a; Kostrikov et al., 2020), learning visual features from raw pixels only using a reward function might fail to learn good features in terms of the performance and sample efﬁciency.
To address this challenge, a number of deep RL approaches (Sermanet et al., 2018; Dwibedi et al.,
; Anand et al., 2019; Laskin et al., 2020b; Mazoure et al., 2020; Stooke et al., 2020; Schwarzer [Figure] et al., 2021) leverage the recent advance of self-supervised learning which effectively extracts highlevel features from raw pixels in an unsupervised fashion.
In Laskin et al.
(2020b)
; Stooke et al.
(2020), they propose to train the convolutional encoder for pairs of images using a contrastive loss (van den Oord et al., 2018).
For training the RL agent, given a query and a set of keys consisting of positive and negative samples, they minimize the contrastive loss such that the query matches with the positive sample more than any of the negative samples (Laskin et al., 2020b; Stooke et al., 2020).
While the parameters of the query encoder are updated through back-propagation using the contrastive loss (van den Oord et al., 2018), the parameters of the key encoder are computed with an exponential moving average (EMA) of the query encoder parameters.
The output representations of the query encoder are passed to the RL algorithm for training the agent.
Schwarzer et al.
(2021) proposes the self-predictive representations (SPR) method that trains the RL agent by leveraging the query and positive sample only, following Grill et al.
(2020) that achieves state-of-the-arts performance without the negative samples in the self-supervised learning.
Especially, it extends a model-free RL agent by adding the transition model that predicts its own latent representation multiple steps into the future (Schwarzer et al., 2021).
The predicted future representations and the representations for future states computed using a target encoder serve as the query and positive sample, respectively.
These approaches have shown compelling performance and high sample efﬁciency on the complex control tasks when compared to existing image-based RL approaches (Kaiser et al., 2019; van Hasselt et al., 2019; Kielak, 2020).
While these approaches can effectively encode the global representations of images with the selfsupervised representation learning, there has been no attention on the local spatial structures present in the consecutively stacked images.
Our key observation is that spatial deformation, i.e. , the change in terms of the spatial structures across the consecutive frames, can provide plenty of local samples for training the RL agent.
A two-frame ﬂow estimation (Dosovitskiy et al., 2015; Ilg et al., 2017;
Jonschkowski et al., 2020), which has been widely used for video processing and recognition in computer vision, can be an appropriate tool in modeling the spatial deformation.
In this work, we propose a novel approach, termed self-supervised spatial representations ( S 3 R), that learns spatial representations for effectively encoding the spatial structures in a self-supervised fashion.
Note that we use the term ‘spatial representations’ to indicate a set of feature maps extracted over frames for inferring locally-varying ﬂow maps.
The spatial representations generated from an encoder are used to predict the ﬂow maps among the input frames by minimizing an image reconstruction loss in a self-supervised manner.
A ﬂow-based warping is then applied to generate future representations.
We further extend our framework by leveraging SPR method (Schwarzer et al., 2021).
As depicted in Figure 1, we impose similarity constraints on the three representations; query representations warped by the estimated ﬂow maps, self-predicted target representations from the transition model (Schwarzer et al., 2021), and target representations of future state.
Note that Shang et al.
(2021) encodes the temporal information by concatenating latent differences of input frames, but the simple substraction operation has certain limitations in effectively capturing the spatial deformation.
Contrary to Amiranashvili et al.
(2018) in which the ﬂow map is directly fed into RL algorithms with a stack of images, our method leverages the ﬂow maps to effectively encode the spatial representations of the images and warp spatial representations at the future state.
Our contribution is summarized as follows.
• Our method learns the spatial representations using the self-supervised ﬂow model for encoding local spatial structures from the consecutive frames used in RL algorithms.
• We propose to impose the similarity constraint on the spatial representations as well as the global representations, providing plenty of supervision for training the encoder of deep RL.
• We compute the future representations through the ﬂow-based warping for imposing the similarity constraint with target representations.
R ELATED W ORK
Self-supervised Representation Learning : The self-supervised representation learning aims to learn general features from large-scale unlabeled images or videos without expensive data annotations.
The contrastive methods have achieved state-of-the-art performance in the self-supervised representation learning.
The contrastive learning aims to bring positive samples closer while sep- arating negative samples from each other (Hadsell et al., 2006).
Wu et al.
(2018) formulate the contrastive learning as a non-parametric classiﬁcation problem at the instance level, and propose to learn visual features with the memory bank and noise contrastive estimation (NCE) (Gutmann &
Hyvärinen, 2010; Mnih & Kavukcuoglu, 2013).
The method in van den Oord et al.
(2018) proposes a probabilistic contrastive loss, called InfoNCE, for inducing representations by leveraging positive and negative samples.
The InfoNCE loss has widely been adopted in Chen et al.
; He et al.
(2020); Hénaff et al.
; Tian et al.
Chen et al.
(2020) present a simple framework for contrastive self-supervised learning without specialized architecture (Bachman et al., 2019; Hénaff et al., 2020) or memory bank (Wu et al., 2018), but it requires a large batch size for using enough negative samples when computing the InfoNCE loss (van den Oord et al., 2018).
He et al.
(2020) propose to build a dynamic dictionary with a queue to avoid the use of large batches when collecting negative samples, and also uses the moving averaged (momentum) encoder for target data (positive and negative samples of query data).
Grill et al.
(2020) use the momentum encoder to produce representations of the targets as a means of stabilizing the bootstrap step.
This enables for learning the representations with only positive samples, which are generated by data augmentation, for a given query without the need to carefully set up negative samples.
The method in Chen & He (2020) further extends this idea by using only stop-gradient operation without using the momentum update.
While these approaches focuses on learning global representations of a single image, our method proposes to learn spatial representations for effectively encoding the spatial structures ( i.e. , ﬂow map) in the consecutive images.
Self-supervised Representation Learning in Deep RL :
Representation learning is crucial for RL algorithms to learn policies with high-dimensional visual observations.
The future prediction conditioned on the past observations and actions serves as auxiliary tasks to improve the sample efﬁciency of model free RL algorithms.
Gelada et al.
(2019) train a transition model to predict representations of future states together with a reward prediction loss.
Guo et al.
(2020) present Predictions of Bootstrapped Latents (PBL) that builds on multi-step predictive representations of future observations for deep RL.
The method in Schwarzer et al.
propose Self-Predictive Representations (SPR) based on an action-conditioned transition model that can predict future representations computed using a target (momentum) encoder.
Our method attempts to predict the future representations through a warping operation using ﬂow maps computed from the spatial representations of the consecutive frames.
Contrastive learning has been used to extract desired latent representations of visual observations used in the RL algorithms.
For training robot agents, Sermanet et al.
(2018) present the time-contrastive networks (TCN) that train viewpoint-invariant representations using a metric learning.
This work was extended in Dwibedi et al.
by embedding multiple frames at each timestep for learning task-agnostic representations such as position and velocity attributes in continuous control tasks.
In Anand et al.
(2019), the representations for RL algorithms are learned by maximizing mutual information (Hjelm et al., 2019) across spatially and temporally distinct features of an encoder of visual observations.
Schwarzer et al.
(2021) leverage the self-supervised learning (Grill et al., 2020) for imposing the similarity constraint between self-predictive and target representations.
Laskin et al.
(2020b) introduce Contrastive Unsupervised representations for Reinforcement Learning (CURL) that learns the representations from visual inputs using the InfoNCE loss (van den Oord et al., 2018).
Stooke et al.
(2020) present Augmented Temporal Contrast (ATC) using image augmentations and
InfoNCE loss (van den Oord et al., 2018) for representation learning, and decouples it from policy learning.
From a different perspective, Hansen et al.
(2020) propose to adapt the policy network through self-supervised representation learning in unseen environments where it is difﬁcult to predict changed rewards.
Our method imposes the similarity constraint on the spatial representations as well as the global representations, thus providing plenty of supervision for training the encoder of deep RL.
Visual Correspondence Learning : Visual correspondence estimation is a long-standing research in the computer vision community.
It aims to establish a pair of corresponding pixels between two (or more) views taken under different locations (stereo matching) or timestep (optical ﬂow).
Recent methods for stereo matching (Žbontar & LeCun, 2016; Chang & Chen, 2018; Zhang et al., 2019) and optical ﬂow estimation (Dosovitskiy et al., 2015; Ilg et al., 2017; Sun et al., 2018) have been advanced largely thanks to the expressive power of deep networks.
Though both approaches share a similar objective of ﬁnding corresponding pixels across views, the optical ﬂow is known to be effective for encoding temporal motion trajectories, while the stereo matching is tailored to predicting 3D depth map in the scene.
The commonly used architecture for two-frame optical ﬂow estimation involves the feature map extraction of two frames, correlation volume computation, a series of convolutions for reﬁnement, and ﬂow regression.
While state-of-the-arts ﬂow estimation methods require using ground truth
ﬂow maps as an explicit supervision (Dosovitskiy et al., 2015;
Ilg et al., 2017; Sun et al., 2018), some unsupervised learning approaches have attempted to infer ﬂow maps with an image reconstruction loss for imposing the constraint that corresponding pixels should have similar intensities (Ren et al., 2017;
Meister et al., 2018; Wang et al., 2018; Jonschkowski et al., 2020).
In our work, we present the self-supervised ﬂow network that learns spatial representations from the consecutive frames used in the RL algorithms.
M ETHOD
We consider the Markov Decision Process (MDP) setting where an agent interacts with environments in a sequence of observations, actions, and rewards.
We denote o k , a k , and r k as the observation, the action of the agent, and the reward received at timestep k .
Since our method is a general framework that leverages the representation learning for training the RL agent, it can be combined with any RL algorithm.
Following the state-of-the-arts RL approaches (Laskin et al., 2020b; Stooke et al., 2020;
Schwarzer et al., 2021) using the self-supervised learning, we adopt the Soft Actor Critic (SAC) method (Haarnoja et al., 2018a) for continuous control task in DeepMind Control Suite benchmark, and Rainbow DQN (van Hasselt et al., 2019) for discrete control task in Atari Games.
The proposed self-supervised spatial representations ( S 3 R) are used as an auxiliary task for training RL agents.
3. S ELF - SUPERVISED F LOW M ODEL FOR S PATIAL R EPRESENTATION
We start with how to generate the spatial representations for capturing spatial deformations from the consecutively stacked frames in a self-supervised manner.
An instance used by the model-free off-policy RL algorithms (Haarnoja et al., 2018a; van Hasselt et al., 2019) is a stack of images, not a single image.
Given an input raw observation o k
=
I k , ..., I k + M } where I k is an image at timestep k , the spatial encoder features e k =
z k , ..., z k + M } are ﬁrst generated by applying an encoder individually to each of the input observations o k .
Note that z ∈ R h × w ×
d is a 3-D volume with a spatial resolution h × w and a feature dimension d .
We apply query encoder and target encoder to o k and o k +1 , respectively, and denote the output of the query encoder e q as z q , and the output of the target encoder e t as z t .While
the existing methods (Sermanet et al., 2018; Dwibedi et al., 2018;
Laskin et al., 2020b; Schwarzer et al., 2021) feeds the stacked frames to the encoder at once, which can be viewed as an early fusion (Karpathy et al., 2014), our method generates the set of the spatial latent representations individually with the encoder.
Later, they are fused using 1 × 1 convolutional layer in a manner similar to late fusion (Simonyan & Zisserman, 2014) (see Figure 2).
A similar strategy was also adopted in Shang et al.
(2021) for encoding the temporal information when training the RL agent.
The set of spatial representations is used to predict the spatial deformations, i.e. , ﬂow maps between two consecutive frames via a supervision by the image reconstruction loss (Godard et al., 2017; 2019; Wang et al., 2020).
We compute a correlation volume V a,b ∈ R h × w × r 2 using a dot product between two latent representations z a and z b (Dosovitskiy et al., 2015) as follows:
[Equation]
where p and q represent 2D feature position in z a and z b , and ¯ r indicates the kernal size for computing correlation, r = 2¯ r + 1 .
Computing the patch similarity in (1) for all combinations of p and q (totally, h 2 · w 2 times) yields huge amount of computations.
Thus, the maximum displacement for computing the patch similarity is limited for q ∈ N ( p ) where N ( p ) represents a square window of size rcentered at p , following visual correspondence literature (Dosovitskiy et al., 2015; Ilg et al., 2017;
Sun et al., 2018).
The correlation volume is fed into a series of convolutions followed by the reﬁnement module, producing a ﬂow map f a → b ∈ R h × w × 2 from I a to I b .
We follow the architecture of ‘FlowNetCorr’ in
Dosovitskiy et al.
(2015) with some modiﬁcations in the spatial resolution, encoder feature dimension and search range r 2 in order to keep the model complexity low.
Note that while ‘FlowNetCorr’ is trained with the supervision by ground truth ﬂow maps, our method relies on the image reconstruction loss for self-supervised learning such that
[Equation]
where I ( p ) indicates an intensity at the pixel corresponding to 2D feature position p .
For computing the loss L r , we resize I a
and I b to the size of the spatial representations, h × w .
We additionally use the Charbonnier regularization loss L reg (Barron, 2019) for producing spatially smooth ﬂow maps.
In Figure 2, we denote F ( z a , z b ) =
f a → b as the self-supervised ﬂow estimation network including the correlation volume computation, the series of convolutions, and the reﬁnement module.
3. R EPRESENTATION L EARNING WITH W ARPING AND S ELF -P REDICTION
Figure 2 illustrates the overall architecture of the proposed S 3 R approach.
Following the prior work on the self-supervised learning (He et al., 2020; Grill et al., 2020; Chen & He, 2020), we use the query encoder E q with the parameters θ q and the target encoder E t with the parameters θ t for encoding the query observation o k and the target observation o k +1 .
While the parameters θ q of the query encoder are updated through back-propagation, the parameters θ t of the target encoder are updated with the query encoder parameters θ q using a stop-gradient operation (Chen & He, 2020) as θ t ← θ q .
Flow Learning and Warping :
By minimizing (2), we ﬁrst compute a set of M + 1 external ﬂow maps { f extk + i +
→ k + i | i = 0 , ..., m } with the self-supervised ﬂow network F such that
[Equation]
Note that the external ﬂow map is predicted from the target feature z
tk
+ i +1 to the query feature z
qk
+ i .
Then, we warp the query features e qk = { z qk , ..., z qk + M } into the future state via the inverse warping (Jaderberg et al., 2015) using M + 1 external ﬂow maps.
The warped query features { z wk +1 , ..., z wk + M +1 } are then fused using 1 × 1 convolution, producing the warped query representation Z wk +1 at the timestep k + 1 (see Figure 2).
We can also predict internal ﬂow maps within the query features e qk = { z qk , ..., z qk + M } as f inta → b =
F ( z qa , z qb ) .
Various combinations of a and b are possible for computing the internal ﬂow maps, and we choose to compute a single ﬂow map f intk → k + M .
We found that this is an appropriate choice in terms of computational efﬁciency and performance as the external ﬂow maps are already used to impose the structural similarity constraint between multiple frames, and is effective in dealing with the case where the external ﬂow between two consecutive frames is relatively small.
More details are presented in the Appendix.
The loss function L f for computing the internal and external ﬂow maps is given as
[Equation]
Representation Learning with Global-level and Pixel-level Similarities : To measure the globallevel similarity between the warped query representation
Z wk +1 and the target representation
Z tk + which is the fusion of target encoder features e tk +
= { z tk +1 , ..., z tk + M +1 } , we use two projection heads and one predictor in a manner similar to Grill et al.
; Chen & He (2020).
We project the two representations
Z wk +1 and Z tk +1 into a smaller latent space by passing them into the query projection head P q with parameters ξ q and the target projection head P t with parameters ξ t , and also apply an additional query prediction head Q q to the query projection.
The target projection head parameters ξ t are updated using the query projection head parameters ξ q with the stop-gradient operation as in the target encoder, i.e. , ξ t ← ξ q .
The prediction loss is computed using the cosine similarity between the warped query representation y wk
+
=
Q q ( P q ( Z wk +1 )) and the observed target representation y tk
+1 =
P t ( Z tk +1 ) , which is deﬁned as
[Equation]
In Figure 2, we denote the module consisting of the query projection and prediction heads and the target projection heads as ‘ Sim ( Z a , Z b ) ’ where Z a and Z b can be the warped representation from query and spatial representations from target network, respectively.
We further extend our method by leveraging a self-prediction module conditioned on an action.
The action-conditioned convolutional transition model of Schwarzer et al.
(2021) is applied in our framework with some modiﬁcations.
We generate the query representation Z qk by applying 1 ×convolution to the query features { z qk , ..., z qk + M } and then feed it into the convolutional transition model H .
Note that we use only a single next prediction Z {+p k +
+}
= H ( Z q k , a k ) of the transition from the query representation
Z qk , unlike Schwarzer et al.
(2021) that recursively generates a sequence of L predictions Z pk +1:
k + L with Z pk + l
+
=
H ( Z pk + l , a k + l ) for l = 0 , ..., L − 1 and Z pk
=
Z qk .
We found that since the spatial representation learning using the ﬂow estimation and warping provide a sufﬁcient amount of supervision, the single self-prediction is enough to train the query encoder.
The self-prediction Z pk +1 is fed into the query projection head P q and the query prediction head Q q such that y pk {++
+}
= Q q ( P q ( Z pk +1 )) .
The prediction loss is also computed using the cosine similarity loss L s ( y pk +1 , y tk +1 ) .
The loss function L sim for measuring the similarity between the three representations y wk +1 , y pk +1 , and y tk +1 is summarized as
[Equation]
where y wk +
=
Q q ( P q ( Z wk +1 )) , y pk
+1 =
Q q ( P q ( Z pk +1 )) , y tk +
= P t ( Z tk +1 ) , and λ is a hyperparameter of loss functions.
We also include L 1 loss
L 1 ( Z pk +1 ,
Z tk +1 ) on the original spatial latent space to impose the pixel-level similarity.
Note that L 1 loss between Z wk +1 and Z tk +1 is not used, as the external ﬂow loss L r ( f ext ) implicitly considers the similarity between the two.
Finally, the query representation Z qk ∈ R h × w × c is fed into the deep RL algorithm, e.g. , c = 64 for
DQN (van Hasselt et al., 2019) and c = 32 for SAC (Haarnoja et al., 2018a).
Loss functions : The ﬁnal loss function is summarized as
[Equation]
(7) where L RL ( Z qk ) indicates the loss of the RL algorithms which use Z qk as an input.
α and β are hyper-parameters that balance the loss functions.
We summarize the overall method in Algorithm 1.
3. I MPLEMENTATION D ETAILS
Self-supervised Flow Model : The input image
I
i is of 84 × 84 for Atari Games and DeepMind
Control (DMControl) Suites.
The query and target encoders generates z qi , z ti +1 ∈ R 7 × 7 × 64 ( i = k, ..., k + 3 ) for Atari Games and z qi , z ti +1 ∈ R 32 × 32 × 32 ( i = k, ..., k + 2 ) for DMControl Suites, respectively.
The search window for computing the correlation volume V is 6 × 6 for Atari games and DMControl Suites.
The correlation volume goes through 3 × 3 convolution layers 3 times.
The decoder is then applied to provide a dense ﬂow map.
The decoder includes three un-convolutional layers, consisting of un-pooling and convolution, and the coarser ﬂow maps and encoder feature maps are concatenated into each un-convolutional layer (Dosovitskiy et al., 2015).
Action Conditioned Transition Model : This basically follows the structure of Schwarzer et al.
, but in order to compute the structural similarity with the features of each frame, we set the input and output channel size to 256, which was originally set to 64 in Schwarzer et al.
The transition model includes two convolutional layers interweaved with ReLU and batch normalization (Ioffe & Szegedy, 2015), with the current representations Z qk
and the action a k of one-hot vector taken to each location being fed to the ﬁrst convolutional layer.
Other Details : The query and target projection heads, P q and P t , are implemented as the multi-layer perceptron (MLP).
For the query prediction head Q q , we reuse the ﬁrst linear layer of the RL head similar to Schwarzer et al.
We used λ = 0 .
1 in (6) and α = 5 , β = 1 in (7) to balance the weight of the losses.
More details are presented in the Appendix.
E XPERIMENTAL R ESULTS 4. E VALUATION ON A TARI G AMES
To compare the performance of the proposed method with state-of-the-arts, we chose Atari
Games introduced in Kaiser et al.
; van Hasselt et al.
(2019) where only 100K environment
[Table]
time steps.
We compared results with state-based SAC and pixel-based
SAC (Haarnoja et al., 2018b), SAC+AE (Yarats et al., 2019), Dreamer (Hafner et al., 2019a), PlaNet (Hafner et al., 2019b), CURL (Laskin et al., 2020b), RAD (Laskin et al., 2020a)
[Table]
and DrQ (Kostrikov et al., 2020)
[Table]
steps, corresponding to two hours of gameplay experiences, are available for training data. This sample-efﬁcient setup, which uses much less environment steps than the standard setup of 50,000K environment steps, has been adopted for evaluating the performance of recent sample-efﬁcient deep RL algorithms (Kaiser et al., 2019; van Hasselt et al., 2019; Kielak, 2020; Laskin et al., 2020b;
[Table]
W
P
[Equation]
R
[Table]
Schwarzer et al., 2021). We compared our results with various RL algorithms including SimPLe (Kaiser et al., 2019) which learns a pixel-level transition model for Atari, Data-Efﬁcient Rainbow (DER) (van Hasselt et al., 2019) which modiﬁes the Rainbow hyperparameters for improving the sample efﬁciency, OTRainbow (Kielak, 2020) which is an over-trained version of the Rainbow for the sample efﬁciency, CURL (Laskin et al., 2020b) which proposes the use of image augmentation with the contrastive loss (van den Oord et al., 2018) for self-supervised representation learning, DrQ (Kostrikov et al., 2020) which uses the modest image augmentation to improve the sample efﬁciency, and SPR (Schwarzer et al., 2021) which trains an agent to predict its own latent state representations into the future.
Following the experimental setup on the above-mentioned approaches, we evaluated on 26 environments of Atari 2600 games by measuring the average return after 100K interaction steps.
As shown in Table 1, the proposed method ( S 3 R) achieved the best performance on 13 out of 26 environments.
We trained our method with 10 random seeds, similar to other methods.
4. E VALUATION ON DMC ONTROL S UITE
Various approaches including ours have been benchmarked on the DMControl Suite where the agent operates from pixels to evaluate challenging visual continuous control tasks (Tassa et al., 2018).
We compared our results with State-SAC which supposes that the agent has access to low-level state based features, Pixel-SAC (Haarnoja et al., 2018b) which directly operates from pixels, SAC+AE (Yarats et al., 2019) which uses a joint learning of SAC with β -VAE (Higgins et al., 2017), VAE (Kingma & Welling, 2013), and regularized autoencoder (Vincent et al., 2008), Dreamer (Hafner et al., 2019a) and PlaNet (Hafner et al., 2019b) which learn a latent space world model, CURL (Laskin et al., 2020b) which uses image augmentation with the contrastive loss (van den Oord et al., 2018), RAD (Laskin et al., 2020a) and DrQ (Kostrikov et al., 2020) which demonstrate that data augmentation can greatly improve the performance of model-free RL algorithms and achieve stateof-the-art performance on DMControl Suite.
Table 2 demonstrates that the self-supervised spatial representations of S 3 R achieved best performance on 4 out of 6 environments for 500K time steps including Cartpole Swingup, Reacher Easy, Walker Walk and Ball in Cup Catch.
We trained our method with 10 random seeds, and the results with 5 random seeds are provided in the Appendix.
4. A BLATION S TUDY
Table 3 measured the average performance over 10 random seeds according to the combinations of several losses on DMControl Suite (Tassa et al., 2018) with 500K time steps.
• ‘F’ using only the ﬂow loss in (4) • ‘F+W’ using ﬂow loss and similarity loss with ‘W’arped query and target representations in (6) • ‘P’ using prediction loss with self-‘P’redicted and target representations in (6) • ‘F+P’ using the ﬂow loss, L1 loss, and the prediction loss in (6).
The network trained with only ‘F’ produces worse performance compared to ‘F+W’, ‘F+P’ and ‘F+W+P’, but still produces comparable performance to state-of-the-arts, implying that even without the warping and transition model, simply guiding the encoder to extract features for the ﬂow prediction helps the RL agent to perform well enough.
The performance of ‘F+W’ and ‘F+P’ is similar, but ‘F+W’ has slightly better performance on step scores with smaller standard deviations.
This implies that the warped query representation to the future state using the estimated ﬂow is capable of providing as useful supervision as the self-predicted representation that uses an action-conditioned transition model to predict future representations.
The performance was further boosted, when using ‘F+W+P’ altogether ( = S 3 R).
To measure only the impact of each loss, data augmentation was not performed.
Note that ‘P’ is similar to SPR method (Schwarzer et al., 2021) using only one prediction.
C ONCLUSION
We have presented the self-supervised spatial representations, termed S 3 R to encode local spatial structures in an unsupervised manner.
The ﬂow maps inferred by the proposed method offer plenty of supervision for learning the spatial representations, and also compute warped predictions at future frame.
S 3 R achieves state-of-the-art performance on Atari benchmark with 100K steps and
DMControl Suites with 100K/500K steps.
We have shown the importance of learning the spatial representations in improving the performance and sample-efﬁciency of image-based RL algorithms.
We hope this can facilitate future works at various aspects for RL based on self-supervised learning.
R EPRODUCIBILITY S TATEMENT
For the reproducibility of the proposed method we opened our code at https://sites.google.
com/view/iclr2022-s3r as stated in the Abstract.
We also provide the full hyperparameters in Section A of Appendix, so others can use the code with the same settings we experimented with
A I MPLEMENTATION D ETAILS
This section provides tables summarizing the hyperparameters used for experiments on Atari Games (Kaiser et al., 2019) and DMControl Suite (Tassa et al., 2018) and the detailed description of our network architecture.
A. H YPERPARAMETERS ON A TARI G AMES AND DMC ONTROL S UITE
[Table]
For the reproducibility of our work
[Table]
we provide full hyperparameters for our experiments on Atari
Games (Kaiser et al., 2019) in Table 4.
We follow the common practices used to set up Rainbow
DQN van Hasselt et al.
in existing methods Laskin et al.
(2020b)
; Schwarzer et al.
(2021) for experimenting on Atari Games.
Table 5 shows a full list of hyperparameters for DMControl suites experiments.
We utilize the similar hyperparameters
[Table]
A. N ETWORK A RCHITECTURE
The detailed network architecture of our method is provided in Table 6.
We describe the architectures of the encoder, FlowNet and transition model.
‘N’, ‘K’ and ‘S’ of convolution operations represent the channel, kernel size, and stride, respectively.
‘LReLU’ and ‘BN’ indicate LeakyReLU
Maas et al.
(2013) and Batch Normalization Ioffe & Szegedy (2015).
We also provide input and output tensors of each layer for better understanding
[Table] [Table]
A. ENVIRONMENTS
S 3 R was implemented in Pytorch (Paszke et al., 2019) with the use of rlpyt (Stooke & Abbeel, 2019) and Mujoco (Todorov et al., 2012) license, and was simulated on 1 Titan RTX GPU. A. D ETAILS OF USING FLOW MAPS
Use of internal ﬂow : When four consecutive frames are used in Rainbow DQN ( M = 3 ), the external ﬂows are speciﬁed using query images o k
( I k ∼ I k +3 )
and target images o k
{++
+}
( I k +1 ∼ I k +4 ) as follows: • ﬂow ( z qk , z tk +1 ) • ﬂow ( z qk +1 , z tk +2 ) • ﬂow ( z qk +2 , z tk +3 ) • ﬂow ( z qk +3 , z tk +4 ). The internal ﬂow is computed using query images o k as
• ﬂow ( z qk , z qk +3 ).
The external ﬂows computed between consecutive two frames are used to warp the query representations.
The internal ﬂow is also predicted using the same self-supervised ﬂow network yet with two distant frames ( z qk → z qk +3 ).
We found that this is often effective in dealing with the case where the ﬂow between two consecutive frames is relatively small.
Namely, the internal ﬂow between k th and ( k + 3) th frames can be complementary when the external ﬂows are rather small and thus ﬂow learning becomes less effective.
The internal ﬂow is computed with a stack of images o k , while the external ﬂows are computed between o k and o k +1 .
This is why we name two ﬂow parts ‘internal’ ﬂow and ‘external’ ﬂow, respectively.
Note that the internal ﬂow is computed only between two distant frames ( k th and ( k + 3) th frames in o k ), considering that the external ﬂow is computed between two consecutive frames as above.
Image and representation warping :
The representation we use contains spatial information.
For instance, suppose the convolutional feature map z of the size 32 × 32 × 64 is generated from an input image I of the size 128 × 128 × 3 .
The feature map can have a lower spatial resolution (by the downsampling operator such as max-pooling or strided convolution) and a larger channel dimension than the input image.
The 64 dimensional vector at a spatial position of the feature map represents the information of corresponding positions of the input image.
Therefore, we can employ the ﬂow used for image warping for warping the feature map (query representations).
This kind of implementation has been commonly used in many computer vision tasks where an image alignment is required.
Visualization of ﬂow maps : For a better understanding of the proposed method, the examples of the ﬂow maps predicted by S 3 R are represented in Figure 3.
[Figure] B P ERFORMANCE CONSISTENCY EVALUATION We consider that the performance depends on the choice of the seed, so we measure the performance by using 10 random seeds and averaging results.
To prove that our performance improvement is consistent and not caused by the noise of the estimation, we observed the change in the average performance according to the number of the random seeds.
Figure 4 and 5 represent the quantitative evaluation on the Atari games (Kaiser et al., 2019) and DMControl Suite (Tassa et al., 2018)
according to the number of the random seed.
We show result of 5 random seeds and 10 random seeds to compare the performance.
In Figure 5, we additionally showed the standard deviation of the performance.
The bar graph represents the average of the performance, and the line graph represents the standard deviation of the performance.
From Figure 4 and 5, it can be seen that our results do not differ signiﬁcantly depending on the number of random seeds.
It can be seen that our performance is signiﬁcantly higher than the SOTA performance regardless of random seed, and is consistent.
C A BLATION STUDY ON D ATA A UGMENTATION
As described in RAD (Laskin et al., 2020a) and DrQ (Kostrikov et al., 2020), image augmentation can improve the data-efﬁciency and generalization of RL methods.
To study the impact of data augmentation when used with the proposed method, we measured the average performance overrandom seeds according to the data augmentation on DMControl Suite (Tassa et al., 2018).
In Table
[Table]
we evaluated the performance of the proposed method when used with crop and translation proposed in Laskin et al.
(2020a).
Slightly different from the result presented in Laskin et al.
(2020a)
[Table]
Cartpole Swingup and Reacher
Easy has the best performance when no augmentation was used, Finger Spin and Cheetach Run has the best performance for translation
[Table]
and Walker Walk and Ball in cup Catch has the best performance for crop.
Since S 3 R learns ﬂow in an end-to-end manner with RL algorithm, it is analyzed that the results are different from those of Laskin et al.
(2020a).
[Table]
, Walk
[Table] [Table]
D A DDITIONAL EVALUATION METRIC
In all experiments, the evaluation on Atari Games (Kaiser et al., 2019) was conducted by measuring the performance with 10 or more random seeds, following the previous studies including SPR (Schwarzer et al., 2021) and CURL (Laskin et al., 2020b).
Recently, Agarwal et al.
(2021) analyzes the problems related to statistical uncertainty in the existing evaluation method of Atari Games (Kaiser et al., 2019).
Accordingly, we add the results of a more complete evaluation using ‘Probability of Improvement’ proposed in Agarwal et al.
(2021) in Figure 6.
This evaluation method estimates how likely an algorithm improves upon another algorithm Agarwal et al.
For instance, ‘P(SimPLe>X)’ indicates the probability (written in the horizontal line) that ‘SimPLe’ is better than another algorithm, called ‘X’, which is listed in the vertical line (e.g.,
S3R, SPR,..., DER).
Namely, the probability that SimPLe is better than SPR, ‘P(SimPLe>SPR)’, has an average value of about 0.3.
In this context, in the six graphs above, the smaller the value of S 3 R, the higher the performance.
Also, in the graph of S 3 R below, ‘P(S3R > X)’, S 3 R ranks the highest as the remaining bars of algorithms are located to the right of the bar of S 3 R. From this analysis, it can be reconﬁrmed that
S 3 R has the most superior performance compared to the state-of-the-art methods on Atari Games (Kaiser et al., 2019) as mentioned in the result of Section 4.1.
E A NALYSIS ON COMPUTATIONAL COST AND PERFORMANCE IMPROVEMENT
Computational cost : The increase in the computational cost for training is unavoidable because S 3 R additionally leverage the ﬂow prediction and warping networks used in the vision task, but we found that the additional computational cost for training is not so signiﬁcant.
For training on DMControl
Suite (Tassa et al., 2018) up to 500K on the same GPU environment, the proposed method takes about 16 hours, whereas the state-of-the-art methods CURL (Laskin et al., 2020b) and SPR (Schwarzer et al., 2021) take about 10 hours and 13 hours, respectively.
Note that the original SPR paper did not provide the code implemented for DM Control Suite, so we conducted the experiments by modifying the original SPR code.
Additionally, the ﬂow prediction and warping networks are used only during training, and the inference process is implemented in the same manner as other methods.
Therefore, the inference time of our method is exactly the same as that of the state-of-the-arts methods (CURL (Laskin et al., 2020b), SPR (Schwarzer et al., 2021), DrQ (Kostrikov et al., 2020))
as long as the same encoder for query images is used.
We further analyzed the performance improvement by our method based on the two benchmarks used in Section 4.
26 Atari Games : Based on the evaluation method used in the existing Atari Games (Kaiser et al., 2019), CURL (Laskin et al., 2020b) recorded the highest mean in 7 games out of 26, and SPR (Schwarzer et al., 2021) recorded the highest mean in 11 games out of 26.
S 3 R has the highest mean in 13 games out of 26.
It can be interpreted that the performance increase of the S 3 R is not small by considering the quantitative aspects of these games.
Also, among the 13 games in which S 3 R has an edge, in particular, in 8 games (Alien, Assault, Gopher, Jamesbond, Krull, Kung Fu Master, Ms
Pacman, and Seaquest), S 3 R records a remarkably higher performance compared to other methods.
Even the performance of certain games is high enough to match that of humans.
This is because the proposed method of capturing the local spatial structure is able to derive an effective representation from the images of the speciﬁc Atari Games with various movements.
However, S 3 R may not be effective for some games.
In particular, S 3 R did not perform well in the task ‘Pong’ in Atari Games (Kaiser et al., 2019).
The biggest reason for this is that there are too few discriminative spatial structures available in the game images.
Therefore, we can be sure that our representation learning method, which effectively captures the spatial structure, will work particularly well for data with much more complex structural features.
In other words, while most of the simple methods suffer from training with data with complex structural features, S 3 R can be a good substitute for addressing this.
A method with relatively low-complexity may be preferred depending on the situation.
However, as mentioned above, it was shown from the Atari benchmark that it is a much better choice to use S 3 R in the tasks that need to capture complex structural features.
DMControl Suite : In the case of DMControl Suite (Tassa et al., 2018), the performance at 100K steps is usually based on when most methods do not converge.
In Table 2, the performance of
S3R recorded the highest mean in 2 tasks out of 6 tasks, and RAD (Laskin et al., 2020a) and DrQ (Kostrikov et al., 2020) also recorded the highest mean in 2 tasks out of 6 tasks, respectively.
It can be interpreted that RAD (Laskin et al., 2020a), DrQ (Kostrikov et al., 2020) and S 3 R are the three methods with the highest convergence speed.
In general, the performance at 500K steps after most methods converge is widely adopted for the evaluation.
In Table 2, S 3 R shows the highest performance in 4 out of 6 tasks compared to stateof-the-arts.
Considering that this performance is the average value obtained by running 10 random seeds, it is a credible assessment.
Also, when compared to the performance improvement rate of other methods, the performance increase of S3R is by no means small.