Year: 2021

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 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.

heading: 1 INTRODUCTION
text: Learning useful representations for downstream tasks is a key component for success in rich observation 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 be 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 a 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-driving car tries to change lanes. However, if such task-irrelevant components are not discarded, they can lead to sudden failure when the car drives through a different environment, say a forest where there are no buildings or billboards. Avoiding brittle behavior is therefore key to efficient deployment of artificial agents in the real world.
There has been a lot of work recently that tries to learn efficiently from pixels. A 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 is 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, as we will show, a lot of the success of data augmentations is an artifact of how the benchmark environments save data, which is not true of the real world [48], thus resulting in failure1. It is important to 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 to model objects, and it becomes hard to discern if methods that perform well are actually good candidates for answering ‘what really matters for robust learning in 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 RL theory for learning with rich observations [17, 4, 9], we 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 an 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, we identify the failure modes of previously proposed objectives and highlight why they result in comparable or worse performance than the considered baseline. Our observations suggest that relying on a particular method across multiple evaluation settings does not work, as the efficacy varies with the exact details of the task, even within the same benchmark (see Figure 1). We note that a finer categorization of available benchmarks based on metrics like density of reward, presence of taskirrelevant components, inherent horizon of tasks, etc., play a 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 a 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.
1It is also hard to pick exactly which data augmentation will work for a particular environment or task [43]. 2For ALE we use the performance of DREAMER after 1M steps, whereas for DMC we consider the
performance after 500k steps.

heading: 2 RELATED WORK
text: Prior work on auxiliary objectives includes the Horde architecture [50], UVFA [44] and the UNREAL 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 environments considered there always included only task-relevant pixel information, thus not pertaining to the hard setting we consider in this work. Representations can be also be fit so as to obey certain state similarities. If these state metrics preserve the optimal policies and are easy 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] or 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 impressive 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, 57]. However, temporal structure can induce positive/negative pairs as well. In such a 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, our work is largely complimentary to these past investigations. Babaeizadeh 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 a particular benchmark category.

heading: 3 METHOD
text: We model the RL problem 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 an agent must act on the basis of rich observations (context) xt 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 a given state is assumed to 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 a 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, our 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]) as the base architecture, and attach the reward and transition
prediction modules to it (See Figure 2). Note that the transition network is over the encoded state st 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 all our experiments. Details about the implementation, network sizes and all hyperparameters is provided in Appendix 3 and Appendix 4 (Table 3) respectively.

heading: 4 EMPIRICAL STUDY
text: 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 our experiments is that the baseline is able to reduce the gap to more sophisticated methods significantly, sometimes even 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.

heading: 4.1 REWARD PREDICTION
text: Figure 3 (left) shows a comparison of ‘with vs without reward prediction’. All other settings are kept unchanged and the only difference is the reward prediction. When the reward model is removed, there remains no grounding objective for the transition model. This results in a representation collapse as the transition model loss is minimized by the trivial representation which maps all observations to the same encoded state leading to degraded performance. This hints at the fact that without a 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 no 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 to not use the standard maximum likelihood based approaches for the transition model loss and instead use a contrastive version of the loss, which has been shown to 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.1.
Linear Reward Predictor. Furthermore, we also compare to the case when the reward decoder is a linear network instead of the standard 1 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.e., in the linear reward predictor case, it might force the representation to only preserve information required
to predict the reward well, which might not be always be enough to predict the optimal actions well. For instance, consider a 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 be discarded if the reward model is directly attached to the encoded state. This idea is similar to 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.

heading: 4.2 TRANSITION PREDICTION
text: 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 is removed, we observe a slight dip in 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 in performance.

heading: 4.3 CONNECTIONS TO VALUE-AWARE LEARNING
text: 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 a 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 as 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 (1-step), since the value function captures the long term nature of the task in hand. To that end, we also test a 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 in 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.

