Original: I NVARIANCE 
Revised: I NVARIANCE T

Original: T HROUGH
Revised: HROUGH

Original: I NFERENCE
Revised: I NFERENCE

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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.
Revised: 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.

Original: Instead of producing invariant visual features through memorization, invariance through inference turns adaptation at deployment-time into an unsupervised learning problem.
Revised: Instead of producing invariant visual features through interpolation, invariance through inference turns adaptation at deployment-time into an unsupervised learning problem.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Results are presented on challenging distractor control suite, a robotics environment with image-based observations.
Revised: Results are presented on challenging distractor control suite, a robotics environment with image-based observations.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Train
Revised: Train

Original: Test
Revised: Test

Original: Let us consider the ability of an intelligent agent to generalize to unseen domains.
Revised: Let us consider the ability of an intelligent agent to generalize to unseen domains.

Original: To have such a discussion, we must ﬁrst consider what generalization means.
Revised: To have such a discussion, we must ﬁrst consider what generalization means.

Original: In much of the learning literature, we typically assume that an agent will accrue experiences of sufﬁcient variation during training, and that these experiences will allow the agent to generalize to novel settings during deployment.
Revised: In much of the learning literature, we typically assume that an agent will accrue experiences of sufﬁcient variation during training, and that these experiences will allow the agent to generalize to novel settings during deployment.

Original: 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.
Revised: 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.

Original: Domain transfer, domain randomization, and meta learning all fall into this category.
Revised: Domain transfer, domain randomization, and meta learning all fall into this category.

Revised: Expanding the support of the training distribution is most often accomplished via artiﬁcial data augmentation.

Revised: 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).

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

Original: While these approaches are powerful, they are often prone to failure in the case where the target domain is not known a priori .
Revised: While these approaches are powerful, they are often prone to failure in the case where the target domain is not known a priori .

Original: 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 signiﬁcantly widen the size of its training support.
Revised: 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 signiﬁcantly widen the size of its training support.

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

Original: We refer to this approach, which aims to improve an agent’s generalization by increasing the richness of its experiences, as in-distribution generalization.
Revised: 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.

Original: A more challenging problem—but  also more realistic— is generalizing out-of-distribution.
Revised: A more challenging problem—but one that is also more realistic—is generalizing out-ofdistribution.

Original: In this setting, we assume the speciﬁcs of the agent’s deployment are not known ex-ante.
Revised: In this setting, we assume the speciﬁcs of the agent’s deployment are not known in advance.

Original: Taking this even further, we assume that access to instrumented reward supervision is unavailable in the deployment environment.
Revised: Taking this even further, we assume that access to instrumented reward supervision is unavailable in the deployment environment.

Original: In this setting, the burden placed on the agent hoping to generalize is signiﬁcant.
Revised: Here, the burden placed on the agent hoping to generalize is signiﬁcant.

Original: Since we lack a reward function, ﬁne-tuning via reinforcement learning is impossible.
Revised: Since we do not have access to reward, ﬁne-tuning via reinforcement learning is impossible.

Original: 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.
Revised: 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.

Original: While difﬁcult, this setting of out-of-distribution generalization is nevertheless evergreen in ﬁelds 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.
Revised: While difﬁcult, this setting of out-of-distribution generalization is nevertheless pervasive in ﬁelds 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.

Original: To make progress in this setting of out-of-distribution generalization, we consider how we can best leverage the information we do have: the agent’s prior experiences during training.
Revised: 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.

Original: Speciﬁcally, if we assume the original task the agent was optimized for during training remains well-deﬁned 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 test domain.
Revised: Speciﬁcally, if we assume the original task the agent was optimized for during training remains well-deﬁned 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.

Original: Without further assumptions or loss of generality, we may recast the out-of-distribution generalization problem into an unsupervisedpolicy adaptation between two MDPs that share a similar latent dynamics and reward structure, but with distinct observations.
Revised: 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.

Revised: This type of unsupervised adaptation in RL has not been explored widely yet to the best of our knowledge.

Revised: 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., 2020b).

Original:  In this paper, we investigate ways to improve generalization under this challenging scenario.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We also demonstrate the ability of our method to generalize across environments with different forward dynamics, so long as the distinct dynamics maintain some semblance of similarity.

Original: 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 .”
Revised: 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 .”

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Generalization in reinforcement learning is a longstanding problem.
Revised: Generalization in reinforcement learning is a longstanding problem.

Original: 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).
Revised: 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).

