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

1 INTRODUCTION . 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 . 2 RELATED WORK . 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 . 3 METHOD . 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 can not 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 .

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