L EARNING R EPRESENTATIONS FOR P IXEL - BASED C ONTROL : W HAT M ATTERS AND W HY ?
Anonymous authors
Paper under double-blind review
A BSTRACT
Learning representations for pixel-based control has garnered signiﬁcant attention recently in reinforcement learning.
A wide range of methods have been proposed to enable efﬁcient 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 difﬁcult setting, incorporating background distractors, as a ﬁrst 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 ﬁner 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.
I NTRODUCTION
Learning useful representations for downstream tasks is a key component for success in rich observation environments [14, 38, 46, 53, 54].
Consequently, a signiﬁcant 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 [41, 35, 11], state similarity metrics like bisimulation or policy similarity [61, 60, 1], and pixel reconstruction losses [28, 19, 24, 23].
On a separate axis, data augmentations have been shown to provide huge performance boosts when learning to control from pixels [34, 32].
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 efﬁcient deployment of artiﬁcial agents in the real world.
C
C C
C C
C C
C
P r o b a b ili t y
o
f
R a n k( i n % ) Methods:
Dataset Categories
Baseline
Baseline-v
Contrastive
Metric
Non-Contrastive
Reconstruction
There has been a lot of work recently that tries to learn efﬁciently 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 [28, 61, 35].
A related line of work learns representations by constructing world models directly from pixels [44, 40, 21, 24].
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 [59, 30, 47], 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 [34, 32] 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 [47], thus resulting in failure 1 .
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 [16, 4, 9], we hypothesize certain key components to be responsible for sample efﬁcient learning.
We test the role these play in previously proposed representation learning objectives and then propose 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 with natural distractors [59, 30, 47], DMC Suite [50], and Atari100K [29] 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 proposed baseline.
Our observations suggest that relying on a particular method across multiple evaluation settings does not work, as the efﬁcacy varies with the exact details of the task, even within the same benchmark (see Figure 1).
We note that a ﬁner 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 efﬁcacy of a method.
We list such categorizations as suggestions for more informative future evaluations.
The ﬁndings 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.
R ELATED W ORK
Prior work on auxiliary objectives includes the Horde architecture [49], UVFA [43] and the UNREAL agent [28].
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 ﬁt 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 efﬁcient representations either by learning the metrics like that in bisimulation [17, 61, 60, 8], by recursively sampling states [10] or by model invariance [51].
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, ﬂip, rotate, random convolution, etc.
Latest works [34, 57] 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.
[47] 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 [41, 11].
The most widely used method to generate positive/negative pairs is through various data augmentations [35, 45, 37, 48, 56].
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 [41].
Other ways of generating positive/negative pairs can be through learnt state metrics [1] or encoding instances [24].
Another popular idea for learning representations is learning world models [21] in the pixel space.
This involves learning prediction models of the world in the pixel space using a pixel reconstruction loss [19, 23, 24].
Other methods that do not explicitly learn a world model involve learning representations using reconstruction based approaches like autoencoders [55].
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 speciﬁcally on representation learning methods that improve sample efﬁciency in pixel-based environments. Henderson et al.
[26] 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 ﬁne technique only in a particular benchmark category.
M ETHOD
We model the RL problem using the framework of contextual decision processes (CDPs), a term introduced in Krishnamurthy et al.
[33] to broadly refer to any sequential decision making task where an agent must act on the basis of rich observations (context) x t to optimize long-term reward.
The true state of the environment s t 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 [61, 47].
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 [55]) as the base architecture, and attach the reward and transition [Figure] prediction modules to it (See Figure 2).
Note that the transition network is over the encoded state s t and not over the observations [36].
Moreover, the transition model is ﬁt 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.
E MPIRICAL S TUDY
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 signiﬁcantly, 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.
4. R EWARD P REDICTION
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 efﬁciently 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 ﬁx 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 signiﬁcantly 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 [Figure] 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 ﬁnal performance [11].
We indeed see a similar trend in the RL case as well.
4. T RANSITION P REDICTION
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 sufﬁcient 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.
4. C ONNECTIONS TO V ALUE -A WARE L EARNING
The baseline introduced above also resembles a prominent idea in theory, that of learning value aware models [16, 4].
Value-aware learning advocates for learning a model by ﬁtting it to the value function of the task in hand, instead of ﬁtting 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 deﬁning 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.7.
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 [39] lead to similar scores as the baseline.
[Table]
[Figure]
C OMPARISON
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 [Figure]
Metric Losses .
Two recent works that are similar to the baseline above are DBC [61] 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 [34] and D R Q [32] algorithms, both of which differ very minimally in their implementations.
We run experiments for three different augmentations— ‘crop’, ‘ﬂip’, 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 ﬂip, 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.
[Table]
Contrastive and Non-contrastive Losses .
A lot of recent methods also deploy different contrastive losses (for example, CPC [41] and InfoNCE [42]) to learn representations, which essentially refers to computing positive/negative pairs and pushing/pulling them apart respectively.
In practise, this can be done for any kind of loss function, such as the encoding function f θ [24], or using random augmentations [35, 37], so on and so forth.
We test a simple modiﬁcation, 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 S IM S IAM [12], B ARLOW T WINS [58], and BYOL [20] getting similar or better performance than contrastive methods like S IM CLR [11].
Complete performances are provided in Appendix 8.5.
[Figure]
Note that in the ALE results (Figure 7), SPR [45] leads to the best results overall, which also deploys a speciﬁc similarity loss for transition prediction motivated by BYOL [20].
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 setting 5 (ALE to DMC).
Learning World Models .
We test D REAMER [24], 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,
D REAMER 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.
Their method, called TIA[18] incorporates several other modules in addition to
D REAMER 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.
[Figure]
Relevant Reconstruction and Sparse Rewards .
Since we consider only dense reward based tasks, using the reward model to ground is sufﬁcient 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 [18], especially when concerned with realistic settings like that of distractors.
To that end, we show that in
[Figure]
the sparse reward case, task-relevant reconstruction 6 is sufﬁcient for robust performance.
We show 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 [35], SimPLe [29], DER [52], and OTR [31] while being quite inferior to D R Q 8 [32, 2] and
SPR [45].
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 ﬁnd that when plotting the same average performance results for only the dense reward environments in 100K, the gap in performance between DER and SPR/D R Q decreases drastically.
Note that both SPR builds over DER but D R Q builds over OTR.
We further delve into understanding the superior performance of SPR and D R Q. In particular, SPR combines a non-contrastive transition prediction loss and data augmentations. By non-contrastive, we refer to a self-supervised loss that does not use contrastive learning ideas.
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 a non-contrastive loss does not lead to gains when run without data augmentations.
Finally, we take the DER agent and add data augmentations to it (from D R Q).
This is shown as DER + A UG in
Figure 7.
We see that this leads to collapse, with the worst performance across all algorithms.
Note that D R Q 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).
D ISCUSSION
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 ﬁner categorization of the evaluation benchmark, and not the whole benchmark itself.
Speciﬁcally, focusing on ﬁner 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 speciﬁc 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.
C ONCLUSION
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 beneﬁts 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 ﬁndings 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, D REAMER 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 constrastive loss does not lead to performance improvements over the baseline in DMC Suite (Figure 4 right plot).
• 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 (‘ﬂip’ and ‘rotate’) show the opposite behavior (RAD ablations on DMC Suite in Table 2).
• Non-contrastive losses do not provide much gains when used alone.
With data augmentations, they lead to more signiﬁcant 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 speciﬁc 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. Appendix M ETHODS We model the setting in this paper using the framework of contextual decision processes (CDPs), a term introduced in Krishnamurthy et al.
[33] 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 deﬁned as the tuple ( X , S , A , P , γ, R ) , with observations or contexts x t ∈ X , states s t ∈ S , actions a t ∈ A , state transition distribution P ( s t +1 | s t , a t ) and reward function R ≡ r ( s t , a t , s t +1 ) , which deﬁnes the nature of the task and dynamics of the system.
Here x t refers to a high-dimensional observation while the true state of the environment s t 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 to use the collected experience to maximize expected sum of rewards, R = (cid:80) ∞ t =0 γ t r t , consider a discount factor γ ∈ [0 , 1) . 1. SAC SAC or Soft Actor-Critic [22] 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 τ
=
( x t , a t , x t +1 , r t ) from replay buffer D , the exact losses are described as follows:
[Equation]
The value function is then ﬁt w.r.t to the entropy regularized (soft) Q function.
[Equation]
Finally the policy is ﬁt to maximize both the Q function and the entropy.
[Equation]
1. SAC-AE
SAC-AE [55] extends SAC by adding an autoencoder for learning the representation s t =
f θ ( x t ) , 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.
1. D REAMER
D REAMER [24, 25] 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 [23]) 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 ( s t +1 | s t , a t , x t ) , and deﬁne targets for the dynamics model.
The pixel space prediction loss is referred to as the pixel reconstruction loss and is deﬁned with a reconstruction model q ( ·| s t ) .
The agent minimizes the expectation ( L DREAMER ) over a ﬁnite horizon summation of reconstruction ( L tO ) , reward prediction ( L tR ) and transition prediction ( L t D ) losses, such that:
[Equation]
s t , a t ) (cid:1)
1. RAD AND D R Q
RAD [34] uses SAC as the main algorithm and adds data augmentations to replay buffer data.
These augmentations include: random cropping, translate, grayscale, cutout, rotate, ﬂip, color jitter, and random convolution.
For ALE, in general, rangom shifts and intensity are used as augmentations by
DrQ [32].
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.
1. CURL
CURL [35] 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 q T Wk , where W is a learned parameter.
Here, q and k are the latent representations of the raw anchors (query)
x q and targets (keys) x k , such that q = f θ ( x q ) and k = sg ( f θ ( x k )) , where sg denotes the stop gradient operation.
1. DER AND OTR
Both DER [52] and OTR [31] refer to a highly tuned version of the Rainbow [27] 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 efﬁcient DQN, overtrained DQN), which most notably does not use a distributional component in it.
1. SPR
SPR [45] or self-predictive representations deploys an auxiliary transition prediction loss which is learned in a non-contrastive self-supervised manner.
In particular, instead of using the actual next state s t
+ k as the target for the predicted next state ˆ s t + k , the authors use a projection network to transform both the targets ( g m ) and the predictions ( g o ).
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 deﬁne a cosine similarity loss L Cosine .
The use of the projector and predictor networks is similar to that of in self-supervised learning methods in vision, such as BYOL [20].
[Equation]
where
[Equation]
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 θ .
1. DBC
DBC [62] 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 speciﬁc 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 (cid:48) ) and minimizing the L1 norm between any two non-identical transitions in the sampled batch as follows:
[Equation]
(6) where s i =
f θ ( x i ) , s i
= sg (cid:0)
f θ (
1. B ASELINE
We use a simple soft actor-critic setup with an embedding function f θ : X → S (similar to SACAE [55]) as the base architecture, and attach the reward and transition prediction modules to it (See
Figure 2).
We deﬁne the transition prediction by ˆ P ( s t , a t ) and the reward prediction by ˆ R ( s t +1 ) such that, ˆ s t +1 =
ˆ P ( s t , a t ) and ˆ r t =
ˆ R ( s t +1 ) .
Note that the transition network is over the encoded state
ˆ s t = f θ ( x t ) .
and not over the observations x t [36].
The overall auxiliary loss function is thus deﬁned as follows:
[Equation]
1. V ALUE - AWARE M ODEL L EARNING
Based on the baseline architecture, the value aware experiments are performed with the SAC critic Q , policy π and the corresponding value aware loss, L V A , as (cid:0) sg ( V t +1 )
− ˆ V t +1 (cid:1) 2 , where
[Equation]
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.7.
B ACKGROUND D ISTRACTOR D ETAILS
We use the distracting control suite with natural distractors as previously introduced in Kay et al.
[30], Zhang et al.
[59] and used in Zhang et al.
[62], Stone et al.
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.
I MPLEMENTATION D ETAILS
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.
[55] 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 conﬁguration 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 ﬂip 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.
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.
[Table]
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
[Table]
https://github.com/denisyarats/drq
DBC https://github.com/facebookresearch/deep_bisim4control
PI
[Table]
SAC https://github.com/google-research/pisac
[Table]
https://github.com/danijar/dreamer
[Table]
https://github.com/kyonofx/tia
[Table]
https://github.com/aravindsrinivas/curl_rainbow
[Table]
https://github.com/mila-iqia/spr
Computation .
All experiments were conducted on either system conﬁguration of:
1. 6 CPU cores of ®Intel Gold 6148 Skylake@2.4 GHz, one ®NVidia V100SXM2 (16G memory)
GPU and 84 GB RAM.
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 ﬁnal version of this paper which corresponds to roughly 4 years of GPU training.
H YPERPARAMETERS
The full set of hyperparameters used for the baseline experiments are provided in Table 3 below.
[Table]
A DDITIONAL B ASELINE A BLATIONS
[Figure]
[Figure]
B ASELINE WITH M ULTISTEP T RAINING
[Figure]
[Figure]
D ATASET C ATEGORIZAIONS AND S PECIFICATIONS
[Table]
[Table]
[Table]
DMC S UITE : C OMPLETE R ESULTS 8. R OLE OF E ACH C OMPONENT IN B ASELINE
[Figure]
[Figure]
8. A RRANGEMENT OF R EWARD AND T RANSITION L OSSES
[Figure]
and conclude that both perform equally well.
We choose reward through transition as it is marginally better
[Figure]
8. R OLE OF A UGMENTATIONS
[Figure]
[Figure]
[Figure]
8. R OLE OF C ONTRASTIVE L OSSES WITH A UGMENTATIONS
[Figure]
8. R OLE OF C ONTRASTIVE AND N ON -C ONTRASTIVE LOSSES
[Figure]
8. R OLE OF R ECONSTRUCTION LOSSES
[Figure]
[Figure]
[Figure]
[Figure]
ALE B ENCHMARK : A TARI 100K
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]