Ignorance is Bliss: Robust Control via Information Gating
Authors: Manan Tomar
Riashat Islam
Matthew E. Taylor
Sergey Levine
Philip Bachman
Venue: NeurIPS 2023 poster (2023)
Keywords: representation learning
mutual information
TL;DR: Learning to gate information (in the input space) in accordance with any given downstream loss

Abstract: Informational parsimony provides a useful inductive bias for learning representations that achieve better generalization by being robust to noise and spurious correlations. We propose *information gating* as a way to learn parsimonious representations that identify the minimal information required for a task. When gating information, we can learn to reveal as little information as possible so that a task remains solvable, or hide as little information as possible so that a task becomes unsolvable. We gate information using a differentiable parameterization of the signal-to-noise ratio, which can be applied to arbitrary values in a network, e.g., erasing pixels at the input layer or activations in some intermediate layer. When gating at the input layer, our models learn which visual cues matter for a given task. When gating intermediate layers, our models learn which activations are needed for subsequent stages of computation. We call our approach *InfoGating*. We apply InfoGating to various objectives such as multi-step forward and inverse dynamics models, Q-learning, and behavior cloning, highlighting how InfoGating can naturally help in discarding information not relevant for control. Results show that learning to identify and use minimal information can improve generalization in downstream tasks. Policies based on InfoGating are considerably more robust to irrelevant visual features, leading to improved pretraining and finetuning of RL models.

Ignorance is Bliss:
Robust Control via Information Gating
Manan Tomar University of Alberta
Riashat Islam McGill University
Matthew E. Taylor University of Alberta
Sergey Levine University of California, Berkeley
Philip Bachman Microsoft Research Montreal
Abstract
Informational parsimony provides a useful inductive bias for learning represen- tations that achieve better generalization by being robust noise and spurious correlations. We propose information gating as a way to learn parsimonious repre- sentations that identify the minimal information required for task. When gating information, we can learn to reveal as little information as possible so that a task remains solvable, or hide as little information as possible so that task becomes vnsolvable. We gate information using differentiable parameterization of the signal-to-noise ratio, which can be applied to arbitrary values in network, e.g., erasing pixels at the input layer or activations in some intermediate layer. When gating the input layer, our models learn which visual cues matter for a given task. When gating intermediate layers, our models learn which activations are needed for subsequent stages of computation. We call our approach InfoGating. We apply InfoGating to various objectives such as multi-step forward and inverse dynamics models, Q-learning, and behavior cloning, highlighting how InfoGating can naturally help in discarding information not relevant for control. Results show that learning to identify and use minimal information can improve generalization in downstream tasks. Policies based on InfoGating are considerably more robust to irrelevant visual features, leading to improved pretraining and finetuning of RL models.
Original Input
Masking Encoder
Masking Decocer
Regular
Downstream
Output Mask
InfoGGated Input
Figure 1: InfoGating at the input layer predicts a mask conditioned on the original input and then processes the masked input to compute representations that minimize downstream loss. The mask is encouraged to remove much of the original input possible without hurting performance the downstream task. The mask applied by taking a convex combination the input and Gaussian noise, weighted by the mask.
Introduction
Pretraining models on large, diverse datasets and transferring their representations to downstream applications is becoming common practice in deep learning. Such representations should capture useful information from available data while ignoring irrelevant features and noise [7] J. For instance, an object recognition model may achieve high accuracy on its training set while strongly relying on "spurious" correlations between background features and important objects. When transferred to
"Corresoondence manan. tomar@gma i....amii.ooommmmm...mmmmmmmmmm. com
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

