Self-Supervised Policy Adaptation during Deployment

1 INTRODUCTION . Deep reinforcement learning ( RL ) has achieved considerable success when combined with convolutional neural networks for deriving actions from image pixels ( Mnih et al. , 2013 ; Levine et al. , 2016 ; Nair et al. , 2018 ; Yan et al. , 2020 ; Andrychowicz et al. , 2020 ) . However , one significant challenge for real-world deployment of vision-based RL remains : a policy trained in one environment might not generalize to other new environments not seen during training . Already hard for RL alone , the challenge is exacerbated when a policy faces high-dimensional visual inputs . A well explored class of solutions is to learn robust policies that are simply invariant to changes in the environment ( Rajeswaran et al. , 2016 ; Tobin et al. , 2017 ; Sadeghi & Levine , 2016 ; Pinto et al. , 2017b ; Lee et al. , 2019 ) . For example , domain randomization ( Tobin et al. , 2017 ; Peng et al. , 2018 ; Pinto et al. , 2017a ; Yang et al. , 2019 ) applies data augmentation in a simulated environment to train a single robust policy , with the hope that the augmented environment covers enough factors of variation in the test environment . However , this hope may be difficult to realize when the test environment is truly unknown . With too much randomization , training a policy that can simultaneously fit numerous augmented environments requires much larger model and sample complexity . With too little randomization , the actual changes in the test environment might not be covered , and domain randomization may do more harm than good since the randomized factors are now irrelevant . Both phenomena have been observed in our experiments . In all cases , this class of solutions requires human experts to anticipate the changes before the test environment is seen . This can not scale as more test environments are added with more diverse changes . Instead of learning a robust policy invariant to all possible environmental changes , we argue that it is better for a policy to keep learning during deployment and adapt to its actual new environment . A naive way to implement this in RL is to fine-tune the policy in the new environment using rewards as supervision ( Rusu et al. , 2016 ; Kalashnikov et al. , 2018 ; Julian et al. , 2020 ) . However , while it is relatively easy to craft a dense reward function during training ( Gu et al. , 2017 ; Pinto & Gupta , 2016 ) , during deployment it is often impractical and may require substantial engineering efforts . 1Webpage and implementation : https : //nicklashansen.github.io/PAD/ In this paper , we tackle an alternative problem setting in vision-based RL : adapting a pre-trained policy to an unknown environment without any reward . We do this by introducing self-supervision to obtain “ free ” training signal during deployment . Standard self-supervised learning employs auxiliary tasks designed to automatically create training labels using only the input data ( see Section 2 for details ) . Inspired by this , our policy is jointly trained with two objectives : a standard RL objective and , additionally , a self-supervised objective applied on an intermediate representation of the policy network . During training , both objectives are active , maximizing expected reward and simultaneously constraining the intermediate representation through self-supervision . During testing / deployment , only the self-supervised objective ( on the raw observational data ) remains active , forcing the intermediate representation to adapt to the new environment . We perform experiments both in simulation and with a real robot . In simulation , we evaluate on two sets of environments : DeepMind Control suite ( Tassa et al. , 2018 ) and the CRLMaze ViZDoom ( Lomonaco et al. , 2019 ; Wydmuch et al. , 2018 ) navigation task . We evaluate generalization by testing in new environments with visual changes unknown during training . Our method improves generalization in 19 out of 22 test environments across various tasks in DeepMind Control suite , and in all considered test environments on CRLMaze . Besides simulations , we also perform Sim2Real transfer on both reaching and pushing tasks with a Kinova Gen3 robot . After training in simulation , we successfully transfer and adapt policies to 6 different environments , including continuously changing disco lights , on a real robot operating solely from an uncalibrated camera . In both simulation and real experiments , our approach outperforms domain randomization in most environments . 2 RELATED WORK . Self-supervised learning is a powerful way to learn visual representations from unlabeled data ( Vincent et al. , 2008 ; Doersch et al. , 2015 ; Wang & Gupta , 2015 ; Zhang et al. , 2016 ; Pathak et al. , 2016 ; Noroozi & Favaro , 2016 ; Zhang et al. , 2017 ; Gidaris et al. , 2018 ) . Researchers have proposed to use auxiliary data prediction tasks , such as undoing rotation ( Gidaris et al. , 2018 ) , solving a jigsaw puzzle ( Noroozi & Favaro , 2016 ) , tracking ( Wang et al. , 2019 ) , etc . to provide supervision in lieu of labels . In RL , the idea of learning visual representations and action at the same time has been investigated ( Lange & Riedmiller , 2010 ; Jaderberg et al. , 2016 ; Pathak et al. , 2017 ; Ha & Schmidhuber , 2018 ; Yarats et al. , 2019 ; Srinivas et al. , 2020 ; Laskin et al. , 2020 ; Yan et al. , 2020 ) . For example , Srinivas et al . ( 2020 ) use self-supervised contrastive learning techniques ( Chen et al. , 2020 ; Hénaff et al. , 2019 ; Wu et al. , 2018 ; He et al. , 2020 ) to improve sample efficiency in RL by jointly training the self-supervised objective and RL objective . However , this has not been shown to generalize to unseen environments . Other works have applied self-supervision for better generalization across environments ( Pathak et al. , 2017 ; Ebert et al. , 2018 ; Sekar et al. , 2020 ) . For example , Pathak et al . ( 2017 ) use a self-supervised prediction task to provide dense rewards for exploration in novel environments . While results on environment exploration from scratch are encouraging , how to transfer a trained policy ( with extrinsic reward ) to a novel environment remains unclear . Hence , these methods are not directly applicable to the proposed problem in our paper . Generalization across different distributions is a central challenge in machine learning . In domain adaptation , target domain data is assumed to be accessible ( Geirhos et al. , 2018 ; Tzeng et al. , 2017 ; Ganin et al. , 2016 ; Gong et al. , 2012 ; Long et al. , 2016 ; Sun et al. , 2019 ; Julian et al. , 2020 ) . For example , Tzeng et al . ( 2017 ) use adversarial learning to align the feature representations in both the source and target domain during training . Similarly , the setting of domain generalization ( Ghifary et al. , 2015 ; Li et al. , 2018 ; Matsuura & Harada , 2019 ) assumes that all domains are sampled from the same meta distribution , but the same challenge remains and now becomes generalization across meta-distributions . Our work focuses instead on the setting of generalizing to truly unseen changes in the environment which can not be anticipated at training time . There have been several recent benchmarks in our setting for image recognition ( Hendrycks & Dietterich , 2018 ; Recht et al. , 2018 ; 2019 ; Shankar et al. , 2019 ) . For example , in Hendrycks & Dietterich ( 2018 ) , a classifier trained on regular images is tested on corrupted images , with corruption types unknown during training ; the method of Hendrycks et al . ( 2019 ) is proposed to improve robustness on this benchmark . Following similar spirit , in the context of RL , domain randomization ( Tobin et al. , 2017 ; Pinto et al. , 2017a ; Peng et al. , 2018 ; Ramos et al. , 2019 ; Yang et al. , 2019 ; James et al. , 2019 ) helps a policy trained in simulation to generalize to real robots . For example , Tobin et al . ( 2017 ) ; Sadeghi & Levine ( 2016 ) propose to render the simulation environment with random textures and train the policy on top . The learned policy is shown to generalize to real robot manipulation tasks . Instead of deploying a fixed policy , we train and adapt the policy to the new environment with observational data that is naturally revealed during deployment . Test-time adaptation for deep learning is starting to be used in computer vision ( Shocher et al. , 2017 ; 2018 ; Bau et al. , 2019 ; Mullapudi et al. , 2019 ; Sun et al. , 2020 ; Wortsman et al. , 2018 ) . For example , Shocher et al . ( 2018 ) shows that image super-resolution can be learned at test time ( from scratch ) simply by trying to upsample a downsampled version of the input image . Bau et al . ( 2019 ) show that adapting the prior of a generative adversarial network to the statistics of the test image improves photo manipulation tasks . Our work is closely related to the test-time training method of Sun et al . ( 2020 ) , which performs joint optimization of image recognition and self-supervised learning with rotation prediction ( Gidaris et al. , 2018 ) , then uses the self-supervised objective to adapt the representation of individual images during testing . Instead of image recognition , we perform test-time adaptation for RL with visual inputs in an online fashion . As the agent interacts with an environment , we keep obtaining new observational data in a stream for training the visual representations . 3 METHOD . In this section , we describe our proposed Policy Adaptation during Deployment ( PAD ) approach . It can be implemented on top of any policy network and standard RL algorithm ( both on-policy and off-policy ) that can be described by minimizing some RL objective J ( θ ) w.r.t . the collection of parameters θ using stochastic gradient descent . 3.1 NETWORK ARCHITECTURE . We design the network architecture to allow the policy and the self-supervised prediction to share features . For the collection of parameters θ of a given policy network π , we split it sequentially into θ = ( θe , θa ) , where θe collects the parameters of the feature extractor , and θa is the head that outputs a distribution over actions . We define networks πe with parameters θe and πa with parameters θa such that π ( s ; θ ) = πa ( πe ( s ) ) , where s represents an image observation . Intuitively , one can think of πe as a feature extractor , and πa as a controller based on these features . The goal of our method is to update πe at test-time using gradients from a self-supervised task , such that πe ( and consequently πθ ) can generalize . Let πs with parameters θs be the self-supervised prediction head and its collection of parameters , and the input to πs be the output of πe ( as illustrated in Figure 1 ) . In this work , the self-supervised task is inverse dynamics prediction for control , and rotation prediction for navigation .

Summary: This paper studies an important problem in vision-based RL: how to adapt a pre-trained policy to an unseen environment in a self-supervised manner. To do this, the authors introduce an auxiliary task branch that can be used to tune the intermediate representation of the policy network on the fly in a self-supervised manner(e.g. inverse dynamic prediction). The experiments in the DeepMind Control suite, CRLMaze, and robot manipulation tasks show the generalization and effectiveness of the proposed method in various vision-based RL problems.