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:  **Introduction:**  In this paper the authors aim to tackle the problem of learning a model that generalizes well on a test distribution that samples outside of the training data distribution.  They aim to do this not via data augmentation but by using an inference model over the latent space of the input.  Such an approach needs to go beyond in-distribution generalization. If the target domain is not know a priori things like domain transfer, domain randomization and meta-learning techniques may fail.  Out-of-distribution generalization is very important but more challenging:  * Prior experience will be critical  * Lack a reward function over the test distribution  * In robotics for instance, this type of learning is critical  The authors propose to: *"recast out-of-distribution generalization problem into an unsupervised policy adaptation between two MDPs that share a similar latent dynamics and reward structure, but with distinct observations"*  Statement of purpose for the paper: *"Harness probabilistic inference to produce invariance by taking advantage of latent structure"*  This work attempts to distinguish itself from other approaches that *"bake policy invariance"* in at training time.   **Problem Formulation:**  Given an MDP: M = <S, O, A, R, P, γ>, S contains ground truth info not always accessible to agent.  In particular two distinct MDPs: A train MDP M_train and a test MDP M_test.  The agent has no access to the reward function at test time.  The problem is defined: *"Often, we are most interested in the case where shift between M_train and M_test is induced by differences in the observation spaces O_train and O_test. That is, the state and reward structures between the train and test MDPs are quite similar, but the observation spaces between the two MDPs are significantly different."* It is noted that recent work in RL has tried to deal with this via data augmentation but this can lead to instability.  **Invariance Through Inference and Algorithm**  The goal here is to learn an encoder mapping semantically similar states in the train and test distribution to latent vectors.  This is to be accomplished given there is no access to rewards at test time and no access to paired trajectories (ie. matched test/train trajectories). The approach is defined by two objectives:  1. a distribution matching objective encourages the latent distribution induced by both MDPs to match.   2. a GAN-style loss to help ensure that the latent code is not using distracting information that would help distinguish the train and test MDPs.  The authors state that *"This adaptation objective is effectively maximizing the mutual information between the prototypical representation kept from during training, and the adapted representations of the new observations"* and further claim that *"our adaptation objective produce invariance by making statistical inference in the latent space."*  As regards the algorithm they also state that: *"The encoder G takes an observation from target environment, and learns to trick the discriminator, while the discriminator predicts whether the input comes from the encoder or source buffer."*  **Experiments**  The authors us the Deepmind Control Suite (DMC) as a training environment and the Distracting Control Suite (DistractingCS) for a test environment.  *"DistractingCS adds three types of distractions to the DeepMind Control Suite through deviations in background, camera pose, and color."*. Results are presented showing that baseline methods augmented with ITI show much slower performance degradation as the distraction intensity is increased.  The authors also compare ITI to Policy Adaptation during Deployment (PAD).  They compare to ITI on DMC for fixed distraction intensity. ITI proves to yield stronger results in this case and it is posited that PAD does not encourage latent train and test distributions to match as ITI does: *We suspect this instability is caused by the large deviations in the latent variable distribution as a result of changes in the target environment. In particular, we posit that the signal from PAD’s inverse dynamics head does not encourage the latent train and test distributions to match, which is a feature specifically baked into Invariance through Inference.*