Light-weight probing of unsupervised representations for Reinforcement Learning
Authors: Wancong Zhang
Anthony GX-Chen
Vlad Sobal
Yann LeCun
Nicolas Carion
Venue: Submitted to ICLR 2023 (2023)
Keywords: machine learning
unsupervised learning
reinforcement learning
computer vision

Abstract: Unsupervised visual representation learning offers the opportunity to leverage large corpora of unlabeled trajectories to form useful visual representations, which can benefit the training of reinforcement learning (RL) algorithms. However, evaluating the fitness of such representations requires training RL algorithms which is both computationally intensive and has high variance outcomes. To alleviate this issue, we design an evaluation protocol for unsupervised RL representations with lower variance and up to 600x lower computational cost. Inspired by the vision community, we propose two linear probing tasks: predicting the reward observed in a given state, and predicting the action of an expert in a given state. These two tasks are generally applicable to many RL domains, and we show through rigorous experimentation that they correlate strongly with the actual downstream control performance on the Atari100k Benchmark. This provides a better method for exploring the space of pretraining algorithms without the need of running RL evaluations for every setting. Leveraging this framework, we further improve existing self-supervised learning (SSL) recipes for RL, highlighting the importance of the forward model, the size of the visual backbone, and the precise formulation of the unsupervised objective.

