Invariance Through Inference
Authors: Takuma Yoneda
Ge Yang
Matthew Walter
Bradly C. Stadie
Venue: ICLR 2022 Submitted (2022)
Keywords: invariance
representation learning
out-of-distribution
domain adaptation
generalization

Abstract: We introduce a general approach, called invariance through inference, for improving the test-time performance of a behavior agent in deployment environments with unknown perceptual variations. Instead of producing invariant visual features through memorization, invariance through inference turns adaptation at deployment-time into an unsupervised learning problem by trying 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 task reward, including changes in camera poses from the challenging distractor control suite. 

Under review as a conference paper at ICLR 2022
INVARIANCE THROUGH INFERENCE
Anonymous authors
Paper under double-blind review
ABSTRACT
We introduce general approach, called Invariance through Inference, for im- proving the test-time performance of an agent deployment environments with unknown perceptual variations. Instead of producing invariant visual fea- tures through interpolation, invariance through inference turns adaptation at deployment-time into unsupervised learning problem. This is achieved in prac- tice by deploying straightforward algorithm that tries to match the distribution of latent features the agent's 's prior experience, without relying on paired data. Although simple, we show that this idea leads to surprising improvements on a variety adaptation scenarios without access deployment-time rewards, in- cluding changes in camera poses and lighting conditions. Results are presented on challenging distractor control suite, a roboties environment with image- image-based observations.
INTRODUCTION
Let consider the ability of intelligent agent to gen- eralize to unseen domains. To have such a discussion, must first consider what generalization means. In much of the learning literature, we typically assume that
Test
Train
an agent will accrue experiences of sufficient variation during training, and that these experiences will allow the agent to generalize to novel settings during deploy- ment. Since the richness of the agent's s experience is of paramount importance the quality of its generalization, there exist a broad family of methods that expand the sup- port distribution of the agent's ' training set. Domain trans- fer, 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-hased control tasks, for example, image observa- tions are cropped, shifted, rotated, and discolored to make learned policies more robust shifts in the input obser- vation 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 al., 2020; Kumar et al., 2021), and care must be taken avoid destabilizing training.
Methods such as data aug- mentation try make the training dis- tribution large at the expense of com- plexity and performance. In spite of these efforts, these methods still often fail adequately cover the target distri- bution.
Figure 1
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 oiito-ddistrbbution, 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 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, in-distribution generalization.

---

Under review as a conference paper at ICLR 2022
A more challenging problem- but one that is also more realistic is generalizing out-of- distribution. In this setting, we assume the specifics of the agent's 's deployment are not known in advance. Taking this even further, assume that access to instrumented reward supervision is unavailable in the deployment environment. Here, the burden placed on the agent hoping to gener- alize is significant. Since do not have access to reward, fine-tuning via reinforcement learning is impossible. Prior experience will help, but a truly outtr--distribution environment of which the agent has no prior knowledge will make the transfer problem non-trivial. While difficult, this setting 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 dif- ferent lighting conditions, environmental objects, physical dynamics, and sensor modalities during deployment.
of
To make progress on nut-of-distributin 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 loss of generality, we may recast the out-of- disttiibuution 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 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, discussed in Section 5. Another approach is to enforce Cycle-GAN (Zhu et al., 2017) style objective to learn cycle- consistent mappings between domains (Zhang et al., 2020). Although it is shown 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 prob- abilistic inference to produce invariance by taking advantage of the latent structure the agent has already discovered the environment. As we will see, this pays great dividends in producing robotic agents that can function effectively in deployment environments that exhibit different light- ing, 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."
RELATED WORK
Generalization in reinforcement learning is a longstanding problem. Recent work has shown that accurately quantifying generalization in this setting poses 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 al., 2021). This line of work showed how to stabilize data- augteneeeetted RL, which lead to significant increase the gemeralization 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 '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).
2

---

Under review as a conference paper at ICLR 2022
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 distribu- tion (Ganin et al., 2016). One common class of techniques for addressing this problem is to learn rep- resentations that maximize domain confusion, thus requiring that the agent learn nomain-invariant features. This idea has been applied successfully in imitation learning (Stadie 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 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 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) parameter- ized via the tuple M (S, O, A, P,n), where S, O, and A are the state, observation, and action spaces, respectively, P: S A+ is the transition function, R S x A h IR is the scalar reward, and y is the discount factor. Recall that the state space S usually represents ground-truth informa- tion about the agent's environment, which we assume the agent can not directly access. Instead, the agent receives stream of observations o E O that convey information about the environment. In this paper, we assume the environment is fully oollyvbbber...e.. 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 x A + [0, 1] that maximizes the expected discounted return J E |> 7² R(st, ae) over infinite horizon, represented as the Q" (0, a) func- tion. In this paper, we assume that the policy consists of an encoder F O + Z, where Z is a compact latent space, and a policy n. : Z x A [0, 1].
We consider setting wherein there are two distinct domains: source domain Mare and target domain Migr 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, interested in the case where shift between Mse and Mig is induced by differences in the observation spaces Oue and Oig.- 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.
most
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 to make use of unpaired trajectory samples from each domain to drive adap- tation. 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 πº 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.
3

---

