Measuring Visual Generalization in Continuous Control from Pixels

Self-supervised learning and data augmentation have significantly reduced the performance gap between state and image-based reinforcement learning agents in continuous control tasks . However , it is still unclear whether current techniques can face the variety of visual conditions required by real-world environments . We propose a challenging benchmark that tests agents ’ visual generalization by adding graphical variety to existing continuous control domains . Our empirical analysis shows that current methods struggle to generalize across a diverse set of visual changes , and we examine the specific factors of variation that make these tasks difficult . We find that data augmentation techniques outperform self-supervised learning approaches , and that more significant image transformations provide better visual generalization . 1 Introduction . Reinforcement Learning has successfully learned to control complex physical systems when presented with real-time sensor data ( Gu et al. , 2017 ) ( Kalashnikov et al. , 2018 ) . However , much of the field ’ s core algorithmic work happens in simulation ( Lillicrap et al. , 2015 ) ( Haarnoja et al. , 2018a ) , where all of the environmental conditions are known . In the real world , gaining access to precise sensory state information can be expensive or impossible . Camera-based observations are a practical solution , but create a representation learning problem in which important control information needs to be recovered from images of the environment . Significant progress has been made towards a solution to this challenge using auxiliary loss functions ( Yarats et al. , 2019 ) ( Zhang et al. , 2020 ) ( Srinivas et al. , 2020 ) and data augmentation ( Laskin et al. , 2020 ) ( Kostrikov et al. , 2020 ) . These strategies often match or exceed the performance of state-based approaches in simulated benchmarks . However , current continuous control environments include very little visual diversity . If existing techniques are going to be successful in the real world , they will need to operate in a variety of visual conditions . Small differences in lighting , camera position , or the surrounding environment can dramatically alter the raw pixel values presented to an agent , without affecting the underlying state . Ideally , agents would learn representations that are invariant to task-irrelevant visual changes . In this paper , we investigate the extent to which current methods meet these requirements . We propose a challenging benchmark to measure agents ’ ability to generalize across a diverse set of visual conditions , including changes in camera position , lighting , color , and scenery , by extending the graphical variety of existing continuous control domains . Our benchmark provides a platform for examining the visual generalization challenge that image-based control systems may face in the real world while preserving the advantages of simulation-based training . We evaluate several recent approaches and find that while they can adapt to subtle changes in camera position and lighting , they struggle to generalize across the full range of visual conditions and are particularly distracted by changes in texture and scenery . A comparison across multiple control domains shows that data augmentation significantly outperforms other approaches , and that visual generalization benefits from more complex , color-altering image transformations . 2 Background . Reinforcement Learning : We deal with the Partially Observed Markov Decision Process ( POMDP ) setting defined by a tuple ( S , O , φ , A , R , T , γ ) . S is the set of states , which are usually low-dimensional representations of all the environment information needed to choose an appropriate action . In contrast , O is typically a higher-dimensional observation space that is visible to the agent . R is the reward function S x A → R , T is the function representing transition probabilities S x S → [ 0 , 1 ] . φ is the emittion function S → O that determines how the observations in O are generated by the true states in S. A is the set of available actions that can be taken at each state ; in the continuous control setting we focus on in this paper , A is a bounded subset of Rn , where n is the dimension of the action space . An agent is defined by a stochastic policy π that maps observations to a distribution over actions in A . The goal of Reinforcement Learning ( RL ) is to find a policy that maximizes the discounted sum of rewards over trajectories of experience τ , collected from a POMDP M , denoted ηM ( π ) : ηM ( π ) = E τ∼π [ t=∞∑ t=0 γtR ( st , at ) ] , where γ ∈ [ 0 , 1 ) is a discount factor that determines the agent ’ s emphasis on long-term rewards . Continuous Control : Many of the challenging benchmarks and applications of RL involve continuous action spaces . These include classic problems like Cartpole , Mountain Car and Inverted Pendulum ( Moore , 1990 ) ( Barto et al. , 1990 ) . However , recent work in Deep RL has focused on a collection of 3D locomotion tasks in which a robot is rewarded for a specific type of movement in a physics simulator . At each timestep , the state is a vector representing key information about the robot ’ s current position and motion ( e.g . the location , rotation or velocity of joints and limbs ) . The action space is a continuous subset of Rn , where each element controls some aspect of the robot ’ s movement ( e.g . the torque of a joint motor ) . There have been several standardized implementations of these environments ( Schulman et al. , 2015 ) ( Brockman et al. , 2016 ) . This paper will use modified versions of those provided in the DeepMind Control Suite ( DMC ) ( Tassa et al. , 2018 ) , built using the MuJoCo physics simulator ( Todorov et al. , 2012 ) . Soft Actor Critic : Soft Actor Critic ( SAC ) ( Haarnoja et al. , 2018a ) is an off-policy actorcritic method that achieves impressive performance on many continuous control tasks . For each training step , a SAC agent collects experience from the environment by sampling from a high-entropy policy ( a ∼ πθ ( o ) ) paramaterized by a neural network with weights θ . It then adds experience to a large buffer D in the form of a tuple ( o , a , r , o′ ) 1 . SAC samples a batch of previous experience from the buffer , and updates the actor to encourage higher-value actions , as determined by the critic network ( Qφ ( o , a ) ) : Lactor = − E o∼D [ min i=1,2 Qφ , i ( o , ã ) − α log πθ ( ã|o ) ) ] , ã ∼ πθ ( o ) ( 1 ) The critic networks are updated by minimizing the mean squared error of their predictions relative to a bootstrapped estimate of the action-value function : Lcritic = E ( o , a , r , o′ ) ∼D [ ( Qφ ( o , a ) − ( r + γ ( min i=1,2 Qφ′ , i ( o ′ , ã′ ) − α log πθ ( ã′|o′ ) ) ) ) 2 ] , ã′ ∼ πθ ( o′ ) ( 2 ) The log πθ terms are a maximum-entropy component that encourages exploration subject to a parameter α , which can either be fixed or learned by gradient descent to approach a target entropy level ( Haarnoja et al. , 2018b ) . The min operation makes use of two critic networks that are trained separately , and helps to reduce overestimation bias ( Fujimoto et al. , 2018 ) . φ′ refers to target networks , which are moving averages of the critic networks ’ weights and help to stabilize learning - a common trick in Deep RL ( Lillicrap et al. , 2015 ) ( Mnih et al. , 2015 ) . Generalization in Reinforcement Learning : Consider a set of POMDPs M = { ( S0 , φ0 , O0 , A0 , R0 , T0 , γ0 ) , ... , ( Sn , φn , On , An , Rn , Tn , γn ) } . Given access to one or several training tasksMtrain ∼M , we would like to learn a policy that can generalize to the entire 1We leave out the terminal boolean d common to most implementations both for notational simplicity and because many of the DMC tasks reset after a fixed amount of time . set . This is analogous to the way we train supervised models to generalize across an entire distribution despite only having labels for a finite number of inputs . In supervised learning , we would partition our labeled data into training and testing sets , and measure the difference in our model ’ s performance between them . We can createMtrain andMtest in a similar way ( Zhang et al. , 2018a ) . This paper focuses on the zero-shot generalization setting ( |Mtrain| = 1 ) , and will normalize by the training return in an effort to compare the generalization of policies with different baseline levels of performance . A policy ’ s generalization er- ror , EG ( π ) , is defined as : EG ( π ) ≡ ( ηMtrain ( π ) − 1|Mtest| ∑|Mtest| i=1 ηMtest , i ( π ) ) / ( ηMtrain ( π ) ) . 3 A Benchmark for Measuring Visual Generalization in Continuous Control from Pixels . 3.1 Continuous Control From Pixels . While SAC performs well on many of the DMC tasks , it has access to state information that would be difficult to replicate in the context of real-world robotics ( O = S ) . Ideally , we could learn a policy from imperfect or high-dimensional observations . This had led to a more challenging benchmark where the same set of tasks are learned from image renderings of the environment . Many successful approaches build around the addition of a convolutional encoder to the SAC agent architecture . The encoder attempts to process a stack of consecutive frames into a low-dimensional representation that is then treated as a surrogate state vector for the actor and critic ( s ) . However , the RL setting ’ s sparse gradient signal makes it difficult to train a high-capacity encoder with a limited number of samples , and baseline implementations typically fail to make progress ( Tassa et al. , 2018 ) . Several recent works have addressed this problem by adding additional objectives to the encoder ’ s gradient update . SAC+AE adds an image reconstruction term that trains a decoder network to recover the original image from the encoder ’ s compressed representation , as in a standard autoencoder ( Yarats et al. , 2019 ) . CURL uses a contrastive loss ( Srinivas et al. , 2020 ) that encourages similar representations across random crops of its input images . DBC uses a bisimulation metric to discard observation information that isn ’ t relevant to control ( Zhang et al. , 2020 ) . All three result in performance at or near the levels of state-based SAC and experiments show that the representation learned by the encoder recovers key state information . Data Augmentation in Deep RL : Another recent line of work applies standard data augmentations from image processing to the training loops of RL algorithms . This simple technique can outperform more complicated approaches like SAC+AE and CURL . There are some subtleties in both the types of data augmentations that are used and how they are applied within the training update . RAD ( Laskin et al. , 2020 ) renders the environment at a higher resolution and randomly crops each observation as it sampled from the replay buffer . Data Regularized Q-Learning ( DrQ ) ( Kostrikov et al. , 2020 ) is a similar approach that includes the option to augment the o′ images separately within each timestep in hopes of computing a lower-variance target for the critic updates . Instead of upscaling , DrQ pads each image by replicating border pixels before randomly cropping back to the original size . Data Regularized Actor-Critic ( DrAC ) ( Raileanu et al. , 2020 ) more explicitly regularizes the outputs of the actor and critic networks to be invariant under each augmentation . It also attempts to learn which transformations to apply automatically ; experiments in all of these works show this to be a surprisingly high-stakes problem , as many augmentation choices have little to no effect on learning ( cf . Laskin et al . ( 2020 ) Figure 2a , Kostrikov et al . ( 2020 ) Appendix E Figure 6 , and Appendix A.3.3 ) . The search finds random crop to be the most helpful in many of the Procgen benchmark games ( Cobbe et al. , 2019a ) and shows that it is a strong default even when outperformed by other methods . Most of the successful augmentations aside from Random Crop are color-based transformations . These include Network Randomization ( NR ) ( Lee et al. , 2019 ) - which uses randomly initialized convolutional layers to create visually distinct filters of the same image - and Color Jitter ( CJ ) , which shuffles the color palette . Instead of benchmarking NR , DrQ and RAD separately , we implement a custom version of Pixel SAC that lets us combine aspects of all three algorithms . The overall structure is similar to the originals : we use a large convolutional encoder that is only updated by the critic ’ s gradients and copy many of the baselines ’ hyperparameters for the sake of comparison . Data Figure 1 : Example “ Walker , Walk ” visual seeds . The random aesthetic changes generated by DMCR allow for significant visual diversity . φ0 uses the default DMC assets . augmentation is applied in a way that is most similar to RAD ( Laskin et al. , 2020 ) . We introduce a new hyperparameter β ∈ [ 0 , 1 ] which controls how we mix the augmented and original batches for the actor and critic updates . β = 1 uses augmented data only ; β = 0 uses the original . We set β = .9 . This setup looks to alleviate the need to perform test time augmentation , even when the augmentation significantly alters the image , as in ( Lee et al. , 2019 ) ( see more analysis in Appendix A.3.1 ) . In addition , we encourage the output of the encoder , eξ to be invariant to our augmentation pipeline z , by adding an explicit regularization term to the encoder loss : Lencoder = Lcritic + λ E o∼D [ ||eξ ( o ) − eξ ( z ( o ) ) ||2 ] ( 3 ) We set λ = 1e-5 , and only let gradients flow through the eθ ( z ( o ) ) pass , for stability . Appendix A.3.2 discusses this idea in more detail . After a comparison ( Appendix A.3.3 ) , we settle on the pad/crop transformation from DrQ ( Kostrikov et al. , 2020 ) as our primary augmentation . We will refer to this algorithm as SAC+AUG . Any other transformations that are used will be added to the acronym in results and figures . For example , SAC+AUG that applies Color Jitter before the DrQ crop will be denoted SAC+CJ+AUG .

Summary: This paper presents a new dataset for continuous control tasks based on the DeemMind control suite, but with varying backgrounds, colors and lighting/camera conditions. Instead of the fixed appearance in the original suite the frames from this dataset contain a lot of varying structure. The purpose of the dataset is to measure the generalization and transfer abilities of pixel-based continuous control agents across different visual variations.