new data with the same objects but a different distribution of backgrounds, performance may suffer. Similarly, a policy for vision-based robot control [23 10 may perform well in its original training environment, but fail catastrophically when transferred to a new setting (28) where all task-relevant features are the same but minor background features differ from training.
A promising principle for encouraging robustness to out-of-distribttion observations is informational parsimony: learning features that contain minimal information (39] about the input while still containing enough information to solve the task(s) of interest. Such regularization is usually based on imposing an information bottleneck (IB) [38] on the network that tries to minimize the amount of information flowing from the input, through the bottleneck, to the output prediction. Although IB approaches have been beneficial in some cases, adding an information bottleneck at the penultimate step of computation does little to prevent overfitting in the preceding steps of computation, which typically comprise the overwhelming bulk of the model. To that end, we consider restricting the information used throughout computation rather than just restricting the information that comes out of a computation.
In this paper, we propose learning input-conditioned functions that gate the flow of information through model. For example, we can consider what happens when one inserts an IB near the beginning of model's computation rather than near the end. In this case, we train the information gating functions to minimize information flow from the input into the rest of the model while still permitting the model to solve the task(s) of interest. We implement this gating using a differentiable parameterization of the signal-to-noise ratio. As the noise level goes up, the signal level goes down, and hence the model displays reduced dependencies on the input. Gating functions can be learnt in conjunction with any downstream loss corresponding to some task of interest. For example, the downstream loss could be standard contrastive loss for self-
learning or an inverse dynamics loss when learning representations for reinforcement learning (RL). In both cases, we can learn gating functions that reveal minimal information used to optimize the loss or learn adversarial gating functions that will remove any information that can be used to optimize the loss. We primarily focus InfoGating representations used for learning control policies in RL. A natural focus of the RL paradigm is on learning a mapping from observations to actions, thus discarding most information not useful for control. We bring forth such notion of only capturing what the agent can affect through InfoGating. Using background distractors and multi-object interactions as a form of noise/irrelevant features, we see that InfoGating is able remove almost all of the irrelevant information from the pixels, leading to better -uttof-distribution generalization in the presence of noise and better in-distribution generalization in the presence of multiple task-irrelevant objects.
Our main contributions are as follows: 1. A general purpose, practical framework for informational parsimony called InfoGating that can restrict information flow throughout& computation to learn robust representations. 2. Qualitative analyses on the properties of the gating functions learned with InfoGating that show they are semantically meaningful and enhance intrepretability. 3. Quantitative analyses of applying InfoGating in the context of various downstream objectives which show clear benefits in terms of improved generalization performance.
2 Related Work
Information Bottleneck. Much prior work in learning robust representations has stemmed from the idea of imposing an information bottleneck (IB) [38 39. Imposing an IB involves maximizing performance on the downstream task, while removing as much information about the input as possible from a network's 's internal representations. Typical approaches based on Variational IB (2) achieve this by learning stochastic representations that are constrained to be close to standard Gaussian prior. Variants of dropout [36] 19 can have a similar effect (i.e., adding noise to representations used in the network) but without explicitly maximizing nois//minimizing information in the representations. While IB is typically applied to later stages of a network, InfoGating can be seen as enforcing an IB at the input itself, which is less explored in the IB literature. Like Variational IB methods, and unlike typical dropout- based methods, InfoGating learns how much noise can be added without hurting performance on the downstream task(s). Information Dropout [I] perhaps the most similar prior work to InfoGating. Information Dropout imposes IBs at multiple points in network's 's computation and the IBs are implemented using techniques based on variational dropout [19. A key distinction between InfoGating and Information Dropout is that InfoGating works post-hoc. With InfoGating a model can consider a computation it has already performed and predict ways in which it could produce the same result more parsimoniously
2

---

Masked Image Modelling. There has been a lot of recent work in self- learning where the input masked and the model is tasked with reconstructing the missing bits [5 4T Most of these works operate with fixed, random masking scheme, instead of learning it directly through the downstream loss. The amount of information masked also fixed a priori whereas we seek to mask as much of the input as possible, while still retaining enough information to solve the task. Some related ideas include learning adversarial augmentations for producing new views of an input during contrastive learning [37], while also learning masks that essentially segment the input space [35].
Representation Learning in Reinforcement Learning. Prior work on learning representations for RL includes temporally contrastive learning [33] 27], one-step inverse models [32]. learning through next state reconstruction [14], bisimulation metries [9], and many more. Some of these objectives have been shown to be useful in learning robust representations 40), and we see our method as complementary to these approaches since InfoGating works with any downstream loss function, without needing knowledge of specific values like reward, next state, etc. For instance, bisimulation and task-informed abstractions [ learn to compress the observation space using reward information, while recent work on learning Denoised MDPs [42] aims at regularizing next state reconstruction using a variational objective. On the other hand, InfoGating remains agnostic to choices of downstream objectives.
Background
This section describes the primary downstream losses use with InfoGating in this paper, namely an InfoNCE based contrastive loss and a multi-step inverse dynamics loss.
Mutual Information Estimation via InfoNCE. Given and as two random variables, their mutual information can be defined as the decrease in uncertainty when observing x given z, compared to just observing x: 7(x, z) H(x) H(x z). where H is the Shannon entropy. InfoNCE [31], based on Noise Contrastive Estimation [3]. computes lower bound on T(z z²), where z and z? are two representations of the input produced by some encoder fx). Specifically, the bound is optimized by discriminating "positive" and "negative" pairs:
LINTONCE = Ez
log evixt, x²) + Ex 2- ext(xt,
where v is a pairwise, scalar-valued function of the representations and Z is a batch of "negative samples". Typically, self-supervised contrastive learning uses data augmentation such as random cropping and color jittering to define two augmented views (z² z²) of a given input as "positives" while views z of other inputs are treated as "negatives". The InfoNCE objective then encourages the representations of positive views of x to be close (i.e., (z z²) is high), while pushing apart the negatives (i.c., (z(z),
is low) [4 5).
Multi-Step Inverse Dynamics Models. Multi-step inverse dynamics models X predict the action(s) that took an agent from some observation X; to some future observation X,+ attempting to learn useful representation of the observations. This resembles learning goal-conditioned policies through relabelling future observations as achieved goals. Although the trivial case of k = 1 (a one-step model) does not capture long term dependencies, recent work has shown that multi-step models are able to capture more information that may be useful for controlling the agent inverse models can be used to learn representations by simply predicting the actions (a.,..., a-+ conditioned on x, and X++k- Actions can be predicted using standard max likelihood objective or with a contrastive reconstruction objective which maximizes the InfoNCE- [(xx,, X(+k), (at,... .... 1)). The learning objective in this case looks as follows:
(22]. Multi-step
lower bound on
InfoNCE 2(4k)
where 2, and Z are representations computed from x; and X and Ytu+ is a representation computed from (ar, -1) The negative samples in this case come from randomly sampling action sequences of the appropriate length from the agent's collected experience. the simplest case, the representation Yra+k-1 may depend only on a..
3

---

4 Information Gating
We present two approaches to InfoGating that we describe as cooperative and adversarial. Cooper- ative InfoGating involves keeping as little information as possible to obtain good performance on the downstream task. Adversarial InfoGating involves removing as little information as possible to preclude good performance on the downstream task. We include both approaches in this paper since identifying minimal sufficient information (cooperative) and identifying any useful information (adversarial) may have different affects depending on the downstream loss and the application domain.
The InfoGating approach has two major components. The first is an encoder of the input x, [xx). The second is an information gating function, ig(x (see Figure Tj. In principle, we can gate information passing through any layer in the representation network that computes the encoding = fxx). The (x) function provides continuous-valued masks (values in [0, 1J) that describe where to erase information from the computation graph for f(x). The shape/size of the output of ig(x) will depend on where we want to gate information. For instance, if we wish to gate information in the input pixel space, (z))) masks are the size of the image. In general, the goal of (g(x)) is to erase as much information from fx) as possible without hurting task performance.
4.1 Cooperative InfoGating
We primarily focus on cooperative InfoGating in this paper and we consider two cooperative variants: gating in the input space or in the feature space.
InfoGating in Input Space. We apply InfoGating to input x by taking simple convex combina- tion of x and random Gaussian noise E. The combination weights are given by ig(x):
ig(x) @ x x (1 ig(x)) e
where x denotes the info-gated input and @ denotes clement-wise multiplication. An all-zero mask corresponds complete noise (i.e., erasing all information from the input), while an all-one mask corresponds to keeping the original input. The function (z( is learnt using the same downstream objective that is used to learn f (x). We encourage masks to remove information from the input by
minimizing their LI norm. This tends to produce sparse masks due to properties of the L1 norm 30) The overall objective for learning with InfoGating is
Lig.t = Lusk (f(x))) + A [lig(x)]l1.
(3)
where Lusk refers any objective through which a useful representation z can be learnt. Note that when doing InfoGating, we minimize the downstream loss for the masked input x (first term), instead of the original input x, while masking out as much of the input as possible through the L1 penalty (second term). The X coefficient is a hyperparameter that controls how much of the input is masked, In principle, any loss function can be used in conjunction with InfoGating and we consider multiple downstream objectives: contrastive learning of dynamic, Q-learning, and behavior cloning. We describe the overall procedure in Algorithm [ which assumes contrastive multi-step inverse dynamics modeling as the downstream objective.
InfoGating in Feature Space. As an alternative directly gating the input pixels, we can also consider a variant where f(x masked instead of the input. Consider the following masking:
Masking
Output Nask vifo)
Downstream Look
Rogular
Endoder
2 (g(x)) (1-1g(x))
where (x)). (4)
The training objective remains the same as in the pixel- level case, i.c.. minimize the downstream loss Lusk (z for the masked representation 2 while masking as much of z as possible (Figure 2) This version is closely re- lated to the deep variational information bottleneck (VIB), which minimizes the KL divergence between the (stochas- tie) representation z and a prior distribution, typically chosen to be a standard Gaussian. Roughly, the values in ix(x) can be seen as specifying how many steps of for- ward diffusion to run in Gaussian diffusion process initiated at z, where values near zero correspond
Figure 2: Feature InfoGating, gates rep- resentations produced by the encoder. The gating network and the encoder are jointly trained to minimize the downstream loss, while the gating network is also encouraged mask as much of as possible.

---

Algorithm InfoGating Pseudocode
I: Input encoder f. masking net ig
2: for (x, X,+2; ar) E loader do
3: Xt, X++h = aug(x:), aug(Xe+k)
# random crop augmentation
4: x/" x ix(x;); ig(*++*)
get infogates
5: x ig()))) o x xg(x;)) oe
infogate current state # infogate future state
6: x14+ = (g(+++)) 3 X1+2 (1 -
ig
de
7: Zt, Ze+: = fxxx), f fx
# get encodings
8: if Cooperative then
9: Run Adam update overall loss:
L InfoNCE(2e, Zetk; a+) [g(())) [
10:
else if Adversarial then
11:
Run Adam update on masking loss:
Lig -nfoNC Zt+k, ae) - 1lig(x;)]lı 1|ig(x:+*)ll=
12:
Run Adam update on encoder loss:
Ly InfoNCH(2,, Z 8t)
13:
end if
end for
rumning more steps of diffusion and thus sampling from distribution that is closer standard Gaussian in terms of KL. divergence [18]. Thus, like VIB methods, we aim to minimize J(x, 2), while maximizing the task performance of However, feature space InfoGating uses different optimization objective and parameterization compared to existing VIB approaches. We leave similar variations of InfoGating where arbitrary layers in computation graph are gated for future work
4.2
Adversarial InfoGating
The versions of InfoGating discussed above are cooperative objectives, where the gating network and the encoder work together lower the overall loss. This leads finding representations that capture the minimal sufficient information for a given task. We can also turn this formulation on its head, yielding adversarial objective. In this case, the gating network is tasked with discovering masks that, when used by the encoder, lead to maximizing its loss. Instead of encouraging the masks erase as much of the input as possible, the adversarial objective encourages the masks remove as little information as possible while minimizing the encoder's performance on the downstream task. We write the adversarial objective for ig())) and (ax
Liy = Lusk (f
[|ág(x)]lı,
Lusk
(5)
This gives rise to a min- max objective w.r.t. the encoder and the gating network while cooperative InfoGating corresponds to min-min objective (see lines 11-12 in Algorithm
5 Experiment Setup
To understand whether InfoGating can consistently focus on the minimal possible information required for control, we test generalization performance when using InfoGating with three different downstream control objectives: 1) contrastive dynamics models (both inverse and forward models),
2) Q-learning (TD-based critic updates), and 3) behavior cloning. We choose these three objectives since they remove irrelevant information to varying degrees by default - for example, dynamics models can capture lot of fask-irrelevant information, while behavior cloning models only contain information that is useful for predicting the optimal policy. Through our experiments. we evaluate whether all three of these objectives can benefit from InfoGating, and to what degree
meant to
'Constructing differentiable upper bounds on compute cost one possibility, similar to DARTS 24]. 'Since the multi-step inverse model produces the best results, we use the same for testing the feature InfoGating and adversarial InfoGating versions well.
5

---