heading: 5 COMPARISON
text: 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 a 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 4 (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 to 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 augmentations 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-relevant 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 as 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.
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
any kind of loss function, such as the encoding function fθ [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 a 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 [11]. 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 similarity 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 4 (right) that there is very little or no improvement in performance as compared to the baseline performance. This again suggests that the same algorithmic idea can have a completely different performance just by changing the evaluation setting5 (ALE to DMC).
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 without 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 learning 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 a decoupling between the distractor background and task relevant components. We illustrate the performance of each of the above algorithms in Figure 5 along with a version where we add full reconstruction loss to the baseline. Complete performances are provided in Appendix 8.6.
Relevant Reconstruction and Sparse Rewards. Since we consider only dense reward based tasks, using the reward model to ground is sufficient to learn good representations. More sophisticated auxiliary tasks considered in past works include prediction of ensemble of value networks, prediction of past value functions, prediction of value functions of random cumulants, and observation reconstruction. However, in the sparse reward case, grounding on only the reward model or on past value functions can lead to representation collapse if the agent continues to receive zero reward for a long period of time. Therefore, in such cases where good exploration is necessary, tasks such as 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 is sufficient for robust performance. We show
4TIA [19] performance data is taken from their publication. 5The 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 transition loss.
6Part Recons. in Figure 6 involves reconstructing the DMC robot over a solid black background.
this in Figure 6 along with performance of baseline and augmentations. Of course, how one should come up with techniques that differentiate between task-relevant and task-irrelevant 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 to human-achieved scores at 100K steps or 400K frames. We consider the categorization proposed by Bellemare et al. [6] based on the nature of reward (dense, human optimal, score exploit and sparse). We implemented two versions of the baseline algorithm, one 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], SimPLe [30], DER [53], and OTR [32] while being quite inferior to DRQ8 [33, 2] and SPR [46]. In comparison to DER, since our implementation of 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 to 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 to 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 a self-supervised loss may not lead to gains when run 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 a 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).

heading: 6 DISCUSSION
text: 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,
7As also evident by TIA’s [19] performance for DMC ball-in-cup catch experiments. 8We use the DRQ( ) version from Agarwal et al. [2] for fair evaluation and denote it as DRQ. 9Note 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 a fair comparison.
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 a 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.

heading: 7 CONCLUSION
text: 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 as 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 to more robust improvements [38, 16].
• Certain augmentations (‘crop’) do well when data is centered while dropping in performance when data is off-center or when cropping fails to remove considerable amounts of taskirrelevant 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 augmentations, they lead to more significant gains. For Atari100k, 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 4 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 a complete failure in performance in Atari100k.
These results elicit the observation that claiming dominance over other methods for an entire benchmark 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.

heading: 1 METHODS
text: 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 an agent must act on the basis of rich observations (context) to optimize long-term reward. Typically, a CDP is defined as the tuple (X ,S,A,P, γ,R), with observations or contexts xt ∈ X , states st ∈ S, actions at ∈ A, state transition distribution P(st+1|st, at) and reward functionR ≡ r(st, at, st+1), which defines the nature of the task and dynamics of the system. Here xt refers to a high-dimensional observation while the true state of the environment st is not directly available (usually low dimensional) and the agent must construct it on its own. Without knowing P andR, the RL agent’s goal is to use the collected experience to maximize expected sum of rewards, R = ∑∞ t=0 γ trt, consider a discount factor γ ∈ [0, 1).

heading: 1.1 SAC
text: SAC or Soft Actor-Critic [23] deploys a soft or 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 a Qφ functions and a separate value function V . For a transition τ = (xt, at, xt+1, rt) from replay buffer D, the exact losses are described as follows:
LQ(φ) = Eτ∼D [( Qφ(xt, at)− (rt + γV (xt+1)) )2] (1)
The value function is then fit w.r.t to the entropy regularized (soft) Q function. V (xt+1) = Ea′∼π [ Qφ̄(xt+1, a ′)− α log πψ(a′|xt+1) ]
(2)
Finally the policy is fit to maximize both the Q function and the entropy. Lπ(ψ) = −Ea∼π [ Qφ(xt, at)− α logψ(a|xt) ] (3)

heading: 1.2 SAC-AE
text: SAC-AE [56] extends SAC by adding an autoencoder for learning the representation st = fθ(xt), 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 fθ. Only the critic gradients train the representation. This idea has been used as standard practice in most papers that follow.