Original: This challenge is further ampliﬁed in visual RL from pixels, where measuring the difference between two problems is especially difﬁcult.
Revised: This challenge is further ampliﬁed in visual RL from pixels, where measuring the difference between two problems is especially difﬁcult.

Original: One promising class of techniques for generalization in image-based RL utilize image augmentation to greatly increase the size of the training distribution.
Revised: One promising class of techniques for generalization in image-based RL utilize image augmentation to greatly increase the size of the training distribution.

Original: These techniques have enjoyed success in vision research (Chen et al., 2020a;
Revised: These techniques have enjoyed success in vision research (Chen et al., 2020a;

Original: He et al., 2020).
Revised: He et al., 2020).

Original: 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., 2021a).
Revised: 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., 2021a).

Original: This line of work showed how to stabilize data-augmented RL, which lead to a signiﬁcant increase in the generalization capabilities of RL algorithms.
Revised: This line of work showed how to stabilize data-augmented RL, which lead to a signiﬁcant increase in the generalization capabilities of RL algorithms.

Original: 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.
Revised: 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.

Original: Several recent works (Hansen et al., 2020b; Wang et al., 2020) consider this setting of adaptation at test time, and are further discussed in Section 3.
Revised: Several recent works (Hansen et al., 2020b; Wang et al., 2020) consider this setting of adaptation at test time, and are further discussed in Section 3.

Original: Invariant Representation Learning is a broad class of techniques for learning representations that are robust across changes in the agent’s environment or reward.
Revised: Invariant Representation Learning is a broad class of techniques for learning representations that are robust across changes in the agent’s environment or reward.

Original: 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.
Revised: 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.

Original: This is closely related to the idea of self-supervised learning (Chen et al., 2020b; Grill et al., 2020).
Revised: This is closely related to the idea of self-supervised learning (Chen et al., 2020b; Grill et al., 2020).

Original: When applied to visual space RL, it has been shown that self-supervised pre-training provides an efﬁcient way to bootstrap useful invariant representations, even in the absence of a ground-truth reward (Schwarzer et al., 2020; Srinivas et al., 2020).
Revised: When applied to visual space RL, it has been shown that self-supervised pre-training provides an efﬁcient way to bootstrap useful invariant representations, even in the absence of a ground-truth reward (Schwarzer et al., 2020; Srinivas et al., 2020).

Original: Domain Shift considers the problem of test-time shift in the underlying input data distribution (Ganin et al., 2016).
Revised: Domain Shift considers the problem of test-time shift in the underlying input data distribution (Ganin et al., 2016).

Original: 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.
Revised: 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.

Original: This idea has been applied successfully in imitation learning (Stadie et al., 2017).
Revised: This idea has been applied successfully in imitation learning (Stadie et al., 2017).

Original: More recent advances have combined Cycle-GANs with imitation learning to produce features that are cycle consistent across domains (Smith et al., 2019).
Revised: More recent advances have combined Cycle-GANs with imitation learning to produce features that are cycle consistent across domains (Smith et al., 2019).

Original: 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).
Revised: 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).

Original: Meta Learning tries to solve generalization by aggregating learning across distributions of tasks (Finn et al., 2017;
Revised: Meta Learning tries to solve generalization by aggregating learning across distributions of tasks (Finn et al., 2017;

Original: Duan et al., 2016; Rakelly et al., 2019).
Revised: Duan et al., 2016; Rakelly et al., 2019).

Original: Often, meta learning pipelines focus heavily on fast inference at test time.
Revised: Often, meta learning pipelines focus heavily on fast inference at test time.

Original: This is most often accomplished through using some form of bi-level optimization (Rothfuss et al., 2018; Houthooft et al., 2018).
Revised: This is most often accomplished through using some form of bi-level optimization (Rothfuss et al., 2018; Houthooft et al., 2018).

Original: Unfortunately, meta learning methods generally require training over entire task distributions.
Revised: Unfortunately, meta learning methods generally require training over entire task distributions.

Original: Worse still, this results in the need for a distribution of reward functions for each task during both training and test time.
Revised: Worse still, this results in the need for a distribution of reward functions for each task during both training and test time.

Original: Deﬁning even a single reward function is often quite difﬁcult, let alone many closely related reward functions.
Revised: Deﬁning even a single reward function is often quite difﬁcult, let alone many closely related reward functions.