Table 1: Multi-step Inverse Dynamics with InfoGating, Comparing performance in the presence of noisy distractors, report returns achieved by policy produced by behavior cloning on top of pretrained repre- sentations in cheetah-run. Extended results are provided in Appendix C The "inv" model is our baseline with pretraining via step inverse dynamics. The "w/ Rand" model adds random info gating during pretraining and the "w/ IG" model adds learned info gating during pretraining (this is our method), The easy/medium/hart settings add different levels of distractor noise, Results are for seeds each, with mean and std. dev, reported.
cuse
IQI.
DRIMI
Tnv
Inw w/ Drepout
Inv w/ VIB
Inv w/ RCAD Inv w/Rand
Tav w/ IG
expert
easy
10.7
0.1
42.6 36.3
11.79 + 3.5
97.1 17.9
21.0 15.6
9.1
medium
29.2 + 28.8
1.0 + 0.6
29,5 31.9
73.3
+ 12.3
38.9 16.9
19.9 + 24,8
7.2 +: 6.12
97.0 +: 5.7
hard
8.8 + 1.9
13.8 + 7.9
2 0.6
5,9 5.7
34.0 + 19.4
1.4 + 0.4
4.2 + 0.7
44,8 + 18.4
overall
16.2
30.3
106.0
medium
easy
42.6 +: 26.4
2.6 +: 0.2
45.6 23.2
28.0 23.6
113.5 + 11.8
29.6 +: 24.5
42.0 +: 31.8
133.0 12.8
medium
31.4 +:
86.6 55.0
2
39,9 39.4
44.0 22.8
2.2
85.4
13.2
92.0 + 35.2
hard
15.2 + 7.7
3.1 + 1.5
10.7 + 6.0
18.8 + 15.2
6.1 1.8
4.2 + 1.6
5.3 + 1.9
3.0 1.43
overall
29.7
30.7
19.4
28.9
54.5
12.6
44.2
76.0
medium- expert
easy
42,4 +: 31.4
12,9 +: 12.7
25.5
45.7 + 7.4
130.4 4 32.0
51.8 + 38.1
10.4 4 9.9
158.0 H: 21.8
medium
39.6 +: 24.7
22.3 + 13.6
14.2 11.2
15.6 +: 8.1
52.2 26.0
5.4 + 3,9
24.8 34.4
89.2 7.0
hard
10.2
5.2 2.6
16.2 + 14.4
+: 38.7
+: 3.4
38.8
overall
30.7
13.4
14.3
25.8
71.2
22.1
12.7
95.3
We test 1) and 2) on the offline visual D4RL domain [25 and 3) on the Kitchen [12] manipulation domain. InfoGating can also be used in non-RL settings like self-supervised learning, as discussed in Appendix A We now describe the details of each use of InfoGating and the corresponding experimental results. Hlper-parameters are listed in Appendix E
5.1 Inverse Dynamics Models
We use multi-step inverse models as our primary downstream loss due to its ability to recover the latent state effectively [17 (22) Consider the contrastive loss based on InfoNCE described in Eq. To apply InfoGating, we mask both the current observation x, and the goal observation Xt+ using the masks from ix( and ig (x++x) respectively. We then process both masked observations through
the encoder to compute the corresponding representations z = fx((") and 2 = fxx(+)) These are then optimized through a loss similar to Eq. 3
L InfoNCE((z"" 19, 2) a.) + A (llig(xx)]11 + [[ig(x+++)1l1)
(6)
Note that both current and future observation masks are penalized through the L1 term. We test the inverse model with InfoGating on datasets consisting of offline observation-action pairings. The datasets contain pixel-based observations with video distractors playing in the background [25). We
first pretrain the representations with the InfoGating objective in Eq. 6 and then perform behavior cloning (BC) using 2-layer MLP over the frozen representations. We have tried replacing BC- based evaluation with, e.g., Q-Learning baaed evaluation, which leads to similar relative scores. We report scores for BC-based evaluation for easier reproducibility and reduced dependence hyperparameters.
We work with three levels of distractor difficulty: easy, medium, and hard. These levels correspond to the amount of noise added to the observations. At evaluation time, a noise-free observation space is used, thus creating an out-of-disttibutton shift in the distractor. We use random cropping as a pre- prepprocessing step for the observations in all our experiments. We compare inverse dynamics with InfoGating to six baselines: control specific methods, i.e. 1) IQL (2 2) DRIML 27. and 3) inverse dynamics without InfoGating (the standard formulation from Eq. T.. and regularization specific methods that work top of inverse dynamics by adding 4) Dropout, 5) VIB bottleneck, 6)
adversarial learning (RCAD) 34 and 6) random masks rather than learnable masks as in InfoGating.
See Table 16 for results.
We observe that InfoGating leads to considerable performance gains over all baselines. Although adding random masking is better than no masking, learnable masks lead to much better performance.

---

