Year: 2021

Abstract:
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 (SR) for effectively encoding such spatial structures in an unsupervised manner. Given the input frames, the spatial latent volumes are first generated individually using an encoder, and they are used to capture the change in terms of spatial structures, i.e., flow maps among multiple frames. To be specific, the proposed method establishes flow 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 flows, 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 significantly boosted by the proposed self-supervised learning of spatial representations. The code is available at https://sites. google.com/view/iclr2022-s3r.

text: 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 (S3R) for effectively encoding such spatial structures in an unsupervised manner. Given the input frames, the spatial latent volumes are first generated individually using an encoder, and they are used to capture the change in terms of spatial structures, i.e., flow maps among multiple frames. To be specific, the proposed method establishes flow 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 flows, 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 significantly boosted by the proposed self-supervised learning of spatial representations. The code is available at https://sites. google.com/view/iclr2022-s3r.

heading: 1 INTRODUCTION
text: 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 efficiency.
To address this challenge, a number of deep RL approaches (Sermanet et al., 2018; Dwibedi et al., 2018; Anand et al., 2019; Laskin et al., 2020b; Mazoure et al., 2020; Stooke et al., 2020; Schwarzer
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 efficiency 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 flow 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 (S3R), 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 flow maps. The spatial representations generated from an encoder are used to predict the flow maps among the input frames by minimizing an image reconstruction loss in a self-supervised manner. A flow-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 flow 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 flow map is directly fed into RL algorithms with a stack of images, our method leverages the flow 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 flow 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 flow-based warping for imposing the similarity constraint with target representations.

heading: 2 RELATED WORK
text: 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 classification 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. (2020); He et al. (2020); Hénaff et al. (2020); Tian et al. (2020). 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., flow 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 efficiency 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. (2021) 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 flow 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. (2018) 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 difficult 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 flow). Recent methods for stereo matching (Žbontar & LeCun, 2016; Chang & Chen, 2018; Zhang et al., 2019) and optical flow 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 finding corresponding pixels across views, the optical flow 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 flow estimation involves the feature map extraction of two frames, correlation volume computation, a series of convolutions for refinement, and flow regression. While state-of-the-arts flow estimation methods require using ground truth flow 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 flow 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 flow network that learns spatial representations from the consecutive frames used in the RL algorithms.

heading: 3 METHOD
text: We consider the Markov Decision Process (MDP) setting where an agent interacts with environments in a sequence of observations, actions, and rewards. We denote ok, ak, and rk 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 (S3R) are used as an auxiliary task for training RL agents.