heading: 1.3 DREAMER
text: DREAMER [25, 26] learns a latent space dynamics model, both for the transition and reward dynamics, and then uses model based updates with a 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 to construct a representation model, p(st+1|st, at, xt), and define targets for the dynamics model. The pixel space prediction loss is referred to as the pixel reconstruction loss and is defined with a reconstruction model q(·|st). The agent minimizes the expectation (LDREAMER) over a finite horizon summation of reconstruction (LtO), reward prediction (LtR) and transition prediction (LtD) losses, such that:
LDREAMER = E [∑ t ( LtO + LtR + LtD )] (4)
LtO = − ln q(xt|st), LtR = − ln R̂(rt|st), LtD = β KL ( p(st+1|st, at, xt) ‖ P̂(st+1|st, at) )

heading: 1.4 RAD AND DRQ
text: 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 as compared to other augmentations. Therefore, when no mention of the exact augmentations is provided, assume that random crop is being used.

heading: 1.5 CURL
text: CURL [36] uses data augmentations (just as RAD) over a standard SAC agent and additionally deploys an auxiliary loss that tries to learn representations such that augmented versions of the same observation or context are pulled together in the representation space while pushing apart augmented versions of two different observations. CURL uses the bi-linear inner-product qTWk, where W is a learned parameter. Here, q and k are the latent representations of the raw anchors (query) xq and targets (keys) xk, such that q = fθ(xq) and k = sg(fθ(xk)), where sg denotes the stop gradient operation.

heading: 1.6 DER AND OTR
text: Both DER [53] and OTR [32] refer to a 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.

heading: 1.7 SPR
text: SPR [46] or self-predictive representations deploys an auxiliary transition prediction loss which is learned in a non-contrastive self-supervised manner by using cosine similarity loss. In particular, instead of using the actual next state st+k as the target for the predicted next state ŝt+k, the authors use a projection network to transform both the targets (gm) and the predictions (go). Furthermore, the predictions are then then passed through a predictor network (q), and the output of this network and of the target projector are used to define a cosine similarity loss LCosine. The use of the projector and predictor networks is similar to that of in self-supervised learning methods in vision, such as BYOL [21].
LCosine = − H∑ h=1 ( yt+h ‖yt+h‖ )T ( ŷt+h ‖ŷt+h‖ ) , (5)
where ŝt+k = P̂ (st, at, . . . , at+k−1) , ŷt+h = q (go (ŝt+h)) , yt+h = gm (st+h) In particular, SPR also uses data augmentations and a convolutional transition dynamics predictor. When not using data augmentations, SPR uses dropout in the encoder network fθ.

heading: 1.8 DBC
text: DBC [63] or Deep Bisimulation for Control is an algorithm based on bisimulation metric losses which tries to 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 (x, a, r, x′) and minimizing the L1 norm between any two non-identical transitions in the sampled batch as follows:
LDBC = ( ‖si − sj‖1 − |ri − rj | − γ W2 ( P̂(si, ai), P̂(sj , aj) ))2 (6)
where si = fθ(xi), si = sg ( fθ(xi) ) and W2 is the 2-Wasserstein distance metric between two transition distributions. Please refer to the paper for further details.

heading: 1.9 BASELINE
text: We use a simple soft actor-critic setup with an embedding function fθ : X → S (similar to SACAE [56]) as the base architecture, and attach the reward and transition prediction modules to it (See Figure 2). We define the transition prediction by P̂(st, at) and the reward prediction by R̂(st+1) such that, ŝt+1 = P̂(st, at) and r̂t = R̂(st+1). Note that the transition network is over the encoded state ŝt = fθ(xt). and not over the observations xt [37]. The overall auxiliary loss function is thus defined as follows:
Lbaseline = ( st+1 − P̂(st, at) )2︸ ︷︷ ︸ Transition prediction loss + ( R(st+1)− R̂(P̂(st, at)) )2︸ ︷︷ ︸ Reward prediction loss
(7)