Original: Recent work has tried to alleviate this burden with unsupervised training (Hsu et al., 2018).
Revised: Recent work has tried to alleviate this burden with unsupervised training (Hsu et al., 2018).

Original:  P ROBLEM F ORMULATION
Revised:  P ROBLEM F ORMULATION

Original: Consider the standard inﬁnite horizon Markov decision process (MDP) (Puterman, 2014)) parameterized via the tuple M = (cid:104) S, O, A, R, P, γ (cid:105) , where S , O , and A are the state, observation, and action spaces, respectively, P : S × A (cid:55)→ S is the transition function, R :
Revised: Consider the standard inﬁnite horizon Markov decision process (MDP) (Puterman, 2014) parameterized via the tuple M = (cid:104) S, O, A, R, P, γ (cid:105) , where S , O , and A are the state, observation, and action spaces, respectively, P : S × A (cid:55)→ S is the transition function, R :

Original: S × A (cid:55)→ R is the scalar reward, and γ is the discount factor.
Revised: S × A (cid:55)→ R is the scalar reward, and γ is the discount factor.

Original: Recall that the state space S usually represents ground truth information about the agents environment which we assume the agent can not  access.
Revised: 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.

Original: The observations space O is the information stream the agent receives  about the environment.
Revised: Instead, the agent receives a stream of observations o ∈ O that convey information about the environment.

Revised: In this paper, we assume the environment is fully-observable.

Revised: In other words, the true state can be perfectly inferred from the corresponding observation, given the right feature extractor.

Original: The goal of reinforcement learning is to solve for the policy π : O × A (cid:55)→ [0 , 1] that maximizes the expected discounted return J = E  (cid:80) ∞
Revised: The goal of reinforcement learning is to solve for the policy π : O × A (cid:55)→ [0 , 1] that maximizes the expected discounted return J = E [ (cid:80) ∞

Original: γ t R ( o t , a t ) over inﬁnite horizon, represented as the Q π ( o, a ) function.
Revised: γ t R ( s t , a t )] over inﬁnite horizon, represented as the Q π ( o, a ) function.

Original: In this paper, we consider a setting wherein there are two distinct MDPs: A train MDP M train and a test MDP M test .
Revised: In this paper, we assume that the policy π consists of an encoder F : O (cid:55)→ Z , where Z is a compact latent space, and a policy π z :

Original: During training, the agent has access to M train and can collect trajectory data ρ train = ( o i , a i , r i +1 , . . . ) , where o i ∈ O , a i ∈ A  , and r i +1 ∈ R is the corresponding reward. During testing, the agent is placed into the previously unseen M test environment.
Revised: Z × A (cid:55)→ [0 , 1] .

Original: Crucially, during the testing phase, we will assume that the agent can only observe ( o test , a test ) ∼ M test i.e., that the agent has no access to the reward function at test time. The goal of the agent is to leverage the data it collected from M train , as well as the unsupervised trajectory data ( o 1 , a 1 , . . ., o t , a t ) to perform fast inference in the testing environment: [Equation]
Revised: We consider a setting wherein there are two distinct domains: a source domain M src and a target domain M tgt .

Original: where o t +1 ∼ π ( o t , a t ; ρ train ) makes clear the dependence of the policy π on data from the training environment ρ train . Note that this objective is for evaluation purposes only, because we do not have access to r test during learning. In the objective, we have used the abuse of notation o t +1 ∼ π ( o t , a t ; ρ train ) instead of the correct o t +1 ∼ P ( o t , a t ) , a t ∼ π ( o t ; ρ train ) , since the next observation is sampled by following along the policy and space was limited.
Revised: 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.

Original: 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 .
Revised: Often, we are most interested in the case where shift between M src and M tgt is induced by differences in the observation spaces O src and O tgt .

Original: That is, the state and reward structures between the train and test MDPs are quite similar, but the observation spaces between the twoMDPs are signiﬁcantly different.
Revised: That is, the state and reward dynamics between the source and target domains are quite similar, but the observation spaces between the two are signiﬁcantly different.

Original: Practically speaking, this can happen quite easily in the presence of distractors, corrupted or malformed inputs, or shifts in sensor readings at test time.
Revised: Practically speaking, this can happen quite easily in the presence of distractors, corrupted or malformed inputs, or shifts in sensor readings at test time.

Original: 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 signiﬁcantly different image observations.
Revised: 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 signiﬁcantly different image observations.