Under review as a conference paper at ICLR 2023
LIGHT-WEIGHT PROBING OF UNSUPERVISED REPRESEN- TATIONS FOR REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Unsupervised visual representation learning offers the opportunity to leverage large corpora of unlabeled trajectories form useful visual representations. which can benefit the training of reinforcement learning (RL) algorithms. However, evaluating the fitness of such representations requires training RI. algorithms which is computationally intensive and has high variance outcomes. To alleviate this issue, we design an evaluation protocol for unsupervised RL representations with lower variance and up to 600x lower computational cost. Inspired by the vision community, we propose two linear probing tasks: predicting the reward observed in given state, and predicting the action of an expert in given state. These two tasks are generally applicable to many RL domains, and we show through rigorous experimentation that they correlate strongly with the actual downstream control performance on the Atari 100k Benchmark. This provides better method for exploring the space of pretraining algorithms without the need of running RL evaluations for every setting. Leveraging this framework, we further improve existing self-supervised learning (SSL) recipes for RL, highlighting the importance of the forward model, the size of the visual backbone, and the precise formulation the unsupervised objective. Code will released upon acceptance.
INTRODUCTION
Learning visual representations is critical step towards solving many kinds of tasks, from supervised tasks such as image classification or object detection, to reinforcement learning (RI.). Ever since the early successes of deep reinforcement learning (Mnih et al. 2015), neural networks have been widely adopted to solve pixel- reinforcement learning tasks such as arcade games (Bellemare et al.) 2013), physical continuous control [Todorov et al. 2012: Tassa et al. 2018), and complex video games (Synnaeve al. 2018 Oh et al. 2016)) However, learning deep representations directly from rewards is a challenging task, since this learning signal is often noisy, sparse and delayed.
With ongoing progress in unsupervised visual representation learning for vision tasks (Zontarr ct al. 202T Chen et al.) 2020a bj Grill et al., 2020; Caron et al. 2020 202I). recent efforts have likewise applied self-supervised techniques and ideas improve representation learning for RI.
Some promising approaches include supplementing the RI. loss with self-supervised objectives (Laskin al. 2020 Schwarzer al. (2021a), or first pre-training the representations on a corpus of trajectories (Schwarzer et al. [022bb] Stooke et al. 2021). However, the diversity in the settings considered, as well as the sell-supervised methods used, make it difficult to identify the core principles of successful self-supervised methods in RL. Moreover, estimating the performance of RL algorithms notoriously challenging (Henderson al., 2018 Agarwal al. (021): it often requires repeating the same experience with a different random seed, and the high CPU-to-GPU ratio is compute requirement of most online RI. methods that is inefficient for typical research compute clusters. This hinders systematic exploration of the many design choices that characterize SSL methods.
In this paper, we strive to provide a reliable and lightweight evaluation scheme for unsupervised visual representation in the context of RL. Inspired by the vision community, we propose to evaluate the representations using linear probing, by training a linear prediction head on top of frozen features. We devise two probing tasks that deem widely applicable: predicting the reward in given state, and predicting the action that would be taken by a fixed policy in a given state (for example that of expert). We stress that these probing tasks are only used means of evaluation. Because

---

Under review as a conference paper at ICLR 2023
Spearman's 0.958
2e-5 (**8)
RL HNN (95 CI of IQM)
RL HNS (shaded)
Reward F1 (line)
75
Reward F1 Mean (%)
RL. HNS (IQM)
0.5
65
0.4
30
60
70
10
Reward Prediction FI (%)
Number of Seeds
Figure I: Left: Correlation between the SSL representations' abilities linearly predict the presence of reward in a given state, versus RL performance using the same representations, measured as the interquartile mean of the numan-normallee a separate SSL pretraining method. A linear line of best fit is shown with 95 confidence interval. We compute Spearman's rank correlation coefficient (Spearman's r) and determine its statistical significance using permutation testing (with = 50000). Right: When comparing two models, the reward probing score can give low variance reliable estimates of RI. performance, while direct RI. evaluation may require many seeds to reach meaningful differences in mean performance.
score (HNS) over 9 Atari games. Each point denotes
very little supervised data is required, they are particularly suitable for situations where obtaining the expert trajectories or reward labels is expensive. Through thorough experimentation, show that the performance of the SSL algorithms (in terms of their downstream RL outcomes) correlates with the performance in both probing tasks with statistically significant (p<0.001) Spearman's rank correlation, making them particularly effective proxies. Given the vastly reduced computational burden of linear evaluations, we argue that it enables much easier and straightforward experimentation of SSL design choices, paving the way for more systematic exploration of the design space.
Finally, we leverage this framework to systematically assess some key attributes of SSL methods. First off, we explore the utility and role of learning a forward model as part of the self-supervised objective. We investigate whether its expressiveness matters and show that equipping with the ability to model uncertainty (through random latent variable) significantly improves the quality of representations. Next, we identify several knobs the self-supervised objective, allowing to carefully tune the parameters in principled way. Finally, we confirm the previous finding (Schwarzer ctt al. 2021b) that bigger architectures, when adequately pre-trained, tend to perform better.
Our contributions can be summarized as follows:
Design of a rigorous and efficient SSL evaluation protocol in the context RL
Empirical demonstration that this evaluation scheme correlates with downstream RL. perfor- mance
Systematic exploration of design choices in existing SSL methods.
2 RELATED WORK
2.1 REPRESENTATION LEARNING
There has recently been a surge interest and advances in the domain of self-supervised learning in computer vision. Some state-of-art techniques include contrastive learning methods SimCLR. MoCov2 (Chen et al. 2020a b); clustering methods Sw.VV (Caron et al. 2020); distillation methods BYOL, SimSiam. OBoW (Gril et al. 2020 Chen and He 2021 Gidaris et al. 2020); and information
maximization methods Barlow Twins and VicReg (Zbontar et al. 202 Bardes et al. 2021)
These advances have likewise stimulated development in representation learning for reinforcement learning. A line of work includes unsupervised losses as an auxiliary objective during RL training to improve data efficiency. Such objective can be contrastive (Laskin et al. 2020 Zhu et al. 2020)
2

---

Under review as a conference paper at ICLR 2023
or non-contrastive (Schwarzer al. 202 Yu et al.
2022).
ST-DIM (Anand et al. 2019). ATC
incorporate temporal information in their contrastive objective, adapting similar techniques from the unsupervised video representation learning (Sermanet et al. 2 Proto-RI. Yaarats et al. 2021a) uses a Aw-V--ike objective to learn representation as well as guide effective exploration during pre-training. Similarly, CRL (Du et al. 2021) trains policy to optimize a SimCLR loss, then shows transfer to RL, imitation learning and image classification. Closer our approach, SGI (Schwarzer et al.) 202Ib) pretrains both an encoder and forward prediction model by minimizing the distance between predictions and target latents using BYOL, and the encoder recycled during RL for improved data efficiency. While different in spirit, many model based methods also train encoder from corpus of trajectory, either by explicit pixel reconstruction Kaiser et al. (2020); |Hffnrr et al. (202T) or embedding space Ye et al.) (202T); Schrittwieser et al. 2020). Self-supervised representations have also been used for imitation learning (Aytar et al. 2018 Pari et al. 202T as well as exploration (Burda et al. 2019a).
(Stooke et al.) 2021) | and BVS-DIN DIM Mengistt et al.
2022)
2.2 REPRESENTATION PROBING IN REINFORCEMENT LEARNING
Some prior work Racah and Pal 2019 Guo et al. 2018: Anand et al. 2019 Higgins et al. 2018 Dittadi et al. 2022) evaluate the quality of their pretrained representations by probing for ground truth state variables such as agent/object locations, game scores or model-specific quantities (eg. ELBO). Das et (2020) propose to probe representations with natural language question-answering Despite the efficiency of these probing methods, their designs are highly domain-specific and require careful handcrafting for each environment. In addition, they fail demonstrate the actual correlation between probing and RL performances, which makes their practical usefulness uncertain. On the other hand, the authors of ATC (Stooke et al. (02T)) propose evaluate representations by finetuning for RI. tasks using the pretrained encoder with weights frozen. Similarly. [Lskin et al. (2021) propose a unified benchmark for SSL. methods continuous control but still require full RI. training. Our work seeks to bridge these two approaches by demonstrating the correlation between linear probing and RL performances, as well as designing probing tasks that are generalizable across environments.
3 A FRAMEWORK TO DEVELOP UNSUPERVISED REPRESENTATIONS FOR RL
In this section, we detail our proposed framework for training and evaluating unsupervised represen- tations for reinforcement learning.
3.1 UNSUPERVISED TRAINIINN
The network is first pre-trained on a large corpus of trajectories. Formally, we define a trajectory Ti of length T3 as sequence of tuples T [(o,,,;) az) t E [I]]. where o, is the observation of the state at time f in the environment and 0g was the action taken in this state. This setting is closely related to Batch RL (Lange al., 2012), with the crucial difference that the reward is not being observed. In particular, should be possible to use the learned representations to maximize any reward (Tuuaii and Ollivier 2021). The training corpus corresponds Dunsup ( J. We note that the policy used to generate this data is left unspecified in this formulation, and is bound to be enviromment-specific Since unsupervised methods usually necessitate a lot of data, this vre-trannig corpus is required to be substantial. In some domains, it might be straightforward to collect large number of random trajectories to constitute Dunsup- In some other cases, like self self-driving, where generating random trajectories is undesirable, expert trajectories from humans can be used instead.
set of such trajectories:
The goal of the pre-training step is to learn the parameters 0 of an encoder ENCa which maps any observation o of the state 8 (for example raw pixels) to a representation e ENC, (o). This representation must be amenable for the downstream control task, for example learning policy.
3.2 EVALUATION
In general, the evaluation of RL algorithms is tricky due to the high variance in performance (Henderson et al. 20188). This requires evaluating many random seeds, which creates a computational burden. We side-step this issue by formulating an evaluation protocol which is light-weight and purely
3

---

Under review as a conference paper at ICLR 2023
supervised. Specifically, we identify proxy supervised tasks that are broadly applicable and relevant for control. We further show in the experiment section that they are sound, in the sense that models' performance the proxy tasks strongly correlates with their performance in the downstream control task interest. Similar to the evaluation protocol typically used for computer vision models, rely linear probing, meaning that train only linear layer on top of the representations, which are kept frozen.
Reward Probing Our first task consists predicting the reward observed in given state. For this task, require a corpus of trajectories Dew 1''..-.. .''} for which the observed rewards are known, ie. T' = [[or, a,, rx) E [1]
In the most general setting, it can be formulated as regression problem, where the goal is to minimize the following loss:
((V)) reward-reg
Σ
Σ
11 (ENCa((O)) rz rzl12
Drew
[T's)
T'.EDrow Drow
(DGAGNETT')
Here, the only learnt parameters v are those of the linear prediction layer hpp
In practice, in many environments where rewards are sparse, the presence or absence of a reward is more important than its magnitude. To simplify the problem in those cases, we can cast it as a binary prediction problem instead (this could be extended to ternary classification the sign of the reward of interest):
C reward-classif
Σ
(T))
Drew
IT'sl
Reward prediction is closely related to value prediction, a central objective in RL that is essential for value-based control and the critic actor-critic methods. The ability to predict instantaneous reward, akin predicting value with a very small discount factor, can be viewed learned representation's 's ability to encode the value function, and has been demonstrably helpful for control, particularly in sparse reward tasks (Jaderberg et al. 2017). Thus, we hypothesize reward prediction accuracy to be good probing proxy task for our setting as well.
lower bound on the
Action prediction Our second task consists in predicting the action taken by an expert in a given state. For this task, require corpus of trajectories Dexp (Ti, Tr,) generated by expert policy. We stress that this dataset may be much smaller than the pretraining corpus since we only require fit and evaluate linear model. The corresponding objective is as follows:
C(v) action-classif
Σ
Σ (o.,..TT.)
CrossEntropy (ENCg((:)))
[eepl
T.EDop
This task closely related to imitation learning, however, we are not concerned with the performance of the policy that we learn by-product.
SELF PREDICTIVE REPRESENTATION LEARNING FOR RL
In work, focus evaluating and improving particular class of unsupervised pretraining algorithms that involves using a transition model to predict its own representations in the future (Schwarzer al. 2021b: Guo al. 2018 Gelada al. 2019). This pretraining modality is especially well suited for RL, since the transition model can be conditioned on agent actions, and can be repurposed for model-based RL after pretraining. Our framework is depicted in Fig this section, we present the main design choices, and we investigate their performance in Section
4.1 TRANSITION MODELS
Our baseline transition model is a 2D convolutional network applied directly to the spatial output of the convolutional encoder (Schwarzer al. 2021b: Schrittwieser al. 2020). The network consists
4

---

Under review as a conference paper at ICLR 2023
Encoder
Online
Predictor
E141
Predictor
Qc-1
Projector
CE loss
Projector
CE loss
Prujector
à;
d++1
Distance
Distance
Projector
Projector
Oc+1
Target
OL+2
Target
Encoder
Encoder
Figure 2: Model diagram. The observations consist of stack of 4 frames, to which we apply data augmentation before passing them to a convolutional encoder. The predictor is a recurrent model outputting future state embeddings given the action. We supervise with an inverse modeling loss (cross entropy loss on the predicted transition action) and an SSL loss (distance between embeddings).
of two 64-channel convolutional layers with 3x3 filters. The action represented as one-hot encoding spatially replicated (in 2D map) and concatenated with the representation input along the channel dimension.
We believe well-established sequence modeling architecture such GRU serve superior transition model. Its gating mechanisms should be better at retaining information from both the imme- diate and distant past, especially helpful for learning dynamics in partially observable environment.
Encoder eo = 60 = ENCo(oo)
RecurrentModdl è fc((r-1:00-1
In addition to the deterministic GRU model above, we also experiment with a GRU variant where introduce stochastic states to allow model to generalize better stochastic environments, such as Atari with sticky actions (Machado et al.) (018). Our model is based on the RSSM from DreamerV2 (Hainer al. 2021), with the main difference being that while pixel reconstruction is used as the SSL objective in the original work, we minimize the distance between predictions and targets purely the latent space. Following DreamerV2, we optimize the latent variables using ttraight-through gradients (Bengio et al.] 2013), and minimize the distance between posterior (2) and prior (z) distributions using KL loss.
Encoder: :
e, ENCa(O:)
RecurrentModel:
h fa(he-1, 21- a1-1)
PosteriorModel:
zt ~ po(zellie,ed et)
PriorPredictor:
~ jo (E(()e))
LatentMerger:
E, = go (hr,
4.2 PREDICTION OBJECTIVES
The objective of self predictive representation learning is to minimize the distance between the predicted and the target representations, while ensuring that they do not collapse to a trivial solution. Our baseline prediction objective is BYOL (Grill et al. 2020), which is used in SGI (Schwarzer et al. 2021b). The predicted representation ê and the target representation ē are first projected to lower dimensions to produce gr+* and V+++. BYO then maximizes the cosine similarity between the predicted and target projections, using linear prediction function q translate from ý to ÿ:
5

---

Under review as a conference paper at ICLR 2023
g((+kk) yo+k
Lo BY (0er4ks (1a4k::0a4k)))
( 122 10
In the case of BYOL, the target encoder and projection module are the exponentially moving average of the online weights, and the gradients are blocked on the target branch.
As alternative prediction objective, experiment with Barlow Twins (Zbontar et al., 2021 Similar to BYOL, Barlow Twins minimizes the distance of the latent representations between the online and target branches; however, instead of using predictor module and stop gradient the target branch, Barlow Twins avoids collapse by pushing the cross-correlation matrix between the projection outputs the two branches to be close to the identity matrix possible. To adapt Barlow Twins, we calculate the cross correlation across batch and time dimensions:
En,((i (Bbaas)
I, BT (jetth, Brstk) = Σ Cis) +
Σ Ct where Cij = Vo, Eore(tr...d K
where À is positive constant trading off the importance of the invariance and covariance terms of the loss, C' is the cross-correlation matrix computed between the projection outputs of two branches along the batch and time dimensions, b indexes batch samples, t indexes time, and i, index the vector dimension of the projection output.
By enabling gradients on both the prediction and target branches, the Barlow objective pushes the predictions towards the representations, while regularizing the representations toward the predictions. In practice, learning the transition model takes time and we want to avoid regularizing the represen- tations towards poorly trained predictions. To address this, we apply a higher learning rate to the prediction branch. We call this technique Barlow Balancing, and implement it in Algorithm I
Algorithm 1: PyTorch-style pseudocode for Barlow Balancing
BarlowLoss = . LLLL* RT (g) j.detach()) (1 p)) "T (j.detach(). 7)
4.3 OTHER SSL OBJECTIVES
SGI's authors (Schwarzereta b02 showed that in the absence of other SSL objectives, pretraining with BYOL prediction objective alone results in representation collapse; the addition of inverse dynamics modeling loss is necessary prevent collapse, while the addition of goal-oriented RL loss results in minor downstream RL performance improvement. In inverse dynamics modeling, the model is trained using cross-entropy model p(a, 0ttk (t-k+1), effectively predicting the transition action between two adjacent states. The goal-oriented loss tries to predict distance to states in the near future from the sampled trajectories (details in Appendix).
RESULTS
5.1 EXPERIMENTAL DETAILS
We conduct experiments the Arcade Learning Environment benchmark (Bellemare et al. 2013). Given the multitude of pretraining setups we investigate, we limit our experiment to Atari games")
Pretraining We use the aualicly-vvailabl DQN replay dataset (Agarwal et al. 2020), which contains data from training a DQN agent for 50M steps with sticky action (Machado et al. 2018). We select 1.5 million frames from the 3.5 to 5 millionth steps of the replay dataset, which constitutes trajectories of a weak, partially trained agent. We largely follow the recipe of SGI (Schwarzer et al. 2021bb) where we jointly optimize the self prediction, goal-conditioned RL, and inverse dynamics modeling
Amidar, Assault, Asterix, Boxing, Demon Attack, Frosthite, Gopher, Krull, Seaquest

---

Under review as a conference paper at ICLR 2023
Demon attack
Gopher
f
5
t = 10
t t=1
t t= 5
t= 10
GRU-det, BYOL
GRU-latent, BYOL
GT
Figure 3: Decoding results, using de-convolutional model to predict the pixel values from frozen state representations. Both games exhibit stochastic behaviours. In Demon attack, both models fail to capture the position of the enemies. In Gopher, the enemy (circled in red) is moving randomly, but thanks to the latent variable, the GRU-latent model is able to predict a possible position, while the deterministic model regresses the mean.
losses for 20 epochs; in some of our experiments we remove one or both of the last two objectives. We use the data-augnn aat--augmentations introduced by Yarats al. (2021b). All experiments are performed on a single MI50 AMD GPU, and the pretraining process took to 8 days depending on the model.
Reward probing We focus the simplified binary classification task of whether reward occurs in given state. We use 100k frames from the 1-1.1 millionth step of the replay dataset, with a 4:1 train/eval split. We train logistic regression model frozen features using the Cyanure (Mairal 2019) library, with the MISO algorithm (Mairal 2015) coupled with QNING acceleration (Lin et al. 2019) for maximum of 300 steps. We do not use any data augmentation. We report the mean Fl averaged across all 9 games. On MI50 AMD GPU, each probing run takes 10 minutes.
Action probing We use the last 100k (4:1 train/eval split) frames of the DQN replay dataset, which correspond to a fully trained DQN agent. We train a linear layer on top of frozen, un-augmented features for 12 epochs with softmax focal loss (Lin et al.) 2017) using SGD optimizer with learning rate 0.2, batch size 256, le-6 weight decay, stepwise scheduler with step size 10 and gamma 0.1. We report the Multiclass (weighted average of scores each class) averaged across all games.
RL evaluation We focus on the Atari 100k benchmark (Kaiser et al. 2020), where only 100k interactive steps are allowed by the agent. This is roughly equivalent to two hours human play, providing approximation for human level sample-efficiency. We follow Schwwrzer al. 2021bb training protocol using the Rainbow algorithm (Hessel et al. 2018) with the following differences: freeze the pretrained encoder (thus only training the Q head), do not apply auxiliary SSL losses while fine-tuning, and finally disable noisy layers and rely instead on e-greedy exploration. This changes made make the RL results reflect closely as possible the performance induced by the quality of the representations. On a MI50 AMD GPU, each run takes between 8 and 12 hours.
We evaluate the agent's performance using human-noneee normalizzed score (HNS), defined as (agentscore- randomsscre)) / (humanscore wwwpsnscooee-rnoonaaaaooooo randomscooreeeee We calculate this per game, per seed by averaging scores over 100 evaluation trajectories at the end of training. For aggregate metries across games and seeds, we report the median and interquartile mean (IQM). For median, we first average the HNS across seeds for each game, and report the median the averaged HNS values. For IQM, we first take the middle 50% of scores across both seeds and games, then report the average. While median commonly reported for Atari Atari100k, recent work has recommended IQM as superior aggregate metric for the RL setting due to its smaller uncertainty (Agarwal et al. 2021); we also follow the cited work to report the 95% bootstrapped confidence intervals for these aggregate metrics.
Unless specified otherwise, the experiments use the medium ResNet-M from Schwarzer et al. 1b02 and the inverse dynamics loss as an auxiliary loss. In BYOL experiments, the target network is exponential moving average of the online network, while in Barlow Twins both networks identical, following the original papers. For additional details regarding model architectures and hyperparameters used during pretraining and RI. evaluation, please refer to Appendix.
7

---

Under review as a conference paper at ICLR 2023
5.2 IMPACT OF TRANSITION MODELS AND PREDICTION OBJECTIVES
Table 1: Fl scores on probing tasks for different transition models and prediction objectives. All standard deviations are the order of le-4
Table 2: Fl scores probing tasks for different
Barlow variants, All standard deviations are on the order of le-4 which we omit below.
Pred Obj Transition
Reward
Action
Pred Obj
Reward
Action
BYOL
Conv-det
64.9
22.7
Barlowos
65.0
26.3
GRU-det
62.2
26.8
Barlow
67.5
26.2
GRU-latent
63.4
23.2
Barlow 65.0
24.7
Barlow
67.7
25.8
Barlowo.7
Conv-det 52.7
24.9
GRU-latent 67.5
26.2
In table we report the mean probing F1 scores for the convolutional, deterministic GRU, and latent GRU transition models trained using either the BYOL or Barlow prediction objective. When using the BYOL objective, the relative probing strengths for the different transition models are somewhat ambiguous: while the convolutional model results better reward probing F1, the GRU models are superior in terms of expert action probing.
Interestingly, we observe that after replacing BYOL with Barlow, the probing scores for the latent model improve, while those of the deterministic models deteriorate. Overall, the particular combina- tion pre-training using the GRU-latent transition model with the Barlow prediction objective results in representations with the best overall probing qualities. Since the deterministic model's 's predictions are likely to regress to the mean, allowing gradients to flow through the target branch the case Bar- low objective can regularize the representations towards poor predictions, and can explain their inferior probing performance. Introducing latent variables can alleviate this issue through better predictions.
We stress that the transition models are not used during probing, only the encoder is. These experi- ments show that having a more expressive forward model during the pre-training has a direct impact on the quality of the learnt representations. In Fig 3. investigate the impact of the latent variable on the information contained in the representations, by training a decoder on frozen features.
In table 2, we show the results from experimenting with different variants of the Barlow objective. We find that using higher learning rate for the prediction branch (Barlowo.7. with 7:3 prediction to target Ir ratio) results in better probing outcome than using equal learning rates (Barlowo.s) or not letting gradients flow in the target branch altogether (Barlow,, here the target encoder is copy of the online encoder). This suggests that while it helpful to regularize the representations towards the predictions, there is potential for them being regularized towards poorly trained ones. This can be addressed by applying a higher learning rate on the prediction branch.
We also demonstrate that using frozen, random target network (Barlowrand) results in good features, and in our experiments it gets the best reward probing performance. This contradicts findings from the vision domain (Grill et al. 2020), but corroborates self-supervised results from other domains such as speech (Chiu et al. 2022)
Random networks have also been shown to exhibit useful inductive biases
Table 3: Representation probing and RL results. Mean binary F1 for reward, mean multiclass F1 for next action (all stdev le-4). RL metrics are aggregated on 10 seeds of 9 games. The 95% CIs are estimated using the percentile bootstrap with stratified sampling (Agarwal et al. 202T).
Backhone Transition
Objectives
Reward Action
Median
1QM
L
GRU-Jat
Barlow.ed. inv
70.3
26.7
L
GRU-Jat
Barlowor.
69.0
27.7
GRU--at
Harlowand: inv
67.7
GRU-Jat
Barlowan. inv
67.4
26.2
M
GRU-Jat
BYOL, goal. inv
63.4
23.2
GeU-det
BYOL, goal,
62.2
Conv-det
BYOL, goal, inv
64.9
22.7
GRU--la
Barlowan
56.2
24.4
Conv-det
B Barlowen, 200l inv
52.7
24.8
None
Stooke et al. 2021
51.3
22,9
None
2019
34.6
21.3
None
VAB-T Stooke et
33.2
20.4
0.4
0.8
0.5
1.0
Human Normalized Score
8

---

Under review as a conference paper at ICLR 2023
for exploration (Burda al., 2019b a). An explanation is that random targets act a regularization that prevent partial collapse by enforcing a wide range of features to encoded by the model.
5.3
IMPACT OF AUXILIARY SSL OBJECTIVES AND ENCODERS
Table 4: Fl scores on probing tasks for different auxiliary objectives. All stdev se le-4.
Table 5: F1 scores probing tasks for different encoders. All stdev s le-4.
SSL Objs
Reward
Action
Pred Obj
Encoder
Reward
Action
BYOL, inv, goal
63.4
23.2
Barlowo.7
Res-M Res-L
67.5
26.2
BYOL, inv
57.3
22.6
69.0
27.7
BYOL
25.9
5.9
Barlowrand
Res-M
67.7
25.8
Barlowo,7, inv, goal
66.5
26.2
Res-I.
70.3
26.7
Barlowo,7. inv
67.5
26.2
Barlowo.7
56.2
24.4
SSL objective Although pretraining with multiple objectives can sometimes result in better down- stream performance, in practice they also make it harder to tune for hyperparameters and debug, therefore is desirable to use the least number of objectives that can result in comparable performance.
In table 4 we show the effects of inverse dynamics modeling (inv) and goal-conditioned RL (goal) objectives on probing performance. The BYOL model experiences partial collapse without the inverse dynamics modeling loss, while the addition of goal loss improves the probing performance slightly. This is in congruence with results reported by Schwarzer et al. (2021b) for the same ablations.
The Barlow-only model performs significantly better than the BYOL-only model in terms of probing scores, indicating that the Barlow objective is less prone to collapse in the predictive SSL setting. Similar to the BOL model, the Barlow model can also be improved with inverse dynamics modeling. while the addition of goal loss has a slight negative impact.
Encoders SGI (Schwarzer et al. 2021b) showed that using bigger encoders during pretraining results in improved downstream RL performance. We revisit this topic from the point of finding out whether the pretrained representations from bigger networks also have better probing qualities. We experiment with the medium (ResNet-M) and large (ResNet-L) residual networks from SGI. In table 5 we show that Barlow models pretrained using the larger ResNet have improved probing scores.
5.4 CORRELATIONS BETWEEN PROBING AND RL PERFORMANCES
If our goal is to use linear probing guide to identify superior pretraining setup for RL, then they are only useful to the extent to which they correlate with the actual downstream RL performance. We perform RL evaluations for representative setups (the best settings from each of table 4 as
well as two contrastive methods: ST-DIM (Anand et al. 2019) and ATC (Stooke et al. 2021) and a reconstruction-batt
method VAE-T Stooke et al. 2021 We report their probing and aggregate RL metrics in table 3 with the confidence intervals of the aggregate RI metrics depicted on the right. We find that the rank correlations between reward and action probing F1 scores and the RL aggregate metrics are significant (Figure []. In summary, our results show the proposed probing scheme is reliable guide for designing pretraining setups that deliver significant downstream RL performance improvements.
6 CONCLUSION
In this paper we have investigated the opportunity to replace costly RL evaluation with lightweight linear probing task to assess the quality of learned representations. Reward and action probing are task-agnostic and should cover most practical applications. Using this methodology to guide us, we have demonstrated the impact of a number of key design choices the pre-training methodology. We hope that these results encourage the research community to systematically explore the design space to further improve the quality supf-supe representations for RL.
"See appendix for details ATC, ST-DIM and VAE-T

---

Under review as a conference paper at ICLR 2023
REFERENCES
Agarwal, R., Schuurmans, D., and Norouzi, M. (2020). An optimistic perspective on offline rein- forcement learning. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 104-114, PMLR.
Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A. C., and Bellemare, M. (2021). Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems, 34.
Anand, A., Racah, E., Ozair, S., Bengio, Y., Côté, M., and Hjelm, R. D. (2019), Unsupervised state representation learning in atari. In Wallach, H. M., Larochelle, H., Beygelzimer, A., d' Alché-Buc, F., Fox, E. B., and Gamettt, R., editors, Advances in Neural Information Processing Systems 32: Annual Conference Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 8766-8779.
Aytar, Y., Pfaff, T., Budden, D., Paine, T. L.., Wang. Z., and de Freitas, N. (2018). Playing hard exploration games by watching youtube. In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, N., and Garnett, R., editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NewrIPS 2018, December 3-8, 2018, Montréal, Canada, pages 2935-2945.
K.,
Bardes, A., Ponce, J., and LeCun, Y. (2021). Vicreg: Varinne-invariiccocovvvaaa regularizattonnnnn for self superviee learning. ArXiv preprint, 1atee200000499
Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M. (2013). The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279.
Bengio, Y., Léonard, N., and Courville, A. (2013). Estimating or propagating gradients through stochastic neurons for conditional computation. ArXiv preprint, abs/ 338.3432..
Burda, Y., Edwards, H., Pathak, D., Storkey, A. Darrell, T., and Efros, A. A. (2019a). Large-scale study of curiosity-driven learning. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net.
Burda, Y., Edwards, H., Storkey, A. J., and Klimov, O. (2019b). Exploration by random network distillation. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net.
Caron, M., Misra, I., Mairal, J., Goyal, P., Bojanowski, P., and Joulin, A. (2020). Unsupervised learning of visual features by contrasting cluster assignments. Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H., editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual.
Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., and Joulin, A. (2021). Emerging Properties in Self- Self-Supervised Vision Transformers. ArXiv preprint, abs/2104.14294
Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. E. (2020a). A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 1597-1607. PMLR.
Chen, X., Fan, H., Girshick, R., and Hc, K. (2020b). Improved basclines with momentum contrastive learning. ArXiv preprint, abs/2003.04297.
Chen, X. and He, K. (2021). Exploring simple siamese representation learning. In Proceedings of the IEFF/CVF Conference Computer Vision and Pattern Recognition, pages 15750-15758.
Chiu, C.-C., Qin, J., Zhang, Y., Yu, J., and Wu, Y. (2022). Self-supervised learning with random- projection quantizer for speech recognition. ArXiv preprint, abs/2202.01855.
10

---

Under review as a conference paper at ICLR 2023
Das, A., Carnevale, F., Merzic, H., Rimell, L., Schneider, R., Abramson, J., Hung, A., Ahuja, A., Clark, S., Wayne, G., and Hill, F. (2020). Probing emergent semantics in predictive agents via question answering. In Proceedings of the 37th International Conference Machine Learning, ICMI, 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 2376 2391. PMLR.
Dittadi, A., Träuble, F., Wüthrich, M., Widmaier, F., Gehler, P., Winther, O., Locatello, F., Bachem, O., Schölkopf. B., and Bauer, S. (2022). The role of pretrained representations for the ood generalization of reinforcement learning agents. (arXiv:2107.05688) arXiv:2107.0566 [cs, stat].
Du, Y., Gan, C., and Isola, P. (2021). Curious representation learning for embodied intelligence. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 10408-10417.
Gelada, C., Kumar, S., Buckman, J., Nachum, O., and Bellemare, M. G. (2019). Deepmdp: Learning continuous latent space models for representation learning. In Chaudhuri, K. and Salakhutdinov. R., editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 2170-2179 PMLR.
Gidaris, S., Bursuc, A., Puy, G., Komodakis, N., Cord, M., and Pérez, P. (2020). Online bag-of- wordsl-words generation for unsupervised representation learning. ArXiv preprint, abs/2012.11552.
Grill, J., Strub, F., Altché, F., Tallec, C., Richemond, P. H., Buchatskaya, E., Doersch, C., Pires, B. Á., Guo, Z., Azar, M. G., Piot, B., Kavukcuoglu, K., Munos, R., and Valko, M. (2020). Bootstrap your own latent A new approach to self-supervised learning. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H., editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual.
Guo, Z. D., Azar, M. G., Piot, B., Pires, B. A., and Munos, R. (2018). Neural predictive belief representations. ArXiv preprint, abs/181 1.06477.
Hafner, D., Lillicrap, T. P., Norouzi, M., and Ba, J. (2021). Mastering atari with discrete world models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net.
Henderson, P., Islam, R., Bachman, P., Pineau, J., Precup, D., and Meger, D. (2018). Deep rein- forcement learning that matters. In Mcllraith, S. A. and Weinberger, K. Q., editors, Proceedings of the Thirty-Seen Secrry--econdaao AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (IAAI-18), and the 8th AAAI Symposium Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pages 3207-3214, AAAI Press.
Hessel, M., Modayil, J., van Hasselt, H., Schaul, T., Ostrovski, G., Dabney, W., Horgan, D., Piot, B., Azar, M. G., and Silver, D. (2018). Rainbow: Combining improvements in deep reinforcement learning. In Mcllraith, S. A. and Weinberger, K. Q.. editors, Proceedings of the Thirty-Secome AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artifi- cial Intelligence (1AA1-18), and the 8th AAAI Symposium on Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pages 3215-3222 AAAI Press.
Higgins, L,, Pal, A., Rusu, A. A., Matthey, L., Burgess, C. P., Pritzel, A., Botvinick, M., Blundell, C., and Lerchner, A. (2018). Darla: Improving zero-shot transfer in reinforcement learning. [cs, stat].
arXiv:1707
Jaderberg, M., Mnih, V., Czamnckii W. M., Schaul, T., Leibo, Z., Silver, D., and Kavukcuoglu. K. (2017). Reinforcement learning with unsupervised auxiliary tasks. In 5th International Conference on Learning Representations, ICLR 20/7, Toulon, France, April 24-26, 2017 Conference Track Proceedings. OpenKeview.net.
Kaiser, L.. Babaeizadeh, M., Milos, P., Osinski, B., Campbell, R. H., Czechowski, K., Erhan, D., Finn, C., Kozakowski, P., Levine, S., Mohiuddin, A., Sepassi, R., Tucker, G., and Michalewski. H. (2020). Model based reinforcement learning for atari. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
11

---

Under review as a conference paper at ICLR 2023
Kingma, D. P. and Welling, M. (2014). Auto-encoding variational bayes. arXiv:1312.6114 [es, stat].
(#rXiv:1312.6114)
Lange, S., Gabel, T., and Riedmiller, M. (2012). Batch reinforcement learning. In Reinforcement learning, pages 45-73. Springer.
Laskin, M., Srinivas, A., and Abbeel, P. (2020). CURL: contrastive unsupervised representations for reinforcement learning. In Proceedings of the 37th International Conference Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 5639-5650 PMLR.
Laskin, M., Yarats, D., Liu, H., Lee, K., Zhan, A., Lu, K., Cang, C., Pinto, L., and Abbeel, P. (2021). Urlb: Unsupervised reinforcement learning benchmark. ArXiv preprint, abs/2110.1510
Lin, H., Mairal, J., and Harchaoui, Z. (2019). An inexact variable metric proximal point algorithm for generic quasi-newton acceleration. SIAM Journal on Optimization, 29(2):1408-144 1443.
Lin, T., Goyal, P., Girshick, R. B., He, K., and Dollár, P. (2017). Focal loss for dense object detection. In IEEE International Conference on Computer Vision, ICCV 2017, Venice, Italy, October 22-29. 2017, pages 2999-3007, IEEE Computer Society.
Machado. M. C., Bellemare, M. G., Talvitie, E.,, Veness, J., Hausknecht, M., and Bowling, M. (2018). Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-562.
Mairal, J. (2015). Incremental intorizzttoonmmnnmizzaiin optimization with application to large-scale machine learning. SIAM Journal on Optimization, 25(2):829-855.
Mairal, J. (2019). Cyanure: An Open-Source Toolbox for Empirical Risk Minimization for Python, C++, and soon more. ArXiv preprint, abs/1912.08166.
Mengistu, M. T., Alemu, G., Chevaillier and Loor, P. D. (2022). Unsupervised learning of state representation using balanced view spatial deep infomax: Evaluation on atari games, In ICAART.
Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. (2015). Human-level control through deep reinforcement learning. nature, 518(7540):329-533
Odena, A., Dumoulin, V., and Olah, C. (2016). Deconvolution and checkerboard artifacts. Distill.
Oh, J., Chockalingam, V., Singh, S. P., and Lee, H. (2016). Control of memory, active perception, and action in minecraft. Balcan, M. and Weinberger, K. Q., editors, Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 JMLR Workshop and Conference Proceedings, pages 2790-2799.JMLL JMLR.org.
Pari, J., Muhammad, N., Arunachalam, S. P., Pinto, L., al. (2021). The surprising effectiveness of representation learning for visual imitation. ArXiv preprint, abs/2112.015//
Perez, E., Strub, F., de Vries, H., Dumoulin, V., and Courville, C. (2018). Film: Visual reasoning with general conditioning layer. In Mcllraith, S. A. and Weinberger, K. Q., editors, Proceedings of the Thirty-Secont AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (1AAI-18), and the 8th AAAI Symposium Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pages 3942-3951. AAAI Press.
Racah, E. and Pal, C. (2019). Supervise thyself: Examining self-supervised representations in interactive environments. ArXiv preprint, abs/1906.1 1951..
Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel. T., et al. (2020). Mastering atari, go. chess and shogi by planning with a learned model. Nature, 588(7839):604460))
12

---

Under review as a conference paper at ICLR 2023
Schwarzer, M., Anand, A., Goel, R., Hjelm, R. D., Courville, A. C., and Bachman, P. (2021a). Data-efficient reinforcement learning with self- seff-predictive representations. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net.
Schwarzer, M., Rajkumar, N., Noukhovitch, M., Anand, A., Charlin, L., Hjelm, R. D., Bachman, P., and Courville, A. C. (2021b). Pretraining representations for data-efficient reinforcement learning. Advances in Neural Information Processing Systems, 34.
Sermanet, P., Lynch, C., Chebotar, Y., Hsu, J., Jang, E., Schaal, S., Levine, S., and Brain, G. (2018). Time-contrastive networks: Self-supervised learning from video. In 2018 IEEE international conference on robotics and automation (ICRA), pages 1134- IEEE.
Stooke, A. and Abbeel, P. (2019). rlpyt: A research code base for deep reinforcement learning in pytorch. (arXiv: (9990019000))) arXiv: 1909.01500 [es].
Stooke, A., Lee, K., Abbeel, P., and Laskin, M. (2021). Decoupling representation learning from reinforcement learning. In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning. ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 9870-9879, PMLR.
Synnaeve, G., Lin, Z., Gehring, J., Gant, D., Mella, V., Khalidov, V., Carion, N., and Usunier, N. (2018). Forward modeling for partial observation strategy games A starcraft defogger. In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R., editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pages 10761-10771.
Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., Casas, D. d. L., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq, A., et al. (2018). Deepmind control suite. ArXiv preprint, abs/1801.000 801.00690.
Todorov, E., Erez, T., and Tassa, Y. (2012). Mujoco: A physics engine for model-based control. In
2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026-5033. IEEE.
Touati, A. and Ollivier, Y. (2021). Learning One Representation to Optimize All Rewards. ArXiv preprint, abs/2103.07945.
Yarats, D., Fergus, R., Lazaric, A., and Pinto, I.. (2021a). Reinforcement learning with prototypical representations. In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 1920-11931. PMLR.
Yarats, D., Kostrikov, I., and Fergus, R. (2021b). Image augmentation is all you need: Regular- izing deep reinforcement learning from pixels. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net.
Ye, W., Liu, S., Kurutach, T., Abbeel, P., and Gao, Y. (2021). Mastering atari games with limited data. Advances in Neural Information Processing Systems. 34.
Yu, T., Zhang, Z., Lan, C., Chen, Z., and Lu, Y. (2022). Mask-based latent reconstruction for reinforcement learning. ArXiv preprint, abs/2201. 12096.
Zbontar, J., Jing, L., Misra, I., LeCun, Y., and Deny, S. (2021). Barlow twins: Self-supervised learning via redundancy reduction. Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 12310-12320, PMLR
Zhu, J., Xia, Y., Wu, L., Deng, J., gang Zhou, W., Qin, T., and Li, H. (2020). Masked contrastive representation learning for reinforcement learning. ArXiv preprint, abs/2010.07470.
13

---

Under review as a conference paper at ICLR 2023
A MODELS AND HYPER-PARAMETERS
A.1 BACKBONES
M and L models are ResNet-M and ResNet-L from SGI (Schwarzer et al. 2021b). The ResNet-M encoder consists of inverted residual blocked with an expansion ratio of 2, with batch normalization applied after each convolutional layer; it uses 3 groups with 32, 64, and 64 channels, and has residual blocks per group: it down-scales the input by factor of 3 in the first group and in the latter 2 groups. This yields a representation of shape 64x7x7 when applied to 84x84-dimensions Atari frames. ResNet-L uses 3 groups with 48, 96, and 96 channels, and has 5 residual blocks per group; uses larger expansion ratio of 4, producing representation shape of 96x7x7 from an 84x84 frame. This enlargement increases the number of parameters by approximately a factor of 5.
S model is the model used in Stooke et al. (2021). It consists of three convolutional layers, with [32, 64, 64] channels. kernel sizes [8, 4, 3). and strides [4, 2, 1], listed from first to last layer.
A.2 TRANSITION MODELS
We experimented with three transition models: convolutional model, deterministic GRU, and latent GRU. Our convolutional model is based on SGI (Schwarzer et al.) (021b). The input into the convolutional transition model is the concatenation of the spatially replicated 2D action map and the representation e, along the channel dimension. The network itself consists of two 64-channel convolutional layers with 3x3 filters, separated by ReLU activation and batch normalization layers.
The deterministic GRU has hidden dimension 600 and input dimension 250. The input a, is prepared by passing the one-hot action vector through a 250 dimensional embedding layer. The initial hidden state én is generated by projecting the representation eo through a 600 dimensional linear layer with ELU activation and dropout. Layer normalization is applied the hidden input all timesteps.
The latent GRU model is based on Dreamerv2's RSSM (Hafner et al. 2021). and is consisted of recurrent model, posterior model, prior predictor, and latent merger. The recurrent model has hidden dimension and input dimension of 600. The initial hidden state hg and input 2g are zero vectors. The flattened stochastic variables 2, and one-hot action vector de are first concatenated and then projected to 600 dimension through a linear layer with ELU activation, before being passed into the recurrent model as input. Layer normalization is applied to the hidden input at all non-zer zerzzrr timesteps.
The posterior model is two-layer MLP with 600 dimensional bottleneck separated by ELU activation. takes the concatenation of representation e, and recurrent hidden output h as input, and outputs a 1024 dimensional vector representing the 32 dimensional logits for 32 latent categorical variables. 2, is sampled from the posterior logits. The prior model is a two-layer MIP with 600 dimensional bottleneck separated by ELU activation. Its output format is same that of the posterior model. é, is sampled from the prior logits. The latent merger is a linear layer that projects the concatenation of he and flattened 2, the same dimension of representation et.
A.3 SSL PROJECTION MODULE
In the case of the deterministic GRU, è is first projected to the same dimension of representation through a linear layer. Henceforth we shall assume that è underwent this step for GRU
The predicted representation é and target representation ē are projected to 1024 dimensional vectors ÿ and ÿ through linear layer. The B BYLL objective involves processing ý with an additional linear layer with output dimension 1024. The Barlow objective involves applying batch normalization to ÿ and ÿ prior to taking the covariance and variance losses.
The inverse dynamics model is two-layer MLP with 256 dimensional bottleneck separated by ReLU activation. It takes the concatenation of ÿr and U4+1 as input, and outputs logits with dimension equivalent to number of actions.
14

---

Under review as a conference paper at ICLR 2023
A.4 ATC, VAE-T, ST-DIM
We the implementation, hyperparameters and architecture from the codebase of (Stooke et al. 2021) and (Stooke and Abbeel 2019) for these models. We change the dataset to the one used in all our experiments We use the dataset described in section to train these models, and train all methods for 58,500 updates. ATC momentum encoder and online branch applied different augmentations of an image to pre-train the encoder. VAE-T from Stooke et al. (2021) uses variational auto-encoder (Kingma and Welling 2014) objective to reconstruct the frame from the next time step given an image at the current time step. ST-DIM (Anand et al. 2019) also uses InfoNCE objective, and in additional to traditional global-global infomax, introduces global-local infomax by using local representations taken from the feature map output of the convolutional encoder and the global pooled feature vector positive pairs. For more details, we refer the reader to the referenced works.
Sreenented-Teeporrt Contrast) uses InfoNCE loss between output of the
A.5 IMAGE RECONSTRUCTION MODEL
We used decoder architecture that mirrors the structure of the ResNet-M encoder. In decoding, instead of transposed convolutions we used upsampling with the nearest value followed by a regular convolution (Odena et al. (016)). We used mean squared error between the reconstructed pixels and the target image as the training criterion. Models were trained and evaluated on the same data as reward and action probing, for 30 epochs using Adam optimizer with learning rate 0.001.
A.6 HYPERPARAMETERS
See tables 6
7 8 9 for hyperparameter values. For ATC, ST-DIM and VAE-T hyperparameters, see
Stooke et 2021))
A.7 IMAGE AUGMENTATION
We use the same image augmentations as used in SGI (Schwarzer et al., 202 which itself used the augmentations in DrQ (Yarats al. (202T)), in both pretraining and tnn-iuning. We specifically apply random crops (4 pixel padding and 84x84 crops) and image intensity jittering.
A.8 GOAL-ORIENTED RL LOSS
Goal-oriented RL loss is taken directly from SGI (Schwarzer et al. 2021b). This objective trains goal-conditional DQN, with rewards specified by proximity to sampled goals. First, goal is sampled to be the state encoding either of the near future in the current trajectory (up to 50 steps in the future), with probability of 20%, of the future state in another trajectory in the current batch. Then, we add Gaussian noise to obtain the final goal g: < an + (1 azgg, where a ~ Uniform (0.5). and n is vector sampled from isotropic Gaussian normalized to have length of 1. Then, in order to obtain the reward of taking action de going from state 8t to 81+1, we first encode the states with the target encoder é (oCtarget(o)). &,
ENCtarget (0t+1). Then, calculate the reward as: 8.19
R(ē,, ( d(ē,, d(ē++1,g). where d(ē.. exp @ 12zlloll 2). We use FiLM (eree et al. 2018 ct al. 2015).
condition the Q-function Q(o,, (r,g) on g. and optimize the model using DQN Mnih
B FORWARD MODEL PROBING
While our principal goal is to demonstrate the correlation between representation probing and offline RL performances, also apply the reward probing technique to predictions in order to evaluate the qualities transition models under different pretraining setups.
In table I0 show the effects of using different transition models during pretraining on prediction probing performance. All models are trained with ResNet-M encoder and inverse loss. Goal loss is also applied to the BYOL models.
15

---

Under review as a conference paper at ICLR 2023
Table 6:
Hyperparameters for pretraining and RI. evaluation.
Parameter
Setting
Pretrain & RL.
Gray-scaling
True
Observation down-sampling
84x84
Frames stacked
4
Action repetitions
4
Sticky Action
True
Reward clipping
[-]. I]
Terminal on loss of life
Truc
Optimizer
Adam
Optimizer: learning rate
0.0001
Optimizer: 31
0.9
Optimizer: 32
0.999
Optimizer: c
0.00015
Minibatch Size
64
Max gradient norm
10
Pretrain
Prediction Depth Epochs
10
20
Goal loss weight
0 or 1
Inverse loss weight
0 or 1
RL
Max frams per episode Update
108K
Distributional Q
Dueling
True 51
Support of Q-distribution
Discount factor
0.99
Priority exponent
0.5
Priority correction
0.4 + 1
Exploration
e-greedy
Training steps
100K
Evaluation trajectories
100
Min replay size for sampling
2000
Replay period every
1 step 2
Updates per step
Multi-step return length
10
network: channels
32,64,64
network: filter size
8x8, 4x4, 3x3
network: stride
4,2.1
network: hidden units Non-linearity Target network: update period
256
ReLU
Table SSL specific hyperparameters.
Table 8: GRU-latent specific hyperparameters.
Parameter
Setting
Parameter
Setting
BYOL
loss weight
kl loss weight
0.1
0.99
kl balance
0.95
Barlow
loss weight
0.002
0.0051
Table 9: Optimal RI. learning rate for different setups. Identified by sweeping through 2, 5e-5. le 4, 2e- and evaluated on games frostbite, assault, gopher, and demon attack.
Encoder
Transition
Objectives
Learning Rate
ResNet-M
Conv-det
BYOL, goal. inv
2c
ResNet-M
GRU-det
BYOL, goal, inv
20 4
ResNet-M
GRU-latent
BYOL, goal, inv
2e 4
ResNet-M
GRU-latent
Barlowon. inv
le-4 4
ResNet-M
GRU-latent
Barlow inv
5e-5
ResNet-L
GRU-latent
Barlowa7. inv
5e-5
ResNet-L
GRU-latent
Barlow inv
1e-4
16

---

Under review as a conference paper at ICLR 2023
Table 10: Mean reward probing F1 scores for pretraining setups with different transition mod-
Table 11: Mean reward probing Fl scores for pretraining setups with different prediction ob- jectives. Evaluated on 5 and 10th predictions. All standard deviations are order of le-4.
els. Evaluated on 5th and 10th predictions. All standard deviations are order of le-4.
Pred Obj
Transition
Pred 5
Pred
Pred Obj
Pred 5
Pred
BYOL
Conv-det
33.1
28.4
BYOL
33.4
28.9
GRU-det
33.0
27.4
Barlowo.s
40.2
30.2
GRU-latent
33.4
28.9
Barlowo.7
39.5
30.2
Barlow
37.4
29.7
Barlowo.7
Conv-det
32.0
27.6
Barlowrand
36.8
27.5
GRUU-det
30.1
25.0
GRU-latent
39.5
30.2
In the deterministic setting, the predictions of the GRU model worse than those of the convolutional model. The introduction of stochasticity appears to fix the underlying issue for predictions, resulting the latent GRU model having the best overall prediction probing performance.
One possible explanation for Conv-det having better predictions than GRU-det is that the spatial inductive bias the convolutional kernels acts as constraint and helps regularize the predictions from regressing to the mean. However, this is more effectively solved by the introduction of latent variables into GRU during training and inference.
In table 11 we show the effects of using different prediction objectives during pretraining on pre- diction probing performance. All models are trained with ResNet-M encoder, GRU-latent transition model, and inverse loss; goal loss is also applied to the BYOL model.
Comparing to the BYOL model, Barlow models generally have higher probing scores predictions. We also note that for Barlow models, regularizing the representations towards the predictions (by setting Barlow Balance < I) improves the qualities of predictions. This is likely because it makes the prediction task casier, making it more likely to learn capable transition model.
This reasoning can also explain why the Barlow model with frozen, random target network achieves superior probing result for representation (table 2 but worse result for predictions compared to the other Barlow versions. Predicting a random target representation is likely more difficult than predicting learned representation, and this may in turn encourage the model rely more learning a powerful encoder and posterior model, and less on learning an accurate transition model.
C FULL RL RESULTS
Table 12: Full RL Results for representative pretraining setups. Setup names are represented (encoder]-[transiitee) model)-(ssl losses). M and refer to ResNet M and ResNet L, CD is convolutional model, GD is deterministic GRU, GL is latent GRU, By and Bt refer to Byol and Barlow, and refer to goal and inverse losses.
Amidar Assault Asterix Boxing DemonAtt Frostbite Gopher Seaquest Krull
Random
5.8
222.4
210.0
0.1
152.1
65.2
257.6
68.4
1598.0
Human
1719.5
742.0
8503.3
12.1
1971.0
4334.7
2412.5
42054.7
2665.5
VAF-T
56.1
526.8
355.3
28.1
830.8
394.0
484.2
311.9
3907.3
ST-DIM
64.7
492.2
335.0
19.8
524.5
269.7
354.2
376.1
2650.6
ATC
100.8
653.8
376.9
21.5
812.3
1002.1
601.5
456.4
3611.8
M-CD-ByGI
169.6
693.1
393.1
54.5
458.1
1058.9
1323.4
461.7
5541.4
M-CD-Bta,Gi GI
206.1
545.6
500.0
21.5
357.7
518.5
880.5
482.4
4216.0
M-GD-ByGI
204.5
552.6
625.2
51.0
723.5
979.7
1299.2
597.1
5006.3
M-GL-ByGI
170.9
392.0
527.2
49.1
1842.9
541.9
1489.7
609.9
4753,9
M-GL-Bta7
97.9
846.0
442.5
$3.9
311.5
461.5
731.0
622.1
4176.4
M-GL-Bta;I
189.9
861.8
426.4
63.2
1048.8
2020.1
857.6
$79.0
5111.4
M-GL-Btrandt
161.8
954.6
569.1
59.6
4373.0
1067.4
1068.8
734.5
5422.6
J-GL-Bttzl
173.5
1072.1
540.0
72.6
1143.9
1633.4
1274.1
578.7
5383.4
L-GL- Gll
136.3
1273.7
506.5
64.0
4112.8
1163.7
1594.3
653.1
5453.6
17

---

Under review as a conference paper at ICLR 2023
D STATISTICAL HYPOTHESIS TESTING OF RANK CORRELATION
Spearman's 's 0.967
0.00 (883)
Spearman's 0,808
F2 0.0006 (***)
RL. HNS (IQM)
RL HNS (IQM)
0.5
0.5
-0.5
60
70
15
25
30
Reward Prediction Fl (%%
Action Prediction F1 (%)
Figure 5: Correlation between the SSL representations' abilities linearly predict (Left) presence of immediate reward and (Right) action, versus RL performance using the same representations, measured as the interquartile mean of the normannorrmalized score (HNS) over 9 Atari games. Each point denotes a separate SSL pretraining method. A linear line of best fit is shown with 95 confidence interval. We compute Spearman's rank correlation coefficient (Spearman's 's r) and determine its statistical significance using permutation testing (with 50000). Compared to Fig. we added one extra model which obtained poor probing results to demonstrate that the correlations holds for a wide range of performance levels.
In Fig. 5 we show the correlations results for both the action and reward predictions. We estimate Spearman's 's rank correlation coefficient (Spearman's 's r) between the linear probing performance and the (interquartile) mean RL using Spearman's 's instead of the Pearson correlation coefficient is because are interested in whether the relative ranking of the models on the linear probing tasks is indicative the relative ranking of the same models when RL is trained top of it. As example, this allows to say if model A out-ranks model B in the reward prediction task, an RL model trained on top of model A's representations will likely out-perform an RL model trained on top of model B representation. However, it does not let us predict by how much model A will out-perform model B.
score (HNS) over 9 Atari games. The reason for
Let denote the difference in ranking between the linear probing performance and the RL perfor- mance, Spearman's 's (denoted as below) is computed as,
6
6
1
where d; is the difference in ranking for the i-th model, and n is the total number of models we have.
We perform statistical hypothesis testing on p with null hypothesis م (no correlation between linear probing performance and RL performance) and alternative hypothesis tion). The null distribution is constructed nonparametrically using permutation testing: we sample random orderings of the observed linear probing performance and RI. performance independently and compute p. This is repeated 50,000 times to generate the null distribution (which is centered at p = we do not expect randomly ordered values to be correlated). We then compare our observed to this distribution and perform one-tailed test for the proportion of samples larger than our observed report p-value.
(positive correla-
D. RANK CORRELATION ON A DIFFERENT DATASET
In Fig. [ we explored the corrllation between the RL performance and the reward probing task, where the dataset used for the reward probing was set of quasi-random trajectories from the DQN dataset, coming from very beginning of the training run of the DQN agent used to collect the data. It is natural to ask whether the correlation results we obtain are sensitive to the specific dataset used. To put this question to the test, re-run the same reward probing task, this time on the "expert" dataset, Le. the last trajectories of the DQN dataset, corresponding to a fully trained agent. The results are shown in Fig 6 The Spearman's correlation coefficient that we obtain is the exact same the one for
18

---

protocol detailed in Under review as a conference paper at ICLR 2023
Spearman's 's 0.952
p 0.00004 (***)
RL HINS (IQM)
0.5
0
-0.5
30
40
50
60
70
Reward Prediction FI (%)
Figure 6: Correlation between SSL representations' abilities to linearly predict presence of immediate reward different probing dataset (expert trajectories instead of random ones), for the set of models using Barlow and BYO objectives.
the random trajectory dataset (even though the reward statistic are different, see Table 14). showing that the correlation result not sensitive the probing dataset used.
D.2 CONFIDENCE INTERVAL OF RL PERFORMANCE AS FUNCTION OF INDEPENDENT RUNS
We further show the confidence interval of the estimated mean RL performance the number of independent runs increase. From our total of 10 independent runs each game, we sample with replacement & < 10 runs (k being number independent runs we "pretend" to have instead of the full 10), independently for each game. We can compute the IQM over this sample to get an estimate for the IQM as if we only have independent runs. We repeat this process 10,000 times construct the 95 confidence interval of the empirical IQM for different k's. Illustrative examples of how much this confidence interval shrinks for different pairs of models shown Fig.
We observe in Fig. 7 the mean RL performance estimates have CIs that eventually separate with many independent runs. This is an unbiased but high variance and computationally intensive estimator the true expected RI. performance. On the other hand, the reward prediction FI score is computationally cheap, low variance and accurate estimator the relative model ranks in mean RL performance. This further corroborates our previous results of positive correlation between reward prediction Fl score and mean RL performance (Fig.
E COMPARISON WITH DOMAIN SPECIFIC PROBING BENCHMARKS
Table 13: Comparison of the correlation of domain specific and domain agnostic (reward prediction) probes with the RL performance.
Spearman's
AtariARI Reward (ours)
0.527
0.058
0.782
0.003
One of the key advantages our probing method is that it is domain agnostic, unlike the previously proposed AtariARI benchmark (Anand et al., (219) which acquires probing labels through the RAM state the emulator, making their method impractical for image-based trajectories.
To better understand how our probing metrics compare with the domain specific ones in terms of correlations with RL performances, we perform the AtariARI probing benchmarks using our pretrained encoders the overlapping games (Boxing, Seaquest, Frostbite, DemonAttack) used in both works. For AtariARI, we first calculate the average probe scores across categories, then average this quantity across the games. For reward probing, apply
19

---

Under review as a conference paper at ICLR 2023
RL HNS (95 (QM)
RL HNS (shaded)
75
0.8
(line)
0.6
700
0.4
65
60
6
8
10
Number Seeds
Resnet-T., GRU-latent, Bariowez
Resnet-M, CRU-Jatent, BYOL & goal inv
1.4
RL HNS (shaded)
Reward (line)
74
Mean 60
RL HNS (95 CI (QM)
RL HNS (shaded)
Reward F1 (line)
70
Reward F7 Mean 0%
1.2
HNS C1
1.0
70
0.4
66
0.8
Reward
66
0.2
0.6
62
10
Number of Seeds
Number Seeds
Resnet-L, GRU-latent, Barlowrand Resnet-M. M, GRU latent, Barlowaard & inv
Resnet-M, CRU-latent, Barlowe. & Inv Resnet-M, Conv-det, BYOL goal inv
Figure 7: Illustrative examples of 95 confidence interval (CT) of the interquartile mean estimate (IQM) of the human normalized score (HMS) for different models. The CI shrinks function of the number of independent runs. CI is estimated using 10,000 bootstrapped samples with replacement. Model names are given below each figure, note the figure colors matches the colors in table The top figure is the same as Fig. [ (Right). For all cases, the reward prediction Fl score gives us an accurate low-variance estimate for how the two models rank relative one another.
section 5.1. For RL performance we use the IQM. We report the correlation between the probing metrics and RL performances across different models.
Our results are summarized in Table We find that the correlation between the average probing Fls and RL performances is stronger for our reward probing method. In particular, our probing method has a significant correlation with RL performances (p < 0.05), while the AtariARI probing method does not.
F PROBING DURING TRAINING
We show evolution of probing performance as training progresses in figure g
G REWARD STATISTICS IN PROBING DATASETS
In table report the percentage of states that have non-zero reward in each of the games, for two different subsets of data:
Checkpoint which correspond quasi-random trajectories from the beginning of the training process of DQN. This is the data used for the reward probing in Fig 1.
Checkpoint 50, which is the last checkpoint of the DQN replay dataset, and corresponds to the fully trained DQN agent, that we assimilate to an expert agent. This data is used for action probing, and for reward probing in Fig
20

---

Under review as a conference paper at ICLR 2023
60.0%
Mean Reward F1
55.0%
45.0%
40.0%
Resnel-M, M, GRU-latent, Barlowrund & inv
Resnet-M, M, Conv-det, BYOL & goal & inv
35.0%
10 11 epoch
18
Figure 8: Average reward probing Fis for two SSL setups during different training epochs. Epoch constitutes untrained model.
Table 14: Percentages of positive rewards checkpoints and 50 the DQN replay dataset for games. Checkpoint is used for reward probing and checkpoint 50 used for expert action probing.
Game
Ckpt 1 %
Ckpt 50 %
Amidar
2.7
5.2
Assault
3.6
6.8
Asterix
5.0
6.0
Boxing
3.5
9.3
DemonAttack
2.1
4.7
Frostbite
4.2
2.9
Gopher
2.8
8.5
Krull
13.2
41.7
Seaquest
1.5
7.5
All the games have a fairly small percentage of positive reward states, and we generally observe higher percentage of reward in checkpoint 50, which is expected since the agent is more capable by then.
0.8
assault
0.8
boxing
0.6
Spoarman 1ho
demon attack trastbite
Spearman rho
gopher
0.4
krull
0.2
suaquest
0.2
0.15
0.30
0.45
0.60
-0.2
10
Rlwward percent
Reward percent
Figure 9: Left: Spearman's correlation coefficient (for the set of models using Barlow and BYOL objectives) between the RL performance each individual game and the reward probing Fl, plotted as a function the percentage of rewards observed in this game. Right: p- p-values associated with each of the Spearman's coefficients.
21

---

Under review as a conference paper at ICLR 2023
G.1 IMPACT OF SPARSITY ON THE CORRELATION
In Fig we plot the Spearman's correlation coefficient between the RL performance on each individual game and the reward probing Fl, a function of the percentage of reward observed in each game (see Table 14p. We do not observe any particular pattern with respect to the sparsity, suggesting that the probing task is not very sensitive the sparsity level of each individual game. Note however that, as usual in the Atari benchmark, it is difficult to draw conclusion from any given individual game, and the statistical significance of our results only emerge when considering the set of games as whole. Indeed, only 3 games achieve individual statistical significance at < 0.01 (Boxing, Seaquest and Assault), while the other do not obtain statistically significant correlations.
H LIMITATIONS
One limitation of the current work is that for the presented probing methods to work one needs a subset of the data either with known rewards, where ideally rewards are not too sparse, or with expert actions. If none of the two is available, our method cannot be used. For the reward probing task, the usefulness of the method also depends the hardness of the reward prediction itself. If the prediction task is too easy, for example because there are rewards at every step, or because the states with rewards are completely different than the without (such that even randomly initialized model would yield features allowing linear separation between the two types of states), then the performance all the models this task are going to be extremely similar, with the only differences coming from random noise. In such a case, the performance of the prediction task cannot be used to accurately rank the quality of the features of each of the models.
For future work we also would like to extend the findings of this paper to more settings, for example different environments.
22


Reviews:
[{"review_id": "y3peFpzM-w", "reviewer": "Reviewer_vEah", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper investigates whether or not light-weight probings to the action between continuous states and reward can measure the pretrained encoder performance on RL tasks. To do this, they pretrained the encoders through self-supervised learning loss with a transition model implemented through recurrent module or recurrent state space modeling with several variations. They tested the pretrained encoder for linear probing and RL tasks. They showed the correlations between the performances on linear probings and RL tasks. They found the linear probing to reward is highly correlated with performance on RL tasks, while the relationship between action linear probing and RL performance is weaker.", "summary_of_the_review": "This paper investigates the linear probing to the action, and reward could be an indicator of RL performance. Mainly they evaluated the SSL methods, BYOL, and Barlow Twins for 9 Atari tasks. It is interesting and could be helpful for others because applying the pretrained encoder to RL tasks is one of the ways to improve the sample efficiency on RL, but the evaluation is very expensive. However, I felt their writing is rushed, so I think minor revision is necessary. ", "strength_and_weaknesses": "Strength\n- The motivation behind this paper is good. Applying the pretrained encoders to RL tasks has been investigated [1,2]. It requires lots of resources. This investigation could be a piece of good evidence to skip the costly evaluation.\n- The many variations and ablations are evaluated. For transition modeling, deterministic recurrent modules and RSSM are validated. For SSL, BYOL and Barlow Twins are tested with various configurations. The ablation studies are reported with and without each objective, such as inverse dynamics modeling and goal-conditioned RL.\n\nWeaknesses\n- They only investigated the SSL methods, not other unsupervised methods such as VAE. I expected they would cover the overall methods from their title, but it is not.\n- They only evaluated nine games. It could not be enough to back up the conclusion.\n- In the equation for reward-reg loss in Reward Probing in section 3.2, should the encoder get $o_{t+1}$ not $o_t$? Because the reward $r_t$ is given when the action is given on the observation $o_t$. \n- In the equation for action-classif loss in Action Prediction in section 3.2, shouldn't the input be $o_t$ and $o_{t+1}$? In Figure 2, the consecutive observations are given to the loss, but not in this equation.\n- In Figure 2, I cannot understand the below sentence. Why is the stacked observation related to data augmentation?\n    - > The observations consist of a stack of 4 frames, to which we apply data augmentation before passing them to a convolutional encoder.\n-  > The action is represented as a 2D one-hot vector and appended to the input to the first convolutional layer.\n    - Why is the action represented through a 2D one-hot vector? The action space is larger than the 2 Dimension.\n- For RSSM, you used the discrete latent. Why didn't you try the continuous latent version [3]? Perhaps, because the discrete version outperforms the continuous latent version in [4], but the discrete latent variable training is more unstable than the continuous latent variable training, so maybe the RSSM with continuous latent variable could outperform the discrete latent version.\n- In BYOL loss equation in section 4.2, should $q(\\hat{y}_{t+k})$ be $q(\\hat{e}_{t+k})$?\n- For Algorithm 1, why did you use the Pseudo code block? It is just a single equation.\n- For the goal-oriented RL loss, please introduce it roughly in the main paper even though the details are in Appendix.\n    - In A.7, there are typo $\\tilde{e}_t+1$.\n- > Similar to the BYOL model, the Barlow model can also be improved with inverse dynamics modeling, while the addition of goal loss has a slight negative impact.\n    - It is interesting. Could you analyze this?\n\n[1] Schwarzer, Max, et al. \"Pretraining representations for data-efficient reinforcement learning.\" Advances in Neural Information Processing Systems 34 (2021): 12686-12699.\n\n[2] Dittadi, Andrea, et al. \"The role of pretrained representations for the ood generalization of rl agents.\" arXiv preprint arXiv:2107.05686 (2021).\n\n[3] Hafner, Danijar, et al. \"Dream to control: Learning behaviors by latent imagination.\" arXiv preprint arXiv:1912.01603 (2019).\n\n[4] Hafner, Danijar, et al. \"Mastering atari with discrete world models.\" arXiv preprint arXiv:2010.02193 (2020).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "G3paIwoJ4X", "reviewer": "Reviewer_LabE", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper develops an evaluation protocol for unsupervised visual pretraining. They learn linear probes to predict expert agent actions and rewards from encoded states. These probes provide a more cost-efficient way of comparing visual representation learning methods for RL. The evaluation protocol is tested on a handful of Atari tasks and the authors show that the performance of networks on the linear probes well correlates with RL performance.", "summary_of_the_review": "The paper tackles an interesting problem, but I think it should include evaluation over a larger set of ", "strength_and_weaknesses": "This paper introduces an interesting idea for the important problem of cost-effective evaluations of visual representations. Currently results are limited to Atari and the evaluation is only performed for a small number of models + self-supervised losses. Because the main goal of this paper is to provide an evaluation protocol that can be used in place of downstream RL performance, it would be helpful to see a much broader range of losses and models as well as more difficult control tasks. It's also unclear how well these evaluation protocols predict downstream performance in the presence of task transfer: e.g., one goal in developing visual representation pre-training methods is to get a good generalizable encoder. Because the evaluation leverages reward information and information about the optimal policy, it doesn't seem like it would predict the fitness of an encoder for new tasks.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "FGDHxiYmX3x", "reviewer": "Reviewer_etov", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper attempts to propose an evaluation protocol for lightweight probing of unsupervised representations and investigates the correlation between RL performance and linear probing from a pretrained representation. Authors are testing this on a very specific class of self-predictive (recurrent) representation models that are being trained with SSL objectives. The authors extensively analyse several design choices of their studied model regarding performance on the two probing tasks of predicting reward or action from a held-out labelled train set. Finally, the paper aims to present some correlation between probing and RL performance on 9 Atari games.\n", "summary_of_the_review": "While this paper attempts to address an important challenge in RL and aims to propose a generally applicable evaluation protocol for unsupervised pretrained RL representations I am not fully convinced that the paper holds up to these promises and several claims in the paper. I am willing to change my opinion in case I missed central parts of the paper but want to encourage the authors to further improve the manuscript based on some suggestions and comments above and recommend rejection.", "strength_and_weaknesses": "Strengths\n\n1. Identification of an important problem setup in RL, that is how to assess which representations and pretraining data is best suited for improving downstream RL performance.\n2. Extensive analysis of various design choices of the model architecture studied\n\nWeaknesses\n1. It is not very clear what the goal or objective of the paper is. Authors say they propose an evaluation protocol for unsupervised RL representations that saves up to 600x computation cost. However, this computational cost saving and protocol seems to be not addressed or described in detail in the main paper and doesn't appear to be the main focus of the paper. Instead, the authors analyse a very specific class of self-predictive (recurrent) representation models that are being trained with SSL objectives. I would have expected to see a more thorough coverage of different unsupervised representation learning methods and more empirical analysis to support these claims (see additional comments in point 3)\n2. It is not very clear what is actually novel about the proposed model and what is based on prior work. The listed contributions are very vague. I would hope the authors can clarify what exactly their contributions are. Also, there exist additional prior works that very extensively studied various light-weight probing tasks on unsupervised representations and how their performance relates to RL performances [1,2]\n3. I am a bit confused by the experiments, especially the correlative analysis and it is not very clear to me if this holds beyond the very particular method and environments. In particular Figure 3 is very confusing and I do not fully understand what the 7 representative setups are supposed to be; how this relates to the 9 rows/models presented in Figure 3 (where do they come from?); and how this relates to the 7 Atari games studied in Figure 1. Without this additional information, it is not clear to me if it is sound to draw a general conclusion about the correlation between probing task performance and RL performance or if there are any other confounders. Also, can authors provide error bars for results in Table 1-5?\n4. Minor: There appears to be a slight mismatch between the title/abstract and the presented experiments and contributions of the main paper. The authors study a very particular narrow setting but the title suggests a generally applicable evaluation protocol for different unsupervised representation learning methods. \n5. Regarding the two probing tasks I wonder how generally applicable they really are. E.g. the reward task seems to be constrained to have labelled data from a very early point in RL training (i.e. rather random policy), whereas the action prediction is limited to labelled data from close to expert trajectories. To be generally applicable I would have hoped to see both probes being trained on the same labelled dataset to compare apples to apples. \n\n[1] Dittadi, Andrea et al. “The Role of Pretrained Representations for the OOD Generalization of Reinforcement Learning Agents”, ICLR 2022.\n\n[2] Higgins, Irina, et al. \"Darla: Improving zero-shot transfer in reinforcement learning.\" International Conference on Machine Learning. PMLR, 2017.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "0t8sTNccIZ", "reviewer": "Reviewer_Vgoa", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a method for evaluating unsupervised representation learning in reinforcement learning.  Using a linear probe on top of frozen, pretrained representations, the paper suggests learning to predict reward values from various states in downstream tasks.  Additionally, the paper uses a linear probe to predict expert actions from learned representations.  They authors show evidence that, for a selection of representation learning approaches, the F1 score of the linear probe correlates strongly with full reinforcement learning on the downstream task.", "summary_of_the_review": "The paper presents an interesting approach to an interesting problem, with the promise of helping evaluate representations much more quickly.  Although the ablation studies are thorough, there could be a broader comparison to other styles of representation learning, which would significantly strengthen the claims that linear probing for reward prediction correlates well with downstream RL training.", "strength_and_weaknesses": "The paper tackles a very difficult and relevant problem, that of evaluating self-supervised representations.  The paper shows evidence that linear probing can give strong indications of eventual RL training performance, which promises to shorten evaluation time and could be impactful in the representation learning for reinforcement learning field.  My main concern with the paper is the lack of diversity in methods used to assess the correlation between linear probes and RL training performance.  All methods compared are ablations of the self-predictive representation approach described in the paper.  While these are important and elucidating experiments, I would like to see a broader set of methods compared, like augmentation-based representations (DrQ or CURL).  Do these correlations hold in these cases as well?  Also, I'm curious about the noise in the linear probe F1 score.  Do the numbers reported in the tables stay the same regardless of random seed?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject