Year: 2021

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

text: 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.”

heading: 2 RELATED WORK
text: 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).

heading: 3 PROBLEM FORMULATION
text: 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.

heading: 4 INVARIANCE THROUGH INFERENCE

heading: 4.1 DISTRIBUTION MATCHING
text: We first deploy a random exploration policy in both the source and target domains to collect samples of observation trajectories
Bsrc = ( . . . , osrct , a src t , o src t+1, . . . ) ∼ Pπ̄(Msrc) (1a)
Btgt = ( . . . , otgtt , a tgt t , o tgt t+1, . . . ) ∼ Pπ̄(Mtgt), (1b)
where Pπ̄(M) denotes the distribution of transitions produced with a random policy in domainM. Given an observation ot, an encoder F produces a corresponding latent representation zt
zsrct = F (o src t ) (2a) ztgtt = F (o tgt t ). (2b)
We seek to match the distribution over ztgtt with that over z src t by adapting the weights of the encoder F , without access to test-time rewards. We employ adversarial training to achieve this. Specifically, our architecture (Fig. 2) includes a discriminator D that tries to distinguish between embeddings from the source domain zsrct 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. This results in matching latent distribution over zsrct and z tgt t . Following the Wasserstein GAN (Arjovsky et al., 2017) formulation, we express our distribution matching loss as
Jadv = EPπ̄(Msrc) [ D ( F̄ (osrct ) )] + EPπ̄(Mtgt) [ 1−D ( F (otgtt ) )] , (3)
where F̄ indicates that the weights of the network are frozen. 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.

heading: 4.2 DYNAMICS CONSISTENCY
text: 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 sufficiently 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). 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.
Algorithm 1 Invariance through Inference Require: Pretrained encoder F , discriminator D, buffers B∗src and Btgt, inverse dynamics network
parameterized with ψinv for i=1:N do . Pre-fill buffers
Sample zt, at, zt+1 ∼ Pπ̄(Msrc;F ) Sample ot, at, ot+1 ∼ Pπ̄(Mtgt) B∗src ← B∗src ∪ (zt, at, zt+1) Btgt ← Btgt ∪ (ot, at, ot+1)
for i = 1 : Tdynamics do . Pretrain dynamics networks Sample zt, at, zt+1 ∼ B∗src ginv ← ∇ψinvLdyn(zt, zt+1, at;ψinv) ψinv ← Optimizer(ψinv, ginv)
for i = 1 : Tadapt do . Adaptation main loop Sample zsrct , a src t , z src t+1 ∼ B∗src
Sample otgtt , a tgt t , o tgt t+1 ∼ Btgt gD ← ∇θD [ D(zsrct ) + (1−D(F (o tgt t ))) ] . Distribution matching (adversarial) loss
gF1 ← ∇θF [ D(zsrct ) + (1−D(F (o tgt t ))) ] gF2 ← ∇θFLdyn(F (o tgt t ), F (o tgt t+1), at;ψ ∗ inv) . Dynamics consistency loss θD ← Optimizer(θD, gD) θF ← Optimizer(θF , gF1 , gF2)
Specifically, let Cinv(zt, zt+1;ψinv) be the inverse dynamics network that predicts the action at associated with the transition from zt to zt+1, and Cfwd(zt, at;ψfwd) be the forward dynamics network that predicts the next latent zt+1from zt and at. We then define the dynamics consistency loss as the error in the inverse dynamics predictions
Ldyn(zt, zt+1, at;ψfwd, ψinv) = ‖Cfwd(zt, at;ψfwd)− zt+1‖22 + ‖Cinv(zt, zt+1;ψinv)− at‖22. (4)
Most importantly, we pretrain the inverse dynamics network on transitions sampled from the source domain Bsrc ∼ Pπ̄(Msrc) (Eqn. 1a) with observations encoded with F
ψ∗fwd, ψ ∗ inv = arg min
ψfwd,ψinv
EPπ̄(Msrc;F ) [ Ldyn(zsrct , zsrct+1, asrct ;ψfwd, ψinv) ] , (5)
where Pπ̄(M;F ) is the distribution over latent trajectories. Under the assumption that the underlying dynamics are shared by the domains, we then utilize these learned dynamics models to further encourage the latent embeddings for the target domain to align with those of the source domain. Specifically, during adaptation, we freeze the weights of the forward and inverse dynamics network, and minimize the following dynamics consistency loss with respect to encoder F using transitions sampled from the target domain Btgt ∼ Pπ̄(Mtgt) (Eqn. 1b)
Jdyn = EPπ̄(Mtgt;F ) [ Ldyn(ztgtt , z tgt t+1, a tgt t ;ψ ∗ fwd, ψ ∗ inv) ] . (6)

heading: 4.3 JOINT OBJECTIVE
text: We adapt our encoder by minimizing a loss that combines both the adversarial loss Jadv (Eqn. 3) and the dynamics consistency loss Jdyn (Eqn. 6). Specifically, we solve for the parameters of the encoder through the following minimax objective
arg min F arg max D
[Jadv + Jdyn] . (7)
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. Algorithm 1 summarizes the whole procedures to sample trajectory, pretrain dynamics networks and adapt encoder via the objective.
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. For this reason we refer to the derived method as Invariance through Inference.

