Invariance Through Inference

We introduce a general approach , called Invariance through Inference , for improving the test-time performance of an agent in deployment environments with unknown perceptual variations . Instead of producing invariant visual features through interpolation , invariance through inference turns adaptation at deployment-time into an unsupervised learning problem . This is achieved in practice by deploying a straightforward algorithm that tries to match the distribution of latent features to the agent ’ s prior experience , without relying on paired data . Although simple , we show that this idea leads to surprising improvements on a variety of adaptation scenarios without access to deployment-time rewards , including changes in camera poses and lighting conditions . Results are presented on challenging distractor control suite , a robotics environment with image-based observations . 1 INTRODUCTION Let us consider the ability of an intelligent agent to generalize to unseen domains . To have such a discussion , we must first consider what generalization means . In much of the learning literature , we typically assume that an agent will accrue experiences of sufficient variation during training , and that these experiences will allow the agent to generalize to novel settings during deployment . Since the richness of the agent ’ s experience is of paramount importance to the quality of its generalization , there exist a broad family of methods that expand the support distribution of the agent ’ s training set . Domain transfer , domain randomization , and meta learning all fall into this category . Expanding the support of the training distribution is most often accomplished via artificial data augmentation . In pixel-based control tasks , for example , image observations are cropped , shifted , rotated , and discolored to make learned policies more robust to shifts in the input observation at test time ( Hansen & Wang , 2020 ; Yarats et al. , 2021 ) . In RL , this type of augmentation can cause value function estimation to become unstable ( Laskin et al. , 2020 ; Raileanu et al. , 2020 ; Kumar et al. , 2021 ) , and so care must be taken to avoid destabilizing training . While these approaches are powerful , they are often prone to failure in the case where the target domain is not known a priori . In this case , if the target domain falls out-of-distribution , the agent will lack the ability to self correct and will fail , even if great pains were taken to significantly widen the size of its training support . Indeed , as we show in Section 5 , even the most robust data augmentation schemes fail when the target observation space differs substantially from that of training . We refer to these approaches Invariance through Interpolation , which aims to improve an agent ’ s generalization by increasing the richness of its experiences , as in-distribution generalization . A more challenging problem—but one that is also more realistic—is generalizing out-ofdistribution . In this setting , we assume the specifics of the agent ’ s deployment are not known in advance . Taking this even further , we assume that access to instrumented reward supervision is unavailable in the deployment environment . Here , the burden placed on the agent hoping to generalize is significant . Since we do not have access to reward , fine-tuning via reinforcement learning is impossible . Prior experience will help , but a truly out-of-distribution environment of which the agent has no prior knowledge will make the transfer problem non-trivial . While difficult , this setting of out-of-distribution generalization is nevertheless pervasive in fields such as robotics , wherein a robot trained from pixel observations in the factory will inevitably encounter new settings with different lighting conditions , environmental objects , physical dynamics , and sensor modalities during deployment . To make progress on out-of-distribution generalization , we consider how we can best leverage the information we do have : the agent ’ s prior experiences during training . Specifically , if we assume the original task the agent was optimized for during training remains well-defined in the target domain , then it is also safe to assume that the agent does have some understanding of the underlying Markov decision process in the target domain . Without further assumptions or loss of generality , we may recast the out-of-distribution generalization problem into an unsupervised adaptation between two MDPs that share underlying dynamics and reward structure , but with distinct observations . This type of unsupervised adaptation in RL has not been explored widely yet to the best of our knowledge . One of the simplest approaches in this vein is to train a policy with a self-supervised auxiliary task such as inverse dynamics , and rely entirely on it during adaptation ( Hansen et al. , 2020 ) . Yet attractive for its simplicity , it often performs relatively poorly , as discussed in Section 5 . Another approach is to enforce Cycle-GAN ( Zhu et al. , 2017 ) style objective to learn cycleconsistent mappings between domains ( Zhang et al. , 2020 ) . Although it is shown to work well in multiple scenarios , it requires access to states , which is not generally available in pixel-based control tasks we consider . In this paper , we investigate ways to improve generalization under this challenging scenario . Rather than producing policy invariance by explicitly baking this property into training , we harness probabilistic inference to produce invariance by taking advantage of the latent structure the agent has already discovered in the environment . As we will see , this pays great dividends in producing robotic agents that can function effectively in deployment environments that exhibit different lighting , object texture , or camera poses . To distinguish our approach from prior works that attempt to produce generalization by baking invariances into the policy at training time , we refer to our method as “ Invarivance Through Inference . ” 2 RELATED WORK . Generalization in reinforcement learning is a longstanding problem . Recent work has shown that accurately quantifying generalization in this setting poses a challenge ( Lee et al. , 2019 ; Cobbe et al. , 2019 ; Zhang et al. , 2018 ) . This challenge is further amplified in visual RL from pixels , where measuring the difference between two problems is especially difficult . One promising class of techniques for generalization in image-based RL utilize image augmentation to greatly increase the size of the training distribution . These techniques have enjoyed success in vision research ( Chen et al. , 2020a ; He et al. , 2020 ) . However , it was only recently shown that such data augmentation leads to instability in value function estimation in RL algorithms ( Hansen & Wang , 2020 ; Hansen et al. , 2021 ) . This line of work showed how to stabilize data-augmented RL , which lead to a significant increase in the generalization capabilities of RL algorithms . The primary shortcoming of these methods is that it is largely impossible to effectively cover the entire potential test distribution by expanding the training distribution , and some test time adaptation is often needed . Several recent works ( Hansen et al. , 2020 ; Wang et al. , 2020 ) consider this setting of adaptation at test time , and are further discussed in Section 3 . Invariant Representation Learning is a broad class of techniques for learning representations that are robust across changes in the agent ’ s environment or reward . Often , the agent will use auxiliary information related to dynamics ( Hafner et al. , 2020 ) or reward prediction ( Jaderberg et al. , 2016 ) to learn a latent representation that is invariant across changing environmental conditions . This is closely related to the idea of self-supervised learning ( Chen et al. , 2020b ; Grill et al. , 2020 ) . When applied to visual space RL , it has been shown that self-supervised pre-training provides an efficient way to bootstrap useful invariant representations , even in the absence of a ground-truth reward ( Schwarzer et al. , 2020 ; Srinivas et al. , 2020 ) . Domain Shift considers the problem of test-time shift in the underlying input data distribution ( Ganin et al. , 2016 ) . One common class of techniques for addressing this problem is to learn representations that maximize domain confusion , thus requiring that the agent learn domain-invariant features . This idea has been applied successfully in imitation learning ( Stadie et al. , 2017 ) . More recent advances have combined Cycle-GANs with imitation learning to produce features that are cycle consistent across domains ( Smith et al. , 2019 ) . Another approach bottlenecks the information on the latent code of a VAE-based architecture to learn features that are more robust across domain shift in a variety of adversarial settings ( Peng et al. , 2018 ) . Meta Learning tries to solve generalization by aggregating learning across distributions of tasks ( Finn et al. , 2017 ; Duan et al. , 2016 ; Rakelly et al. , 2019 ) . Often , meta learning pipelines focus heavily on fast inference at test time . This is most often accomplished through using some form of bi-level optimization ( Rothfuss et al. , 2018 ; Houthooft et al. , 2018 ) . Unfortunately , meta learning methods generally require training over entire task distributions . Worse still , this results in the need for a distribution of reward functions for each task during both training and test time . Defining even a single reward function is often quite difficult , let alone many closely related reward functions . Recent work has tried to alleviate this burden with unsupervised training ( Hsu et al. , 2018 ) . 3 PROBLEM FORMULATION . Consider the standard infinite horizon Markov decision process ( MDP ) ( Puterman , 2014 ) parameterized via the tupleM = 〈S , O , A , R , P , γ〉 , where S , O , and A are the state , observation , and action spaces , respectively , P : S ×A 7→ S is the transition function , R : S ×A 7→ R is the scalar reward , and γ is the discount factor . Recall that the state space S usually represents ground-truth information about the agent ’ s environment , which we assume the agent can not directly access . Instead , the agent receives a stream of observations o ∈ O that convey information about the environment . In this paper , we assume the environment is fully-observable . In other words , the true state can be perfectly inferred from the corresponding observation , given the right feature extractor . The goal of reinforcement learning is to solve for the policy π : O × A 7→ [ 0 , 1 ] that maximizes the expected discounted return J = E [ ∑ ∞ γ tR ( st , at ) ] over infinite horizon , represented as the Qπ ( o , a ) function . In this paper , we assume that the policy π consists of an encoder F : O 7→ Z , where Z is a compact latent space , and a policy πz : Z ×A 7→ [ 0 , 1 ] . We consider a setting wherein there are two distinct domains : a source domain Msrc and a target domainMtgt . Crucially , when the agent is placed in the previously unseen target domain , we assume that it can only access its observations and the actions it took , but not reward . Often , we are most interested in the case where shift betweenMsrc andMtgt is induced by differences in the observation spaces Osrc and Otgt . That is , the state and reward dynamics between the source and target domains are quite similar , but the observation spaces between the two are significantly different . Practically speaking , this can happen quite easily in the presence of distractors , corrupted or malformed inputs , or shifts in sensor readings at test time . Consider a robot that is trained from image observations in a clean environment , and then expected to perform the same task in the wild , where changing lighting conditions and environmental conditions might lead to significantly different image observations . As a result of this distributional shift in the observations on which the policy is conditioned , deploying an agent trained in the source domain directly in the target domain ( i.e. , zero-shot transfer ) results in poor performance . Given an agent pretrained on a source domain , our goal is to adapt it to the target domain . As we do not have knowledge about the correspondence of observations in the two domains , this setting essentially requires us to make use of unpaired trajectory samples from each domain to drive adaptation . The proposed method , Invariance through Inference , aims to only adapt the encoder F in a manner that minimizes distributional shift in Z between the domains and , in turn , enables the policy πz to transfer to the unseen target domain , by utilizing only unpaired transitions . We achieve this by jointly optimizing two objectives : distribution matching and dynamics consistency .

Summary: This paper proposes to learn the invariances relating to MDPs with different observations (e.g. invariances in the underlying dynamics of the systems) using a generative model of the observations in an unsupervised manner. The inference in this generative model is accomplished using a GAN framework with auxiliary losses for the inverse dynamics of the system.