Since random masking provides similar benefits to data augmentation. some performance improve- ment over the standard inverse model expected. However, learning masks clearly does more than just data augmentation. Moreover, InfoGating in the input space also performs more robustly than methods that add regularization in the model weights (Dropout), an intermediate layer (VIB) or at the output (adversarial samples). This result highlights how valuable removing information in the input space can be for certain applications (as can be visualized in Figure [. Finally, these results also indicate that minimizing the downstream objective (in this case the inverse model loss is not always sufficient for learning robust representations. By incorporating the InfoGating objective, we add an inductive bias that encourages the model to focus on the most minimal, useful information without additional supervision. This inductive bias towards informational parsimony empirically produces more powerful representations.
5.2 InfoGating for Stabilizing Q-Learning
Our prior experiments focused on pretraining representations and then learning regression- -basd policies, while keeping the representation fixed. In this section, we ask if the same conclusions hold we train the representation end-to-end with Q-Learning loss instead. We use the DrQ 200 loss as the downstream objective for this variant of InfoGating. Our InfoGating formulation for Q-Learning essentially amounts to doing standard DrQ training with observations that are masked by the gating network. The gating network learning is driven by the Q-Learning loss directly as follows:
Lig.0 = (Qo(x an) r(x,, ar) - (xx++11(++1) 2
(7) where Qo and Q/ denote the current and target Q networks respectively, while r(x,, ar) is the obtained reward. In Table 2 all three distractor settings, while with InfoGating we see significant improvements in performance. This result shows how using minimal information in the input space can be beneficial in stabilizing TD-based critic training.
[1ig(x))]1,,
we observe that the base DrQ algorithm is quite prone to failure for
Table 2: InfoGating for Q-Learning. We use DrQ as the base Q-Learning algorithm and add InfoGating to it. The experimental setup is the same as in Table T6
Results are for an expert policy level. medium
case
easy
hard
overall
DrQ
5.7 + 2.3
29.8 + 20.5
3.4 + 0.8
12.9
DrQ + IG
63.4 22.6
52.0 = 6.2
5.3 # 1.9
40.2
5.3 Finetuning General Representations with Behavior Cloning
Recent work 29 26 has investigated learning general representations from large datasets involving diverse object interactions and different varieties of tasks. A natural question arises when using such pretrained representations for downstream tasks how should the representation be fine-tuned so that only the relevant object and task features for the given task are used for learning the task's 's policy? InfoGating can be seen as a natural way to extract only information pertaining the downstream task. Having tested InfoGating extensively on noisy environments, we now study whether there are similar benefits when there is no explicit noise present in the environments, but there are multiple objects/pixel components which could act as potential distractors. We choose five different tasks from the Kitchen 12 environment to test this hypothesis.
Specifically, given pretrained features, we test whether fine-tuning them using InfoGating is more powerful than fine-tuning with only behavior cloning (BC) loss. We test two pretraining variations here: one corresponding to ImageNet features and the other corresponding to CLIP features. Both are 1) trained on ImageNet data, alongside language clippings for CLIP and 2) fine-tuned using a behavior cloning (BC) policy with a 2-layer MLP attached on top of the pre-trained encoding. Table shows normalized success rates (each out of maximum score of 100) for both BC and BC with InfoGating. We consistently see that InfoGating is able to mask out most pixel-level information except the robot gripper and the objects being manipulated (see Figure 3 for visualizations of the leant masks), thus leading to strictly better success rates than the baseline BC policy.
6 Ablations
We study various design choices involved in InfoGating and include the most important ablations in this section. Further ablation results are provided in Appendix D
7

---

Table 3: Finetuning with Bchavior Cloning Kitchen-v3. We compare fine-tuning performance of BC with without InfoGating. Representations are pre-trained either on ImageNet labels or through CLIP training. In both cases, InfoGating learns to remove irrelevant objects in the environment surroundings and leads to consistently higher performance.
ImageNet
CLIP
BC w/ IG
BC
BC
w/ IG
BC
knobl-on
17.6 1 3.1
7.3 1 4.2
13.3 1 1.49
11.3 1 0.94
Idoor-open
9 9 5.3
5 4 1.9
11.0 4.8
4.6 + 1.49
light-on
19.6 7.7
11 = 3.6
24.0
12.8
11.0
3.0
sdoor-open
61.6 5.0
36.6
+ 16.6
61.6 - 8.7
42,0
5.9
micro-open
13.0 + 7.8
4.3 + 2.9
12.0 + 5.6
5.5 2.1
overall
24.1
12.8
24.3
14.8
Figure Visualizing Masks Produced by InfoGating. Masks improve during training on the Kitchen-sdood -dooo-----o------ open-v3 task (left to right corresponds to increasing training steps). InfoGating is able to learn semantically meaningful masks, even beyond the planar control domains such as the chectah-run task.
InfoGating in the Feature Space. We compare the feature InfoGating version with the Variational Information Bottleneck (VIB), while using the same backbone encoders for both methods (see Table 4). Both VIB and feature InfoGating lead to some performance improvement over the base architecture, but do come close the score for the pixel InfoGating version. This is potentially due to not removing background distractor information in the first layer itself. Interestingly, feature InfoGating almost matches the pixel InfoGating performance for the hard distractor case. Note that there exist variations such as change in body color, camera zoom, etc in the hard distractor case, that remain even after pixel-level InfoGating. We suspect that feature InfoGating is able to mask out such features but is hurt by the background distractor instead, hence falling to the same performance the pixel-level variant.
Table 4: InfoGating in the Feature Space. We compare the Variational Information Bottleneck with InfoGating feature space. The experimental setup is the same in Table T6 Results for expert policy level.
easy
medium
hard
overall
VIB
97.7 + 32.2
86.2 + 24.8
11.0 + 5.2
64.9
feat. IG
71.4 + 44.5 76.7 + 21.3
58.9 + 28.3
69.0
Adversarial InfoGating. We test the adversarial InfoGat- ing algorithm using multi-step inverse model loss, just how
InfoGated Observation
Original Observation
as we did for cooperative InfoGating. the adversarial masks tend to hide the agent body. Some additional image content also erased, since a precise silhouette of the agent's 's pose would be highly predictive of the agent's pose. To test how useful this kind of mask- ing is, we simultaneously train separate encoder over the reverse the adversarial masks. the adversarial process hides the robot pose successfully, then the reverse mask should contain maximal information about the agent and minimal information about the background. We in- deed see that the adversarial InfoGating model performs similarly to the cooperative version, thus verifying that it an equally viable approach for learning parsimonious representations (see Table 5).
Figure Left. Original distractor based image. Right. Leartt adversarial mask over the original image.
Effeet of Mask Sparsity. We test what range of À values leads to improved performance (Figure 5) As is expected, when À is too high, the entire input masked and no useful information is captured.

---

Table 5: Adversarial InfoGating, We compare differences in the representations learnt with Cooperative
Adversarial InfoGating. The experimental setup the same as in Table Results are for an expert policy level. overall
case
easy
medium
hard
cooperative
176.2 +: 9.1 adversarial
97.0 +
44.8 18.4
105.9
202.3 +: 2.64
84.8 41.2
4.6 1.7
96.8
Similarly, when À is too low, none of the input is masked and we recover the performance of the base model, one that directly minimizes the downstream loss.
Mixing Original and Masked Inputs to En-
coder. Since the masking network and the en- coder process different kinds of input (unmasked
1e=1
ing if the encoder can be better regularized by forcing to be able to process both masked and unmasked images. We can also encourage the encoder to be robust to changes the input/- mask mapping learned by the masking network by occasionally swapping masks between im- ages in a batch while training the encoder. We test both of these separately (Table (6) and ob- serve that training the encoder on masked and unmasked inputs provides strong improvements the base InfoGating model. This also has the added benefit of not using the masking net- work at evaluation time, and only using the coder, which has now been implicitly trained to be invariant to the masked and unmasked inputs. Note that we use this mixed input InfoGating as the default for all our experiments.
and masked input, respectively), is worth ask-
performance
300
X=
5e-2
A
200
1.0
100
mask scaling X
Figure 5: Scaling À leads to different masks which reveal different amounts of information. For small X, the masked observation still contains distracting/non-saliii information, thus hurting performance. Similarly, when masking is aggressive, much information lost and performance goes down.
Table 6: Mixing Masking and Unmasked Inputs. Experiment setup follows Table version is marked in gray Results are for an expert policy level.
Default InfoGating
easy
medium
hard
overall
w/ mix input
176.2 1: 9.1
97.0 1: 5.7
44.8 1: 18.4
106
shuffic mask
170.0 + 31.6
54.6 + 22.3
9.8 + 10.2
78.1
w/o mix input
119.4 + 15.6
18.4 + 17.5
10.5 + 7.4
49.4
Limitations and Future Work
Although InfoGating helps learn robust representations, it does not recover the same performance for all distractor levels as in the case when no distractors are present. Ideally, we should be able to learn masks that fully remove the background information and thus recover the distractor-fret performance. This might be down to two reasons. First, the masks may leave room for noise information to escape around the edges of the object/agent they mask. Second, the encoder that processes the masked images should be robust to slight variations in the masking patterns between training and evaluation samples. practice, using a UNet for more accurate masks and training the encoder on a mix of masked and unmasked inputs helped remedy these issues quite well, but there still remains room for improvement. This motivates the idea of efficiently InfoGating at multiple layers, without having separate masking networks for each InfoGating layer.
Furthermore, our initial exploration of InfoGating can be viewed as a step towards learning object- centric representations without any explicit notion of objects forced into the model architecture or learning objective. For example, settings where an agent's actions may affect multiple objects, and we want predict future observations conditioned on the agent's actions, the masks produced by InfoGating will need to reveal some information about (i.e., "look at") each object. The pursuit of informational parsimony should discourage masks from revealing more of each object than necessary.

---

8 Conclusion
We present InfoGating, a wide-ranging method to learn parsimonious representations that are robust to irrelevant features and noise. We describe two different approaches to InfoGating: cooperative and ad- versarial, while demonstrating that the gating can be learnt both the input space or any intermediate feature space. We apply InfoGating multiple downstream objectives including multi-step inverse dynamics models, Q-Learning, and behavior cloning. InfoGating produces semantically meaningful masks that improve interpretability, and leads to consistently better performing representations, in terms both out out-of-distribution (visual D4RL distractor noise) and in an-distribution (irrelevant Kitchen objects) generalization.
Acknowledgments
The authors would like to thank Charlie Snell, Colin Li, and Dibya Ghosh for providing feedback on earlier draft, and Katie Kang for suggesting the paper title. Part of this work has taken place in the Intelligent Robot Learning (IRL) Lab at the University of Alberta, which is supported in part by research grants from the Alberta Machine Intelligence Institute (Amii); Canada CIFAR Al Chair, Amii; Compute Canada; Huawei; Mitacs; and NSERC.
References
[1] Alessandro Achille and Stefano Soatto. Information dropout: Learning optimal representations through noisy computation. IEEE transactions on pattern analysis and machine intelligence, 40(12):2877-2905 2018.
[2] Alexander Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.004100 2016.
[3] Mahmoud Assran, Mathilde Caron, Ishan Misra, Piotr Bojanowski, Florian Bordes, Pascal Vincent, Armand Joulin, Mike Rabbat, and Nicolas Ballas, Masked siamese networks for label-efficient learning. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXXI, pages 456-473, Springer, 2022.
[4] Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. Advances in neural information processing systems, 32, 2019.
[5] Ting Chen, Simon Korblithh, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1,597-1607. PMLR, 2020.
[6] Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 15750-15758, 2021.
[7] Thomas G. Dietterich, George Trimponias, and Zhitang Chen. Discovering and removing exogenous state variables and rewards for reinforcement learning. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference 2018, Stockholmssüsssan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 1261-1269. PMLR, 2018. /ttp:tt///prceeeeeee htm
Machine Learning. ICML
m1 mlr.pree wwt.ppess//bb//ieeeeeeee
[8] Yonathan Efroni, Dipendra Misra, Akshay Krishnamurthy, Alekh Agarwal, and John Lang- ford. Provable rl with exogenous distractors via multistep inverse dynamics. arXiv preprint 100.v221 2021.
[9] Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing, 40(6):1662-1714 201
[10] Chelsea Finn, Ian J. Goodfellow, and Sergey Levine. Unsupervised learning for physical inter- action through video prediction. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Is- abelle Guyon, and Roman Garnett, editors, Advances in Neural Information Processing Systems
10

---

29: Annual Conference on Neural Information Processing Systems 2016. December 5-10, 2016. Barcelona, Spain, pages 64-72, 2016. URL /tL/:ptp///ttcceeingggssso newraisss................. cc/
html
[11] Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning, pages 3480-3491. PMLR, 2021.
[12] Abhishek Gupta, Vikash Kumar, Corey Lyneh, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. arXiv preprint arXiv: 1910. 1956, 2019.
[13] Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the thirteenth international con- ference artificial intelligence and statistics, pages 297-304. JMLR Workshop and Conference Proceedings, 2010.
[14] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.0160 2019.
[15] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16000-16009, 2022.
[16] Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint 190. /2261, 2019.
[17] Riashat Islam, Manan Tomar, Alex Lamb, Yonathan Efroni, Hongyu Zang, Aniket Di- dolkar, Dipendra Misra, Xin Li, Harm van Seijen, Remi Tachet des Combes, al. Agent- controller representations: Principled offline rl with rich exogenous information. arXiv preprint arXiv:2211.00164, 2022.
[18] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances neural information processing systems, 34:21696-21707, 2021.
[19] Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparame- terization trick. Advances in neural information processing systems, 28, 2015.
[20] Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004.13649, 2020.
[21] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-leamiig. arXiv preprint arXiv:2110.06100 2021.
[22] Alex Lamb, Riashat Islam, Yonathan Efroni, Aniket Didolkar, Dipendra Misra, Dylan Foster, Lekan Molu, Rajan Chari, Akshay Krishnamurthy, and John Langford. Guaranteed discovery of controllable latent states with multi-step inverse models. arXiv preprint arXiv:2207.0822. 2022.
[23] Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. Mach. Learn. Res., 17:39:1-39:40, 2016. URL /ttppttp://miir.or
papers/v17/15-525 html
[24] Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: differentiable architecture search.
In International Conference on Learning Representations (ICLR), 2019.
[25] Cong Lu, Philip J Ball, Tim GJ Rudner, Jack Parker- Parker-Holder, Michael A Osborne, and Yee Whye Teh. Challenges and opportunities in offline reinforcement learning from visual observations. arXiv preprint (xxxx:22266047777 2022.
[26] Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. Vip: Towards universal visual reward and representation via value-implicit pre-training, arXiv preprint arXiv:2210.00030, 2022.

---

[27] Bogdan Mazoure, Remi Tachet des Combes, Thang Long Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. Advances in Neural Information Processing Systems, -6966 2020.
[28] Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S. Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real-world environments through meta-reinforrement learning. In 7th International Conference Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL https //apenreview net forum? i Hdd-HyztsoCsYY
[29] Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A universal visual representation for robot manipulation. arXiv preprint arXiv:2203. /2601, 2022.
[30] Andrew Y Ng. Feature selection, 1 1 vs. 1 2 regularization, and rotational invariance. In Proceedings of the twenty-firs international conference on Machine learning, page 78, 2004.
[31] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, arXiv preprint orXiv:1807.0377 2018
[32] Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity Curiosity-driven driven exploration by self-supervised prediction. In International conference on machine learning, pages 2778- 2787. PMLR, 2017.
[33] Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine, and Google Brain. Time-contrastive networks: Self-supervised learning from video. In 2018 IEEE international conference on robotics and automation (ICRA), pages 1134-1141. IEEE, 2018.
[34] Amrith Setlur, Benjamin Eysenbach, Virginia Smith, and Sergey Levine. Adversarial unlearning: Reducing confidence along adversarial directions. arXiv preprint arXiv:2206.01367, 2022.
[35] Yuge Shi, N Siddharth, Philip Torr, and Adam Kosiorek. Adversarial masking for self- supervised learning. In International Conference Machine Learning, pages 20026-20048 PMLR, 2022.
[36] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-198, 2014.
[37] Alex Tamkin, Mike Wu, and Noah Goodman. Viewmaker networks: Learning views for unsupervised representation learning. arXiv preprint arXiv:2010.07432, 2020.
[38] Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 ieee information theory workshop (itw), pages 1-5. IEEE, 2015.
[39] Naftali Tishby, Fernando C Pereira, and William Bialek. The information bottleneck method. arXiv preprint physics/0004057, 2000.
[40] Manan Tomar, Utkarsh A Mishra, Amy Zhang, and Matthew E Taylor. Learning representations
for pixel-based control: What matters and why? arXiv preprini arXiv:2111.07775. 2021.
[41] Pascal Vincent, Hugo Larochelle, Isabelle Lajoie, Yoshua Bengio, Pierre-Antoine Manzagol, and Léon Bottou. Stacked denoising autoencoders: Learning useful representations in deep network with a local denoising criterion. Journal of machine learning research, 11(12), 2010.
[42] Tongzhou Wang, Simon S Du, Antonio Torralba, Phillip Isola, Amy Zhang, and Yuandong Tian. Denoised mdps: Learning world models better than the world itself. arXiv preprint arXiv:2206. 15477, 2022.
12

---

A InfoGating with SimSiam
Random masking has been used as an SSL objective learn useful features for downstream object detection/classifcation Can InfoGating be used as a general strategy to recover similarly useful features? We test this by applying InfoGating with self-supervised representation learning algorithms, in particular SimSiam
[6].
The SimSiam 66] objective uses cosine similarity loss between the representation of view z and a predicted output of the representation of the second view p p(z2). where the function called the predictor. This is place of using the InfoNCE loss define contrastive pairs for z and 2 Both methods lead to similar performing representations with the main difference being that SimSiam does not require negative examples. The overall SimSiam loss can be written as:
SimSSam Z1, Zg) D(p1. 22) 2 D(p2, 21) /
where D denotes the cosine similarity function:
P
D z) =
|pl2
|z||2"
Note that D not a symmetric quantity as it uses a stop gradient operation on its second argument Having described how SimSiam works, we can now go into how InfoGating is applied alongside it. Given two augmented views of the input image and X2, we mask both views to get 10 and 2g then process them through the encoder to compute the SimSiam loss L(2)" 2."). Finally, we also add the original SimSiam loss, i.e. one over the unmasked inputs to the overall objective:
L = LsimSium (z Z2 ig + LSimSiam (Z),
+ A (Ilig(x)) IllI (x2g(1))11))
We test this version of InfoGating CIFAR-I0, while evaluating performance Corrupted CIFAR- 10 [16]. Figure shows that InfoGating leads to improved evaluation scores comparison to the base SimSiam method, thus showcasing better robustness to test-time induced corruptions.
SimSiam
SimSiam
InfoGating
Performance
60
40
20
Speckle-Noise Motion-Blur
Pixelate
Frost
Contrast
Brightness
Figure 6: Generalizing Corruptions through InfoGating. We directly compare SimSiam performance w/ and w/o InfoGating on different subsets of the corrupted CIFAR-10 dataset. Adding InfoGating improves robustness to test-time corruptions.
B Visualizing InfoGating Masks
Learning info-gates directly from the downstream loss allows for adapting the mask based on the task hand. On the other hand, fixed random masking scheme offers limited benefits, largely pertaining to enhanced data augmentation. For instance, using InfoNCE loss derived from augmented views of the input with no distractors leads to learning accurate masks that capture the agent pose almost perfectly (see Figure 77. However, when distractor noise is added, the same InfoNCE loss leads to masks that are spread across the input, failing to remove background noise. Switching the loss to multi-step inverse dynamics loss leads to masks that focus on the agent pose and successfully remove the background noise. Similarly, using the forward dynamics loss as the downstream objective, masks do not clearly follow the agent pose in the inverse model case. Similarly, when the
13

---

InfoNCE w/a Distractor
InfeNCE w/ Distractor
Forward w/
Inverse w/
Distractor
Distractor
Figure Effect of Losses InfoGating Masks. InfoNCE w/o Distractor (mean 208.0) learns almost perfect masks, while InfoNCE Distractor (mean score: 10.4) learns masks that are spread across input, thus failing to learn a robust enough representation. Forward dynamics w/ Distractor (mean score: 60.4) and Inverse dynamics w/ Distractor (mean score: 176.2) lead to much more accurate masks and thus lead to better performing representations well.
Figure 8: Visualizing Masks Learnt by InfoGating. STL-10 images alongside the masked image. The masking network gradually learns trace edge boundaries both foreground and background objects.
downstream objective is self-supervised loss that focuses on all parts of the input (as opposed to only information useful to control agent), we that the learnt masks capture the edge boundaries of both background and foreground objects (see Figure S). These different variations highlight the importance of choosing the right downstream loss well as learning InfoGating masks in conjunction with the given loss.
C Extended Results
Throughout the paper, we use contrastive variant of the multi-step inverse dynamics loss. This is implemented by first encoding the current and future/goal observations (which are info-gated) x;" and x into corresponding encodings 2 and 2 We then concatenate the encodings along with the action to form a single encoding triplet. This is then passed through 2-layer MLP to output logits for the energy value corresponding to the given (z ig Z ar) triplet. The actions actually taken by the agent make up for 'positive' triplet and the model is trained output low energy value. Similarly, is replaced by random actions å, to form a 'negative' triplet, for which the model trained output high energy value.
We can consider a similar version that uses forward contrastive loss instead. To that end, we simply encode the current action and observation pair (x"", a ) to 2,") while encoding the future/goal observation x to a similar size vector zitk- We then train these by applying standard InfoNCE over the two encodings. The forward model is usually not enough to encode sufficient information for control and therefore lacks the guarantees that multi-step inverse model has. We this by info-gating both the current and future observation pair, just as in the inverse model. Next we process both of the masked observations through the encoder, and then concatenate the action a with the z
14

---

representation, before projecting it to the same dimension as the goal representation z Finally, we penalize the info- gafo-gates for both current and goal observations as done previously to obtain the following objective:
C InfoNCE(2,, 2,4x) (llig(x))l1 [ig[( II)),
where z.-6 = g(2 a.) is projection function that maps the state and action embedding to space of the same size " Note that in the InfoNCE loss, the negatives now come from sampling different future observation representations from the batch.
Table 7:
InfoGating with Forward Dyanmics Models. Experiment setup follows Table 16 "fwd" is the baseline forward dynamics model and "Twd IG" adds InfoGating. For reference, training forward dynamics model without distractors has a return of 173.6 13.5.
case
easy
medium
hard
overall
fwd
6.0 + 4.6
7.6 + 5.9
4.2 + 1.6
5.9
fwd + IG
60.4 + 27.9
62.1 + 29.0
18.9 + 5.2
47.1
D Extended Ablations
Sharing Parameters for Masking Network and Encoder. Instead of using two different encoders for the masking network and the encoding network, we ask if the same kind of masks can be recovered when sharing parameters between the two encoders (see Table (). In such case, the network first outputs the mask by processing the original image through its encoder, then masks the image, and then passes the masked image through its encoder again generate embedding vector. Such shared parameter setup by default ensures that the encoder has chance to see both masked and unmasked images, which could help avoid overfitting to the distribution masks output by the mask encoder.
Table 8: Sharing Masking Network and Encoder Parameters. Experiment setup follows Table 16 Default InfoGating version is marked gray.
case
casy
medium
hard
overall
unshared
176.2 + 9.1
97.0 + 5.7
44.8
18.4
106.0
shared
104.3 + 22.6 44.2 + 44.0
22.6 + 6.3
57.0
We generally observe that the performance of InfoGating suffers when sharing parameters. We suspect this result may depend strongly on hyperparameters (e.g. model capacity, noise scale, etc.).
E Hyperparameter Details
We use a UNet architecture all throughout this paper for implementing InfoGating on the pixel-level. We use a warm-up consisting of 5k gradient steps where the InfoGating network is not trained, while the downstream encoder is. This is done so that the learnt masks are not affected by initial gradient errors in the downstream loss. Once the warm-up period is over, the loss in Equation 3Jss deployed as usual. The details for the masking network (which follows a UNet architecture) and the encoder and Table 10 respectively. Note that these network architcctures correspond to the visual D4RL locomotion tasks. We add additional layers based on differences in input image sizes (84 84) for the Kitchen (256 x 256) and CIFAR/STL--10 experiments (32 x 32/ 966 96).
described in Table and
15

---

Table 9: UNet Architecture.
Down Sampling (InfoGating Network)
3 x 3 conv2d 32, stride pad 1, GroupNorm, ReLU
3 x 3 conv2d 32, stride 1, pad 1, GroupNorm, ReLU
3 x 3 conv2d 64, stride 1, pad 1, GroupNorm, ReLU
3 3 conv2d 64, stride 1, pad 1, GroupNorm, ReLU
3 x 3 conv2d 64, stride 1, pad 1, GroupNorm, ReLU
MLP
FC 128, ReLU
FC 128, ReLU
FC 1600, ReLU
Up Sampling (InfoGating Decoder)
x 3 conv2d 64, stride 1, pad 1, GroupNorm, ReLU
3 x 3 conv2d 64, stride 1, pad 1, GroupNorm, ReLU
3 x 3 conv2d 32, stride 1, pad 1, GroupNorm, ReLU
3 x 3 conv2d 32, stride 1, pad 1, GroupNorm, ReLU
3 x 3 conv2d 32, stride 1, pad 1, GroupNorm, ReLU
Table 10: Encoder Architecture.
6 x 6 conv2d 128, stride 6, pad 0, ReLU
x I conv2d 128, stride 1, pad 0, ReLU
x 3 conv2d 128, stride 1, pad 0, ReLU
4 x 4 conv2d 256, stride 2, pad O, ReLU
FC 256, LayerNorm, ReLU
Rebuttal
In-distribution Distractor (Train)
Held-out Distractor (Eval)
Masked Observation
Un-Masked Observation
Figure 9: InfoGating masks for held-out distractors (during evaluation time), and for in-iistrbbution distractors (during training). The IG masks generalize quite well to unseen distractors.
16

---

Table 11: Visual D4RL Locomotion Training Details.
batch size
128
À
0.1
IG warm-up
5k steps
cropping padding
4
Trame_stack
3
action_repeal
2
buffer_size
100000
learning_rane
le-4
eval_episodes
10
Table 12: Kitchen Training Details.
batch size
32
X
schedule(0.1,, 3.0, 2k steps)
IG warm-up
1k steps none
cropping padding
frame_stack
num_ num_demos
trajectories
learning_ree rate
le-3
IG learning rate
le-4
action reppaat
eval_trajectoriie
50
Table CIFAR/STI--10 Training Details.
batch size
32
momentum
0.9
reighh_decca
1e-6
À
0.06
IG warm-up
100 steps
cropping scale
(0.2, 1.0)
learning_rand
0.3
IG learning rate
0.3
num_ num_epochs
60
Table 14: Noisy Evaluation. Comparing InfoGating performance with multi-step inverse dynamics when training on noisy distractors and evaluating held-out out-of-distribution distractors (instead of noise-free evaluation as in Table the main paper). We also evaluate performance unmasked observations (right-last column).
JQL
DRIML
Inv Dropout
Iav VIB expert
Inv w/ RCAD
Inv w/ IG
Inv w/ (unmusked eval)
33.8 1: 6.6
1..0
68.4 12.2
56,9
41.9 11.9
92.3 + 19.4
206.4 =: 5.8
medium
12.6.:
12.0
2..4
48.9
60.2
hard
17.4 2.9
10.6 + 2.1
17.1 + 3.9
17.3 3.9
9.9 + 3.8
17.6 4 4.5
12.5 6.4
medium
easy
47.5 12.9
26.1 4: 0.4
54.4
84.2 + 10.0
25.1 12.2
$3.1 7.8
78.0 4: 27.5
medium
13.0 4.8
16.14:3.77
22.3
18.8 6.4
4: 1.1
48.1 8.3
130.0 21.2
4.2
17

---

Table 15: Walker Results. Comparing InfoGating performance with multi-step inverse dynamics when training on noisy distractors and evaluating on held-out distractors (instead of noise-free evaluation as in Table 1). case
IQI.
DRIMI.
w/ Dropout
w/ VIB
Inv w/ RCAD
Inv w/
expert
57.2 + 3.9
30.0 + 0.7
239.2 + 24.7
178.2 + 39.3
185.5 + 26.9
207.0 + 45.0
medium
151.0 + 24.0
50.0 + 6.4
177.3 + 28.4
240.0 + 30.6
231,4 + 12.7
242.0 + 20.6
medium-expert
135.0 +: 20.4
57.9 # 12.4
210.5 + 11.2
187.1 + 26.8
207.9 + 25.2
215.2 +: 24.2
overall
114.4
45.9
209.0
201.7
207.6
221.4
Table 16: Reliance on A. Comparing InfoGating performance when using three different À networks one network. All experimental settings are the same Table of the main paper. For the multiple À version, do not train on mix of unmasked and masked observations.
case Inv w/ IG (A 0.1, 0.2, 0.3) Inv w/ IG (A 0.1)
easy
165.9 13.1
176.2 9.1
medium
95.3 = 9.3
97.0 + 5.7
hard
55.6 = 5.3
44.8 + 18.4
18


Reviews:
[{"review_id": "ZHrw5OroNR", "reviewer": "Reviewer_L82J", "summary": "In this paper, the authors propose InfoGating as a way to learn parsimonious representation that could achieve better generalization by being robust to noise and spurious correlations. Representations that identify the minimal information required for certain task. Those representations attempt to be robust to out-of-distribution observations. \n\nIn contrast with previous works that applied a similar concept, i.e., information bottleneck, in later stage of the computation, InfoGating is applied to the input layer, which allows models to learn which information is key to solve the task. For example, what pixels matter for a given task.\n\nThere are two approaches proposed for InfoGating namely cooperative and adversarial. The former aims to identifying minimal sufficient information. The latter to identifying any useful information.\n", "questions": "Q1. Can you please explain further what is the metric used on the experiments? \n\nQ2. In the experiments with inverse dynamic models, How consistent is the masking between two consecutive frames?\n\n", "limitations": "I found the limitations are well covered. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "S1. The method is simple and well explained, easy to read. It sounds correct.\n\nS2. I found the experimental setup very interesting and well choose. I have seen methods motivated by similar concepts in self-supervised learning and efficient transformer [ATS: Adaptive token sampling for efficient vision transformers. Fayyaz et al. ECCV-22], but not for RL tasks.\n\nS3. Ablation results on feature space vs pixel InfoGating shows the importance of removing distractors early in the pipeline. Supporting the motivation of this work.\n\nS4. Appendix, Visualization and additional results are good and provide good insides to readers.", "weaknesses": "W1. Since the InfoGaiting is presented as a general method to obtain more robust representations that can deal with out-of-distribution observations, I was expecting to see experiments in more dedicated settings like NICO++ [NICO++: Towards Better Benchmarking for Domain Generalization]\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and positive feedback on our paper. We appreciate the time and effort you have put into evaluating our work.\n\n**Benchmarking on vision-based datasets**\n\n> Since we focus on RL settings, we did not test on vision-based domain generalization, but this could be interesting in future work. Our new experiments evaluate on multiple domains (multiple background distractors, please see Table 1 in the rebuttal PDF), going beyond the leave-one-out evaluation strategy, as pointed out in the NICO++ paper.\n\n**Q1. Can you please explain further what is the metric used on the experiments?** \n\n> We evaluate using the return obtained after training a linear policy on representations learned with any given method. The representations are trained with distractor-based observations while the policy is evaluated without distractors. Our new experiments show that these conclusions still hold if the policy is evaluated with distractors.\n\n**Q2. In the experiments with inverse dynamic models, How consistent is the masking between two consecutive frames?**\n\nWe observe masks to be consistent across consecutive frames. The masking network learns to focus on the robot contours, as shown in Figure 7 in the Appendix. It would be interesting to see what happens when inferring masks simultaneously across multiple frames, since temporal consistency could be leveraged to further reduce the number of visible pixels.\n\nWe hope that our rebuttal has addressed your questions and concerns, and we appreciate your consideration in revising the evaluation of our paper.\n"}, {"review_id": "lPtrS4BKEH", "reviewer": "Reviewer_8eSS", "summary": "This paper introduces a novel masking technique called InfoGating for learning masks in contrastive loss settings with InfoNCE.  The proposed approach is simple, well-motivated, and is evaluated in several RL setting including inverse dynamics models, Q-learning, and behavior cloning.  Some ablation studies on regularization of the mask, as well as variants of the objective are also considered.", "questions": "* How does infogating satisfy or violate the IB principle in deep learning?  Is it possible for the model to result in higher information in later layersdepending on how much gating happens in features of the network? Does this violate any benefits of the IB principle, for example better generalization?\n* How does IG perform compared with other methods for masking/information bottleneck on Q-learning and behavior cloning settings?  \n* What is the connection between IG and adversarial masking?", "limitations": "The proposed approach is not compared with other masking approaches for traditional augmentation and image classification settings.", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The proposed approach is a novel way for learning masking and is applied for RL tasks, whereas traditional approaches are typically evaluated for classification, or SSL tasks.  Evaluation of the proposed approach covers multiple RL frameworks including Q learning, behavior cloning, and inverse dynamics.\n* The proposed method is stated very clearly including all hyperparameters for reproducibility, clear pseudocode for inverse dynamics, and a clear description of the approach with intuition.", "weaknesses": "* My primary concern with this paper is that the proposed approach does not have sufficient evaluation not situate itself well with prior works.  There are two notabl method comparisons, VIB (2017), and RCAD (2022).  Neither of these approaches, however are masking-based approaches which is the main line of this work from my point of view.  These methods are also not included in comparisons on behavioral cloning and Q-learning experiments.  I would like to see additional comparisons on Q-learning and behavioral cloning, and perhaps some experiments comparing other masking techniques which from my understanding have been used for RL such as Q-learning: https://proceedings.neurips.cc/paper_files/paper/2022/hash/a0709efe5139939ab69902884ecad9c1-Abstract-Conference.html, https://proceedings.neurips.cc/paper_files/paper/2022/hash/802a4350ca4fced76b13b8b320af1543-Abstract-Conference.html.  Although these approaches approach masking from a slightly different angle, my understanding is that they are still aimed at better observation learning.  There are also works orthogonal to this that focus on sparse convolutions: https://openaccess.thecvf.com/content_CVPR_2020/html/Verelst_Dynamic_Convolutions_Exploiting_Spatial_Sparsity_for_Faster_Inference_CVPR_2020_paper.html.  I think contrasting with these approaches may also distinguish the paper.\n* Additionally, cheetah run is only one of many tasks within the continuous control suite.  It would be good for the authors to clarify why this task in particular, and if IG performs well across multiple additional tasks.  Further, the reported cheetah run numbers are a fair bit lower than those traditionally reported for the return (usually 300+).  Is there commentary on why the numbers are much lower in this setting, and why learning curves are not provided?\n* There are some followup questions about whether or not information gating without controlling  for consistency in all layers is good.  In particular, whether this may violate properties of the IB principle which say that the information will only decrease throughout the network.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and constructive feedback on our paper. We appreciate the time and effort you have put into evaluating our work.\n\n**Comparison with Masking-based methods** \n> We ran new tests using mask-based latent reconstruction (MLR). Even after extensive hyper-parameter tuning, we were not able to avoid collapse in representation. While we were able to produce meaningful results after slightly modifying the method, the results in each case were inferior to our method. In our RL setting, our method is best seen as a form of regularization. This is a somewhat different goal from the Dynamic Convolutions work.\n\n**Results beyond Cheetah domain**\n> Our new tests include results in the Walker domain. We primarily focused on Cheetah because the vd4rl domain only included distractor datasets for all three policy levels for Cheetah. For the Walker results, we collected our own distractor datasets. The results are provided in Table 2 in the rebuttal PDF. We largely observe that InfoGating performs better than the other five baselines.\n\n**InfoGating and the IB principle**\n> Our method strictly removes information from the input, and information can’t be added as one progresses through a compute graph (without taking additional inputs) due to the data processing inequality.\n\n**Other IB methods in the Q-learning and Behavior Cloning Setting**\n> Q-learning - Since the Q-learning setup is quite similar to the inverse dynamics, i.e. they both use the same encoder, we expect the dropout results to be similar to that of the inv w/ dropout results in Table 1 of the main paper.\n\n> Behavior Cloning - In our preliminary experiments, we added dropout to the policy network for both behavior cloning and behavior cloning w/ IG versions. We did not see any significant improvement in performance for both versions, while keeping the model architectures fixed. It is quite likely that since the heavy lifting is done by the representation encoder, adding dropout to the policy network does little to improve performance.\n\n> Because of the bottleneck on compute resources, we could not run more experiments to validate these preliminary observations. We are happy to include these in the final version of the paper if the reviewer feels strongly about these results.\n\n**Connection between IG and Adverarial Masking**\n\n> The “cooperative” version of InfoGating is quite different from Adversarial Masking as in ADIOS. In particular, the IG masks in this setting seek to reveal the minimal info required for solving a task. The “adversarial” version of IG is similar to a single-mask version of ADIOS. Extending IG to use multiple masks is straightforward.\n\nWe hope that our rebuttal has addressed your questions and concerns, and we appreciate your consideration in revising the evaluation of our paper.\n"}, {"review_id": "LWOLQvcacr", "reviewer": "Reviewer_xdoN", "summary": "The authors proposed a mutual information-based encoder that generates masks to gate inputs in order to either pass on minimal information for downstream tasks, or to remove any useful information that could be used to optimize the downstream loss in an adversarial setting. This is a general model that can be applied to a variety of tasks, and the authors focus on the application to reinforcement learning tasks.", "questions": "For the minimax objective, are there any nested loops when updating the two sets of parameters?\nIs there code available?", "limitations": "The authors appropriately addresss ethical or social impacts of their research.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Through experiments, the authors show that the idea of directly removing visual information in a learnable way through dynamic masking improves the downstream policy learning in terms of performance and training stability when there are visual distractors. \n* InfoGating is very general and could be applied to input or any intermediate representations, with good interpretability\n", "weaknesses": "* training a minimax objective is usually not very stable, and the authors are not very clear in algorithm 1 or the appendix about the details of their minimax training. See questions.\n* The idea of masking inputs directly is reasonable, and experiments show that it improves performance and training stability. But it seems not so obviously advantageous on intermediate layers. The authors discussed in Section 6 and 7, but there's no clear statement or experiments towards this line. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and positive feedback on our paper. We appreciate the time and effort you have put into evaluating our work.\n\n**Minimax objective details**\n\n> We did not use nested loops or different numbers of gradient updates for optimizing the image and mask encoders in the (minimax) adversarial objective. For adversarial losses, we simply switch the sign prior to backprop depending on which parameters we’re updating. We did not encounter stability issues.\n\n**Masking intermediate features**\n\n> We agree that we do not provide strong evidence that masking intermediate features is particularly useful in our RL settings. But, we believe it could be useful in other settings, e.g., for sparsifying attention to reduce compute in LLMs. See, e.g., works like “DARTS: Differentiable Neural Architecture Search” (Liu et al., ICLR 2019).\n"}, {"review_id": "uN4w4OmO4a", "reviewer": "Reviewer_TDKb", "summary": "The authors hypothesise that gating information propagation in neural networks will lead to better generalisation. To achieve this they propose a system that performs gates information using a multiplicative differentiable operation, which they call InfoGating. To show the validity of their approach they compare to related models on several tasks such as contrastive dynamics control, Q-Learning and behaviour cloning.", "questions": "1. What is the relation between lambda and the overall performance? This is the main hyperparmeter of the system yet there is no graph showing how manipulating it affects performance.\n2. If the authors have access to ground truth masks for some or all of the datsets, could they compare the masks produced by their approach with the ones produced by InfoGating?", "limitations": "Main technical issue I see is that the reliance on the lambda parameter is very similar to the reliance of a $\\beta$ parameter in the $\\beta$-VAE literature, which creates a tradeoff that is not always optimal.\n\nIt would have been interesting to test the information gating approach with other more principled techniques like object-centric learning which already segment the images into relevant parts (or at least they try to, eg [1]). Then info gating would be only responsible for selecting the appropriate objects for the downstream task.\n\nThis approach is also very similar to work in meta-learning which draws on the parallel between gating and neuromodulation in Neuroscience [2]. It would have been interesting to have both a conceptual and experimental comparison between the two.\n\n### References\n\n[1] Locatello, F., Weissenborn, D., Unterthiner, T., Mahendran, A., Heigold, G., Uszkoreit, J., ... & Kipf, T. (2020). Object-centric learning with slot attention. *Advances in Neural Information Processing Systems*, *33*, 11525-11538.\n\n[2] Beaulieu, S., Frati, L., Miconi, T., Lehman, J., Stanley, K. O., Clune, J., & Cheney, N. (2020). Learning to continually learn. *arXiv preprint arXiv:2002.09571*.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The method is well motivated.\n2. The tasks are relevant to the issue the model is seeking to solve.\n3. The comparison models appear to be relevant and well justified.\n4. They perform relevant ablations of their design choices.", "weaknesses": "While the experimental results show improved generalisation, it is not clear how general this approach will prove to be. So, while I agree that masking irrelevant parts of the input so as to limit distractor information is a good strategy (and one that humans definitely use when interacting with the world), I am less clear that this specific approach fully captures how this process works and will generalise effectively beyond to more complex datasets.\n\nTo the author’s credit, they do point out limitations of their approach.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and constructive feedback on our paper. We appreciate the time and effort you have put into evaluating our work.\n\n**Question 1: relation between lambda and the overall performance / reliance on the lambda parameter**\n\n> We illustrate the effect of lambda on performance in Figure 5 in the paper which shows how test performance varies with lambda. As lambda increases performance improves until further increasing lambda causes too much information loss. We have added new results in Table 3 of the rebuttal PDF, where we show that InfoGating is considerably robust to $\\lambda$ value. Particularly, we trained InfoGating with three masking networks, each with a different lambda value. At each update step, a single masking network is selected at random. We see that the performance is similar to the one obtained when training with a single lambda value, hinting at how InfoGating is not too sensitive to the masking sparsity.\n\n**Question 2: ground truth masks**\n\n> There are no ground truth masks for the settings we consider, and it’s not clear how to produce such masks. However, Figure 7 in the Appendix shows how our masks generally capture the agent pose while removing most other information.\n\nWe will note the parallel between gating and neuromodulation in the paper by adding a few lines about the conceptual similarities. Thank you for the pointer.\n\nWe hope that our rebuttal has addressed your questions and concerns, and we appreciate your consideration in revising the evaluation of our paper.\n"}, {"review_id": "rzqmlFEyhH", "reviewer": "Reviewer_WnVz", "summary": "This paper empirically investigates a few approaches to modulating the amount of information used by a neural network learner in a variety of control-adjacent learning problems. The proposed approach, InfoGating, learns an input-conditioned continuous-valued mask that is applied to the same input or features thereof. The learner is trained to optimize a task loss and minimize the information allowed through by the mask. Experiments on representation learning via multi-step inverse dynamics modeling demonstrate the relative benefit of InfoGating over other approaches to information parsimony such as the variational information bottleneck. Experiments on Q-learning and fine-tuning pre-trained visual representations for behavior cloning demonstrate the benefit of InfoGating over naive baselines.", "questions": "- (Q1) What do the learned masks look like for i) inputs whose underlying state is trained on, but with held-out distractors; ii) inputs whose underlying state is held-out, but with trained-on distractors; and iii) inputs whose state and distractors are both held-out? This would go a long way in explaining the benefits of (or diagnosing weaknesses with) the proposed approach.\n\n- (Q2) Why does baseline inverse dynamics modeling in Table 1 do so poorly, given the protocol of train-on-noisy, test-on-noise-free? Concretely, why doesn't adding the noise to the training data confer benefits a la data augmentation, since it specifies input transformations that the output (actions) should be invariant to? Is it because the considered noise (frames from continuously playing videos) has temporal consistency?", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "### Originality\nThe particular formulation of InfoGating as using input-conditioned masking amortized via a separate neural network is, to my knowledge, original. The authors remark that InfoGating is closely related to previously proposed ideas in the space of information-based regularization of neural networks.\n\n### Quality\nThe scope of the experiments chosen to validate the proposed idea is perhaps the main strength of this work. Three distinct experimental testbeds (dynamics modelling + probing, Q-learning, fine-tuning pre-trained representations for behavior cloning) are considered, as are a few variants of InfoGating (input gating vs. feature gating, cooperative masking vs. adversarial masking).\n\n### Clarity\nThe proposed methods are simple and clearly described. However, while some experimental details (architectures and hyperparameters) are presented in the supplementary material, this seems insufficiently detailed to guarantee reproducibility, so I would encourage the authors to release code.\n\n### Significance\nThe considered problem is important and the proposed methods are simple and intuitively reasonable.", "weaknesses": "I list my perceived weaknesses in rough decreasing order of importance. \n\n- (W1) The experimental design leaves a lot to be desired. First and foremost, for the main experiment (multi-step inverse dynamics modeling + behavior cloning probe on half-cheetah), the authors report performance on \"a noise-free observation space\" (line 221), which I take to mean the original DeepMind Control Suite background or similar. I would instead expect evaluation to use held-out noise (distractions from the Distracting Control Suite) to properly assess the generalization capabilities of the models. As-is, contribution claim 3 (line 62) is unsubstantiated, since, to my knowledge, we typically do not use \"generalization\" to mean noisy train -> noiseless test (I would instead characterize this as denoising), and neither of the other testbeds include comparisons to other information-based regularization techniques. \n\n- (W2) The design choice of operationalizing information regularization via a learned input-conditioned mask pushes the burden of generalization to the masking network. This seems to be the main distinguishing aspect of InfoGating over other forms of information regularization, yet the impact of this design is neither evaluated in isolation nor in toto.\n\n- (W3) There are several key statements that are unclear or confusing.\n    - \"Although IB approaches have been beneficial in some cases, adding an information bottleneck at the penultimate step of computation does little to prevent overfitting in the preceding steps of computation, which typically comprise the overwhelming bulk of the model\" (line 34). It's not clear to me that IB approaches predominantly bottleneck just before computation output. In particular, methods that use informational/Gaussian/variational dropout seem to contradict this. I would appreciate citations supporting the quoted statements.\n    - \"Roughly, the values in $ig(x)$ can be seen as specifying how many steps of forward diffusion to run in a Gaussian diffusion process initiated at $z$, where values near zero correspond to running more steps of diffusion and thus sampling from a distribution that is closer to a standard Gaussian in terms of KL divergence\" (line 174). This sounds like an interesting connection, but it seems to be only true if the mask specifies a constant value. Otherwise, different pixels are noised to different extents, which doesn't seem to correspond correctly to the standard modelling assumptions in diffusion models.\n\n- (W4) The material in the half-page Background section is completely orthogonal to the proposed idea. I would move a condensed version to the experimental section since the content really is just context for the experiments. Relatedly, Algorithm 1 is InfoGating applied specifically to multi-step inverse dynamics modeling, yet one of the contribution claims is that InfoGating is general purpose.\n\nAddressing W1 and showing good performance on more meaningful out-of-distribution evaluation would cause me to increase my score substantially.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and constructive feedback on our paper. We appreciate the time and effort you have put into evaluating our work.\n\n**W1: Evaluations on noiseless environments and W2: Generalizations of masking network**\n\n> Yes, “noise-free observation space” indicates the default DM Control Suite backgrounds. We agree that our main results with background distractors should include results where both training and test episodes include distractors. We performed these tests and present them in Table 1 in the PDF accompanying this rebuttal. We also performed tests in response to your point about offloading generalization duties to the masking network. Specifically, we include results for using the image encoder trained on IG-masked images with unmasked images at test time. Note that this encoder is trained with a mix of masked and unmasked images during training (so unmasked images are not “out of distribution”), and that all other methods use unmasked images (so this comparison is “fair” in this sense).\n\n> From Table 1 in the rebuttal PDF, the image encoder trained via inverse dynamics with InfoGating outperforms the baselines when testing with unmasked images. Depending on how many distractors are seen during training relative to how many are held out for use in testing, we also see some signs of potential overfitting in the mask encoder. I.e., Inverse Dynamics + IG with masked images at test time can suffer when only a small number of distractors are seen during training, even though the same image encoder performs quite well when working with unmasked images during testing.\n\n> We thank the reviewer for pointing out these gaps in our experiments and believe that the expanded results provide stronger support for the claimed contributions in our submission. These results suggest an interpretation of our method as performing task-adaptive data augmentation. I.e. the mask encoder applies cutout to the training images with the goal of cutting out as much as possible without making the task intractable. From this perspective, it is natural to use the IG-trained image encoder with unmasked images at test time, which makes generalization by the mask encoder less critical to our method.\n\n**IB approaches predominantly bottleneck just before**\n\n> We cannot provide a specific citation for our statement that existing IB-like methods predominantly implement bottlenecks at late stages of the overall computation (eg, the last linear layer as in the original work by Alemi et al.), since it is merely an observation on our part. We cite variational information dropout as a counterpoint to our observation, and will edit our claims to be more clear that we’re stating an opinion about how the potential gain from applying some sort of IB early or late in the compute graph is not well reflected in the balance of early/late IB usage in prior publications. We will also clarify that this observation only pertains to “actively info minimizing” forms of IB, and was not intended to include methods based on dropout where the degree of info minimization is set a priori.\n\n**Connections to diffusion models**\n\n> Regarding similarities between InfoGating and diffusion models, we note that the forward diffusion processes in standard DDMs are pixel-wise independent and we posit that one could train a model where the reverse process proceeds at different rates for each pixel. Generating data for training a reverse process with variable per-pixel diffusion rates would be simple due to the pixel-wise independence of the forward process, though one may need some extra bookkeeping to track how far through the forward/reverse process each pixel is. Of course, ease of implementation is not equivalent to ease of getting good results.\n\n**W4: paper structure**\n\n> We can rewrite Algorithm 1 in a more general form in the main paper. We chose to tailor it to inverse dynamics in the submitted draft since we focused a bit more on this setting in the paper.\n\n**Q1: Generalization of IG masks**\n\n> Figure 1 in the rebuttal PDF visualizes the mask for in-distribution and held-out distractors, showing that the the InfoGating masks generalize well to out-of-distribution distractors. Note that since we train on previously collected offline datasets, it is not possible to infer if an underlying state is in-distribution or not.\n\n**Q2: Baseline inverse dynamics performance**\n\n> We include a single, static distractor dataset while training and so your point on \"a kind of data augmentation effect coming into play\" would be valid if we trained on multiple distractor datasets (creating an implicit invariance to the distractors). The temporal consistency in the background is also a potential reason for the Walker experiments, which uses the single distractor video.\n\nWe hope that our rebuttal has addressed your questions and concerns, and we appreciate your consideration in revising the evaluation of our paper."}]

Author rebuttal:
We thank the reviewers for their feedback. We have individually responded to points made by each reviewer and provided further experiments in support of our response (please refer to the attached rebuttal PDF). Below is a list of additional experiments we have included:

1. **Evaluations of all baselines and InfoGating on multiple distractors (including unmasked evaluations for IG):** We show that even when evaluations are not done on noise-free observations, InfoGating is considerably better performing.

2. **Results on the walker domain, for all baselines and InfoGating.**

3. **Visualizations of InfoGating masks on in-distribution and held-out distractors.**

4. **InfoGating with multiple masking networks**, each with a different $\lambda$ value, showing that IG relies very loosely on the particular  value of $\lambda$.
 
We hope our response and the additional results can help address the reviewers' questions and concerns. For criticisms not addressed in this rebuttal, we invite the reviewers to point out those which seem most pressing during the interaction period. Thank you!

Meta-review: This paper presents information gating as a mechanism to use minimum task-relevant information to improve robust learning. The method is applied Q-learning. It boosts performance by concentrating on the informative content, enhancing the robustness of the reinforcement learning model.

The reviewers' feedback acknowledges the strong points of the paper, particularly its clear presentation and its application to reinforcement learning. They appreciate the insights from your ablation results and the overall experimental setup. All reviewers recommend unanimously for acceptance.

Whilst, reviewers emphasise the need for more extensive evaluation, suggesting direct comparisons with relevant existing methods and a broader range of tasks to demonstrate generalisability. They've also highlighted the need to discuss the relation of the method with attention.

The AC had a chance to read the paper closely and went through all the reviews. Besides the strong points mentioned above, the paper lacks in that the method is not discussed and compared with attention methods. The proposed InfoGating has an alternating optimisation, but that requires more ablation to be convincing.  It holds true that there may not be a widely recognized practice in the field of RL to incorporate attention for information gating, but the AC strongly encourages the authors to consider these insights for refinement. Additionally, incorporate the experiments from the rebuttal phase in the paper.
Decision: Accept (poster)