Original: Recent work has tried to handle this  shift between training and testing environments by greatly expanding the support of the training distribution during training (Hansen & Wang, 2020). This is most often accomplished via artiﬁcial data augmentation. Image observations are cropped, shifted, rotated, and discolored to make learned policies more robust to shifts in the input observation at test time. 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.
Revised: 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.

Original: While data augmentation techniques have greatly improved the robustness of RL algorithms to observation shift at test time, potential issues linger. In Figure 2, we see a cartoon illustrating the primary worry. No matter how much the support over the input observation space is expanded, the size of the target testing observation space is very small in comparison to the overall size of all observation spaces.
Revised: Given an agent pre-trained on source domain, our goal is to adapt it to the target domain.

Original: Thus, there is still a fairly reasonable chance that even extreme data augmentation measures will fail to adequately cover the test MDP’s observation space.
Revised: 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.

Original: Indeed, the results in Section 5 show this to often be the case—even the most robust data augmentation schemes fail when the test MDP’s observation space differs substantially from that of training. In the next section, we will develop a technique, Invariance Through Inference, that  enablesquick test-time adaptation onto the  target observation space by utilizing only unsupervised transition pairs.
Revised: 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 .

Original:  I NVARIANCE THROUGH I NFERENCE 4. A LOSS FUNCTION FOR I NVARIANCE THROUGH I NFERENCE
Revised:  I NVARIANCE THROUGH I NFERENCE 4. D ISTRIBUTION M ATCHING

Original: Following the previous section, assume we are in a setting with two MDPs, M train and M test .
Revised: We ﬁrst deploy a random exploration policy in both the source and target domains to collect samples of observation trajectories

Original: We assume these MDPs share similar states, dynamics, and reward structures.
Revised: [Equation]

Original: However, the observation spaces incur a signiﬁcant distributional shift.
Revised: where P ¯ π ( M ) denotes the distribution of transitions produced with a random policy in domain M .

Original: Nevertheless, since the underlying state and dynamics that are relevant to the task are similar, there exists some encoding z = G θ ( o, a ) s.t.
Revised: Given an observation o t , an encoder F with parameters θ produces a corresponding latent representation z t

Original: [Equation]
Revised: [Equation]

Revised: We seek to match the distribution over z tgt t with that over z src t by adapting the weights of the encoder

Original: In other words, we want to learn a latent encoder that maps semantically similar observations to close latent vectors, learning toignore irrelevant distractions in the observation space.
Revised: F , without access to test-time rewards. We employ adversarial training to achieve this. Speciﬁcally, our architecture (Fig. 2) includes a discriminator D that tries to distinguish between embeddings from the source domain z src t and those from the target domain z tgt t . At the same time, we adapt the parameters of the encoder to produce a latent embedding that is indistinguishable to the discriminator.

Original: In many real world deployment scenarios, test-time rewards R test are unavailable.
Revised: This results in matching latent distribution over z src t and z tgt t . Following the Wasserstein GAN (Arjovsky et al., 2017) formulation, we express our distribution matching loss as [Equation] where ¯ F indicates that the weights of the network are frozen.

Original: Thus, we have hope to train this  encoder without access to test-time rewards.
Revised: The encoder tries to minimize this objective while the discriminator acts as an adversary and seeks to maximize it, resulting in a GANlike minimax game. 4. D YNAMICS CONSISTENCY Theoretically, the use of adversarial training can result in an encoder that maps source and target domain observations to latent embeddings that have identical distributions. However, it is possible for a sufﬁciently expressive encoder to map the same target observations to a random perturbation of observations in the source domain while still matching the distributions (Zhu et al., 2017).

Original: Given this constraint, we note that the shared underlying dynamics are the most feature-rich and informative source of mutual information between the two environments.
Revised: To mitigate this, we take advantage of the shared underlying dynamics as a feature-rich and informative source of mutual information between the two domains.

Revised: Speciﬁcally, let C inv ( z t , z t +1 ) be the inverse dynamics network that predicts the action a t associated with the transition from z t to z t

Revised: +1 .

Original: We would therefore like our encoded latents to minimize an inverse dynamics model loss 
Revised: We then deﬁne the dynamics consistency loss as the error in the inverse dynamics predictions

Original: [Equation]
Revised: [Equation]

Original: However, simply training this model with data fromboth MDPs will not ensure that observations with similar ground-truth states are actually aligned. Instead, the encoder might simply learn to handle each MDP independently.
Revised: Most importantly, we pre-train the inverse dynamics network on transitions sampled from the source domain B src ∼ P