heading: 1.10 VALUE-AWARE MODEL LEARNING
text: Based on the baseline architecture, the value aware experiments are performed with the SAC critic Q, policy π and the corresponding value aware loss, LV A, as ( sg(Vt+1)− V̂t+1 )2 , where
Vt+1 = Eat+1∼π [Q(st+1, at+1)− α log π(at+1|st+1)] V̂t+1 = Eat+1∼π [Q(ŝt+1, at+1)− α log π(at+1|ŝt+1)] ,
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.

heading: 2 BACKGROUND DISTRACTOR DETAILS
text: 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 task-irrelevant 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.

heading: 3 IMPLEMENTATION DETAILS
text: We use the SAC algorithm as the RL 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 as shown in RAD and CURL i.e. random cropping (84× 84) from observations of (100× 100) dimensions for DMC and random shifts of ±4 pixels for ALE. Random flip and rotate are applied in a similar way as well. For the reconstruction loss augmentations, a 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 https://github.com/Showmax/kinetics-downloader SAC-AE https://github.com/denisyarats/pytorch_sac_ae RAD https://github.com/MishaLaskin/rad CURL-DMC https://github.com/MishaLaskin/curl DrQ https://github.com/denisyarats/drq DBC https://github.com/facebookresearch/deep_bisim4control PI-SAC https://github.com/google-research/pisac Dreamer https://github.com/danijar/dreamer TIA https://github.com/kyonofx/tia CURL-ALE https://github.com/aravindsrinivas/curl_rainbow SPR https://github.com/mila-iqia/spr
Computation. All experiments were conducted on either system configuration of:
1. 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.2GHz, one ®NVIDIA V100 Volta (16GB HBM2 memory) GPU and 84 GB RAM
The average completion time of the baseline experiments was 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 is expected to 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.

heading: 4 HYPERPARAMETERS
text: The full set of hyperparameters used for the baseline experiments are provided in Table 3 below.

heading: 5 ADDITIONAL BASELINE ABLATIONS

heading: 6 BASELINE WITH MULTISTEP TRAINING
text: 0.0 0.1 0.2 0.3 0.4 0.5 0.0
0.2
0.4
0.6
Cartpole Swingup
0.0 0.1 0.2 0.3 0.4 0.5 0.0
0.1
0.2
Cheetah Run
0.0 0.1 0.2 0.3 0.4 0.5 0.00
0.25
0.50
0.75
Finger Spin
0.0 0.1 0.2 0.3 0.4 0.5 0.00
0.05
0.10
0.15
0.20 Reacher Easy
0.0 0.1 0.2 0.3 0.4 0.5 0.0
0.2
0.4
Walker Walk
0.0 0.1 0.2 0.3 0.4 0.5 0.00
0.01
0.02
0.03
0.04 Hopper Hop
Environment Steps (×106)
No rm
al ize
d Sc
or e
Reconstruction horizon=1 Reconstruction horizon=2 Reconstruction horizon=4 Reconstruction horizon=6
Figure 13: Baseline w/ multi-step training and reconstruction Performance of the proposed baseline on six DMC tasks with multi-step training and added reconstruction loss for each prediction. This resembles the architecture of DREAMER [25] corresponding to baseline. Full reconstruction tries to capture task-irrelevant information leading to decrease in performance. Further, the performance degrades irrespective of the choice of the horizon.

heading: 7 DATASET CATEGORIZAIONS AND SPECIFICATIONS

heading: 8.1 ROLE OF EACH COMPONENT IN BASELINE

heading: 8 DMC SUITE: COMPLETE RESULTS

heading: 8.3 ROLE OF AUGMENTATIONS

heading: 8.2 ARRANGEMENT OF REWARD AND TRANSITION LOSSES

heading: 8.4 ROLE OF CONTRASTIVE LOSSES WITH AUGMENTATIONS

heading: 8.6 ROLE OF RECONSTRUCTION LOSSES

heading: 8.5 ROLE OF CONTRASTIVE AND SSL W/O CONTRASTIVE LOSSES

heading: 8.8 VALUE AWARE LEARNING

heading: 8.7 ROLE OF AUGMENTATIONS WITH RECONSTRUCTION LOSSES

heading: 9 ALE BENCHMARK: ATARI 100K

Review summaries:
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. 

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. 

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. 

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.