Under review as a conference paper at ICLR 2022
target buffer
source buffer
Z
Ol
D
source target
1+1
F
Z++1
Cdyn
a
Our proposed Invariance through Inference architecture. The encoder F takes an obser- vation of target domain, and learns fool the discriminator, while the discriminator D predicts whether the input an encoded target observation or a latent sample from source buffer. This adver- sarial training encourages the distribution of encoder outputs similar to the latent embedding sampled from the source buffer. Cdyn is the pretrained forward and inverse dynamics networks used only to guide the encoder during adaptation.
Figure 2:
INVARIANCE THROUGH INFERENCE
4.1 DISTRIBUTION MATCHING
We first deploy random exploration policy in both the source and target domains to collect samples of observation trajectories
Bid (..., Ot 1111. (M(M.)) (Mg))
(la)
Byg (..., Of ar
(1b)
where Ps (M) denotes the distribution of transitions produced with random policy in domain M. Given an observation encoder F produces a corresponding latent representation 2
2t
F
(2a)
F..
(2b)
We seek to match the distribution over z; with that over 23th by adapting the weights the encoder F. without access test-time rewards. We employ adversarial training achieve this. Specifically, our architecture (Fig. 2) includes a discriminator D that tries to distinguish between embeddings from the source domain = and those from the target domain z lgt At the same time, we adapt the parameters of the encoder to produce latent embedding that is indistinguishable to the discrimina- tor. This results in matching latent distribution over z and 2!th Following the Wasserstein GAN (Arjovsky et al., 2017) formulation, we express our distribution matching loss as
Jasy = Ep, (Mac) (F("() + (P. - D (F((
(3)
where F' indicates that the weights of the network are frozen. The encoder tries to minimize this objective while the discriminator acts an adversary and seeks to maximize it, resulting in GAN- like minimax game.
4.2 DYNAMICS CONSISTENCY
Theoretically, the use of adversarial training can result 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 the shared underlying dynamics as a feature-rich and informative source of mutual information between the two domains.
4

---

Under review as a conference paper at ICLR 2022
Algorithm Invariance through Inference
Require: Pretrained encoder F, discriminator D, buffers 8 and Big, inverse dynamics network parameterized with i/inv for i=l:N do
D Pre-fill buffers
Sample 2,, at. 24+1 Pa (Mscc:))
Sample oc, de, 0++1
Pa (Migt)
_ f- Bic U (2t, at 21+1)
Bigi Bigs U (ot, ae, 0++1)
for i = 1 Taymamics do
Di Pretrain dynamics networks
Sample Z,, a1; Zc+1 Bic
Fiov f V Ldyn (Zt, Zt+1, Oe; Vinv
Vinv t Optimizeee(eeee Vinn ginv)
for i Tatapt do
Adaptation main loop
Sample 2 aare
=++1 B
Sample of at I+1 ~ Bygt
JD +- Voo (2 STC)
+ (1 [ (FF())))
Distribution matching (adversarial) loss
DF Vor (( (te) + +(1 ((F(())
gF + Ver Cayn (F(o)"). F(ot+t) ((o azi mmm)
B Dynamics consistency loss
0D + Optimizer(@p.. ID)
0F + Optimizer (0, 8r 9F, 9F%
: vinv) be the inverse dynamics network that predicts the action a, asso- ciated with the transition from 2 to 2t+1, and (Fe dzs t'twa) be the forward dynamics network that predicts the next latent 24+1 24+1 from 2 and ag. We then define the dynamics consistency loss as the error in the inverse dynamics predictions
Specifically, let Cinv (2,, 24i1:
Ldyn (Zes 244 Ozj Vrwds )iin)) |[Cred Z2, Oei (twa) 224 + ||Cinv 1 Vine) a, l12.
(4)
Most importantly, we pretrain the inverse dynamics network on transitions sampled from the source domain Bare Ps arc (Eqn. la) with observations encoded with F
fwa 2/."
arg min (Pe(M
(2 1, at snc, 1/rwds Vinv)]
Dirud. Vlaw
where (M F) is the distribution over latent trajectories. Under the assumption that the underly- ing 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 Big (M(Mi (Eqn. 1b)
Jayn (px(MM:F) (2y 1244 4g 1st
wiiad Viim))]
4.3 JOINT OBJECTIVE
We adapt our encoder by minimizing a loss that combines both the adversarial loss Jadv (Eqn. 3) and the dynamics consistency loss Jayn (Eqn. 6). Specifically, we solve for the parameters of the encoder through the following minimax objective
arg min arg max (Jaa + Jayn]
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 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 adapta- tion objective produce invariance by making statistical inference the latent space. For this reason refer to the derived method as Invariance through Inference.
5

---

Under review as a conference paper at ICLR 2022
Figure 3: Samples from the modified DistractingCS with intensities increasing from (left-most col- umn) zero (right-most column) one for the (top row) color, (middle row) camera pose, and (bottom row) background distractions.
EXPERIMENTS
We want to understand the impact of test-time inference on an agent's ability to perform out-of- distribution generalization. This section will compare Invariance through Inference with two state- of-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), baseline that, like our method, adapts the policy without access 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 that test-time ference 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 Distract- ing 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 the background image, deviations in color, and changes to the camera pose relative 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 be- ginning 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 inten- sity 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 single distraction during testing (rather than all three) in order to better understand the impact of each type of distraction on overall performance. Figure shows an example distractions across intensities Walker-walk domain.

---

Under review as a conference paper at ICLR 2022
Episode Return
1000
1000
1000
Dr248
SVEA
SVEA
500
500
500
SAC+ ITI
SAC
0.5
1.0
Distraction Intensity
Distraction Intensity
Distraction Intensity
(a) Background
(b) Colors
(c) Camera Posc
Figure 4: The gain of applying Invariance through Inference various distracted baselines. Dashed lines denote the performance of the baseline agent in the target environment (i.e., zero-shot), while solid lines represent the performance gains of the base agent with ITT (our method).
INVARIANCE THROUGH INFERENCE ON DMCONTROL
This section studies the impact of test-time adaptation 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, 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: Episode return in the target (test) environments (mean and standard deviation) before (zero- shot) and after (+ITI) adaptation for SAC, SVEA, and DrQ-v2 with background distraction at an intensity setting of 1.0. The performance of each baseline in the source (training) environments be found in the Appendix.
SAC
SVEA
DrQ-12
Domain
Zero-shot
+ITI
Zero-shot
+ITI
Zero-shot
+ITI
Dalliin.cup-catcch155
50
227 222
490 4 376
987
88
1 425
caripole-balance 434 1 275
585 1 295
446 4 330
627 1 258
273 1 107
322 117
cartpoll-swinguup 182 1 147
369
243
269 365
612 213
82 3:
247 136
cheetab-run
169 + 65
248 = 53
317 = 137
378 +55
18
393 = 125
finger-spin
113 162
192
391 467
943
207 328
769 206
finger-turn.easy
163 163 99
146
278 - 180
268 - 241
914 44
reacher-easy
+:
381 + 76
75 + 77
624 + 305
58 4222
685 + 211
walker-stand
330
364 115
917 + 138
999 + 12
630 + 197
868 + 151
walker-walk
242 142
291 134
866 + 45
924 + 45
326 + 195
770 + 140
Table presents the results for the different DistractingCS domains in the presence of background distractions with 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 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 performance degradation.
Figure visualizes the performance of the different methods, averaged over the set of DistractingC'S domains, 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 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.

---

Under review as a conference paper at ICLR 2022
Relative Impry. (%)
80
80
60
60
60
40
40
40
DrQ+2 III
20
20
20
SAC 11]
0-
0
1000
2000
1000
2000
1000
2000
Steps
Steps
Steps
(a) Background
(b) Colors
(c) Camera Pose
Figure 5: Relative improvement (compared zero-shot) function of adaptation steps when applying ITI to different baseline policies. As in Figure 4, each point represents the mean over nine domains and five random seeds. The results correspond to an intensity value of 1.0.
Table 2: Comparison against PAD
Distraction
Zero-shot
+PAD
+ITI
None
835 + 230
Background
213 247
279 + 271
425 + 292
Colors
230 - 263
271 + 300
402 + 339
Camera Pose
319 265
326 + 259
412 + 275
5.2 COMPARISONS WITH PAD
Similar to our approach, PAD pretrains the agent in a clean environment, and then adapts the agent via unsupervised transfer, without assuming access the target environment's reward func- tion (Hansen et al., 2020). To evaluate the robustness of PAD to distractions, we consider Distract- ingCS with fixed distraction intensity of 1.0. Table compares the performance as the difference between the episode returns before and after adaptation along with the episode returns in the clean environment. should be noted that PAD requires pretrain policy along with inverse dynamics prediction objective. Thus we trained SAC with the auxiliary objective specifically for this experi- ment.
Across all environments, we see that PAD struggles to adapt to distractions at test time. We sus- pect 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 dy- namics head does not encourage the latent train and test distributions to match, which is a feature specifically baked into Invariance through Inference.
5.3 FURTHER DISCUSSION
Table 3: Ablations with variants of Invariance through Inference that remove inverse, forward dy-
namics, or adversarial objectives. DrQ-v2 used as a pretrained policy. We compute episode returns from nine domains and five random seeds, and the results correspond to an intensity value of 1.0.
+ITI w/o dyn.
+ITI
Distraction
Zero-shot
+ITI
w/o adv.
Background
228 232
602 + 300
615 289
176 + 221
Colors
234 + 245
536 = 320
534 + 327
117 = 96
Camera Pose
345 287
417 + 284
407 272
208 + 235
In order to better understand the contribution of the different objectives to test- time generalization, we perform series of ablations which we omit either the dynamics con- sistency the adversarial objectives. In these experiments, use a pretrained DrQ-y2 network
Ablation Studies
8

---

Under review as a conference paper at ICLR 2022
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 minp Jayn (Eqn. 6) results in a significant decrease in performance. We theorize that the dynamies consistency objective helps 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 objectíve, 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 pre- served despite the distractions, which then diminishes the net effect of the dynamics consistency objective. For example, one can think the original points in the latent space simply being shifted and affine-transformed. In this case, solely matching the distributions can effectively undo the trans- formations, resulting in 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 consid- ered, 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 nothing particu- larly special about this choice, random exploration may be subbed out with any exploration strategy. One interesting strategy to pre-fill the buffer with an exploration strategy that relies on unsuper- vised 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.
CLOSING REMARKS
We introduced Invariance through Inference, new method for leveraging unsupervised data to im- prove 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 small intrinsic dimensionality. Ex- panding this class of methods to work on longer horizon multi-stage tasks is an intriguing possibility.
REFERENCES
Martin Arjovsky, Soumith Chintala, and Léon Bottou. 3XXiv:7011.08755 2017.
Wasserstein GAN,
arXiv preprint
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the International Conference on Machine Learning (ICML), pp. 1597-1601 2020a.
Xi Chen, Ali Ghadirzadeh, Mårten Björkman, and Patric Jensfelt. Adversarial feature training for generalizable robotic visuomotor control. Proceedings of the IEEE International Conference on Robotics and Automation (ICRA), pp. 1142-1148, 2020b.
Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generaliza- tion in reinforcement learning. In International Conference on Machine Learning, pp. 1882-1289, 2019.
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RL?: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv:161 orXiv:/611.02770) 1.02779, 2016.

---

Under review as a conference paper at ICLR 2022
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leaming for fast adaptation of deep networks. In Proceedings of the International Conference on Machine Learning (ICML), pp. 1135-1135, 2017.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adverssaral training of neural net- works. Journal of Machine Learning Research,
2016.
Jean-Bastion Grill, Florian Strub, Florent Altché, C. Tallec, Pierre H. Richemond, Elena Buchatskaya, C. Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, B. Piot, K. Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent: A new approach to self-supervised learning. arXiv preprint arXiv:2006.07733. 2020.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Picter Abbeel, and Sergey Levine. Soft actor-critic algo- rithms and applications. arXiv preprint www.::::12..1000000 2018.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with discrete world models. arXiv preprint arXiv:2010.02/93; 2020.
Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmen- tation. arXiv preprint arXiv:2011.13388 2020.
Nicklas Hansen, Rishabh Jangir, Yu Sun. Guillem Alenyå, Picter Abbeel, Alexei A Efros, Lerrel Pinto, and Xiaolong Wang. Self-supervised policy adaptation during deployment. arXiv preprint Stti:20007.043099 2020.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep Q-learning with ConvNets and vision transformers under data augmentation. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross B. Girshick. Momentum contrast for unsupervised visual representation learning. Proceedings of the IEEE Conference on Computer Vision and Pullern Recognition (CVPR), pp. 9726-9735, 2020.
Rein Houthooft, Richard Y Chen, Phillip Isola, Bradly C Stadie, Filip Wolski, Jonathan Ho, and Pieter Abbeel. Evolved policy gradients. arXiv preprint
04002, 2018.
Kyle Hsu, Sergey Levine, and Chelsea Finn. Unsupervised learning via meta-learning. arXiv preprint arXiv:1810.02334, 2018.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukeuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:1611.05397 2016.
Aviral Kumar, Rishabh Agarwal, Aaron Courville, Tengyu Ma, George Tucker, and Sergey Levine. Value-based deep reinforcement learning requires explicit regularization. ICML 2021 Rein- forcement Learning for Real Life Workshop, 2021.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Re- inforcement learning with augmented data. arXiv preprint arXiv:2004. /4990, 2020.
Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. A simple randomization technique for generalization in deep reinforcement learning. arXiv preprint arXiv:1910.05396. 2019.
Xue Bin Peng, Marcin Andrychowicz. Wojciech Zaremba, and Pieter Abbeel. Sim-to-real trans- fer of robotic control with dynamics randomization. In Proceedings of the IEEE International Conference on Robotics and Automation (ICRA), May 2018.
Martin L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. John
Wiley & Sons, August 2014.
Roberta Raileanu, Max Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Auto- matic data augmentation for generalization deep reinforcement learning. arXiv preprint arXiv:2006. 12862, 2020.
10

---

Under review as a conference paper at ICLR 2022
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off off-policy reta-reiinorrement learning via probabilistic context variables. In Proceedings of the Interna- tional Conference on Machine Learning (ICML), pp. 5331-5340, 2019.
Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. ProMP: Proximal meta-policy search. arXiv preprint arXiv:1810. 067844 2018.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with momentum predictive representations. arXiv preprint arXiv:2007.05929, 2020.
Laura Smith, Nikita Dhawan, Marvin Zhang, Pieter Abbeel, and Sergey Levine. AVID: Learning multi-stage tasks via pixel-level translation of human videos. arXiv preprint arXiv:1912.0444,, 2019.
Aravind Srinivas, Michael Laskin, and Picter Abbeel. Curl: Contrastive unsupervised representa- tions for reinforcement learning, arXiv preprint 0rXiv:2004.04136 2020.
Bradly C Stadie, Pieter Abbeel, and Ilya Sutskever. Third-person imitation learning. arXiv preprint urXiv:1703.000 2017.
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting con- trol suite A challenging benchmark for reinforcement learning from pixels. arXiv preprint arXiv:2101.02722, 2021.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy Lillicrap. and Martin Ried- miller. DeepMind Control Suite. arXiv preprint 1e.Xivv801..cm 006900 2018.
Dequan Wang, Evan Shelhamer, Shaoteng Liu, Bruno Olshausen, and Trevor Darrell. Tent: Fully test-time adaptation by entropy minimization. arXiv preprint arXiv:2006. 10726, 2020.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous con-
trol: Improved dagm-nudee reinforcement learning. arXiv preprint arXIv:2107.006455 2021.
Amy Zhang, Nicolas Ballas, and Joelle Pineau. dissection of overfitting and generalization in continuous reinforcement learning. arXiv preprint arXiv:1806.07937 2018.
Qiang Zhang, Tete Xiao, Alexei A. Efros, Lerrel Pinto, and Xiaolong Wang. Learning cross-domain correspondence for control with dynamics cycle-consistency. arXiv preprint 1rXiv:2012.09811, 2020.
Jun Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to- to-imag translation using cycle-consistent adversarial networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2017.

---

Under review as a conference paper at ICLR 2022
A APPENDIX
The following provide a more detailed experimental evaluation of Invariance through Inference on the Distracting Control Suite.
B PERFORMANCE IN ORIGINAL DOMAINS
Table 4 presents the average reward for the baseline SAC, SVEA, and DrQ-v2 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 only used in making comparison against PAD.
Table 4: Performance in the original (clean) domains.
Domain
SAC
SAC+Inv
SVEA
DrQ-v2
ba ballLin.cup-catchh
452 + 303
999 + 7
1007 + 4
1007 + 3
cartpole-balance
1022 + 8
988 26
996 21
969 + 123
cartpole-swingup
735 + 167
885 - 24
892 15
874 21
cheetah-run
309 + 26
4
148:100
finger-spin
615 +15+:3
971 + 64
1000
997 36
finger-turneeasy
138 = 28
658 141
317
14
reacher-easy
381 39
723 383
812 293
988 + 28
walker-stand
438 + 111
997 1
1006 4
996 31
walker-walk
393 + 117
915 22
980 16
C ABLATIONS
Tables 5, 6, and provide 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.
Table Ablation with Background distraction
+ITI w/o inv., fwd..
+ITI w/o inv.
+ITT w/o fwd.
+ITT
Domain
Zero-shot
+ITT
w/o adv.
walker-walk
326 + 196
749 = 133
778 142
777 + 145
768 146
268 367
walker-stand
623 = 233
866 = 153
883
110
859 + 144
873 - 129
402 347
cartpole-swingup
82 - 35
231 = 128
395 + 218
288 + 163
246 - 152
117 + 43
ball.in.cup-catch
88 1 39
394 + 387
381 +: 416
400 1: 414
383 1 381
93 1 42
finger-spin
208 + 327
783 : 214
742 - 190
772 225
769 223
74 + 160
reacher-easy
98
726 + 149
713
732 1 104
694 169
91 1
cheetab-run
98 - 90
411 - 191
397 - 152
3 147
419 157
12 - 19
cartpole-balance
271 101
336 = 126
315 - 98
367 84
297 = 121
264 - 80
finger-turneeasy
261 = 244
920 = 41
929 57
899 + 52
909 = 74
256 - 264
D ARCHITECTURES
This section describes the architectures of encoder F, discriminator D, inverse dynamics Cim- and forward dynamics Civd- 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 latent for the actor, and use it as an encoder F. In all 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, 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 three layer
took the entire
12

---

Under review as a conference paper at ICLR 2022
Table 6: Ablation with Color distraction
+ITI w/o inv., fwd.
+ITI
+ITT w/o fwd.
+ITI
Domain
Zero-shot
+ITI
w/o inv.
w/o adv.
walker-walk
80 1 43
481 313
468 305
495 1 330
472 319
26 1
walker-stand
278 150
543 231
603
571 + 259
548 268
145
cartpolo-swingup
152 = 84
552 = 377
520 359
507 + 339
434 - 395
94 56
bell.in.up-catch
239 = 374
812 225
857 = 182
840 + 251
843 = 196
131 - 46
finger-spin
319 + 354
612
561 + 376
591 + 351
603 = 358
143 - 172
reacher-easy
138 + 135
490 = 116
186 392
486 + 362
145 405
140
cheetah-run
193 - 188
122 = 282
416 293
121 + 277
106 - 285
1
cartpole-balance
481 351
602 = 366
576 359
593 + 351
583 - 349
216 94
tinger-turn.easy
194 200
313 = 290
323 - 337
297 + 316
304 - 323
170 + 57
Table 7: Ablation with Camera Pose distraction
+ITI
+ITI w/o inv.
+ITI w/o fwd.
+ITI
Domain
Zero-shot
+ITI
w/o inv., fwd.
w/o adv.
walker-walk
293 + 195
375 + 154
369 + 168
389 + 161
366 + 164
63 71
walker-stand
621 + 202
704 + 105
617 + 155
661 + 146
679 + 93
330 + 252
cartpole-swingup
286 +: 51
234 123
211 + 122
241 + 140
281 42
172 132
cup-
327 227
462 = 307
429 321
566 + 314
400 337
tinger-spin
29 + 25
252 216
222 - 209
275 + 194
251 151:2 206
32 63
zeacher-easy
917 + 101
933 -
925
950-:22
940 + 90
732 - 159
cheetah-run
55 + 20
142
154
141
144 -44+11
24 + 28
cartpole-balance
288
307 100
3 83
379 116
380 64
217 47
finger-turn.easy
287 -
346 = 219
359 122
377 + 200
421 = 180
160 97
multi-layer perceptron (MLP) with and ReLU activations. Inverse dynamics network Civ is five layer MLP and ReLU activations. takes the concatenated latents as its input. Forward dynamics network Crwa takes action and latent, and encode them separately followed by concatenation and further layers. The action fed
linear layer followed by L.,, another linear layer and ReLU with hidden dimension 100. The latent is fed 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.
E HYPERPARAMETERS
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 in (Hansen et al., 2021), and PAD in (Hansen et al., 2020). The dimension of the latent 2, differs depending on the encoder the base policy, but all of the encoders in our experiments produce latent with dimension 100.
SAC and SVEA be found
Table Hyperparameters for dynamics pretraining
Hyperparameter
Valuc
Steps (Tau)
100, 000 256
Batch size
Optimizer
RMSProp(a
0.99, x 10-3)
Learning rate (forward dynamics)
0.001
Learning rate (inverse dynamics)
0.001
13

---

Under review as a conference paper at ICLR 2022
Table Hyperparameters adaptation
Hyperparameter
Value
Capacity of buffers (Nuut)
1, 000, 000 256
Batch size
Discriminator updates per step
Gradient clipping
0.01
Optimizer
RMSProp(a 0.99, 1.0 x 10-*)
Learning rate (encoder)
1.0 10 (for DrQ-v2)
1.0 10 (otherwise)
Learning rate (discriminator)
1.0 10
(for DrQ-v2)
1.0 10
(otherwise)
Learning rate (inverse dynamics)
1.0 10
14


Reviews:
[{"review_id": "z9vAp8t9bBo", "reviewer": "Reviewer_X3P2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper devotes a lot of space to motivating the setup and justifying its value. I do believe that this setup is relevant in practice in robotics applications. There are indeed many cases in which sim2real faces generalization problems. Unfortunately, the value of this approach is not actually demonstrated in a sim2real application, but in transfer between simulated applications.\n\nWhile the goal of the paper seems very worthy, and the proposed algorithm seems to provide empirical benefits, the theoretical presentation and the justification of the approach are severely lacking, and I would say incorrect. In particular:\n- Why is a distinction between S and O even made? The presentation goes beyond the standard RL presentation making a distinction between the state and the observation, and makes a point of saying that the observation doesn't uniquely identify the state. And that the transition function P is a function of states (and actions) and not observations. To then ignore such distinction and make the transition function only a function of the observations, negating the explanation about partial observability.\n- Why is Eq. (5) just adding the two losses? Is this dimensionally correct? How do we know they are balanced? It seems that the relative scaling could be arbitrary by simply changing the scale of the actions.\n- What's \\tilde{z} in Eq. (3)? This is never defined.\n- In Eq. (4),  \\bar{G} is said to be the frozen generator. However, this frozen generator is put inside the expectation over \"outside of the distribution samples\". Isn't the frozen generator used for the training samples, therefore inside the distribution? There's no proper definition of what \"distr\" and \"non-distr\" means.\n- At the beginning of section 4.1, the case is made for a single generator being used both for training and testing data, and that this generator should ignore the distractors. However, Eq. (4) and the rest of the paper uses different generators for training and testing.\n- Figure 4.2 Are you sure this is a D? Shouldn't this be a P_\\phi?\n- Most importantly, the use of the GAN to match the distributions over z is completely unjustified, see more below.\n\nThe need for the GAN: It is unclear what the GAN portion of the architecture is bringing to the table. Since the inverse dynamics head is locked and learned on the training data, making it work on the test data should result in a z that aligns well with the latent space discovered at training time. The paper contains an ablation section showing that removing the GAN makes the system not work at all. How is this possible? The authors hypothesize that \"the inverse dynamics head has a tendency to simply memorize dynamics in the training and testing domains without collapsing the commonalities\". But the authors had said that, after training the inverse dynamics head on the training data, this head is frozen and remains locked during adaptation to test data. How can it possibly memorize the test dynamics then? And, even though I don't understand how this overfitting can happen, if it did, shouldn't the solution be regularization of the inverse dynamics network? The account of this architecture is inconsistent and the GAN seems completely unwarranted. Very little detail about the whole process is given and there's not supplementary information or code available.\n\nMinor:\nTypo: \"training environment the Distracting\"\nThe meaning of the values in Table 1 are unclear. The legend has no information and the text says it's the reward difference between before and after adaptation, but then there's a column for before and another for after, so the numbers themselves are not the difference. What is it?", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper tackles an important problem, but contains technical inconsistencies, lacks rigor and provides very little detail about the proposal as a whole. The use of a GAN, a major part of their architecture, is completely unjustified, and we are told that it's necessary because the system doesn't work if we remove it... but there's no explanation as to why. No code is provided, no supplementary details are provided.\n\nEmpirical results seem good, but there's little support as to why this is the case. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "cumwF2UKClH", "reviewer": "Reviewer_wVDQ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths \n\n* **Novelty**. As far as I am aware the authors propose a novel method for solving the problem of generalisation to new (but similar) domains in RL. \n\n* **Relevance**. This is an important problem to solve for the real world deployment of reinforcement learning agents.\n\n* **Empirical results**. The results provided indicate that the authors' method performs well.\n\n\n## Weaknesses\n\n* **Clarity**. Throughout the paper there were notation and wording choices which hinder clarity, in my opinion. I list them here:\n    1. The phrase \"ex-ante\" is unnecessarily fancy, for lack of a better word. A more commonly used term would be better. \n    2. The statement \"... with different forward dynamics, so long as the distinct dynamics maintain *some semblance of similarity*\" is very vague. Could it not be replaced with something more concrete? \n    3. Regarding notation in Sec 4.1. \n         A. In some cases G_\\theta has two arguments, the observation and action, and in other cases it only has the observation. Which is correct?\n         B. (Related to A), should the text after eqn 2, also mention a mapping similar observations *and actions* to close latents?\n         C. Perhaps it would be helpful to write \\hat{P}_\\phi rather than P_\\phi in eqn 2, to help distinguish P_\\phi from the P in the MDP?\n    4. \\tilde{z} and D are not introduced before eqn 3. \n    5. D_{dist} and D_{non-dist} in eqn 4 are never explicitly described. \n    6. In eqn 5, it is not clear why J_inv is within the arg max over D, given that it does not depend on D (i.e. D does not appear in eqn 2)?\n    7. Section 4.2 would be much more clear if written in an algorithm environment rather than as paragraphs of text. \n    8. In figure 4, should 2. not contain P_\\phi rather than D?  \n       \n\n* **Experimental evaluation.**\n    1. It seems to me that the experimental evaluation in this paper is not very comprehensive, only showing two main experiments and a single ablation. I would expect to see a wider range of experiments in an ICLR paper introducing a new method. Otherwise, how can readers know whether the method is generally applicable? \n    2. Regarding the modifications to DistractingCS, it would be great to see the effects of the various combinations of 2 distractions.\n    3. Why was the comparison to PAD not also made in Figure 5? It would be informative to compare how the degree of distraction impacts PAD. \n    4. How many repetitions were made to get the confidence intervals in Table 1? \n    5. The proposed future work of exploring the effect of exploration strategy seems like it should be included in this paper, in order to paint a full picture of the proposed method. \n\n## Other comments/questions for the authors\n\n1. A relevant but missing citation/discussion in the related work (specifically in **Invariant representation learning**) seems to be World Models (Ha and Schmidhuber, 2018).\n\n2. (nit) Regarding the slight abuse notation for eqn 1, there seems to be enough space to write things correctly. Furthermore, the explanation for the abuse of notation takes more extra space than the correct notation would.\n\n3. Regarding the discussion of data augmentation techniques in the paragraph above Sec 4. While it is true that even extreme data augmentation techniques would fail to cover the test MDP observation space, it is not clear that data augmentation does not ultimately play the same role as *invariance through inference* albeit in a less direct manner. That is, data augmentation should force the policy learned to be invariant to the noise from the data augmentations and thus be more likely to generalise to the test MDP setting.\n\n4. Could the authors elaborate on their choice of the GAN framework as opposed to VAEs or some other deep generative model?\n\n5. How much of a challenge was it to get the adversarial training to converge in the various experiments presented in the text. GANs are well known to be difficult to train and often require many tricks to converge reliably. This was not touched on in the manuscript and seems an important detail for the general applicability of this work. \n\n\n## References\n\nDavid Ha, Jürgen Schmidhuber:\nWorld Models. CoRR abs/1803.10122 (2018)\n", "summary_of_the_paper": "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.", "summary_of_the_review": "While the authors' proposed method seems like a promising solution to a challenging and important problem, this paper is let down by a lack of clarity in presentation, and a narrow range of experimental evaluation. \n\n++++ Post-revision update ++++\n\nThe authors have largely rewritten the paper which has resulted in a much clearer presentation. They have also expanded on their experimental results and clarified that the scope of their experimental work is larger than I initially understood. With these two points in mind, I have increased my score from 5 to 6. Unfortunately, I do still feel that the experimental evaluation could be more comprehensive, which is why I have not increased my score further. Additionally, I have not had the time I would like to go through the heavily revised paper, thus I have decreased my confidence. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "Y7tlO0dwBu8", "reviewer": "Reviewer_HvyM", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I quite like the aim of the proposed method, and the results are impressive.  The fundamental principle of using the discriminator to obtain invariant latent representation of an out-of distribution state space of the MDP seems like a very sensible and practical idea.  The only problem with the paper is that is very high level and I can't quite figure out how it all works and comes together.  Even with Figure 3 and Algorithm overview in section 4.2, it's not clear what happens when.  Seems there are different phases to training...but I can't quite understand what happens when.  When is P_\\phi fixed and when it is trainable?  What is a_t in Equation 2 when training on MDP_test?  What's z and \\tilde{z} in 3?  Do D_{non_dist} and D_{distr} refer to data from MDP_{train} and MDP_{test}?  How big does the source buffer storing training data need to be?  Does it store all training data?  ", "summary_of_the_paper": "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.  ", "summary_of_the_review": "I don't think it would require a huge amount of work, but a bit more meticulous explanation of the loss functions and the training process is needed.  At the moment it seems that only those intimately familiar with the method and code (that is not provided) would be able to implement it.  The direction seems good, results are very encouraging, just need to flesh out the actual method a bit more, otherwise it feels there isn't enough information to reproduce.   ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ImGJVTP3vNB", "reviewer": "Reviewer_jjRF", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths:**\n\n * The approach itself is novel and relies on leveraging the model's ability to infer invariance via an adversarial loss on latent representations from source and target distributions. This provides a more principled way to model environment invariance and can possibly scale well in comparison to data driven approaches. \n\n * The approach does not rely on rewards from the target distribution.  However, it is unclear whether this is typically the case for the other baselines.\n\n* The results showing decreased degradation for distraction on DMC (Fig. 5) look good and are consistent across baselines and distraction dimensions.  The results against PAD are also significantly in favour of ITI.\n\n* A nice ablation demonstrating the effect of the adversarial objective.\n\n**Weaknesses:**\n\n* The evaluation domain seems to be narrow and focuses on robotics control domains or environments that simulate that.  How well does this approach work on other domains?  Are there any theoretical bounds or existing results that might give us confidence that this approach scales well across any out-of-distribution evaluation scenario?\n\n * The paper focuses on comparisons to data augmentation approaches to modeling out-of-distribution data.  Are there other approaches that might be considered?  Could combinations of this approach and data augmentation lead to even more robust adaptation to out-of-distribution estimation?\n\n * A more formal description of an adaptation and the criteria for being out-of-distribution would be very helpful.  In general, in sections 1 and 3 clearer explanation of the problem and definitions would lead to more clarity around the problem that this model proposes to solve.  For instance, what are the failure cases of generalization that this approach aims to solve and what in general are common features of a model with robust generalization?  What if we simply have access to large amounts of more varied data, do the gains of ITI remain strong in such scenarios? \n\n * First sentence of the last paragraph of intro, it is stated \"this challenging scenario\".  Can you be precise and explicit here?\n\n * The *Algorithm Overview* section seems like it could us more detail. An algorithm box could be useful here.  Do we use training samples only to train the encoder?\n\n* Both the ablations and comparison to PAD are done over a fixed set of hyper-parameters.  It would be helpful if these conclusions were better supported across more variation and/or datasets.\n\n* In the experiments section you emphasize that these methods \"expand support for the training distribution\", could you contrast these methods with each other?  Do they have access to reward of the target distribution?\n\n* Table 1 doesn't seem to describe what is being compared.  Could you add more description here?\n\n* When discussing the comparison to PAD, it 's stated that the hypothesis is that it does not encourage latent train and test distributions to match as ITI does, could this be described in more detail wrt to PAD itself?  Just a brief description would help lend insight as to why ITI performs better.\n\n", "summary_of_the_paper": "\n**Introduction:**\n\nIn 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.\n\nOut-of-distribution generalization is very important but more challenging:\n * Prior experience will be critical\n * Lack a reward function over the test distribution\n * In robotics for instance, this type of learning is critical\n\nThe 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\"*\n\nStatement 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.\n\n\n**Problem Formulation:**\n\nGiven 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.\n\n**Invariance Through Inference and Algorithm**\n\nThe 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:\n 1. a distribution matching objective encourages the latent distribution induced by both MDPs to match. \n 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.\n\nThe 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.\"*\n\n**Experiments**\n\nThe 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.\n\nThe authors also compare ITI to Policy Adaptation during Deployment (PAD).  They compare to ITI on DMC for fixed distraction intensity.\nITI 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\nchanges 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.*\n", "summary_of_the_review": "\nWhile I believe that the ideas of this work are novel and very interesting and provide a potential pathway away from data augmentation method for out-of-distribution generalization I have two issues with the paper: 1) the evaluation domain may be too narrow to provide confirmation of the approach with confidence over other methods, 2) more clarity and details of the approach would be helpful in understanding this work, how it contrasts with others, and where it exceeds those. \n\nPost Rebuttal:\n\nMy initial two main concerns with this work: 1) how well the approach generalizes and 2) the clarity of the paper. The updated draft addresses the second concern very well and so I consider that resolved. The sheer volume of updates here is impressive. For the first point, the authors have carried out a large number of new experiments which strengthens the case for control domains but to me the overall claims of the paper seems to be more general than that. That said, I believe this is a worthy contribution given the efforts put forth by the authors and will also increase my score to a 6.  Great work!\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Meta Review for Invariance Through Inference

The motivation of this work is to address the problem of learning a model that generalizes well on a test distribution that samples outside of the training data distribution. Reviewer X3P2 wrote a good summary of the paper:

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.

Reviewers, including myself, recognize the novelty of the approach, in particular appreciate the authors' motivation to provide a more principled way to model environment invariance that can possibly scale well in comparison to data driven approaches. However, the initial round of feedback is generally negative, in particular, most reviewers raise concerns regarding the lack of clarity in presentation, and also have issues with the narrow range of experimental evaluation. Clearly, this is promising work, but possibly had to be rushed for submission.

To the authors' credit, they devoted substantial efforts to completely revamp their paper, addressing many of the issues head on. The resulting updated manuscript is almost a complete rewrite of the paper. All reviewers acknowledge (and praise) the effort from the authors' to improve the paper, and 3 out of 4 reviewers had improved (or maintained) their scores from rejection to a 6. But as the paper is a complete revamp, reviewers did not have the time to assess the entire rewrite of the paper (it's like the need to review a paper from scratch), so the confidence is reduced.

While X3P2 did not change their score, they did lead a discussion amongst myself and other reviewers, and they spent the time to take a detailed look at the completely revised draft. Here are the comments from that discussion, for full transparency:

---

*The paper has been completely revamped, to the point in which the presented technique is actually different (the dynamics loss now includes a new forward term). The changes are overall welcome since it significantly improves the clarity of the presentation. Experiments still show promise.*

*I still have problems with the theoretical aspect of it, though. I think that it is unclear why the proposed system is working and fails to provide the minimal system that works.*

*Equation (4) is dimensionally incorrect. It's summing squared error over actions with squared error over latents. Both of these are arbitrary units that can lead to the forward or backward losses dominating. It's also unclear why both of these losses are necessary and not just the inverse one.*

*Equation (7) is similarly dimensionally incorrect. Again, units are arbitrary and for all we know the joint loss could be ignoring the dynamics loss or the adversarial loss.*

*The use of a GAN and the corresponding loss is unjustified. As the authors acknowledge, if the dynamic loss is very small, then the system should already work. They argue that finding that parameterization without the adversarial loss is challenging. There's a difference between using the right loss and finding the right way to optimize it, but here it seems that the right loss is being modified for optimization purposes. Is the adversarial loss something we really want to minimize or just something that helps find the best dynamics loss? What if we remove the adversarial loss after being close to convergence? What if we use multiple restarts or other techniques to help with the optimization of only the dynamics loss? My take from the theory is that the adversarial loss term shouldn't be needed and that the challenging optimization problem should be addressed (rather than modifying the loss).*

*The new ablation experiments are also confusing: If the dynamics loss is the actual driver, and the adversarial loss only helps with finding a good solution, how come that we get almost equally good results when we remove the dynamic loss? Matching the latent distribution shouldn't be enough to have aligned latents. Maybe there's something about the architecture of F that matches the ground truth, so that matching the distribution aligns the latents. This hints towards the adversarial loss actually playing an important role beyond helping with the optimization problem. This is not supported at all by the theory, since matching distributions should result in arbitrary latents and potentially performance of a random system. In fact, given the problems with units, the joint loss might be dominated by the adversarial loss, which would explain this result.*

*It seems like there's something here, but I think more work is necessary to really understand which pieces are necessary in this system and whether there's some sort of adaptation between the experimental setup and F that would explain why distribution matching results in latent alignment, which is not expected (Zhu et al. 2017). Also, the units problem makes the ablation results even harder to interpret: Maybe the dynamics loss is playing a small role in the joint loss, and that's why removing it completely has a small impact.*

---

After much assessment, while I do find this work to be interesting and potentially highly impactful (since they introduce an alternative approach to data-centric one OOD), the final manuscript's assessment is still borderline (the reviewers all mentioned that while they recognize the improvement, they list issues from preventing their full endorsement), and X3P2 still found several issues with the revision (which I do believe can be addressed in due time). While I'm fully confident that with additional work, this paper could have the potential to be an impactful one, I am currently on the side of not recommending it for acceptance for ICLR 2022.

Note to the PC's, that this is a borderline decision. If the PC's want to flip the decision to an accept, and think the post rebuttal issues are small enough, I'll be fine with that. But in any case, I look forward to seeing a further improved version (of the revamped manuscript) published in a journal or presented at a future conference. Good luck!
Decision: Reject