Learning Representations for Pixel-based Control: What Matters and Why?
Authors: Manan Tomar
Utkarsh Aashu Mishra
Amy Zhang
Matthew E. Taylor
Venue: ICLR 2022 Submitted (2022)
Keywords: Reinforcement Learning
Representation Learning
Pixel-based Control

Abstract: Learning representations for pixel-based control has garnered significant attention recently in reinforcement learning. A wide range of methods have been proposed to enable efficient learning, leading to sample complexities similar to those in the full state setting. However, moving beyond carefully curated pixel data sets (centered crop, appropriate lighting, clear background, etc.) remains challenging. In this paper, we adopt a more difficult setting, incorporating background distractors, as a first step towards addressing this challenge. We present a simple baseline approach that can learn meaningful representations with no metric-based learning, no data augmentations, no world-model learning, and no contrastive learning. We then analyze when and why previously proposed methods are likely to fail or reduce to the same performance as the baseline in this harder setting and why we should think carefully about extending such methods beyond the well-curated environments. Our results show that finer categorization of benchmarks on the basis of characteristics like the density of reward, planning horizon of the problem, presence of task-irrelevant components, etc., is crucial in evaluating algorithms. Based on these observations, we propose different metrics to consider when evaluating an algorithm on benchmark tasks. We hope such a data-centric view can motivate researchers to rethink representation learning when investigating how to best apply RL to real-world tasks.

Under review as a conference paper at ICLR 2022
LEARNING REPRESENTATIONS FOR PIXEL-BASED CONTROL: WHAT MATTERS AND WHY?
Anonymous authors Paper under double-blind review
ABSTRACT
Learning representations for pixel-based control has garnered significant attention recently in reinforcement learning. A wide range of methods have been proposed to enable efficient learning, leading to sample complexities similar to those in the full state setting. However, moving beyond carefully curated pixel data sets (centered crop, appropriate lighting, clear background. etc.) remains challenging. In this paper, we adopt more difficult setting, incorporating background distractors, as first step towards addressing this challenge. We present simple baseline approach that can learn meaningful representations with no metric-basee learning, no data augmentations, no world-model learning, and no contrastive learning. We then analyze when and why previously proposed methods are likely to fail or reduce to the same performance the baseline in this harder setting and why we should think carefully about extending such methods beyond the well curated environments. Our results show that finer categorization of benchmarks on the basis of characteristics like density of reward, planning horizon of the problem, presence of task-irrelevant components, etc., is crucial in evaluating algorithms. Based these observations, we propose different metrics to consider when evaluating an algorithm on benchmark tasks. We hope such data-centric view can motivate researchers to rethink representation learning when investigating how to best apply RL to real-world tasks.
INTRODUCTION
Learning useful representations for downstream tasks a key component for success in rich ob- servation environments [14, 39, 47, 54, 55). Consequently, a significant amount of work proposes various representation learning objectives that can be tied to the original reinforcement learning (RL) problem. Such auxiliary objectives include the likes of contrastive learning losses [42, 36, 11], state similarity metrics like bisimulation or policy similarity [62, 61, 1], and pixel reconstruction losses [29, 20, 25, 24]. On a separate axis, data augmentations have been shown to provide huge performance boosts when learning to control from pixels [35, 33]. Each of these methods has been shown to work well for particular settings and hence displayed promise to part of a general purpose representation learning toolkit. Unfortunately. these methods were proposed with different motivations and tested on different tasks, making the following question hard to answer:
What really matters when learning representations for downstream control tasks?
Learning directly from pixels offers much richer applicability than when learning from carefully constructed states. Consider the example of a self-driving car, where it is nearly impossible to construct complete state description such that it describes the position and velocity of all objects of interest, such as road edges, highway markers, other vehicles, etc. In such real world applications, learning from pixels offers a much more feasible option. However, this requires algorithms that can discern between task-relevant and task-irrelevant components in the pixel input, i.e., learn good representations. Focusing on task-irrelevant components can lead to brittle or non-robust behavior when put in slightly different environments. For instance, billboard signs over buildings in the background has no dependence on the task in hand while a self- diff-drving car tries to change lanes. However, if such task irrelevant components are not discarded, they can lead sudden failure when the car drives through different environment, say a forest where there are no buildings or billboards. Avoiding brittle behavior is therefore key efficient deployment of artificial agents in the real world.

---