Original: Critically, we do not assume access to paired trajectories, which can be difﬁcult to achieve in practice.
Revised: [Equation]

Original: If we had access to such paired data, we could simply train an encoder that minimizes the distance to both paired observations. At minimum, we would like to ensure that the latent distributions over observations embeddings are similar for both MDPs, since that would hopefully force the encoder to ﬁlter out superﬂuous distractions that separate the two MDPs. Towards this goal, we deﬁne two objectives. First, a distribution matching objective encourages the latent distribution induced by both MDPs to match.
Revised: Eqn

Original: [Equation]
Revised: [Equation]

Original: Second, we employ 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
Revised: 1a) with observations encoded with F

Original: [Equation]
Revised: [Equation]

Original: where  ¯ G denotes the encoder whose weights are frozen after training.
Revised: where P ¯ π ( M ; F ) is the distribution over latent trajectories.

Original: The ﬁnal objective can now be written as
Revised: Under the assumption that the underlying dynamics are shared by the domains, we then utilize these learned inverse dynamics model to further encourage the latent embeddings for the target domain to align with those of the source domain. Speciﬁcally, during adaptation, we freeze the weights of the inverse dynamics network, and minimize the following dynamics consistency loss with respect to encoder F using transitions sampled from the target domain

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: We adapt our encoder by minimizing a loss that combines both the adversarial loss J adv (Eqn. 3) and the dynamics consistency loss J dyn (Eqn. 6).

Original: In summary, in contrast to those methods that produce invariance through memorization, our adaptation objective produce invariance by making statistical inference in the latent space. For this reason we  the derived method
Revised: Speciﬁcally, we solve for the parameters of the encoder through the following minimax objective

Original: Invariance
Revised: [Equation]

Original: 4. A LGORITHM O VERVIEW

Original: Constructing a pipeline to minimize Equation 5 is fairly straightforward.

Original: The architecture is fully represented in Figure 3. In practice, the algorithm proceeds as follows. Start by collecting trajectories in the clean training environment andadding the corresponding encoded observations into a replay buffer. These initial samples are collected under an exploratory policy.
Revised: 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.

Original: In practice, we ﬁnd random exploration is sufﬁcient for many robotics environments. Now that the buffer is nonempty, we can draw samples from it to  pretrainthe inverse dynamics network. With this in hand, we can start adapting the encoder towards the test environment.
Revised: Algorithm 1 summarizes the whole procedures to sample trajectory, pretrain dynamics network and adapt encoder via the objective.

Original: Since we want to ensure the observation embeddings for the train and test environment follow a similar distribution, we use the optimization framework from Generative Adversarial Networks (GANs) to constrain the encoder appropriately. Our encoder, which is analogous to a generator, is encouraged to  produce similar embeddings on both train and test environment observations.
Revised: In summary, in contrast to those methods that produce invariance through interpolation, our adaptation objective produce invariance by making statistical inference in the latent space.

Original: The discriminator learns to classify whether a given input comes from source environment buffer or target environment. Finally, we need a structure that can help guide the encoder toproduce a meaningful embedding for each observation. To wit, we add an inverse dynamics head onto our architecture as a type of auxiliary task, under the assumption that train and test time environments share a high degree of overlap in their environmental dynamics. This head is pre-trained with transitions sampled from the source buffer, and has its weights ﬁxed during adaptation.
Revised: For this reason we refer to the derived method as Invariance through Inference.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We want to understand the impact of test-time inference on an agent’s ability to perform out-ofdistribution generalization.
Revised: We want to understand the impact of test-time inference on an agent’s ability to perform out-ofdistribution generalization.

Original: To wit, this section will compare Invariance through Inference with two state-of-the-art methods for data-augmented reinforcement learning: SVEA (Hansen et al., 2021b) and DrQ-v2 (Yarats et al., 2021).
Revised: To wit, this section will compare Invariance through Inference with two state-of-the-art methods for data-augmented reinforcement learning: SVEA (Hansen et al., 2021b) and DrQ-v2 (Yarats et al., 2021).

Original: Recall from the introduction that these methods vie for increased generalization capabilities by expanding the support of the training distribution.
Revised: Recall from the introduction that these methods vie for increased generalization capabilities by expanding the support of the training distribution.

Original: We expect this form of generalization to be less performant than inference at test time.
Revised: We expect this form of generalization to be less performant than inference at test time.