heading: 3.1 SELF-SUPERVISED FLOW MODEL FOR SPATIAL REPRESENTATION
text: 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 ok = {Ik, ..., Ik+M} where Ik is an image at timestep k, the spatial encoder features ek = {zk, ..., zk+M} are first generated by applying an encoder individually to each of the input observations ok. Note that z ∈ Rh×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 ok and ok+1, respectively, and denote the output of the query encoder eq as zq, and the output of the target encoder et as zt.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., flow 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 Va,b ∈ Rh×w×r 2
using a dot product between two latent representations za and zb (Dosovitskiy et al., 2015) as follows:
Va,b(p, q) = ∑
o∈[−r̄,r̄]×[−r̄,r̄]
< za(p+ o), zb(q + o) >, (1)
where p and q represent 2D feature position in za and zb, and r̄ indicates the kernal size for computing correlation, r = 2r̄+ 1. Computing the patch similarity in (1) for all combinations of p and q (totally, h2 ·w2 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 r2 centered 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 refinement module, producing a flow map fa→b ∈ Rh×w×2 from Ia to Ib. We follow the architecture of ‘FlowNetCorr’ in Dosovitskiy et al. (2015) with some modifications in the spatial resolution, encoder feature dimension and search range r2 in order to keep the model complexity low. Note that while ‘FlowNetCorr’ is trained with the supervision by ground truth flow maps, our method relies on the image reconstruction loss for self-supervised learning such that
Lr(fa→b) = ∑ p |Ia(p)− Ib(p+ fa→b)|+ Lreg, (2)
where I(p) indicates an intensity at the pixel corresponding to 2D feature position p. For computing the loss Lr, we resize Ia and Ib to the size of the spatial representations, h× w. We additionally use the Charbonnier regularization loss Lreg (Barron, 2019) for producing spatially smooth flow maps. In Figure 2, we denote F(za, zb) = fa→b as the self-supervised flow estimation network including the correlation volume computation, the series of convolutions, and the refinement module.

heading: 3.2 REPRESENTATION LEARNING WITH WARPING AND SELF-PREDICTION
text: Figure 2 illustrates the overall architecture of the proposed S3R 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 Eq with the parameters θq and the target encoder Et with the parameters θt for encoding the query observation ok and the target observation ok+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 first compute a set of M + 1 external flow maps {fextk+i+1→k+i|i = 0, ...,m} with the self-supervised flow network F such that
fextk+i+1→k+i = F(z t k+i+1, z q k+i) for i = 0, ...,M. (3)
Note that the external flow map is predicted from the target feature ztk+i+1 to the query feature zqk+i. Then, we warp the query features e q k = {z q k, ..., z q k+M} into the future state via the inverse warping (Jaderberg et al., 2015) using M + 1 external flow maps. The warped query features {zwk+1, ..., zwk+M+1} are then fused using 1× 1 convolution, producing the warped query representation Zwk+1 at the timestep k + 1 (see Figure 2).
We can also predict internal flow maps within the query features eqk = {z q k, ..., z q k+M} as f inta→b = F(zqa, z q b ). Various combinations of a and b are possible for computing the internal flow maps, and we choose to compute a single flow map f intk→k+M . We found that this is an appropriate choice in terms of computational efficiency and performance as the external flow maps are already used to impose the structural similarity constraint between multiple frames, and is effective in dealing with the case where the external flow between two consecutive frames is relatively small. More details are presented in the Appendix. The loss function Lf for computing the internal and external flow maps is given as
Lflow = Lr(f intk→k+M ) + M∑ i=0 Lr(fextk+i+1→k+i). (4)
Representation Learning with Global-level and Pixel-level Similarities: To measure the globallevel similarity between the warped query representation Zwk+1 and the target representation Z t k+1 which is the fusion of target encoder features etk+1 = {ztk+1, ..., ztk+M+1}, we use two projection heads and one predictor in a manner similar to Grill et al. (2020); Chen & He (2020). We project the two representations Zwk+1 and Z t k+1 into a smaller latent space by passing them into the query projection head Pq with parameters ξq and the target projection head Pt with parameters ξt, and also apply an additional query prediction head Qq 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 ywk+1 = Qq(Pq(Z w k+1)) and the observed target representation ytk+1 = Pt(Z t k+1), which is defined as
Ls(ywk+1, ytk+1) = − < ywk+1, y t k+1 >
‖ywk+1‖2‖ytk+1‖2 . (5)
In Figure 2, we denote the module consisting of the query projection and prediction heads and the target projection heads as ‘Sim(Za, Zb)’ where Za and Zb 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 modifications. We generate the query representation Zqk by applying 1 × 1 convolution to the query features {zqk, ..., z q k+M} and then feed it into the convolutional transition model H. Note that we use only a single next prediction Zpk+1 = H(Z q k , ak) of the transition from the query representation Zqk , unlike Schwarzer et al. (2021) that recursively generates a sequence of L predictions Zpk+1:k+L with Z p k+l+1 = H(Z p k+l, ak+l) for l = 0, ..., L − 1 and Z p k = Z q k . We found that since the spatial representation learning using the flow estimation and warping provide a sufficient amount of supervision, the single self-prediction is enough to train the query encoder. The self-prediction Zpk+1 is fed into the query projection head Pq and the query prediction head Qq such that ypk+1 = Qq(Pq(Z p k+1)). The prediction loss is also computed using the cosine similarity loss Ls(ypk+1, ytk+1).
The loss function Lsim for measuring the similarity between the three representations ywk+1, y p k+1, and ytk+1 is summarized as
Lsim = Ls(ywk+1, ytk+1) + Ls(y p k+1, y t k+1) + λL1(Z p k+1, Z t k+1), (6)
where ywk+1 = Qq(Pq(Z w k+1)), y p k+1 = Qq(Pq(Z p k+1)), y t k+1 = Pt(Z t k+1), and λ is a hyperparameter of loss functions. We also include L1 loss L1(Zpk+1, Ztk+1) on the original spatial latent space to impose the pixel-level similarity. Note that L1 loss between Zwk+1 and Z t k+1 is not used, as the external flow loss Lr(fext) implicitly considers the similarity between the two.
Algorithm 1: Self-Supervised Spatial (S3R) Representations Eq , Et: Query encoder, Target encoder F, H: Flow model, Transition model Pq , Qq , Pt: Query projection head, Query prediction, Target projection head Initialize replay buffer and parameters of Eq , Et, Pq , Qq , Pt, F, H; while Training do • Spatial Representation with Flow Learning Generate zqk+i, z t k+i+1 with Eh, Eg for i = 0, ...,M
Learn external flow and internal flow with (4) • Representation Learning with Warping and Self-Prediction Warp a set of query features zqk+i with external flow f ext k+i+1→k+i for i = 0, ...,M . Generate warped query representation Zwk+1 by fusing a set of warped query features. Generate query representation Zqk by fusing a set of query features z q k+i for i = 0, ...,M . Generate predicted target representation Zpk+1 from Z q k using transition model H. Generate target representation Ztk+1 by fusing a set of target features z t k+i+1 for i = 0, ...,M . Zqk goes into RL MLP head.
Compute global and local similarity of (6) from Zwk+1, Z p k+1, and Z t k+1 using ‘Sim’ in Figure 2. Optimize parameters of Eq , Pq , Qq , F, H by minimizing Ltotal = Lflow + αLsim + βLRL in (7). Update parameters of Et and Pt with Eq and Pq;
end
Finally, the query representation Zqk ∈ Rh×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 final loss function is summarized as
Ltotal = Lflow + αLsim + βLRL(Zqk), (7)
where LRL(Zqk) indicates the loss of the RL algorithms which use Z q k as an input. α and β are hyper-parameters that balance the loss functions. We summarize the overall method in Algorithm 1.

heading: 3.3 IMPLEMENTATION DETAILS
text: Self-supervised Flow Model: The input image Ii is of 84 × 84 for Atari Games and DeepMind Control (DMControl) Suites. The query and target encoders generates zqi , z t i+1 ∈ R7×7×64 (i = k, ..., k + 3) for Atari Games and zqi , z t i+1 ∈ R32×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 flow map. The decoder includes three un-convolutional layers, consisting of un-pooling and convolution, and the coarser flow 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. (2021), 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. (2021). The transition model includes two convolutional layers interweaved with ReLU and batch normalization (Ioffe & Szegedy, 2015), with the current representations Zqk and the action ak of one-hot vector taken to each location being fed to the first convolutional layer.
Other Details: The query and target projection heads, Pq and Pt, are implemented as the multi-layer perceptron (MLP). For the query prediction head Qq, we reuse the first linear layer of the RL head similar to Schwarzer et al. (2021). 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.

heading: 4 EXPERIMENTAL RESULTS

heading: 4.1 EVALUATION ON ATARI GAMES
text: To compare the performance of the proposed method with state-of-the-arts, we chose Atari 2600 Games introduced in Kaiser et al. (2019); van Hasselt et al. (2019) where only 100K environment
steps, corresponding to two hours of gameplay experiences, are available for training data. This sample-efficient 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-efficient deep RL algorithms (Kaiser et al., 2019; van Hasselt et al., 2019; Kielak, 2020; Laskin et al., 2020b;
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-Efficient Rainbow (DER) (van Hasselt et al., 2019) which modifies the Rainbow hyperparameters for improving the sample efficiency, OTRainbow (Kielak, 2020) which is an over-trained version of the Rainbow for the sample efficiency, 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 efficiency, 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 (S3R) achieved the best performance on 13 out of 26 environments. We trained our method with 10 random seeds, similar to other methods.

heading: 4.2 EVALUATION ON DMCONTROL SUITE
text: 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 S3R 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.

heading: 4.3 ABLATION STUDY
text: 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 flow loss in (4) • ‘F+W’ using flow 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 flow 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 flow 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 flow 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 (= S3R). 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.

heading: 5 CONCLUSION
text: We have presented the self-supervised spatial representations, termed S3R to encode local spatial structures in an unsupervised manner. The flow maps inferred by the proposed method offer plenty of supervision for learning the spatial representations, and also compute warped predictions at future frame. S3R 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-efficiency of image-based RL algorithms. We hope this can facilitate future works at various aspects for RL based on self-supervised learning.

heading: 6 REPRODUCIBILITY STATEMENT
text: 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.

heading: B PERFORMANCE CONSISTENCY EVALUATION
text: 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 significantly depending on the number of random seeds. It can be seen that our performance is significantly higher than the SOTA performance regardless of random seed, and is consistent.

heading: C ABLATION STUDY ON DATA AUGMENTATION
text: As described in RAD (Laskin et al., 2020a) and DrQ (Kostrikov et al., 2020), image augmentation can improve the data-efficiency and generalization of RL methods. To study the impact of data augmentation when used with the proposed method, we measured the average performance over 10 random seeds according to the data augmentation on DMControl Suite (Tassa et al., 2018). In Table 7, 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), 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, and Walker Walk and Ball in cup Catch has the best performance for crop. Since S3R learns flow in an end-to-end manner with RL algorithm, it is analyzed that the results are different from those of Laskin et al. (2020a).

heading: D ADDITIONAL EVALUATION METRIC
text: 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. (2021). 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 S3R, the higher the performance. Also, in the graph of S3R below, ‘P(S3R > X)’, S3R ranks the highest as the remaining bars of algorithms are located to the right of the bar of S3R. From this analysis, it can be reconfirmed that S3R 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.

heading: E ANALYSIS ON COMPUTATIONAL COST AND PERFORMANCE IMPROVEMENT
text: Computational cost: The increase in the computational cost for training is unavoidable because S3R additionally leverage the flow prediction and warping networks used in the vision task, but we found that the additional computational cost for training is not so significant. 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 flow 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. S3R has the highest mean in 13 games out of 26. It can be interpreted that the performance increase of the S3R is not small by considering the quantitative aspects of these games. Also, among the 13 games in which S3R has an edge, in particular, in 8 games (Alien, Assault, Gopher, Jamesbond, Krull, Kung Fu Master, Ms Pacman, and Seaquest), S3R 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 specific Atari Games with various movements.
However, S3R may not be effective for some games. In particular, S3R 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, S3R 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 S3R 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 S3R 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, S3R 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.

Review summaries:
This paper proposes a representation learning method that leverages unsupervised signals like flow and forward models and the constraints between them and apply it to the state representation problem of RL.   The method is adding a number of auxiliary rewards to the torso which flow and a latent transition model that are constrained to agree with a contrastive loss. The flow has an architecture inspired by FlowNet and trained with warping reconstruction and a spatial smoothness term. A latent transition model is trained contrastively. Then constraints are added to make the latent of the next step agree with both of these predictions.   The method is run on 13 Atari tasks and compared there with competitive baselines and it largely outperforms them. An ablation without baselines is performed a few dm control suite tasks  

In this paper, the authors focus on the problem of representation learning for deep reinforcement learning. To this end, they propose an approach to learning structured representations via establishing flows between latent volumes. Similar to SPR, they predict future representations with a transition model conditioned on actions. They evaluate their approach in two domains: Atari and DeepMind Control Suite and show improved performance over the state of the art. 

This paper focuses on self-supervised visual representation learning for RL applications. As opposed to previous literature that focused on learning global representations for the current observation, the authors propose learning "structured" representations based on flow maps that encode local structure. To do so, they combine: - a self-supervised flow model (from previous literature), trained using an image reconstruction loss; - a classic action-conditioned 1-step transition model using cosine similarity in latent space - a flow based warping (also from previous literature) to provide a second next-step latent prediction task. They evaluate their method on various standard RL tasks in two domains (Atari, DM Control Suite) and show improvements over SOTA in ~half of the tasks.