Under review as a conference paper at ICLR 2022
Methods: Baseline
Baseline.v0
Contrastive
Metric
SSL w/ Contrastive
Reconstruction
Probability of 1 %)
100
80
60
20
C1
C2
C3
C4
C5
C6
C7
C8
Dataset Categories
Figure 1: Comparing pixel-based RL methods across finer categorizations of evaluation benchmarks. Each category 'Cx' denotes different data-centric properties of the evaluation benchmark (e.g-. C1 refers discrete action, dense reward, without distractors, category and the algorithms are provided Table 4 and Table 5. Baseline-vo refers to applying the standard deep RL agent (e.g.. Rainbow DQN [53] and SAC [23]): Baseline refers to adding reward and transition prediction to baseline-vO, as descrihed in Section 3; Contrastive includes algorithms such as PI-SAC [38] and CURL [36]; Metric denotes the metric losses such DBC [62]; SSI. w/o Contrastive includes algorithms such as SPR [46]; Reconstruction' includes DREAMER MER [25] and TIA [19]. For a given method, we always consider the best performing algorithm. Every method leads to varied performance across data categories, making a comparison which is average across all categories highly uninformative.
with data randomly cropped [33, 35J. Exact descriptions each
There has been a lot of work recently that tries to learn efficiently from pixels. dominant idea throughout prior work has been that of attaching an auxiliary loss to the standard RL objective, with the exact mechanics of the loss varying for each method [29, 62, 36]. A related line of work learns representations by constructing world models directly from pixels [45, 41, 22, 25]. We show that these work well when the world model simple. However, as the world model gets even slightly more complicated, which is true of the real world and imitated in simulation with the use of video distractors [60, 31, 48], such approaches can fail. For other methods, it is not entirely clear what component/s in auxiliary objectives can lead to failure, thus making robust behavior hard to achieve. Another distinct idea is of using data augmentations [35, 33] over the original observation samples, which seem to be quite robust across different environments. However, we will show, lot of the success of data augmentations is an artifact of how the benchmark environments save data, which not true of the real world [48], thus resulting in failure'. It is important note that some of these methods are not designed for robustness but instead for enhanced performance on particular benchmarks. For instance, the ALE [7] benchmark involves simple, easy model objects, and becomes hard to discern if methods that perform well are actually good candidates for answering 'what really matters for robust learning the real world.'
Contributions. In this paper, we explore the major components responsible for the successful application of various representation learning algorithms. Based on recent work in RI. theory for learning with rich observations [17, 4, 9], hypothesize certain key components to be responsible for sample efficient learning. We test the role these play in previously proposed representation learning objectives and then consider exceedingly simple baseline (see Figure 2) which takes away the extra "knobs" and instead combines two simple but key ideas, that of reward and transition prediction. We conduct experiments across multiple settings, including the MuJoCo domains from DMC Suite [51] with natural distractors [60, 31, 48]. and Atari100k [30] from ALE [7]. Following this, identify the failure modes of previously proposed objectives and highlight why they result in comparable or worse performance than the considered bascline. Our observations suggest that relying on particular method across multiple evaluation settings does not work, the efficacy varies with the exact details of the task, even within the same benchmark (see Figure 1). We note that finer categorization of available benchmarks based on metrics like density of reward, presence of task- irrelevant components, inherent horizon of tasks, etc., play crucial role in determining the efficacy of a method. We list such categorizations as suggestions for more informative future evaluations. The findings of this paper advocate for more data-centric view of evaluating RL algorithms [13], largely missing in current practice. We hope these insights can lead to better representation learning objectives for real-world applications.
It is also hard to pick exactly which data augmentation will work for a particular environment or task [43]. For ALE use the performance of DREAMER after 1M steps, whereas for DMC we consider the nerformance after 500k steps.
2

---

Under review as a conference paper at ICLR 2022
2 RELATED WORK
Prior work on auxiliary objectives includes the Horde architecture [50], UVFA [44] and the UN- REAL agent [29]. These involve making predictions about features or pseudo-rewards, however only the UNREAL agent used these predictions for learning representations. Even so, the benchmark en- vironments considered there always included only task-relevant pixel information, thus not pertaining to the hard setting consider in this work. Representations can be also be fit as to obey certain state similarities. If these state metrics preserve the optimal policies and are casy to learn/given a priori, such a technique can be very useful. Recent works have shown that we can learn efficient representations either by learning the metrics like that in bisimulation [18, 62, 61, 8]. by recursively sampling states [10] by model invariance [52]. Data augmentations modify the input image into different distinct views, each corresponding to a certain type of modulation in the pixel data. These include cropping, color jitter, flip, rotate, random convolution, etc. Latest works [35, 58] have shown that augmenting the states samples in the replay buffer with such techniques alone can lead to impres- sive gains when learning directly from pixels. Recently. Stone et al. [48] illustrated the advantages and failure cases of augmentations. Contrastive learning involves optimizing for representations such that positive pairs (those coming from the same sample) are pulled closer while negative pairs (those coming from different samples) are pushed away [42, 11]. The most widely used method to generate positive/negative pairs is through various data augmentations [36, 46, 38, 49, 57J. However, temporal structure can induce positive/negative pairs as well. In such case, the positive pair comes from the current state and the actual next state while the negative pair comes from the current state and any other next state in the current batch [42]. Other ways of generating positive/negative pairs can be through learnt state metrics [1] or encoding instances [25]. Another popular idea for learning representations is learning world models [22] in the pixel space. This involves learning prediction
models of the world in the pixel space using a pixel reconstruction loss [20, 24, 25]. Other methods that do not explicitly learn a world model involve learning representations using reconstruction based approaches like autoencoders [56].
Quite a few papers in the past have analysed different aspects of the RL problem. Engstrom et al. [15] and Andrychowicz et al. [3] have focused on analysing different policy optimization methods with varying hyperparameters. Our focus is specifically on representation learning methods that improve sample efficiency in pixel-based environments. Henderson et al. [27] showed how RL methods in general can be susceptible to lucky seedings. Recently, Agarwal et al. [2] proposed statistical metrics for reliable evaluation. Despite having similar structure, work is largely complimentary to these past investigations. Babacizadch et al. [5] analysed reward and transition but only focused on the Atari 200M benchmark and pixel reconstruction methods. In comparison, our work is spread across multiple evaluation benchmarks, and our results show that reconstruction can be a fine technique only in particular benchmark category.
METHOD
We model the RL problem using the framework of contextual decision processes (CDPs), a term introduced in Krishnamurthy al. [34] to broadly refer to any sequential decision making task where an agent must act on the basis of rich observations (context) Ie to optimize long-term reward. The true state of the environment st is not available and the agent must construct it on its own, which is required for acting optimally on the downstream task. Furthermore, the emission function which dictates what contexts are observed for given state is assumed only inject noise that is uncorrelated to the task in hand, i.e. it only changes parts of the context that are irrelevant to the task [62, 48). Consider again the example of people walking on the sides of road while a self-driving car changes lanes. Invariance to parts of the context that have no dependence on the task, e.g. people in the background, is an important property for any representation learning algorithm since we cannot expect all situations to remain exactly the same when learning in the real world. Detailed description of the setup and all the prior methods used is provided in Appendix 1.
We explore the utility of two main components, that of reward and transition prediction, in learning representations. A lot of prior work has incorporated these objectives either individually or in the presence of more nuanced architectures. Here, aim is to start with the most basic components and establish their importance one by one. Particularly, we use a simple soft actor-critic setup (taking inspiration from SAC-AE (56]) the base architecture, and attach the reward and transition
3

---

Under review as a conference paper at ICLR 2022
3++1
Transition Loss
Reward Loss
Encoded State
fo
8(+1
LT
LR
Stop Gradient
fo
Plola (elee ae)
R(
St
B++1
Pz
Actor
Critic
Reward Prodiction
Prodiction
Observation
C Aster
Covette
Action
Figure 2: (Left) Baseline for control over pixels. We employ two losses besides the standard actor and critic losses, one being a reward prediction loss the other latent transition prediction loss. The encoded state sy is the learnt representation. Gradients from both the transition/reward prediction and the critic are used to learn the representation, whereas the actor gradients are stopped. In the ALE setting, the actor and critic losses replaced by Rainbow DQN loss [53]. (Right) Natural Distractor in the background for standard DMC setting (left column) and custom oef-center setting (right column). More details about the distractors can be found in Appendix 2.
prediction modules to it (See Figure 2). Note that the transition network over the encoded state 8t and not over the observations [37]. Moreover, the transition model is fit between the encoded state and the reward model. Unless noted otherwise, we call this architecture as the baseline for our experiments. Details about the implementation, network sizes and all hyperparameters' is provided in Appendix and Appendix 4 (Table 3) respectively.
4 EMPIRICAL STUDY
In this section, we analyze the baseline architecture across six DMC tasks: Cartpole Swingup. Cheetah Run, Finger Spin, Hopper Hop, Reacher Easy, and Walker Walk. A common observation in experiments is that the baseline is able to reduce the gap to more sophisticated methods significantly, sometimes outperforming them in certain cases. This highlights that the baseline might serve as a stepping stone for other methods to build over. We test the importance of having both the reward and transition modules individually, by removing each of them one by one.
4.1 REWARD PREDICTION
Figure 3 (left) shows comparison of 'with vs without reward prediction' All other settings are kept unchanged and the only difference the reward prediction. When the reward model is removed, there remains grounding objective for the transition model. This results in representation collapse as the transition model loss is minimized by the trivial representation which maps all observations to the same encoded state leading degraded performance. This hints at the fact that without valid grounding objective (in this case from predicting rewards), learning good representations can be very hard. Note that it is not the case that there is reward information available to the agent, since learning the critic does provide enough signal to learn efficiently when there are no distractions present. However, in the presence of distractions the signal from the critic can be extremely noisy since it is based on the current value functions, which are not well developed in the initial stages of training. One potential fix for such a collapse is approaches for the transition model loss and instead use a contrastive version of the loss, which has been shown learn general representations in the self-supervised learning setting. We test this later in the paper and conclude that although it does help prevent collapse, the performance is still heavily inferior to when we include the reward model. Complete performances for individual tasks are shown in Appendix 8.
the standard maximum likelihood based
Linear Reward Predictor. Furthermore, we also compare to the case when the reward decoder is a linear network instead of the standard layer MLP. We see that performance decreases significantly in this case as shown in Figure 3 (middle), but still does not collapse like in the absence of reward prediction. We hypothesize that the reward model is potentially removing useful information for predicting the optimal actions. Therefore, when it is attached directly to the encoded state, i.c., in the linear reward predictor case, it might force the representation only preserve information required
4

---

Under review as a conference paper at ICLR 2022
Reward
w/o Reward
Non-linear Rew Dec
Linear Rew Dec
Trarsition
w/o Transition
Nornalized Score
Nornalzze Score
Normalteed Score
0.3
0.3
0.2
0.2
0.2
0.1
0.1
0.1
0.0
0.0
Environment Steps (x10")
Environment nt (x10)
Enviranment (x10)
Figure 3: Baseline Ablations. Average normalized performance across six standard domains from DMC. Mean
and std err. for runs. Left plot: Baseline with without reward prediction Middle plot: Baseline with non-linear vs linear reward predictor/decoder. Right plot: Baseline with vs without transition prediction.
to predict the reward well, which might not be always be enough to predict the optimal actions well. For instance, consider robot locomotion task. The reward in this case only depends on one variable, the center of mass, and thus the representation module would only need to preserve that in order to predict the reward well. However, to predict optimal actions, information about all the joint angular positions and velocities is required, which might he discarded if the reward model is directly attached to the encoded state. This idea is similar why contrastive learning objectives in the self-supervised learning setting always enforce consistency between two positive/negative pairs after projecting the representation to another space. It has been shown that enforcing consistency in the representation space can remove excess information, which hampers final performance [11]. We indeed see a similar trend in the RL case as well.
4.2 TRANSITION PREDICTION
Similarly, Figure 3 (right) shows a comparison of 'with vs without transition prediction'. The transition model loss enforces temporal consistencies among the encoded states. When this module removed, we observe a slight dip performance across most tasks, with the most prominent drop in cartpole as shown in Appendix 8.1 (Figure 16). This suggests that enforcing such temporal consistencies in the representation space is indeed an important component for robust learning, but not a sufficient one. To examine if the marginal gain is an artifact of the exact architecture used, we explored other architectures in Appendix 8.2 but did not observe any difference performance.
4.3 CONNECTIONS TO VALUE-AWARE LEARNING
The baseline introduced above also resembles a prominent idea in theory, that of learning value aware models [17, 4]. Value-aware learning advocates for learning model by fitting it to the value function of the task in hand, instead of fitting it to the true model of the world. The above baseline can be looked at doing value aware learning in the following sense: the grounding to the representation is provided by the reward function, thus defining the components responsible for the task in hand and then the transition dynamics are learnt only for these components and not for all components in the observation space. There remains one crucial difference though. Value aware methods learn the dynamics based on the value function (multi-step) and not the reward function (I-step), since the value function captures the long term nature of the task in hand. To that end, we also test more exact variant of the value-aware setup where we use the critic function as the target for optimizing the transition prediction, both with and without a reward prediction module (Table 1). Complete performances are provided Appendix 8.8. We see that the value aware losses perform worse than the baseline. A potential reason for this could be that since the value estimates are noisy when using distractors, directly using these value function estimates as a target does not help in learning a stable latent state prediction. Indeed, more sophisticated value aware methods such as in Temporal Predictive Coding [40] lead to similar scores as the baseline.
Table 1:
Truly value-aware objectives. We report average final score after 500K steps across six standard
domains from DMC.
Baseline
Value-aware (w/ reward) Value-aware (w/o reward)
Average Scores
0.42 + 0.02
0.36 0.03
0.23 0.03
'DBC [62] performance data is taken from their publication.
5

---

Under review as a conference paper at ICLR 2022
RAD Flp
RAD Crop
Baselrn Crop
Baselie (SSL wio Contrastivel
PHSAG (Conéractiee) Contrastbre Row
0.4
0.5
0.3
0.2
0.4
0.2
0.1
0.2
0.0
0.5 0.0
0.0
0.0
0.2
0.4
0.2
0.5
Environment Stops (x10f)
Enviromenn Steps 6xt0
Envitonment Steps
Figure 4: Baseline Ablations. Average normalized performance across six standard domains from DMC. Mean and std emr for 5 runs. Left plot: Baseline vs state metric losses (DBC [62]). The performance of baseline is compared with bisimulation metries employed by DBC è. Middle plot: Data Augmentations. Cropping removes irrelevant segments while flip and rotate do not, performing similar to the baseline. Bascline with random crop performs equally as good as RAD. Right plot: Contrastive and SSL w/o Contrastive. We replace the transition loss of the baseline with a contrastive version (Contrastive + Rew). Further, we consider simple contrastive (PI-SAC (38]) and SSL w/o contrastive (variant of SPR [46] for DMC) losses as well.
5
COMPARISON
Until now, we have discussed why the two modules we identify as being vital for minimal and robust learning are actually necessary. Now we ask what other components could be added to this architecture which might improve performance, as has been done in prior methods. We then ask when do these added components actually improve performance, and when do they fail. More implementation details are provided in Appendix 3.
Metric Losses. Two recent works that are similar to the baseline above are DBC [62] and MiCO [10], both of which learn representations by obeying distance metric. DBC learns the metric by estimating the reward and transition models while MiCO uses transition samples to directly compute the metric distance. We compare baseline's performance with DBC as shown in Figure (left). Note that without metric, DBC is similar to the baseline barring architectural differences such as the use of probabilistic transition models in DBC compared deterministic models in the baseline. Hence, we observe that the performance "without metric" exceeds that of "with metric".
Data Augmentations, A separate line of work has shown strong results when using data augmen- tations over the observation samples. These include the RAD [35] and DRQ [33] algorithms, both of which differ very minimally in their implementations. We run experiments for three different augmentations- 'crop', 'flip', and 'rotate". The "crop" augmentation always crops the image by some shifted margin from the center. Interestingly, the image of the robot is also always centered, thus allowing 'crop' to always only remove background or task-irrelevant information and never remove the robot or task relevrelevant information. This essentially amounts to not having background distractors and thus we see that this technique performs quite well as shown in Figure 4 (middle). However, augmentations that do not explicitly remove the distractors, such as rotate and flip, lead to similar performance the baseline. This suggests that augmentations might not be helpful when distractor information cannot be removed, or when we do not know where the objects of interest lie in the image, something true of the real world. We test this by shifting the robot to the side, thus making the task-relevant components off-center and by zooming out i.e. increasing the amount of irrelevant information even after cropping. We see that performance of 'crop' drops drastically in this case, showcasing that most of the performance gains from augmentations can be attributed to how the data is collected and not to the algorithm itself. Additional ablations are provided in Appendix 8.3.
Table 2: RAD additional ablations. We report average final score after 500K steps across Cheetah Run and Walker Walk domains from DMC. This illustrates that the performance of augmentations is susceptible to quality of data. Also, for the "Zoomed Out" setting, it is worth noting that both crop and flip settle to the same score.
Standard
Off-center
Zoomed Out
RAD Crop
0.34 + 0.14
0.30 + 0.08
0.23 1 0.10
RAD Flip
0.27 + 0.08
0.29 + 0.07
0.23 = 0.07
Contrastive and SSL w/o Contrastive Losses A lot of recent methods also deploy contrastive losses (for example, CPC [42]) to learn representations, which essentially refers to computing positive/negative pairs and pushing together/pulling apart respectively. In practice, this can be done for

---

Under review as a conference paper at ICLR 2022
any kind of loss function, such the encoding function fe [25], or using random augmentations [36, 38], so on and so forth. We test a simple modification, that of using the contrastive variant of the transition prediction loss than the maximum likelihood version. We see that, in Figure 4 (right), the contrastive version leads to inferior results than the baseline, again suggesting that contrastive learning might not add lot of performance improvement, as has been witnessed in the self-supervised literature with methods like SIMSIAM [12]. BARLOW TWINS [59], and BYOL [21] getting similar or better performance than contrastive methods like SIMCLR [1I]. Complete performances are provided in Appendix 8.5.
Note that in the ALE results (Figure 7), SPR [46] leads to the best results overall, which also deploys a specific sim- ilarity loss for transition prediction motivated by BYOL. [21]. We follow the same setup and test a variant of the baseline which uses the cosine similarity loss from SPR and test its performance on DMC based tasks. We again show in Figure (right) that there is very little or no im-
Ranel ne
TVA wi
Full Rocon
DREAMER w/o di2.
DRLAMER ett did
0.6
Normalized Score
0.4
provement in performance as compared to the baseline performance. This again suggests that the same algorith- mic idea can have completely different performance just by changing the evaluation setting' (ALE to DMC).
0.2
0.0
0.0
0.1
0.2
0.3
0.4
0.5
Environm ent Staps (X10%)
Learning World Models. We test DREAMER [25]. a state of the art model-based method that learns world models though pixel reconstruction on two settings, with and with- out distractors. Although the performance in the "without distractors" case is good, we see that with distractors, DREAMER fails on some tasks, while performing inferior to the baseline in most tasks. This suggests that learning world models through reconstruction might only be a good idea when the world models are fairly simple to learn. If world models are hard to learn, as is the case with distractors, reconstruction based Icarning can lead to severe divergence that results in no learning at all. We also compare against the more recently introduced method from Fu et al. [19]. Their method, called TIA[19] incorporates several other modules in addition to DREAMER and learns decoupling between the distractor background and task relevant components. We illustrate the performance of each of the above algorithms in Figure along with a version where we add full reconstruction loss to the baseline. Complete performances are provided in Appendix 8.6.
Figure 5: Pixel reconstruction. Average normalized performance across six DMC do- mains with distractors. Baseline achieves better performance than SOTA methods like DREAMER and TIA
Relevant Reconstruction and Sparse Rewards. Since consider only dense reward based tasks, using the reward model to ground is sufficient to learn good repre- sentations. More sophisticated auxiliary tasks considered in past works include prediction of ensemble of value net-
Baseline
RAD Fiip DREAMER TIA
Part Recons
RAD Crop
1.0
Normalized Score
0.8
works, prediction of past value functions, prediction of value functions of random cumulants, and observation re- construction. However, in the sparse reward case, ground- ing on only the reward model or on past value functions can lead to representation collapse if the agent continues to receive zero reward for long period of time. There- fore, in such cases where good exploration is necessary, tasks such observation reconstruction can help prevent collapse. Although this has been shown to be an effective technique in the past, we argue that full reconstruction can harm the representations in the presence of distractors. Instead, we claim that reconstruction of only the task relevant components in the observation space results in learning good representations [19], especially when concerned with realistic settings like that of distractors. To that end, we show that in the sparse reward case, task-relevant reconstruction 6 sufficient for robust performance. We show
0.6
0.4
0.2
0.0
0.1
0.4
0.5
Environment ent Steps (110)
Figure 6: Reconstruction and augmenta-
tions for sparse settings. Normalized per- formance for ball-in-cup catch domain from DMC.
"TIA [19] performance data is taken from their publication.
'The SPR version without augmentations actually uses two separate ideas for improvement in performance, a cosine similarity transition prediction loss and a separate convolution encoder for the transition network, making it hard to attribute gains over the base DER [53] to just tramsition loss.
6 Part Recons. in Figure involves reconstructing the DMC robot over solid black background.
7

---

Under review as a conference paper at ICLR 2022
Mean
Median
Median
Rew+
Rew+Trans
Rew Only
Only
SPR*
SPR**
SPR
SPR DrQ CURL OTR Aug
DrQ CURL OTR
DER+Aug
DER
DER
SimPLe
SimPLe
0.25 0.50
0.75
0.15 0.30 0.45
0.06 0.12 0.18 0.24
0.05 0.10 0.15
Human Normalized Score on all games
Human Normalized Score dense games
Figure 7: Atari 100K. Human normalized performance (mean/median) across 25 games from the Atari 100K benchmark. Mean and 95% confidence interval for 5 runs. Left plot: Comparison for all 25 games. Right plot: Comparison for only dense reward games (7 games from Table 7).
this in Figure along with performance of baseline and augmentations. Of course, how one should come up with techniques that differentiate between task-relevant and task-irrelevany components in the observations, remains an open question 7. Additional ablations are provided in Appendix 8.6.
Atari 100K. We study the effect of techniques discussed thus far for Atari 100K benchmark, which involves 26 Atari games and compares performance relative human-achieved scores at 100K steps or 400K frames. We consider the categorization proposed by Bellemare et al. [6] based the nature of reward (dense, human optimal, score exploit and sparse). We implemented two versions of the baseline algorithm,
with both the transition and reward prediction modules and the other with only reward prediction. Our average results over all games show that the baseline performs comparably to CURL [36], SimPI.e [30], DER [53], and OTR [32] while being quite inferior to DR [33, 2] and SPR [46]. In comparison to DER, since our implementation
the baseline is over the DER code, similar performance to DER might suggest that the reward and transition prediction do not help much in this benchmark. Note that the ALE does not involve the use of distractors and so learning directly from the RL head (DQN in this case) should be enough encode information about the reward and the transition dynamics in the representation. This comes as a stark contrast to the without distractors case in DMC Suite, where transition and reward prediction still lead better performance. Such differences might be attributed to the continuous vs discrete nature of DMC and ALE benchmarks. More surprisingly, we find that when plotting the same average performance results for only the dense reward environments in 100K, the gap in performance between DER and SPR/DRQ decreases drastically. Note that both SPR builds over DER but DRQ builds over OTR.
We further delve into understanding the superior performance of SPR and DRQ. In particular, SPR combines a cosine similarity transition prediction loss with data augmentations. To understand the effect of each of these individually, we run SPR without data augmentations, and call it SPR "9 We see that SPR* leads to performance similar to the baseline and the DER agent, suggesting that such self-supervised loss may not lead to gains when without data augmentations. Finally, we take the DER agent and add data augmentations to it (from DRQ). This is shown as DER + AUG in Figure 7. We see that this leads to collapse, with the worst performance across all algorithms. Note that DRQ builds over OTR and performs quite well whereas when the same augmentations are used with DER, which includes a distributional agent in it, we observe a collapse. This again indicates that augmentations can change data in fragile manner, sometimes leading to enhanced performance with certain algorithms, while failing with other algorithms. Segregating evaluation of algorithms based on these differences is therefore of utmost importance. We show the individual performance on all 25 games in Appendix 8.5 (Table 7).
6 DISCUSSION
The above description of results on Atari 100K point to a very interesting observation, that evaluation of different algorithms is very much correlated with a finer categorization of the evaluation benchmark,
7 As also evident by TIA's [19] performance for DMC ball-in-cup catch experiments. We use the Q(e version from Agarwal et al. [2] for fair evaluation and denote as DR Note that this is different from the SPR without augmentations version reported in Schwarzer et al. [46]
since that version uses dropout as well which is not fair comparison.
8

---

Under review as a conference paper at ICLR 2022
and not the whole benchmark itself. Specifically, focusing on finer categorizations such as density of reward, inherent horizon of the problem, presence of irrelevant and relevant task components, discreteness vs continuity of actions etc. is vital in recognizing if certain algorithms are indeed better than others. Figure 1 stands as a clear example of such discrepancies. These observations pave the way for better evaluation protocol for algorithms, one where we rank algorithms for different categories, each governed by a specific data-centric property of the evaluation benchmark. Instead of saying that algorithm X is better than algorithm Y in benchmark Z, our results advocate for an evaluation methodology which claims algorithm X to be better than algorithm Y in dense reward, short horizon problems (considered from benchmark Z), i.e. less emphasis on the benchmark itself and more on certain properties of a subset of the benchmark. Having started with the question of what matters when learning representations over pixels, our experiments and discussion clearly show that largely it is the data-centric properties of the evaluation problems that matter the most.
7 CONCLUSION
In this paper we explore what components in representation learning methods matter the most for robust performance. We focused on the DMC Suite with distractors as the main setting while also extending our observations to DMC Suite without distractors and the Atari 100k benchmark. Our results show that a much simpler baseline, one involving a reward and transition prediction modules can be attributed to much of the benefits in DMC Suite with distractors. We then analysed why and when existing methods fail to perform as good or better than the baseline, also touching on similar observations on the ALE simulator. Some of our most interesting findings are follows:
Pixel reconstruction is a sound technique in the absence of clutter in the pixels, but suffer massively when distractors are added. In particular, DREAMER and adding a simple pixel reconstruction loss leads to worse performance than the baseline in DMC Suite (Figure 5).
Contrastive losses in and of itself do not seem to provide gains when there is a supervised loss available in place of it. We observe that replacing the supervised state prediction loss of the baseline by the InfoNCE contrastive loss does not lead to performance improvements over the baseline in DMC Suite (Figure 4 right plot). On the other hand, using contrastive losses with data augmentations can lead more robust improvements [38, 16].
Certain augmentations ('crop') do well when data is centered while dropping in performance when data off-center or when cropping fails to remove considerable amounts of task- irrelevant information. Other augmentations ('flip' and 'rotate') show the opposite behavior (RAD ablations on DMC Suite in Table 2).
SSL w/o contrastive losses do not provide much gains when used alone. With data augmen- tations, they lead to more sigmificant gains. For Ataril00k, Figure 7 shows that SPR, a state of the art non contrastive method leads to similar performance as the base DER agent when used without data augmentations (denoted by SPR**). Using the SPR inspired loss in DMC Suite also did not lead to gains over the baseline (in Figure right plot).
Augmentations are susceptible to collapse in the presence of distributional Q networks. Figure 7 shows that 'crop' and 'intensity" augmentations added to the DER agent lead to complete failure in performance Atari 100k.
These results elicit the observation that claiming dominance over other methods for entire bench- mark may not be an informative evaluation methodology. Instead, focusing the discussion to a more data-centric view, one where specific properties of the environment are considered, forms the basis of a much more informative evaluation methodology. We argue that as datasets become larger and more diverse, the need for such an evaluation protocol would become more critical. We hope this work can spur further discussion in categorizing evaluation domains in more complex scenarios, such as with real world datasets and over a wider class of algorithmic approaches.
9

---

Under review as a conference paper at ICLR 2022
REFERENCES
[1] Agarwal, R., Machado, M. C., Castro, P. S., and Bellemare, M. G. Contrastive behavioral similarity embeddings for generalization reinforcement learning. International Conference on Learning Representations, 2021.
[2] Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A., and Bellemare, M. G. Deep rein- forcement learning at the edge of the statistical precipice. arXiv preprint arXiv:2108, 13264, 2021.
[3] Andrychowicz, M., Raichuk, A., Stanczyk, P., Orsini, M., Girgin, S., Marinier, R., Hussenot, L.,, Geist, M., Pietquin, O., Michalski, M., et al. What matters in on-policy reinforcement learning? a large-scale empirical study. arXiv preprint arXiv:2006.05990 2020.
[4] Ayoub, A., Jia, Z., Szepesvari, C.,, Wang, M., and Yang, L. Model-based reinforcement learning with sargetdd regression. In International Conference on Machine Learning, pp. 463 (63-474. PMLR, 2020.
[5] Babaeizadeh, M., Saffar, M. T., Hafner, D., Erhan, D., Kannan, H., Finn, C., and Levine, S. trade-off's of image prediction in visual model-based reinforcement learning. 2020.
[6] Bellemare, M., Srinivasan, S., Ostrovski, G., Schaul, T., Saxton, D., and Munos, R. Unifying count-based exploration and intrinsic motivation. Advances in neural information processing systems, 29: 1er1147-14477 2016.
[7] Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
[8] Biza, O., Platt, R., de Meent, J.-W., and Wong, L. L. Learning discrete state abstractions with deep variational inference. arXiv preprint arXiv:2003.04300 2020.
[9] Castro, P. S. Scalable methods for computing state similarity in deterministic markov decision processes. Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 10069-10076, 2020.
[10] Castro, P. S., Kastner, T., Panangaden, P., and Rowland, M. Mico: tions via sampling-based state similarity for markov decision processes. CoRR, 0bs/2106.08229, 2021. URL https: /////aaxvvv//// org/ abs /2106. 08229.
Learning improved representa-
[11] Chen, T., Korbblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual representations. International conference on machine learning, pp. 1599-1607. PMLR, 2020.
[12] Chen, X. and He, K. Exploring simple siamese representation learning. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 15750-15758, 2021.
[13] Co-Reyes J. D., Sanjeev, S., Berseth, G., Gupta, A., and Levine, S. Ecological reinforeement learning. arXiv preprins arXiv:2006. /2478, 2020.
[14] Du, S., Krishnamurthy, A., Jiang, N., Agarwal, A., Dudik, M., and Langford, J. Provably efficient rl with rich observations via latent state decoding. International Conference on Machine Learning, pp. 1665-1674. PMLR, 2019.
[15] Engstrom, L., Ilyas, A., Santurkar, S., Tsipras, D., Janoos, F., Rudolph, L., and Madry, A. Implementation matters in deep rl: A case study on ppo and trpo. In International conference on learning representations, 2019.
[16] Fan. J. and Li, W. Robust deep reinforcement learning via multi-viww information bottleneck. arXiv preprint arXiv:2102. 2021.
[17] Farahmand, A.m Barreto, A., and Nikovski, D. Value aware loss function for model-base based reinforcement learning. In Artificial Intelligence and Statistics, pp. 1486-1494, PMLR, 2017.
10

---

Under review as a conference paper at ICLR 2022
[18] Ferns, N., Panangaden, P., and Precup, D. Bisimulation metries for continuous markov decision processes. SIAM J. Comput., 4066):1622-1714 December 2011. ISSN 0097-5397. doi: 10.1137/10080484X URL htt.:p/://oo...r//..1111///00/0
[19] Fu, X., Yang. G., Agrawal, P., and Jaakkola," T. Learning task informed abstractions. In Meila, M. and Zhang. (eds.), Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 3480-3491. PMLR, 18-24 Jul 2021. URL http: /tt::/prrooo http///ooceedddege iggs. www. press /
/ v139/ fu2 22..bb.b.... html.
[20] Gelada, C., Kumar, S., Buckman, J., Nachum, O., and Bellemare, M. G. DeepMDP: Learning continuous latent space models for representation learning. In Chaudhuri, K. and Salakhutdinov, Machine Learning. volume 97
R. (eds.), Proceedings of the 36th International Conference of Proceedings of Machine Learning Research, pp. 2170-2179. PMLR, 09-15 Jun 2019. URL https: http//ppoeeeigggeee mlr. prw.sppessseee / vvtp:v//////// http//99//////d// 9a .tml..
[21] Grill, J.-B.. Strub, F., Altché, F., Tallec, C., Richemond, P. H., Buchatskaya, E., Doersch, C., Pires, B. A., Guo, D., Azar, M. G., et al. Bootstrap your own latent: new approach to self-supervised learning. arXiv preprint arXiv:2006.07755 2020.
[22] Ha, D. and Schmidhuber, J. World models. arXiv preprint arXiv: 1803. 10122 2018.
[23] Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Dy, J. and Krause, A. (eds.), Proceedings of the 35th International Conference Machine Learning Research, pp. 1861-1870. PMLR, 10-15 Jul 2018. URL https:// proceedings mll. wwwps/esssssee. v80 haarnojal88 html.
Machine Learning, volume 80 of Proceedings of
[24] Hafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D., Lee, H., and Davidson, J. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555-2565. PMLR, 2019.
[25] Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination. International Conference on Learning Representations, 2020. URL https: wwpppppevviieeeeee net forum? wewett/foummmdd1mmme idd=s100TCADDS
[26] Hafner, D., Lillicrap, P., Norouzi, M., and Ba. J. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021. URL https: //openrevil ew net/ forum? (ommm2ddd00aaay2bb....
[27] Henderson, P., Islam, R., Bachman, P., Pineau, J., Precup, D., and Meger, D. Deep reinforcement learning that matters. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
[28] Hessel, M., Modayil, J., Van Hasselt, H.. Schaul, T., Ostrovski, G., Dabney, W., Horgan, D., Piot, B., Azar, M. and Silver, D. Rainbow: Combining improvements in deep reinforcement learning. In Thirty-second AAAI conference on artificial intelligence, 2018.
[29] Jaderberg, M., Mnih, V., Czarnecki, W. M., Schaul, T., Leibo, J. Z., Silver, D., and Kavukcuoglu, K. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint 2016.
[30] Kaiser, L., Babaeizadeh, M., Milos, P., Osinski, B., Campbell, R. H., Czechowski, K., Erhan, D., Finn, C., Kozakowski, P., Levine, S., et al. Model-based reinforcement learning for atari. arXiv preprint arXiv: 1903.00374 2019.
[31] Kay, W., Carreira, J., Simonyan, K., Zhang, B., Hillier, C., Vijayanarasimhan. S., Viola, F., Green, T., Back, T., Natsev, P., et The kinetics human action video dataset. arXiv preprint arXiv: 7777006550000 2017.
[32] Kielak, K. Importance of using appropriate baselines for evaluation of data-efficiency in deep reinforcement learning for atari. arXiv preprint arXiv:2003. 2020.

---

Under review as a conference paper at ICLR 2022
[33] Kostrikov, I,, Yarats, D., and Fergus, R. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. CoRR, abs/2004.13649 2020. URL https: //axxiv..
[34] Krishnamurthy, A., Agarwal, A., and Langford, J. Pac reinforcement learning with rich observations. arXiv preprint arXiv: 2016.
[35] Laskin, M., Lee, K., Stooke, A., Pinto, L., Abbeel, P., and Srinivas, A. Reinforcement learning with augmented data. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M. F, and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 19884-19895. Curran Associates, Inc., 2020. URL httpp::pproceeaaggg neurips ccc pappappr///// 21e.1111600222000000000000000000
2020/ / fwwe/ee1eeeeeee
pdf..
[36] Laskin, M., Srinivas, A., and Abbeel, P. Curl: Contrastive unsupervised representations for reinforcement learning. Proceedings of the 37th International Conference on Machine Learning. Vienna, Austria, PMLR 119, 2020. arXiv:2004.04136.
[37] Lee, A. X., Nagabandi, A., Abbeel, P., and Levine, S. Stochastic latent actor-critic: Deep rein- forcement learning with a latent variable model. Larochelle, H., Ranzato, M., Hadsell, R., Bal- can, M. F., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 741-752. Curran Associates, Inc., 2020. URL https::ttpt.tct http://///cccccceeeeeeeeee neurips cc/paper/2020/ http://eeeee//0/t5//000000 42 000.0/02222220000001101000000000
[38] Lee, K.-H,, Fischer, I., Liu, A., Guo, Y., Lee, H., Canny, J., and Guadarrama, S. Predictive information accelerates learning in rl. arXiv preprint arXiv:2007.12401 2020.
[39] Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., al. Human-level control through deep reinforcement learning. nature, 5 518(7540):529-533 2015.
[40] Nguyen, T., Shu, R., Pham, T., Bui, H., and Ermon, S. Temporal predictive coding for model- based planning in latent space. arXiv preprint arXiv:2/06.07155. 2021.
[41] Oh, J., Guo, X., Lee, H., Lewis, R., and Singh, S. Action-conditional video prediction using deep networks in atari games. arXiv preprint arXiv: 150.00850, 2015.
[42] Oord, A. d., Li, Y., and Vinyals, O. Representation learning with contrastive predictive coding. arXiv preprint 1807.03777 2018
[43] Raileanu, R., Goldstein, M., Yarats, D., Kostrikov, I., and Fergus, R. Automatic data augmen- tation for generalization in deep reinforcement learning. arXiv preprint 20X1v:2006.12862 2020.
[44] Schaul, T., Horgan, D., Gregor, K., and Silver, D. Universal value function approximators. In International conference machine learning, pp. 1312-1320. PMLR, 2015.
[45] Schmidhuber, J. Formal theory of creativity, fun, and intrinsic motivation (1990-2010). IEEE Transactions on Autonomous Mental Development, 2(3):230-247, 2010.
[46] Schwarzer, M., Anand, A., Goel, R., Hjelm, R. D., Courville, A., and Bachman, P. Data-efficient reinforcement learning with self-predictive representations. arXiv preprint arXiv:2007.05929 2020.
[47] Silver, D., Schrittwieser, J., Simonyan, K., Antonoglou, I., Huang, A., Guez, A., Hubert, T., Baker, L., Lai, M., Bolton, A., al. Mastering the game of go without human knowledge. nature, 550(7676):354-359 2017.
[48] Stone, A., Ramirez, O., Konolige, K., and Jonschkowski, R. The distracting control suite-a chal- chal- lenging benchmark for reinforcement learning from pixels. arXiv preprint arXiv:2101.0272. 2021.
[49] Stooke, A., Lee, K., Abbeel, P., and Laskin, M. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning, pp. 9870-9879. PMLR, 2021.
12

---

Under review as a conference paper at ICLR 2022
[50] Sutton, R. S., Modayil, J., Delp, M., Degris, T., Pilarski. P. M., White, A., and Precup. D. Horde: scalable real-time architecture for learning knowledge from unsupervised sensorimotor interaction. In The 10th International Conference Autonomous Agents and Multiagent Systems-Volume
Volume 2, pp. 761-768, 2011.
[51] Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., Casas, D. d. L., Budden, D., Abdolmaleki, A.. Merel, J., Lefrancq, A., al. Deepmind control suite. arXiv preprint arXiv: 2018.
[52] Tomar, M., Zhang, A., Calandra, R., Taylor, M. E., and Pineau, Model-invariant state abstractions for model-based reinforcement learning, arXiv preprini urXiv:2102.09850 2021.
[53] van Hasselt, H. P., Hessel, M., and Aslanides, J. When to use parametric models in reinforcement learning? In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL http::::ttt/tseeeeee http//peooceenn.seee/eeee//e neurips iss cc/paper/2019 8b75449c6ee 72 fd92d- Paper sepper.pff.. pdf.
file/1b742ae21555
[54] Wahlström, N., Schön, B., and Deisenroth, M. P. From pixels torques: Policy learning with deep dynamical models. arXiv preprint arXiv: a 2015.
[55] Watter, M., Springenberg, J. T., Boedecker, J., and Riedmiller, M. Embed to control: A locally linear latent dynamics model for control from raw images. arXiv preprint arXiv:1506.0755 07365, 2015.
[56) Yarats, D., Zhang, A., Kostrikov, I., Amos, B., Pineau, J., and Fergus, R. Improving sample efficiency in model-free reinforcement learning from images. arXiv preprint arXiv:1910.01741, 2019.
[57] Yarats, D., Fergus, R., Lazaric, A., and Pinto, Reinforcement learning with prototypical representations. arXiv preprint arXiv:2102.1 1271, 2021.
[58] Yarats, D., Kostrikov, I., and Fergus, R. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations, 2021. URL https:////// /wwpnreeeette ew, net. / forum? d=GY 6- Tva..
[59] Zbontar, J., Jing, L., Misra, I., LeCun, Y., and Deny, S. Barlow twins: Self-supervised learning via redundancy reduction. arXiv preprint arXiv:2103.032300 2021.
[60] Zhang, A., Wu, Y., and Pineau, J. Natural environment benchmarks for reinforcement learning. arXiv preprint arXiv: 181i: 060032, 2018
[61] Zhang, A., Lyle, C., Sodhani, S., Filos, A., Kwiatkowska, M., Pineau, J., Gal, Y., and Precup,
D. Invariant causal prediction for block MDPs. In III, H. D. and Singh, A. (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 11214-12224 PMLR, 13-18 Jul 2020. URL https:/ // proceedings mlr. www.s/peessaee vl1 9/ zhang20t html.
[62] Zhang, A., McAllister, R., Calandra, R., Gal, Y., and Levine, S. Learning Invariant Representa- tions for Reinforcement Learning without Reconstruction, arXiv, jun 2020. ISSN 23318422 URL http:://aaane http::::aaxxx..oo////////// wttp://////a/bb//bbbb/e /2w...000...00 10742.
[63] Zhang, A., Sodhani, S., Khetarpal, K., and Pineau, J. Learning Robust State Abstractions for Hidden-Parametae Block MDPs. pp. 1-22, 2020. URL http:://///.. http:://axvvvvv//e abs/2007 2207.. 07206.
13

---

Under review as a conference paper at ICLR 2022
Appendix
METHODS
We model the setting in this paper using the framework of contextual decision processes (CDPs), a term introduced in Krishnamurthy et al. [34] to broadly refer to any sequential decision making task where agent must act the basis of rich observations (context) to optimize long- -erm-term reward. Typically, a CDP is defined as the tuple (2', S. P, y, R), with observations or contexts Te E X, states 8t E 8, actions a, A, state transition distribution P(5 &t, at) and reward function R = r(st, Qt St+1), which defines the nature of the task and dynamics of the system. Here F, refers high-dimensional observation while the true state of the environment 9; is not directly available (usually low dimensional) and the agent must construct it on its own. Without knowing P and R, the RL agent's goal is use the collected experience maximize expected sum of rewards, R Ei=o r'r,, consider a discount factor [0, 1).
1.1 SAC
SAC or Soft Actor-Critic [23] deploys soft entropy regularized actor critic framework, wherein the policy entropy is added as an additional loss to the standard Q function maximization loss. In particular, SAC uses Q₂ functions and separate value function V. For transition (T,, ar Tc) from replay buffer D, the exact losses are described as follows:
Lo(Ø) E,-D [Qo( at) - (ra + (V(L)))))
The value function is then fit w.r.t to the entropy regularized (soft) function.
((I++1) = En'na (5(F++1, - a log ny (a' [m(+1)
(2)
Finally the policy is fit to maximize both the Q function and the entropy.
(x(1)) = -EaNT [Qllte ur) log (u/ar))
(3)
1.2 SAC-AE
SAC-AE [56] extends SAC by adding autoencoder for learning the representation s; fo( which is implicit in the SAC description through the use of deep convolutional networks. A key observation made in this work is that actor gradients should not be used to learn the representation fo. Only the critic gradients train the representation. This idea has been used as standard practice in most papers that follow.
1.3 DREAMER
DREAMER [25, 26] learns latent space dynamics model, both for the transition and reward dynamics, and then uses model based updates with Recurrent State Space Model (RSSM [24]) to optimize the policy. In learning the transition dynamics model, prediction in the latent state as well as in the pixel space is used construct representation model, p(se+1 (s,, d,, xz). and define targets for the dynamics model. The pixel space prediction loss is referred to as the pixel reconstruction loss and defined with a reconstruction model q(- q(-lsx). The agent minimizes the expectation (Cdreamer) over a finite horizon summation of reconstruction (C)), reward prediction (Ct² and transition prediction (CD) losses, such that:
CDREAMER = E
Σ (Cb + C/ + 2b)
(4)
% In Ing((1)))), 16= R(rell)))
C 3 KI. P P(st1199, a:))
14

---

Under review as a conference paper at ICLR 2022
1.4 RAD AND DRQ
RAD [35] uses SAC as the main algorithm and adds data augmentations to replay buffer data. These augmentations include: random cropping, translate, grayscale, cutout, rotate, flip, color jitter, and random convolution. For ALE, in general, rangom shifts and intensity are used as augmentations by DrQ [33]. It is shown that random crop achieves by far the best performance compared other augmentations. Therefore, when no mention of the exact augmentations is provided, assume that random crop being used.
1.5 CURL
CURL [36] data augmentations (just RAD) standard SAC agent and additionally deploys an auxiliary loss that tries to learn representations such that augmented versions of the same observation context pulled together in the representation space while pushing apart augmented versions of two different observations. CURL uses the bi-linear inner-product q" Wk, where W is learned parameter. Here, q and & are the latent representations of the raw anchors (query) Tgg and targets (keys) Tı, such that q fo(x,) and l sg(Jo(zx)), where sg denotes the stop gradient operation.
1.6 DER AND OTR
Both DER [53] and OTR [32] refer to highly tuned version of the Rainbow [28] agent. The hyperparameters are tuned so as to get superior performance over Rainbow on Atari100k, which uses lot less samples than the standard training of 200M steps. There also exist similar versions for the DQN agent (data efficient DQN, overtrained DQN), which most notably does not use a distributional component in it.
1.7 SPR
SPR [46] self-predictive representations deploys an auxiliary transition prediction loss which is learned in a con-contrastive self-supervised manner by using cosine similarity loss. In particular, instead of using the actual next state 5 as the target for the predicted next state $ the authors use projection network to transform both the targets (gm) and the predictions (go). Furthermore, the predictions are then then passed through predictor network (g), and the output of this network and of the target projector are used to define cosine similarity loss LCosine- The use of the projector and predictor networks is similar to that in self-supervised learning methods in vision, such as BYOL [21].
H Σ Am
Yi+h
T
Øc-n
LCosine
(5)
|
[1
where
814k P. (St, az 0+4k-1),
Dith (go ($r+h)),
Be+A 9m (8tth)
In particular, SPR also data augmentations and convolutional transition dynamics predictor. When not using data augmentations, SPR uses dropout in the encoder network fe.
1.8 DBC
DBC [63] or Deep Bisimulation for Control is an algorithm based on bisimulation metric losses which tries pull observations having the same long term reward closer in the representation space.
The specific implementation uses an encoder, reward prediction model and a probabilistic latent transition model. In particular, the encoder is trained by sampling batches of experiences (v, a, (,rxx) x²) and minimizing the LI norm between any two non-identical transitions in the sampled batch as follows:
CDBC (1s,- -35--s1::
ry| ()
(6)
where 8; = fo(z)) 3; = sg(fo(z;) and W2 is the 2- Wasserstein distance metric between two transition distributions. Please refer the paper for further details.
15

---

Under review as a conference paper at ICLR 2022
1.9 BASELINE
We use simple soft actor-critic setup with an embedding function fo X' - S (similar to SAC- AE [56]) as the base architecture, and attach the reward and transition prediction modules to it (See Figure 2). We define the transition prediction by (ss ac) and the reward prediction by R(s++1) such that, 81+1 = P( (e) and r; = R(s++1). Note that transition network over the encoded state & = fa()) and not over the observations x,
[37]. The overall auxiliary loss function is thus
defined as follows:
Lbaseline
(5t+1 - ((s,, at))
(R(5 ((P(()t,))))
Transition prodiction loss
Roward prediction loss
1.10 VALUE-AWARE MODEL LEARNING
Based the baseline architecture, the value aware experiments are performed with the SAC critic Q. policy т and the corresponding value aware loss, Lv as (sg(V++1) Ve+1) where
V Eor+ ~ur (Q(st-1,0++1) log π (04+11 V++1 Ea, 1~ (06++1,0441)
log क (
Here, sg denotes the stop gradient operation. We test the approach both with and without the reward prediction module (R in equation equation 7) and report the results in Table 1. Complete performance plots are provided in Appendix 8.8.
16

---

Under review as a conference paper at ICLR 2022
2 BACKGROUND DISTRACTOR DETAILS
We use the distracting control suite with natural distractors as previously introduced in Kay et al. [31], Zhang et al. [60] and used in Zhang et al. [63], Stone et al. [48]. The background distractor frames are picked up sequentially from videos randomly sampled from the Kinetic dataset ('driving car' class) and were used during both training and evaluation. The choice of video for any two runs of a single DMC task as well as during training and evaluation can vary randomly. The selected sequence of frames are masked alongside the observations obtained from the DMC suite. The resulting observations have two sequences, one corresponds to task-relevant robot movement while the other corresponds to faskvirreeevant background video streams. We consider the addition of distractors to two different settings: centered and off-centered, as shown in Figures 8 and 9 respectively.
Figure 8: Distractor Robot Centered configuration We show the sequence of observations with distractor backgrounds Cheetah Run, Hopper Hop and Walker Walk tasks. is to be noted that there be than one video streams playing sequentially for an environment.
Figure 9: Distractor on Robot Off-centered configuration The background distractor configuration is same the robot-centered setting, but the robot's size and position are changed. This setting increases difficulty for methods dependent on well positioning of the robots.
17

---

Under review as a conference paper at ICLR 2022
3 IMPLEMENTATION DETAILS
We use the SAC algorithm as the RI. objective added with a reward and a deterministic state transition model both of which are MLPs. In this section, we present in detail the architecture, source code links and the computational usage.
Architecture. The encoder architecture is taken from the open-source implementation of Yarats et al. [56] which uses the same encoder for the actor and critic to embed image observations. Only critic gradients are used to train the encoder. Both the reward and transition prediction losses are equally weighed in the gradient update and there is no extra hyperparameter to balance between the two.
The ablation on augmentation uses the same configuration of cropping shown in RAD and CURL i.e. random cropping (84 x 84) from observations of (100 x 100) dimensions for DMC and random shifts of #4 pixels for ALE. Random flip and rotate are applied in similar way as well. For the reconstruction loss augmentations, deconvolution layer is introduced to reconstruct the original pixel-state from the latent state vectors.
Source Codes. For the methods used for comparison, the corresponding code was directly borrowed from their open-source implementations or their performance was taken from Agarwal et al. [2]. The natural distractor background was added from the open source implementation of Kinetic dataset to the other methods for comparison. The respective URLs are given below.
Method
Open-Source URL
Kinetic dataset
SAC-AE
RAD
CURL-DMC
https://ggthhbb..c/m//sssssssssa
DrQ
DBC
PI-SAC
Dreamer
httpss//ggthaab.omm////m///////a
TIA
http:://ggthhubb.cmm//yyooo/////
CURL-ALE
SPR
Computation. All experiments were conducted on either system configuration of:
6 CPU cores of @Intel Gold 6148 Skylake@2.4 GHz, one @NVidia V100SXM2 (16G memory) GPU and 84 GB RAM.
2. 6 CPU cores of @Intel Xeon Gold 5120 Skylake@2.2GGI) one @NVIDIA V100 Volta (16GB HBM2 memory) GPU and 84 GB RAM
The average completion time of the baseline experiments around 20 hours. The time taken was less (around 12 hours) for the SAC or baseline-v0 experiments. For the value aware experiments, RAD and CURL average time of completion was around 15 hours. The time required for Dreamer, baseline with reconstruction losses, SPR variant on DMC and PI-SAC is approximately 28 hours. On an average, each experiment expected be done in 24 hours. A total of around 1500 experiments were performed for the final version of this paper which corresponds to roughly 4 years of GPU training.
18

---

Under review as a conference paper at ICLR 2022
4 HYPERPARAMETERS
The full set of hyperparameters used for the baseline experiments are provided Table 3 below.
Table 3: Hyperparameters for Baseline and related ablations.
Hyperparameter
Values
Observation shape
(84, 84, 3)
Latent dimension
50
Replay buffer size
100000
Initial steps
1000
Stacked frames
Action repeat
finger, spin: walker, walk 8 cartpole, swingup
otherwise
SAC: Hidden units (MLP)
1024
Transition Network: Hidden units (MLP)
128
Transition Network: Num Layers (MLP)
6
Reward Network: Hidden units (MLP)
512
Reward Network: Num Layers (MLP)
Evaluation episodes
10
Optimizer
Adam
(81, B2) (J. Qo)
(.9, (999)
(3,, B2) - (a) Learning rate (fo, Tigh Q+)
(.5, .99))
2e-4 cheetah, run
le-3 otherwise
Learning rate (a)
le-4
Batch Size
128
Q function EMA T
0.005
Critic target update freq Convolutional layers
Number of filters
32
Non-linearity
ReLU
Encoder EMA
0.005
Discount
.99
19

---

Under review as a conference paper at ICLR 2022
5
ADDITIONAL BASELINE ABLATIONS
Cartpole Swingup
Cheetah Run
Finger Spin
1.00
0.75
0.4
0.50
0.2
0.25
0.25
0.00
0.0
0.00
Normalized
0.5
0.0
Reacher Easy
Walker Walk
Hopper Hop
0.6
0.15
0.4
0.10
0.2
0.05
0.0
0.0
0.00
0.0
Envirenment Steps (x10*)
Baseline w/ Distractor
Baseline w/o Distractor
SAC w/ Distractor
SAC wio Distractor
Figure 10: Baseline w/ distractor Baseline w/o distractor SAC w/ and w/o distractor Performance of the proposed baseline on six DMC tasks the presence and absence distractors along with performnce of SAC (Baseline-v() on the same settings. The minor difference in the performance of baseline the two task conditions clearly illustrates the learning of background/distrratoor invariant representations. This is in stark contrast with several experimental analysis showing the overfitting by other algorithms. Additionally, the performance gain as compared to SAC is visible both with without distractor cases.
Cartpole Swingup
Cheetah Run
Finger Spin
1.00
0.75
0.4
0.75
0.50
0.2
0.25
0.25
0.00
0.0
0.00
Normalized
Reacher Easy
Walker Walk
Hopper Hop
1.00
1.00
0.75
0.75
0.2
0.50
0.50
0.1
0.25
0.25
0.00
0.00
0.0
0.0
Environment Steps (x10*)
Baseline Crap
Baseline Flip
RAD Crop
RAD Flip
Figure 11: Baseline w/ augmentations Baseline-v0 w/ augmentations Performance of the baseline with crop and fli augmentations on task data. The comparison conducted based on the results RAD [35] with similar augmentations. Performance curves for each of the six DMC tasks is shown. With cropped data, baseline performs equally well as RAD due to removal of task irrelevant features by cropping. Whereas for flipped data, where none of the task irrelevant features excluded, baseline surpasses RAD.
20

---

Under review as a conference paper at ICLR 2022
BASELINE WITH MULTISTEP TRAINING
Cartpole Swingup
Cheetah Run
Finger
0.6
1.00
0.75
0.6
0.4
0.50
0.2
0.2
0.25
0.0
0.0
0.00
Normalized
0.5
0.0
Reacher Easy
Walker Walk
Hopper Hop
0.6
0.4
0.06
0.3
0.2
0.2
0.02
0.0
0.0
0.00
0.0
Envirenment Steps (x10*)
Baseline horizon-1
Baseline horizon-2
Baseline horizon-4
Baseline horizon-6
Baseline w/ multi-step training Performance of baseline with recurrent latent state space model and reward prediction network (MLP) shown for fixed horizon each six DMC tasks. Increasing horizon degrades performance across all the tasks.
Figure 12:
Cartpole Swingup
Finger Spin
0.6
0.75
0.50
0.2
0.25
0.0
0.00
Normalized
Reacher Easy
Walker Walk
Hopper Hap
0.20
0.04
0.4
0.03
0.10
0.02
0.05
0.01
0.00
0.0
0.1
0.2
0.3
0.44
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Environment Steps (x10*)
Reconstruction horizon-1
Reconstruction harizon-2
Reconstruction horizon-4
Reconstruction horizon-6
Figure 13: Baseline w/ multi-step training and reconstruction Performance of the proposed haseline six DMC tasks with multi-step training and added reconstruction loss each prediction. This resembles architecture of DREAMER [25] corresponding to haseline, Full reconstruction tries to capture task-irrelevant information leading to decrease in performance. Further, the performance degrades irrespective of the choice of the horizon.
21

---

Under review as a conference paper at ICLR 2022
7 DATASET CATEGORIZAIONS AND SPECIHICATIONS
Table 4: Dataset Categorizations. Dense+ refers to games under the 'Human Optimal' 'Score Exploit' category from the categorization provided in [6] and Table 7. The performance of baseline, state metric, contrastive, SSL w/ contrastive and reconstruction in cach of the categories are shown in Figure 1. The algorithm corresponding each method shown in Appendix 7 (Table 5), Also, since augmentations only modify the data, they should included the evaluation categorizations.
C5
C6
C8
Environment
ALE
ALE
ALE
DMC
DMC
DMC
DMC
DMC
Reward Structure
Dense
Dense
Dense+
Dense
Dense
Dense
Sparse
Dense
Distractor
No
No
Yes
Yes
Yes
Yes
No
Augmentations
None
Shift, Intensity
None
None
Crop
Flip
None
None
Table 5: Categorizations and Methods. The algorithm corresponding each of the methods is shown. The
algorithms considered DER [53], DRQ [33], SAC [23], RAD [35], PI-SAC [38], DBC [62], CURL [36]. SPR ([46) and SPR**, DREAMER [25, 26] and [19]. The performance of each the methods in each of the categories are shown in Figure I. For ALE, DREAMER scores after 1M steps were used for comparison.
Mcthods
C1
C2
C3
C4
C5
C6
C7
C8
Bascline
Baseline-v0
DER
DRQ
DER
SAC
RAD
RAD
SAC
SAC
State Metric
DBC
DBC
Contrastive
CURL
CURL
CURL
PI-SAC
PI-SAC
PI-SAC
PI-SAC
SSL w/o Contrastive
SPR**
SPR*"
SPR
SPR**
Reconstruction
DREAMER
DREAMER
TIA
TIA
DREAMER
Table 6: Categorizations and Environments. The environments included in each of the categories is shown. The environments considered from both DMC and ALLE Atari-100k benchmarks. Dense+ ALE refers games under the 'Human Optimal' and 'Score Exploit category from categorization provided in [6] and Table 7. The performance each the methods in each of the categories are shown in Figure 1.
Environments
C1
C2
C3
C4
C5
C6
C7
C8
Saripole
/
Cheetah-Run
Hopper-Hopp
1
1
Walker-Wall
1
Reacher-Easy
1
Finger-Spin
Ball-in-Cup-Catch
Dense Atari100k
Dense+ Atari 100k
22

---

Under review as a conference paper at ICLR 2022
8
DMC SUITE: COMPLETE RESULTS
8.1
ROLE OF EACH COMPONENT IN BASELINE
Cartpole Swingup
Cheetah Run
Finger Spin
0.4
0.8
0.75
0.3
0.6
0.2
Normalized Score
0.1
0.2
0.00
0.0
0.0
0.0
0.1
0.2
0.3
0.4
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.3
0.4
0.5
Reacher Easy
Walker Walk
Hopper Hop
0.06
0.3
0.4
0.04
0.2
0.02
0.0
0.00
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Environment Steps (x10%)
Reward w/ Transition
w/
Reward w/o Transition
Transition
w/o Transition
Figure 14: Baseline component ablations. The importance of the two components of the proposed bascline: Transition and Reward, analyzed based the effect of their absence the performance achieved for each
DMC tasks. The comparison shows that both are necessary achieve performance overall. While absence of both results in some learning, addition transition collapses the representations, On the other hand, reward proves to be non-negligible component.
Cartpole Swingup
Cheetah Run
Finger Spin
0.8
0.6
0.3
0.6
0.2
0.1
0.2
0.0
Normalized
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Reacher Easy
Walker Walk
Hopper Hop
0.06
0.3
0.4
0.2
0.04
0.2
0.1
0.02
0.00
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Eniironnenn Steps (x10%)
Baseline Non-Linear Reward Decoder
Baseline Linear Reward Decoder
Figure 15: Baseline w/ non-linear vx w/ linear reward decoder. Performance across six DMC tasks with linear and non-linear reward decoders with the baseline architecture. Reward decoder formulation plays significant role in the representation learning objective. As linear reward decoder discards fuskiirreeeaant features quickly, left features are not sufficient to act optimally. On the other hand, non-linear reward decoder comparatively sophisticated in recognizing iask-irreevaant features and features required to act optimally.
23

---

Under review as a conference paper at ICLR 2022
8.2
ARRANGEMENT OF REWARD AND TRANSITION LOSSES
Cartpole Swingup
Cheetah Run
Finger Spin
0.8
0.6
0.3
0.4
Normalized Score
0.2
0.1
0.2
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Reacher Easy
Walker Walk
Hopper Hop
0.4
0.6
0.06
0.04
0.2
0.02
0.0
0.00
0.0
0.1
0.2
0.3
0.4
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Enviranment Steps (x10%)
Reward through Transittion
Independent Reward Transittion
Figure 16: Arrangement of Reward and Transition in baseline architecture. Performance across six DMC tasks with reward prediction through transition and independent reward + transition prediction modules. The average performance is similar for both the cases. Having identified the importance of the reward and transition models, the remaining question workable possibilities for this, one being attaching both independently to the encoded state, i.c., R(s, a, s s'), and the other being attaching the transition model to the encoded state and then attaching the reward model to the output of the transition model, i.e. learning reward through transition R(s, s'). We test both of these and conclude that both perform equally well. We choose reward through transition it is marginally better.
how actually assemble these the overall architecture. There
8.3
ROLE OF AUGMENTATIONS
Cartpole Swingup
Run
Finger Spin
1.00
0.75
0.4
0.50
0.2
0.25
0.25
Normalized Score
0.00
0.0
0.00
0.0
0.1
0.2
0.3
0.1
0.3
0.4
0.5
0.0
Easy
Walker
Hopper Hop
1.00
1.00
0.75
0.75
0.2
0.50
0.1
0.25
0.25
0.00
0.0
0.1
0.2
0.5
Environment Steps (x105)
Baseline
RAD Crop
RAD Flip
RAD Rotate
Figure Baseline Augmentations. Performance proposed baseline and RAD [35] with crop/flip/rotate augmentations on the data from the task. Cropping removes the distractor data and excels in getting the best performance in all the six DMC tasks. Baseline performs close to RAD flip and rotate which consider the complete pixel data. This indicates that RAD crop takes advantage of the centered position of the objects and well curated structure the data. This eventually introduces a considerable inductive bias.
24

---

Under review as a conference paper at ICLR 2022
Cartpole Swingup
Cheetah Run
Finger Spin
1.00
1.00
0.75
0.75
0.50
0.50
Normalized Score
0.25
0.25
0.00
0.0
0.00
0.0
0.1
0.2
0.3
0.4
0.5
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.5
Reacher Easy
Walker Walk
Hopper Hop
1.00
1.00
0.75
0.75
0.2
0.50
0.50
0.25
0.25
0.00
0.00
0.0
0.0
0.1
0.2
0.3
0.4
0.5
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Environnmnnt Steps (x10%)
RAD Crop
RAD Flip
Small Robots RAD Crop
Small Robots RAD Flip
Figure 18: RAD w/ standard DMC setting vs w/ zoomed out DMC setting. We conducted an ablation on performance of RAD with crop and flip data augmentations for diverse camera positions for six DMC tasks. modified camera seltings zooms out robot in frame. drop in performance of RAD [35] in modified setting validates that the success of cropping is due to the well-centered and appropriate settings of DMC.
8.4
ROLE OF CONTRASTIVE LOSSES WITH AUGMENTATIONS
Cartpole Swingup
Cheetah Run
Finger Spin
1.00
0.75
0.4
0.75
0.50
0.50
0.2
0.25
Normalized Score
0.00
0.00
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Reacher Easy
Walker Walk
Hopper Hop
1.00
1.00
0.75
0.00
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Environment Steps (x110f)
Baseline
RAD
PI-SAC
CURL (Contrastive)
Crop)
(wl Random Crop)
(Contrastive)
Figure 19: Baseline vs RAD vs PI-SAC vs CURL. We conducted an ablation on performance of SAC with contrastive losses DMC task with data augmentations. Bascline and RAD [35] with random crop compared with PI-SAC [38] and CURL [36]. While CURL uses the contrastive loss of CPC [42], PI-SAC uses Conditional Entropy Bottleneck (CEB) objective. PI-SAC performance is significantly better as compared to CURL.
25

---

Under review as a conference paper at ICLR 2022
8.5
ROLE CONTRASTIVE AND SSL w/o CONTRASTIVE LOSSES
Cartpole Swingup
Cheetah Run
Finger Spin
1.00
0.75
0.75
0.3
0.50
0.50
0.25
0.1
Normalized
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Reacher Easy
Walker Walk
Hopper Hop
0.15
0.4
0.10
0.2
0.05
0.00
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Environment Steps (x119)
Reward w/ Transition
Reward w/ Contrastive
(Contrastive)
(Contrastive)
(Cosine Similarity)
Figure 20: Baseline contrastive vs SSL w/o contrastive. Performance of the baseline is compared with the
addition of contrastive loss instead of the transition loss of the baseline for each of the six DMC tasks. Further, we compare simple contrastive losses like PI-SAC [38] (without augmentations), CURL [36] SSL w/o contrastive losses like an extension of SPR [46] (which used cosine similarity loss) DMC. Contrastive and SSL. w/o Contrastive approaches do not seem to be very effective in the absence of data augmentations.
8.6
ROLE OF RECONSTRUCTION LOSSES
Cartpole Swingup
Cheetah Run
Finger Spin
0.8
0.75
0.6
0.4
0.50
0.4
Normalized Score
0.2
0.2
0.25
0.00
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Reacher Easy
Walker Walk
Hopper Hop
1.00
1.00
0.20
0.75
0.75
D.50
0.50
0.10
0.05
0.00
0.0
0.1
0.2
0.3
0.4
0.5
0.D
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
D.3
0.4
0.5
Environment Steps (x11%)
Baseline
Dreamer Distractor
Dreamer Distractor
TIA w/ Distractor
Figure 21: Baseline vs Methods with Reconstruction Losses. Performance of the proposed baseline with the state of the art algorithms employing reconstruction losses, DREAMER [25] and TIA [19]. DREAMER without distractors marks the upper limit for the maximum achievable performance. The performance of baseline, having a relatively simple architecture, considerably better than DREAMER and TIA.
26

---

Under review as a conference paper at ICLR 2022
Cartpole Swingup
Cheetah Run
Finger Spin
1.00
0.8
0.4
0.75
0.4
0.50
0.2
Normalized Score
0.25
0.0
0.0
0.00
0.0
0.1
0.3
0.4
0.5
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.5
Reacher Easy
Walker Walk
Hopper Hop
1.00
0.6
0.15
0.75
0.4
0.10
0.50
0.2
0.25
0.05
0.0
0.00
0.00
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.3
0.4
0.5
Environmmnt Steps (x10%)
Baseline
Full Reconstruction
Partial Reconstruction
Figure 22: Baseline Baseline with Reconstruction losses. The performance of the variants of baseline with additional reconstruction losses in analyzed for each of the six DMC tasks. The variants include full reconstruction and partial/relevant reconstruction losses. While the former same as the approach employed by DREAMER [25], the later only tries to reconstruct the task-relevan features, which is provided explicitly. We observe that the performance achieved by including the partial reconstruction loss is the ceiling for the performance which be achieved by employing reconstruction losses. However, still observe failure in the case of Reacher Easy task.
Cartpole Swingup
Cheetah Run
Finger Spin
1.00
0.75
0.75
0.50
0.4
0.50
Normalized Score
0.25
0.2
0.25
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Reacher Easy
Walker Walk
Hopper Hop
1.00
1.00
0.75
0.75
0.50
0.50
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Environment Steps (x10%)
Partial Reconstruction
Crop
Dreamer Distractor
Distractor
Figure 23: Baseline vs Augmentations vs Reconstruction. Based on the ablation studies with full and relevant reconstruction losses added to the baseline, the relevant reconstruction variant is contrasted with RAD [35], DREAMER [25] and TIA [19]. RAD with cropping removes illustrates that the representations learned by RAD with cropping is equivalent to reconstructing only the relevant pixels. Full reconstruction by DREAMER constructs the lower margin of the plots and TIA improves moderately on DREAMER by adding dccoupling between task relevant and irrelevant fcatures. However, the efficacy of these methods significantly compromised compared to performance achievable by actually learning task-relevant features.
of the distractor data and the performance
27

---

Under review as a conference paper at ICLR 2022
8.7
ROLE AUGMENTATIONS WITH RECONSTRUCTION LOSSES
Cartpole Swingup
Cheetah Run
Finger Spin
1.00
0.4
0.75
0.75
0.3
0.50
0.25
0.1
0.25
0.00
0.00
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Normaizze
Reacher Easy
Walker
Hopper Hop
1.00
1.00
0.06
0.04
0.25
0.00
0.00
0.0
0.1
0.2
0.3
0.4
0.0
0.1
0.2
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Enviranment Steps (x110))
Baseline Random
Baseline Reconstruction)
Baseline
Reconstruction Random Crop)
Figure 24: Augmentations with Baseline vs Reconstruction losses. The role of augmentations the final performance of reconstruction loss on the distractor observations is analyzed. While augmentations tend to boost the performance for all algorithms including Baseline, RAD (Baseline-vO) and PI-SAC (Contrastive), they do not show much improvement reconstruction losses.
8.8
VALUE AWARE LEARNING
Cartpole Swingup
Cheetah Run
Finger
1.00
0.8
0.4
0.75
0.4
0.2
0.50
0.25
0.0
0.0
0.00
Normalized
0.5
Reacher Easy
Walker Walk
Hopper Hop
0.4
0.6
0.06
0.2
0.4
0.1
0.2
0.02
0.0
0.0
0.00
0.5
0.0
Environment Steps (x10*)
Baseline
Value through Transition w/ Reward
Value through Transition w/o Reward
Figure 25: Baseline Value Aware Learning. Performance of the proposed baseline and truly value aware learning on each the six DMC tasks. We correlate the relation between the baseline and value aware learning. We, then, conduct ablations on the transition loss where we replace 'reward through transition' with "value through transition'. The above plots show the performance of 'value through transition' with/without an added reward loss.
28

---

Under review as a conference paper at ICLR 2022
9 ALE BENCHMARK: ATARI 100K
Table 7: Individual performance scores for games from Atari 100k benchmark categorized based on Bellemare
al. [6]. We observe that augmentations fail to work with DER [53] but perform reasonably well with OTR [32] evident from the performance of DRQ [33]. Further, to analyse the reason behind the performance boost of SPR [46]. we perform experiment by removing dropout and augmentations and denote it SPR"" The
performance drop SPR* as compared with SPR significant.
Game
DER
DER+Aug
OTR
CURL
DrQ
SPR**
SPR
Baseline
Human Optimal
Assault
431.2
155.5
351.9
600.6
452.4
647.9
571.0
474.8
Asterix
470.8
254.4
628.5
734.5
603.5
435.5
977.8
638.4
BattleZone
10124.6
6083.9
4060.6
14870.0
12954.0
12208.0
16651.0
9124.0
Boxing
0.2
9.3
1.2
6.0
-0.7
35.8
1.2
Breakout
1.9
4.48
9.8
4.9
16.1
5.612
17.1
2.9
ChopperCommand
861.8
270.7
1033.3
1058.5
780.3
695.2
974.8
675.2
Crazy Climber
16185.3
5052.6
21327.8
12146.5
20516.5
16441.4
42923.6
23912.6
Demon Attack
508.0
628.8
711.8
817.6
1113.4
397.6
545.2
635.8
Jameshond
301.6
145.6
112.3
471.0
236.0
404.5
365.4
228.1
Pong
-19.3
6,4
1.3
-16.5
-8.5
-0.5
-5.9
-18.8
Exploit
Kangaroo
779.3
1391.3
605.4
872.5
940.6
3749.4
3276.4
653.4
Krull
2851.5
838.4
3277.9
4229.6
4018.1
2266.6
3688.9
2775.3
Kung Fu Master
14346.1
4004.9
5722.2
14307.8 8
8939.0
15469.7
13192.7
10987.6 66
Road Runner
9600.0
2274.9
2696.7
5661.0
8895.1
4503.6
14220.5
6408.4
Seaquest
354.1
114.3
286.9
384.5
301.2
448.4
583.1
518.9
Up N Down
2877.4
2240.1
2847.6
2955.2
3180.8
2413.3
28138.5
2543.2
Dense Reward
Alien
739.9
223.2
824.7
558.2
771.2
835.1
801.5
699.7
Amidar
188.6
42.8
82.8
142.1
102.8
138.5
176.3
115.654
Bank Heist
51.0
53.99
182.1
168.9
205.3
380.9
Frostbite
866.8
325.4
231.6
1181.3
331.1
1248.7
1821.5
1740.9
Hero
6857.0
339.4
6458.8
6279.3
3736.3
6076.2
7019.2
4464.4
Ms Pacman
1204.1
393.6
941.9
1465,5
960.5
1346.9
1313.2
999.2
Qbert
1152.9
1530.3
509.3
1042.4
854.4
1841.75
669.1
1007.9
Sparse Reward
Freeway
27.9
0.6
25.0
26.7
9.8
29.5
24.4
27.54
Private Eye
97.8
54.8
100.0
218.4
-13.6
-49.5
124.0
100.0
29


Reviews:
[{"review_id": "fjSVLTtT56C", "reviewer": "Reviewer_61FY", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "Strengths\n* The authors presented large-scale empirical evaluations and ablation studies to analyze various components for pixel-based control with distractors.\n\nWeaknesses:\n* Many comparisons presented in this paper are in fact system-wise (such as most comparisons shown in Sec. 5), but the authors use these system-wise comparisons to derive conclusions on the family of methods. I find this unconvincing. For a concrete example, CURL is only one implementation of contrastive learning, and it's known to be an ineffective one such that the original CURL implementation actually works better without the contrastive objective. I don't believe conclusions drawn from CURL experiments can general to contrastive learning. There are other variants such as [a] that consider InfoNCE between the past and future and use momentum encoders similar to SPR, achieving significantly better results on DMC. [44] came after [a] and similarly considered predictive information using InfoNCE. [b] presents an even more complicated implementation of contrastive learning. Just to name a few examples. The same concern applies to other components (such as data augmentation, image reconstruction, etc) and the ablation studies, meaning that your conclusions only apply to specific implementations and that it's not obvious that they can generalize to the whole family of methods.\n* Lack of details for experiment set-up: Experiment setup of many figures are unclear to me. For example, are all experiments in Appendix 5 with distractors or not?\n* Lack of details for the proposed dataset categorization (table 3): I found it very difficult to follow. For example, since not every DMC task has sparse reward, it's not clear to me what kind of tasks C7 contains and how they are different to C4.\n* For one, there isn't a clear story line showing in the contributions paragraph in the introduction section and the conclusion section. The authors propose a new baseline method while analyzing existing methods and studying many different components. There seem to be multiple threads in this paper without having a main story.\n\nQuestions:\n* Have you noticed your reacher experiments can only work with data augmentation? It's not clear to me whether it's with distractors or not. Do you want to discuss this?\n* For the statement in conclusion \"Contrastive losses in and of itself does not provide gains when there is a supervised loss available in place of it.\", which supervised loss are you referring to?\n* What does \"non-contrastive loss\" even mean? Any loss that is not contrastive can be referred as \"non-contrastive loss\".\n* I was wondering if the authors are aware of performance issues with CURL, no matter it is an implementation flaw or algorithm design choice problem. Why do you choose to compare to CURL as a contrastive baseline?\n* Contrastive InfoNCE approximate reconstruction distribution [38, a]. Can you provide more insights for why they could perform differently despite the learning objective is similar?\n\n[a] \"Predictive Information Accelerates Learning in RL.\" NeurIPS 2020.\n\n[b] \"DRIBO: Robust Deep Reinforcement Learning via Multi-View Information Bottleneck.\" ICLR 2022 submission. (Of course, I don't expect the authors to cite this paper. This is just an example.)", "summary_of_the_paper": "This paper proposes a baseline RL approach for pixel-based control with distraction in background, which adds reward and latent state transition predictions to actor-critic agent (or DQN agent for Atari experiments). The authors further consider many components proposed in previous contributions on pixel-based control (such as contrastive learning, data augmentation, image reconstruction, etc) and conduct ablation studies on top of the proposed baseline approach. They also directly analyze existing methods for pixel-based control with distraction in background. ", "summary_of_the_review": "Overall, there are good values in the large-scale studies presented in this paper. However, I'm not convinced by conclusions as the authors try to generalize behavior of specific implementation to a family of methods. If I were to implement a new agent, I don't feel like I can believe these conclusions so that makes me question what knowledge this paper can add to the community. Furthermore, many details are either missing or not made clear, and the main story isn't very strong. Therefore, I don't think this paper is ready for publication in the current status.", "strength_and_weaknesses": "", "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": "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": "KHLSgGFw2EX", "reviewer": "Reviewer_RjsF", "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": "[Strengths]\n\n* Representation learning is an important aspect of effective reinforcement learning from observations (here \"pixels\"). Categorization of environments based on characteristics in which different algorithms perform best has certainly helped in the quick development of RL practice. The C1-C8 chategorization and results of Figure 1 are no different and could be of general interest. It is nice to see that for the most part the results confirm existing \"lore\" e.g. world-model reconstruction methods work well without distractors but perform significantly worse with distractors, getting closer to the proposed baseline (difference between C4 and C8)\n\n* Data-augmentation is used to learn invariances with respect to non-informative (for the task) transformations of the input. This is well-known in the representation learning literature and I appreciate the investigation of these properties in the context of RL (discussion in Section 5).\n\n* I agree in principle about the potential utility of predicting rewards to improve the quality of transition model representations, and am happy to see SAC-AE with reward transitions performing on par with more involved baselines.\n\n[Weaknesses]\n\n* It unclear if the main contribution of this work is to introduce an improved variant of the SAC-AE baseline (with reward and transition prediction) to improve robustness against distractors or instead investigate the performance of different approaches to RL in tasks with different characteristics. The paper could do well to focus on a single aspect and develop it more thoroughly. Section 4.2 and the introduction of transition prediction is a prime example of an element that feels forced and not supported by the experiments (Fig. 15 in particular does not inspire particular confidence).\n  \n* To elaborate on above: the paper suffers from tension between attempting to highlight how claiming superiority of a method across all environments is misleading, while at the same time spending much of its time arguing in favor of a specific method that is supposed to \"learn meaningful representations\" (which is never shown). Much like the other previously proposed methods, `baseline` performance itself wildly varies across domains. The work does not investigate the architectural modifications proposed particularly deeply: there are no comparisons to previous approaches that incorporate reward prediction. \n\n* The authors bring up contextual decision processes (CDP) without ever utilizing the framework to illustrate a concept or derive a result. In principle, I understand the choice of CDPs as the focus is on distractors. Are there other motivations behind this choice? Could you not extend the standard POMDP formulation to a case where observations are noise corrupted by noise that is not informative for the task (the so-called distractors?)\n\n* The prescriptive claims in Section 7 should be better contextualized given the results provided in the paper (e.g. \"Pixel reconstruction suffers massively when distractors are added, as shown in ... where ...\"). The degree to which each of these claims is experimentally supported varies greatly. The Section could benefit from providing a summary of experiments done in support of each claim.\n\n[Questions]\n\n* In general, there are only a few domains in which `baseline` appears to beat `baseline-v0` (C4 and C6 and C7), all of which with distractors. Is it correct to state that everything being equal, introducing the proposed modifications to the `baseline` decreases its performance in a setting without distractors, but helps otherwise? The results certainly seem to suggest this is the case.\n\n* The paper is missing a description of what types of distractors are considered and how they have been implemented. Are distractors kept the same (once chosen) across a given training dataset? If not, i.e you sample distractors, why are you treating them any different than data-augmentation (that is guaranteed to retain task information)? This last case would also explain the apparent advantage of introducing the additional reward predictors: you are in effect allowing the representation to learn to ignore backgrounds for example, and building a reward prediction model on this better, \"distilled\" representation. Methods trained in environments without distractors can choose to also rely on non-relevant information for their predictions, which is inherently a worse representation to learn a reward from, hence the difference. Are the distractors instead only introduced at test time?", "summary_of_the_paper": "This work investigates representations for pixel-based reinforcement learning. The authors empirically show how different approaches perform in a variety of domains categorized by their reward structure, presence of distractors and use of data-augmentation techniques. The work further argues in favor of an \"exceedingly simple\" baseline method consisting of SAC-AE modified with transition and reward prediction modules. ", "summary_of_the_review": "In summary, the paper executes a fairly extensive experimental exploration of different RL algorithms in different domains. As is often the case for these types of works, some prescriptive claims made are not fully supported by experimental evidence (e.g Section 4.2). This happens more often when the paper itself argues in favor of a specific algorithmic detail (`baseline` here). Some of the results show evidence for aspects of representation learning in RL that could be of general interest (though they are not all novel in themselves, only summarized together). I am also of the opinion that not including other baselines using reward prediction does not inspire confidence that this approach to its incorporation is any better than the others that exist. I am willing to reconsider my current evaluation.", "strength_and_weaknesses": "", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "1WQmlOh_d9-", "reviewer": "Reviewer_7euW", "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": "Strengths:\n\n1.\tThe paper presents an interesting discussion which is useful to the field of representation learning, visuomotor control etc. that is currently of interest to the community. It is interesting to see the success of a what can be considered fairly standard baseline with reward and transition predictions achieve competitive performance.  \n2.\tIt contains a fairly extensive empirical study of several state of the art and quite recent representation learning techniques such as DREAMER, Task informed abstractions, Deep Bisimulation for Control, CURL, SPR etc.  \n3.\tThe discussion on which losses, decodings and augmentations benefit the representation is interesting and useful. Figure 1 is also a convincing way to tell the story that naïve benchmarking of algorithms does not contain the full picture, and to properly evaluate algorithms, one needs to take a data centric view. This can help kick off a wider discussion on how best to evaluate and benchmark representation learning models.   \n\nWeaknesses:\n\n1. The concept of using reward prediction as a way to help reinforcement learning/representation learning is not very new. For example, [1] contains a discussion on how offline pretraining with reward prediction can help downstream RL/imitation learning. Also relevant are the concept of 'successor features' from [2] helping generalization over tasks where dynamics and transitions are similar, the analysis of reward prediction on representation learning from [3]. This class of techniques is probably closer to the proposed baseline, and comparing performances to some of these would have allowed a more pointed comparison about the reward prediction feature and the novelty of the method. \n\n2. While the main focus of the paper seems to be the discussion itself, and not claiming superiority with the proposed baseline, the main tasks for evaluation are all somewhat visually similar (i.e. DMC agents with distractors). While the distractor backgrounds are likely of different varieties, the underlying task arrangement is still similar: a moving agent on top of a background. It would have been interesting to evaluate how the baseline and other methods perform on more complex tasks such as egocentric visual control.\n\n3. There is not a lot of discussion with regards to generalization of the learned representations. Compared to unsupervised representation learning, one can intuitively understand the advantages of task related supervision such as having access to rewards etc. which in turn benefit the representation. But it is also often hard to handcraft specific tasks, rewards etc. in a way that is beneficial for a model. How well do these representations and associated policies generalize? If several downstream tasks in a given environment are similar to each other (i.e., the task relevant features are similar), can a single trained ‘baseline’ policy generalize? Have the authors considered the case of representation learning where a generalizable representation is learnt in a pre-trained, task-free and reward-free way (one that focuses on agent dynamics, world modeling etc.) and then is finetuned for individual tasks?  \n\nOther comments:\n\n1. The gist of figure 3(a) seems to be the fact that reward prediction plays a major role in creating effective, task-relevant representations which, conceptually is easy enough to understand. I am curious whether such a representation, while beneficial for the task at hand, is detrimental for generalization. For instance, if another way of representation learning (one perhaps not so closely tied to task-specific reward) results in higher sample complexity for individual tasks but better generalization, there might be cases where that is also valuable. Is state transition a more ‘generalizable’ construct than rewards?\n\n2. It is not clear why DREAMER w/o dist in Figure 6, and DREAMER/TIA in Figure 7 have a constant profile of scores. Are they not learning anything?\n\n3. The last part of Section 5 states that performance of a non-contrastive loss often depends on the existence of data augmentations, whereas SPR still maintains strong performance without augmentation. Shouldn’t contrastive loss be the one that is more dependent on data augmentation for sufficient positive-negative pairs? \n\n4. Section 5 states that reconstruction-based techniques like DREAMER suffer when distractors are present. What would be the case if reconstruction is combined with sufficient data augmentations – i.e., different kinds of distractors? Also, what is the intuition for why TIA performs worse than the baseline even with explicit decoupling between task-relevant and task-irrelevant features? \n\n[1] Yang, M., & Nachum, O. (2021). Representation matters: Offline pretraining for sequential decision making. arXiv preprint arXiv:2102.05815.  \n[2] Lehnert, L., Littman, M. L., & Frank, M. J. (2020). Reward-predictive representations generalize across tasks in reinforcement learning. PLoS computational biology, 16(10), e1008317.  \n[3] Hlynsson, H. D., & Wiskott, L. (2021). Reward prediction for representation learning and reward shaping. arXiv preprint arXiv:2105.03172.", "summary_of_the_paper": "This paper presents an approach for learning representations from pixel data that are amenable for control tasks. The proposed approach is a simple baseline that does not require data augmentation, world models, contrastive losses etc. but only contains two simple sub-tasks that are supposed to contribute heavily towards an effective representation: reward prediction and state transition prediction. Along with evaluating this proposed baseline, the paper also compares it to several prior works on representation learning: i.e., several approaches such as data augmentation, distance metric losses, contrastive losses, relevant reconstruction etc. It is shown that the proposed simple baseline either outperforms several of these methods or at least is very close in performance. Finally, the paper presents an interesting discussion about how evaluating an algorithm is not just about the dataset and the chosen benchmark task, but requires a more nuanced point of view of several factors such as reward sparsity, action continuity/discreteness, relevance and irrelevance of features to the task, and so on. The findings of the paper are not just about the effectiveness of the proposed method, but a more overarching view of which types of representation learning methods work in what conditions. ", "summary_of_the_review": "The paper contains a discussion on how to best learn representations from pixel data for control tasks. The authors discuss the insights of  various techniques in the literature and show empirically that simpler constructs that focus on task relevance, system dynamics etc. have the potential to outperform a lot of methods in the literature. While the paper's evaluation focuses on fairly simple tasks, I feel that the empirical results, the extensive analysis and the discussion of strengths and weaknesses of other classes of methods is useful for the community. ", "strength_and_weaknesses": "", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "-fi8m07gAFy", "reviewer": "Reviewer_KvBB", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "Strengths \n- The paper performed various experiments analyzing the contribution of each component used in representation learning for RL. \n- The results suggest, which I agree, that we need to analyze based on the categories of the tasks rather than relying on the average or overall performance of the benchmark.\n- Some experiments (e.g., crop and part reconstruction) are quite interesting as it shows the effect of the object centering. \n\nWeaknesses\n- One of the main claim of the paper is that reward and temporal prediction are vital components providing the most robust performance. I’m not convinced if this can be taken as a general finding about representation learning in RL. It seems that the author’s argument is highly biased mostly to the dense reward setting. For other settings, the result can be different. Specifically, the authors use 8 categories. The proposed baseline seems to work for these categories. However, the result is highly dependent on how the 8 categories are picked because there exist a much larger number of categories in the space of tasks and we don’t know whether the baseline would be robust generally and globally in this space. For example, the reward prediction may not work well for sparse reward tasks (which is currently considered only in one category out of 8). Also, the temporal prediction may not be useful much for fully observable tasks. Also, for unsupervised RL, the result would be pretty different. \n\n- While I found some experiment results are interesting (e.g., cropping ), many of the findings from the paper were actually not surprising rather somewhat already known or expected. For example, in dense reward settings, the fact that it does not perform well after removing reward prediction is pretty obvious. Also, the fact that pixel reconstruction does not work well on image observations with distractors has been shown in many previous works (e.g., TPC [1] and TIA and the related works of these papers). \n\n- For Value-Aware Learning, the paper does not discuss much about the result.\n\n- It would be more interesting if the work also embraces broader classes (e.g., other categories, unsupervised RL, etc.) of tasks to discuss what matters and why about representation learning in RL. \n\n[1] Temporal Predictive Coding For Model-Based Planning In Latent Space", "summary_of_the_paper": "The authors claim that the usefulness of the learning components (e.g., contrastive objective) proposed in previous representation learning for RL, is highly dependent on the specifics of the task category, and thus show that they don't provide robustness across different and diverse task categories. They also claim that the reward and transition prediction are the most vital (and minimal) components providing robustness. Based on this observation, they propose to evaluate the representation learning methods by considering the specific properties of task categories instead of evaluating on a whole benchmark level such as DMC and Atari. ", "summary_of_the_review": "The argument about the vitality and importance of the reward and temporal prediction seems to be an over-claim, or the experiment results do not support this claim enough. Also, many findings of the paper are quite what can be expected and not much surprising. ", "strength_and_weaknesses": "", "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": {}}]

Author rebuttal:

Meta-review: Meta Review for Learning Representations for Pixel-based Control: What Matters and Why?

In this work, the authors presented large-scale empirical evaluations and ablation studies to analyze various components (e.g. contrastive objectives, model-based approaches, data augmentation) for pixel-based control with distractors. Reviewer 7euW wrote a great summary for this paper:

This paper presents an approach for learning representations from pixel data that are amenable for control tasks. The proposed approach is a simple baseline that does not require data augmentation, world models, contrastive losses etc. but only contains two simple sub-tasks that are supposed to contribute heavily towards an effective representation: reward prediction and state transition prediction. Along with evaluating this proposed baseline, the paper also compares it to several prior works on representation learning: i.e., several approaches such as data augmentation, distance metric losses, contrastive losses, relevant reconstruction etc. It is shown that the proposed simple baseline either outperforms several of these methods or at least is very close in performance. Finally, the paper presents an interesting discussion about how evaluating an algorithm is not just about the dataset and the chosen benchmark task, but requires a more nuanced point of view of several factors such as reward sparsity, action continuity/discreteness, relevance and irrelevance of features to the task, and so on. The findings of the paper are not just about the effectiveness of the proposed method, but a more overarching view of which types of representation learning methods work in what conditions. 

Along with myself, most reviewers (including the critical 61FY) agree that there is great value in the large-scale studies presented in this paper. Furthermore, I personally like how it links a large body of recent work in this topic together in one study. The reviews were mixed (6, 6, 3, 3), and the negative reviews (the 3's) generally had issues with not the study or experiments, but the conclusions the authors drew from them. In the words of 61FY (who managed to have a good discussion with the authors):

*I'm not convinced by conclusions as the authors try to generalize behavior of specific implementation to a family of methods. If I were to implement a new agent, I don't feel like I can believe these conclusions so that makes me question what knowledge this paper can add to the community. Furthermore, many details are either missing or not made clear, and the main story isn't very strong. Therefore, I don't think this paper is ready for publication in the current status.*

Although I really appreciate the effort and detail that went into this nice work, based on the current assessments from the 4 reviewers, I can't recommend it for acceptance in its current state. I feel though, that with a change of narrative, or even with a re-examination of the experimental results, the authors can turn the paper around into a highly impactful paper. The description of all of the methods explored, and experiments performed alone makes a wonderful survey of the field with sufficient impact, so I think the authors are *almost* there in publishing a highly impactful work that can make the community look deeper into pixel-based control methods (with distractors). I hope to read an updated version of this work in the future published at a journal or presented at a future conference. Good luck!
Decision: Reject