Original: In our experiments, we will also compare against Policy Adaptation during Deployment (PAD)
Revised: In our experiments, we will also compare against Policy Adaptation during Deployment (PAD)

Original: (Hansen et al., 2020a), a baseline that, like our method, adapts the policy without access to the reward at test time.
Revised: (Hansen et al., 2020a), a baseline that, like our method, adapts the policy without access to the reward at test time.

Original: To further probe the generalization abilities of test-time inference, we conduct an experiment wherein we vary the intensity of environmental distractions.
Revised: To further probe the generalization abilities of test-time inference, we conduct an experiment wherein we vary the intensity of environmental distractions.

Original: The upshot here is that test-time inference signiﬁcantly increases the robustness of our policy to distractions during deployment.
Revised: The upshot here is that test-time inference signiﬁcantly increases the robustness of our policy to distractions during deployment.

Original: We will conclude with some general discussion and remarks regarding the design tradeoffs involved in test-time inference.
Revised: We will conclude with some general discussion and remarks regarding the design tradeoffs involved in test-time inference.

Original: Setup We conduct experiments using ﬁve domains from the DeepMind Control Suite (DMC) (Tassa et al., 2018).
Revised: Setup We conduct experiments using ﬁve domains from the DeepMind Control Suite (DMC) (Tassa et al., 2018).

Original: We use DMC as a clean training environment the Distracting Control Suite (DistractingCS) (Stone et al., 2021) for the distracted test environment.
Revised: We use DMC as a clean training environment the Distracting Control Suite (DistractingCS) (Stone et al., 2021) for the distracted test environment.

Original: DistractingCS adds three types of distractions to the DeepMind Control Suite through deviations in background, camera pose, and color.
Revised: DistractingCS adds three types of distractions to the DeepMind Control Suite through deviations in background, camera pose, and color.

Revised: Algorithm 1 Invariance through

Revised: Require: Pre-trained encoder F , discriminator D , buffers B ∗ src and B tgt , inverse dynamics network parameterized with ψ inv for i=1:

Revised: N do (cid:46) Pre-ﬁll buffers

Revised: Sample z t , a t , z t +1 ∼ P ¯ π ( M src ; F )

Revised: Sample o t , a t , o t +1 ∼ P ¯ π ( M tgt ) B ∗ src ← B ∗ src ∪ ( z t , a t , z t +1 )

Original: Modiﬁcations to DistractingCS
Revised: B tgt ← B tgt ∪ ( o t , a t , o t +1 ) for i=1:T dynamics do (cid:46) Pre-train dynamics network Sample z t , a t , z t +1 ∼ B ∗ src g inv ← ∇ ψ inv L dyn ( z t , z t +1 , a t ; ψ inv ) ψ inv ← Optimizer ( ψ inv , g inv ) for i=1:T adapt do (cid:46) Adaptation main loop Sample z src t , a src t , z src t +1 ∼ B ∗ src Sample o tgt t , a tgt t , o tgt t +1 ∼ B tgt g D ← ∇ θ D (cid:2) D ( z src t ) + (1 − D ( F ( o tgt t ))) (cid:3) (cid:46) Distribution matching (adversarial) loss g F 1 ← ∇ θ F (cid:2) D ( z src t )

Revised: + (1 − D ( F ( o tgt t ))) (cid:3)

Revised: g F 2 ← ∇ θ F L dyn ( F ( o tgt t ) , F ( o tgt t +1 ) , a t ; ψ ∗ inv ) (cid:46) Dynamics consistency loss θ D ← Optimizer ( θ D , g D ) θ F ← Optimizer ( θ F , g F 1 , g F 2 )

Revised: Modiﬁcations to DistractingCS

Original: The default conﬁguration of DistractingCS changes distractions at the start of every episode (e.g., different background images are used at every episode).
Revised: The default conﬁguration of DistractingCS changes distractions at the start of every episode (e.g., different background images are used at every episode).

Original: However, we are interested in measuring an agent’s ability to perform inference across several episodes on the same target environment.
Revised: However, we are interested in measuring an agent’s ability to perform inference across several episodes on the same target environment.

Original: Thus, we modify DistractingCS to sample a distraction once in the beginning of learning, and then use the same distraction across all learning epochs.
Revised: Thus, we modify DistractingCS to sample a distraction once in the beginning of learning, and then use the same distraction across all learning epochs.

Original: This also ensures consistent evaluation across algorithms.
Revised: This also ensures consistent evaluation across algorithms.