heading: 5 EXPERIMENTS
text: We want to understand the impact of test-time inference on an agent’s ability to perform out-ofdistribution generalization. This section will compare Invariance through Inference with two stateof-the-art methods for data-augmented reinforcement learning: SVEA (Hansen et al., 2021) and DrQ-v2 (Yarats et al., 2021). Recall from the introduction that these methods vie for increased generalization capabilities by expanding the support of the training distribution. We expect this form of generalization to be less performant than unsupervised adaptation at test time. In our experiments, we will also compare against Policy Adaptation during Deployment (PAD) (Hansen et al., 2020), a baseline that, like our method, adapts the policy without access to the reward at test time.
To further probe the generalization abilities of test-time inference, we conduct an experiment wherein we vary the intensity of environmental distractions. The upshot here is that test-time inference significantly increases the robustness of our policy to distractions during deployment. We will conclude with some general discussion and remarks regarding the design tradeoffs involved in test-time inference.
Setup We conduct experiments using nine domains from the DeepMind Control Suite (DMC) (Tassa et al., 2018). We use DMC as the training (source) environment and the Distracting Control Suite (DistractingCS) (Stone et al., 2021) as the distracted test (target) environment. DistractingCS adds three types of distractions to the DeepMind Control Suite in the form of changes to the background image, deviations in color, and changes to the camera pose relative to training.
Modifications to DistractingCS The default configuration of DistractingCS changes distractions at the start of every episode (e.g., different background images are used at every episode). However, we are interested in measuring an agent’s ability to perform inference across several episodes on the same target environment. Thus, we modify DistractingCS to sample a distraction once in the beginning of learning, and then use the same distraction across all learning epochs. This also ensures consistent evaluation across algorithms. In accordance with this change, we also modify the intensity benchmark from DistractingCS. In our experiments, intensity measures the deviation between an environment distraction and the train environment’s default value. For example, intensity may measure how far the distracting color is from the default. 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. Figure 3 shows an example of distractions across intensities on Walker-walk domain.

heading: 5.1 INVARIANCE THROUGH INFERENCE ON DMCONTROL
text: This section studies the impact of test-time adaptation on DMControl performance. We begin by pretraining soft actor-critic (SAC) (Haarnoja et al., 2018), SVEA, and DrQ-v2 in a non-distracting training-time environment. After training, we evaluate the learned policies on test environments with distractions of various intensities. This evaluation is zero-shot, i.e., there is no additional training in the test environment.
Table 1 presents the results for the different DistractingCS domains in the presence of background distractions with an intensity level of 1.0. Specifically, we compare the test-time performance of SAC, SVEA, and DrQ-v2 in each domain with the episode rewards that we achieve when using our Invariance through Inference algorithm to adapt the encoder. The baseline algorithms employ image augmentation, which provides some robustness to variations at test time. Even then, however, we find that Invariance through Inference adaptation improves the test-time generalization of all three baseline policies in most domains, often resulting in significant performance gains. In cases where Invariance through Inference does not improve performance, the resulting reward is comparable to the baseline policy, i.e., Invariance through Inference does not result in a performance degradation.
Figure 4 visualizes the performance of the different methods, averaged over the set of DistractingCS domains, as a function of the intensity of the distractions. Since the baseline methods are trained with image augmentation, they do exhibit some robustness to distraction. However, we see this robustness rapidly diminishes as the distraction intensity increases. In particular, large changes to camera pose or the image background proved challenging for standard augmentation procedures. Comparatively, Invariance through Inference makes it much smoother and slower degradation of performance. This supports our hypothesis that adaptation powered by unsupervised learning can significantly widen the generalization abilities of learning algorithms.

heading: 5.2 COMPARISONS WITH PAD
text: Similar to our approach, PAD pretrains 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., 2020). To evaluate the robustness of PAD to distractions, we consider DistractingCS with a fixed distraction intensity of 1.0. Table 2 compares the performance as the difference between the episode returns before and after adaptation along with the episode returns in the clean environment. It should be noted that PAD requires to pretrain a policy along with inverse dynamics prediction objective. Thus we trained SAC with the auxiliary objective specifically for this experiment.
Across all environments, we see that PAD struggles to adapt to distractions at test time. 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.

heading: 5.3 FURTHER DISCUSSION
text: Ablation Studies In order to better understand the contribution of the different objectives to testtime generalization, we perform a series of ablations in which we omit either the dynamics consistency or the adversarial objectives. In these experiments, we use a pretrained DrQ-v2 network
for the algorithm’s base policy, and then perform adaptation across all distractions with an intensity value of 1.0. The results in Table 3 show that the adversarial training is critical to adapt the latent representation in the target domain. Performing adaptation using only the dynamics consistency objective, i.e., arg minF Jdyn (Eqn. 6) results in a significant decrease in performance. We theorize that the dynamics consistency objective helps to refine the local consistency in the latent space when the space in the target domain is close to that of the source domain. If the latent spaces significantly differ, however, we suspect that the gradients may negatively affect convergence.
Compared to the adversarial objective, ablating the dynamics consistency objective has surprisingly little effect on test-time generalization. It may be that the local structure of the latent space is preserved despite the distractions, which then diminishes the net effect of the dynamics consistency objective. For example, one can think of the original points in the latent space simply being shifted and affine-transformed. In this case, solely matching the distributions can effectively undo the transformations, resulting in a representation that is consistent with the original, without the need to explicitly align local structure. We suspect that this may be the case in the domains that we considered, and thus the effect of dynamics consistency objective gets obfuscated.
Pre-Filling the Replay Buffer In our experiments, we pre-filled both the original latent buffer and the target observation buffer with data collected via random exploration. There is nothing particularly special about this choice, random exploration may be subbed out with any exploration strategy. One interesting strategy is to pre-fill the buffer with an exploration strategy that relies on unsupervised pretraining. Given the recent success of these strategies on improving exploration, it seems likely this could further improve performance. We leave this for future work.

heading: 6 CLOSING REMARKS
text: We introduced Invariance through Inference, a new method for leveraging unsupervised data to improve the test-time adaptation of reinforcement learning systems. 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. 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.
Future work in this area might develop better techniques for the initial exploration phase. Many of the environments we considered were quite reversible, with a small intrinsic dimensionality. Expanding this class of methods to work on longer horizon multi-stage tasks is an intriguing possibility.

heading: A APPENDIX
text: The following provide a more detailed experimental evaluation of Invariance through Inference on the Distracting Control Suite.

heading: B PERFORMANCE IN ORIGINAL DOMAINS
text: Table 4 presents the average reward for the baseline SAC, SVEA, and DrQ-v2 on the non-distracted source domains. The method labeled SAC+Inv denotes a soft actor-critic agent that is trained using inverse dynamics as an additional auxiliary objective that is only used in making a comparison against PAD.

heading: C ABLATIONS
text: Tables 5, 6, and 7 provide a per-domain ablation summary for background, color, and camera pose distractions, respectively. As with the results in Table 3, we use DrQ-v2 as the pretrained policy and present the mean reward and standard deviation for five random seeds.

heading: D ARCHITECTURES
text: This section describes the architectures of encoder F , discriminator D, inverse dynamics Cinv and forward dynamics Cfwd. Encoder architectures follow the originally presented design choices of each base policy except for DrQ-v2. We take the part of the original network that produces a latent for the actor, and use it as an encoder F . In all of SAC (of the version used in (Hansen et al., 2021)), SVEA and PAD, this shared latent is set to have dimension 100. For DrQ-v2, we took the entire network architecture from SVEA. The discriminator consists of a linear layer with hidden dimension 100 followed by Layer Normalization (LN) (Ba et al., 2016) and tanh activation, and a three layer
multi-layer perceptron (MLP) with and ReLU activations. Inverse dynamics network Cinv is five layer MLP and ReLU activations. It takes the concatenated latents as its input. Forward dynamics network Cfwd takes action and latent, and encode them separately followed by concatenation and further layers. The action is fed it to a linear layer followed by LN, another linear layer and ReLU with hidden dimension 100. The latent is fed to three layer MLP where the first activation is LN and others are ReLU. The both encoded inputs are concatenated and fed to 4 layer MLP with ReLU activations followed by LN and tanh activation. In all layers except for those explicitly mentioned, hidden dimension is set to 1, 024.

heading: E HYPERPARAMETERS
text: This section details the hyperparameter settings that were used for the experimental evaluation. The hyperparameters relevant to pretraining dynamics network are listed in Table 8, and those relevant to adaptation are listed in Table 9. For the base policies, we followed the hyperparameters and architecture choices on their original paper. Those information on SAC and SVEA can be found in (Hansen et al., 2021), and PAD in (Hansen et al., 2020). The dimension of the latent zt differs depending on the encoder of the base policy, but all of the encoders in our experiments produce a latent with dimension 100.

Review summaries:
 **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.* 

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.

The paper proposes a method for learning invariant latent representation of the observations from MDP processes that share some aspects of their dynamics but differ in states.  This is useful for generlising reinforcement learning agents to wider range of variability of test conditions (i.e. test data can be out of distribution of the train data).  The invariant latent representation is derived through minimisation of mutual information between latent encoding of the in-distribution and the out-of-distribution experiences by "fooling" a GAN-like discriminator tasked at differentiating the two.  

In this paper, the transfer of a reinforcement model (RL) setting from an idealized (training) environment to a more realistic environment with distractors in the observations is considered. Instead of augmenting the training environment with more data so as to make the system more resilient to variations and distractors, the system is adapted at test time to be invariant to the specific distractors found in the environment. Experiments in simulation show the benefits of the proposed approach. Crucially, the agent is not able to access any reward data at test time.