Original: In accordance with this change, we also modify the intensity benchmark from DistractingCS.
Revised: In accordance with this change, we also modify the intensity benchmark from DistractingCS.

Original: In our experiments, intensity measures the deviation between an environment distraction and the train environment’s default value.
Revised: In our experiments, intensity measures the deviation between an environment distraction and the train environment’s default value.

Original: For example, intensity may measure how far the distracting color is from the default.
Revised: For example, intensity may measure how far the distracting color is from the default.

Original: Finally, we modify the environments to only apply a single distraction during testing (rather than all three) in order to better understand the impact of each type of distraction on overall performance.
Revised: Finally, we modify the environments to only apply a single distraction during testing (rather than all three) in order to better understand the impact of each type of distraction on overall performance.

Original: 5. I NVARIANCE THROUGH I NFERENCE O N DMC ONTROL
Revised: 5. I NVARIANCE THROUGH I NFERENCE O N DMC ONTROL

Original: This section studies the impact of test-time adaptation on DMControl performance.
Revised: This section studies the impact of test-time adaptation on DMControl performance.

Original: We begin by pretrained SAC, SVEA, and DrQ-v2 in a non-distracting training-time environment.
Revised: We begin by pretrained SAC, SVEA, and DrQ-v2 in a non-distracting training-time environment.

Original: After training, we evaluate the learned policies on distracted test environments of varying intensities.
Revised: After training, we evaluate the learned policies on distracted test environments of varying intensities.

Original: This evaluation is zero-shot, there is no additional training on the test environment.
Revised: This evaluation is zero-shot, there is no additional training on the test environment.

Original: Figure 5 visualizes the performance of the different methods as a function of the intensity of the distractions.
Revised: Figure 3 visualizes the performance of the different methods as a function of the intensity of the distractions.

Original: Since the baseline methods are trained with image augmentation, they do exhibit some robustness to distraction.
Revised: Since the baseline methods are trained with image augmentation, they do exhibit some robustness to distraction.

Original: However, we see this robustness rapidly diminishes as the distraction intensity increases.
Revised: However, we see this robustness rapidly diminishes as the distraction intensity increases.

Original: In particular, large changes to camera pose or the image background proved challenging for standard augmentation procedures.
Revised: In particular, large changes to camera pose or the image background proved challenging for standard augmentation procedures.

Original: Comparatively, Invariance through Inference saw much smoother and slower degradation of performance.
Revised: Comparatively, Invariance through Inference saw much smoother and slower degradation of performance.

Original: This supports our hypothesis that adaptation powered by unsupervised learning can signiﬁcantly widen the generalization abilities of learning algorithms.
Revised: This supports our hypothesis that adaptation powered by unsupervised learning can signiﬁcantly widen the generalization abilities of learning algorithms.

Original: 5. C OMPARISONS WITH PAD
Revised: 5. C OMPARISONS WITH PAD

Original: Similar to our approach, PAD pre-trains the agent in a clean environment, and then adapts the agent via unsupervised transfer, without assuming access to the target environment’s reward function [Figure] Hansen et al. (2020a).
Revised: Similar to our approach, PAD pre-trains the agent in a clean environment, and then adapts the agent via unsupervised transfer, without assuming access to the target environment’s reward function (Hansen et al., 2020a).

Original: To evaluate the robustness of PAD to distractions, we consider DistractingCS with a ﬁxed distraction intensity of 0 . 6 .
Revised: To evaluate the robustness of PAD to distractions, we consider DistractingCS with a ﬁxed distraction intensity of 0 . 6 .

Original: Table 1 compares the performance as the difference between the reward before and after adaptation along with the reward in the clean environment.
Revised: Table 1 compares the performance as the difference between the reward before and after adaptation along with the reward in the clean environment.

Revised: [Figure]

Original: Across all environments, we see that PAD struggles to adapt to distractions at test time.
Revised: Across all environments [Table] we see that PAD struggles to adapt to distractions at test time.

Original: We suspect this instability is caused by the large deviations in the latent variable distribution as a result of changes in the target environment.
Revised: We suspect this instability is caused by the large deviations in the latent variable distribution as a result of changes in the target environment.

Original: 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 speciﬁcally baked into Invariance through Inference.[Table] [Table] [Table]
Revised: 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 speciﬁcally baked into Invariance through Inference.

Original: [Table]

Original: [Table]
Revised: [Table]

Original: Background 377 .

Original: 53 ± 324 .528 . 14 ± 374 .396 . 90 ± 271 .736 . 34 ± 354 .
Revised: 5. F URTHER D ISCUSSION

Original: Color 437 .

Original: 48 ± 344 .631 .

Original: 84 ± 354 .452 .

Original: 31 ± 337 .

Original: [Table]
Revised: [Table]

Original: Unsupervised Pipeline Ablations To verify that the adversarial training scheme actually contributes to the performance gain during adaptation, we run Invariance through Inference while setting either the adversarial training losses to 0 .
Revised: Pipeline Ablations To verify that the adversarial training scheme actually contributes to the performance gain during adaptation, we run Invariance through Inference while setting either the adversarial training losses to 0 .

Original: This corresponds to an approach that relies purely on inverse dynamics models to drive the latent code.
Revised: This corresponds to an approach that relies purely on inverse dynamics models to drive the latent code.

Original: In these experiments, we use a pre-trained DrQ-v2 for the algorithm’s base policy, and then run adaptation across all distractions with intensity 0 .
Revised: In these experiments, we use a pre-trained DrQ-v2 for the algorithm’s base policy, and then run adaptation across all distractions with intensity 0 .

Original: The results in Table 2 show that the adversarial training is critical for the unsupervised pipeline to learn properly.
Revised: The results in Table 2 show that the adversarial training is critical for the unsupervised pipeline to learn properly.

Original: We theorize that the inverse dynamics head has a tendency to simply memorize dynamics in the training and testing domains without collapsing the commonalities, i.e. it memorizes rather than learns.
Revised: We theorize that the inverse dynamics head has a tendency to simply memorize dynamics in the training and testing domains without collapsing the commonalities, i.e. it memorizes rather than learns.

Original: The adversarial training component insures that the learning is projected onto a latent space that is agnostic of the underlying
Revised: The adversarial training component insures that the learning is projected onto a latent space that is agnostic of the underlying

Original: MDP.
Revised: MDP.

Original: [Table]
Revised: [Table]

Original: Pre-Filling the Replay Buffer
Revised: Pre-Filling the Replay Buffer

Original: In our experiments, we pre-ﬁlled both the original latent buffer and the target observation buffer with data collected via random exploration.
Revised: In our experiments, we pre-ﬁlled both the original latent buffer and the target observation buffer with data collected via random exploration.

Original: There is nothing particularly special about this choice, random exploration may be subbed out with any exploration strategy.
Revised: There is nothing particularly special about this choice, random exploration may be subbed out with any exploration strategy.

Original: One interesting strategy is to pre-ﬁll the buffer with an exploration strategy that relies on unsupervised pre-training.
Revised: One interesting strategy is to pre-ﬁll the buffer with an exploration strategy that relies on unsupervised pre-training.

Original: Given the recent success of these strategies on improving exploration, it seems likely this could further improve performance.
Revised: Given the recent success of these strategies on improving exploration, it seems likely this could further improve performance.

Original: This is left to future work.
Revised: This is left to future work.

Original:  C LOSING R EMARKS
Revised:  C LOSING R EMARKS

Original: We introduced Invariance through Inference, a new method for leveraging unsupervised data to improve the test-time adaptation of reinforcement learning systems.
Revised: We introduced Invariance through Inference, a new method for leveraging unsupervised data to improve the test-time adaptation of reinforcement learning systems.

Original: Empirical results demonstrate that as discrepancies between training and deployment environments become more intense, Invariance through Inference is more capable of adaptation than data augmentation techniques.
Revised: Empirical results demonstrate that as discrepancies between training and deployment environments become more intense, Invariance through Inference is more capable of adaptation than data augmentation techniques.

Original: The problem of test-time adaptation in visual reinforcement learning using unsupervised test-time trajectories is relatively new, but has thus far shown great promise.
Revised: The problem of test-time adaptation in visual reinforcement learning using unsupervised test-time trajectories is relatively new, but has thus far shown great promise.

Original: Future work in this area might develop better techniques for the initial exploration phase.
Revised: Future work in this area might develop better techniques for the initial exploration phase.

Original: Many of the environments we considered were quite reversible with a small intrinsic dimensionality.
Revised: Many of the environments we considered were quite reversible with a small intrinsic dimensionality.

Original: Expanding this class of methods to work on longer horizon multi-stage tasks is an intriguing possibility
Revised: Expanding this class of methods to work on longer horizon multi-stage tasks is an intriguing possibility