Title: Ignorance is Bliss: Robust Control via Information Gating
Venue: NeurIPS 2023
Accepted: yes
Keywords: representation learning

mutual information
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:

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.

+

Masking

Decoder

Masking

Encoder

InfoGated Input

Regular

Encoder

Downstream

Loss

1 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, 8]. 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

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-distribution 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 a 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 a model. For example, we can consider what happens when one inserts an IB near the beginning of a 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 a standard contrastive loss for self-supervised 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 on 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 a 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 to remove almost all of the irrelevant information from the pixels, leading to better out-of-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 a 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 internal representations. Typical approaches based on Variational IB [2] achieve this by learning stochastic representations that are constrained to be close to a 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 noise/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 [1] is perhaps the most similar prior work to InfoGating. Information Dropout imposes IBs at multiple points in a network'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.

Masked Image Modelling . There has been a lot of recent work in self-supervised learning where the input is masked and the model is tasked with reconstructing the missing bits [3, 15, 41]. Most of these works operate with a fixed, random masking scheme, instead of learning it directly through the downstream loss. The amount of information masked is 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 metrics [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 [11] 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.

3 Background

This section describes the primary downstream losses we 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 x and z 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 : I ( x , z ) = H ( x ) -H ( x | z ) , where H is the Shannon entropy. InfoNCE [31], based on Noise Contrastive Estimation [13], computes a lower bound on I ( z 1 , z 2 ) , where z 1 and z 2 are two representations of the input x produced by some encoder f ( x ) . Specifically, the bound is optimized by discriminating 'positive' and 'negative' pairs:

where ψ 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 1 , z 2 ) 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 1 , z 2 ) is high), while pushing apart the negatives (i.e., ψ ( z 1 , z -) is low) [4, 5].

Multi-Step Inverse Dynamics Models . Multi-step inverse dynamics models [8] predict the action(s) that took an agent from some observation x t to some future observation x t + k , attempting to learn a 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 [22]. Multi-step inverse models can be used to learn representations by simply predicting the actions ( a t , ..., a t + k -1 ) conditioned on x t and x t + k . Actions can be predicted using a standard max likelihood objective or with a contrastive reconstruction objective which maximizes the InfoNCE-based lower bound on I (( x t , x t + k ) , ( a t , ..., a t + k -1 )) . The learning objective in this case looks as follows:

where z t and z t + k are representations computed from x t and x t + k , and y t : t + k -1 is a representation computed from ( a t , ..., a t + k -1 ) . The negative samples in this case come from randomly sampling action sequences of the appropriate length from the agent's collected experience. In the simplest case, the representation y t : t + k -1 may depend only on a t .

4 Information Gating

We present two approaches to InfoGating that we describe as cooperative and adversarial . Cooperative 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 , f ( x ) . The second is an information gating function, ig ( x ) (see Figure 1). In principle, we can gate information passing through any layer in the representation network that computes the encoding z = f ( x ) . The ig ( x ) function provides continuous-valued masks (values in [0 , 1] ) 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, ig ( x ) masks are the size of the image. In general, the goal of ig ( x ) is to erase as much information from f ( x ) 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 an input x by taking a simple convex combination of x and random Gaussian noise ϵ . The combination weights are given by ig ( x ) :

where x ig denotes the info-gated input and ⊙ denotes element-wise multiplication. An all-zero mask corresponds to complete noise (i.e., erasing all information from the input), while an all-one mask corresponds to keeping the original input. The function ig ( x ) 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 L1 norm. This tends to produce sparse masks due to properties of the L1 norm [30]. The overall objective for learning with InfoGating is

where L task refers to 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 ig (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 λ 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 1, which assumes contrastive multi-step inverse dynamics modeling as the downstream objective.

InfoGating in Feature Space . As an alternative to directly gating the input pixels, we can also consider a variant where f ( x ) is masked instead of the input. Consider the following masking:

The training objective remains the same as in the pixellevel case, i.e., minimize the downstream loss L task ( z ig ) for the masked representation z ig , while masking as much of z as possible (Figure 2). This version is closely related to the deep variational information bottleneck (VIB), which minimizes the KL divergence between the (stochastic) representation z and a prior distribution, typically chosen to be a standard Gaussian. Roughly, the values in ig ( x ) can be seen as specifying how many steps of for-

+

Masking

Encoder

Regular

Encoder

Output Mask

Downstream

Loss

ward 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 [18]. Thus, like VIB methods, we aim to minimize I ( x , z ) , while maximizing the task performance of z . However, feature space InfoGating uses a different optimization objective and parameterization compared to existing VIB approaches. We leave similar variations of InfoGating where arbitrary layers in a 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 to lower the overall loss. This leads to finding representations that capture the minimal sufficient information for a given task. We can also turn this formulation on its head, yielding an 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 to erase as much of the input as possible, the adversarial objective encourages the masks to remove as little information as possible while minimizing the encoder's performance on the downstream task. We write the adversarial objective for ig ( x ) and f ( x ) as:

This gives rise to a min-max objective w.r.t. the encoder and the gating network while cooperative InfoGating corresponds to a min-min objective (see lines 11-12 in Algorithm 1).

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 a lot of task-irrelevant information, while behavior cloning models are meant to 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 ‡ .

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. Hyper-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 as described in Eq. 1. To apply InfoGating, we mask both the current observation x t and the goal observation x t + k using the masks from ig ( x t ) and ig ( x t + k ) respectively. We then process both masked observations through the encoder to compute the corresponding representations z ig t = f ( x ig t ) and z ig t + k = f ( x ig t + k ) . These are then optimized through a loss similar to Eq. 3:

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 a 2-layer MLP over the frozen representations. We have tried replacing BCbased evaluation with, e.g., Q-Learning-based evaluation, which leads to similar relative scores. We report scores for BC-based evaluation for easier reproducibility and reduced dependence on 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-distribution shift in the distractor. We use random cropping as a pre-processing step for the observations in all our experiments. We compare inverse dynamics with InfoGating to six baselines: control specific methods, i.e. 1) IQL [21], 2) DRIML [27], and 3) inverse dynamics without InfoGating (the standard formulation from Eq. 1), and regularization specific methods that work on 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 improvement over the standard inverse model is 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), in 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 1). 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-based policies, while keeping the representation fixed. In this section, we ask if the same conclusions hold if we train the representation end-to-end with a Q-Learning loss instead. We use the DrQ [20] 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:

where Q θ and Q ′ θ denote the current and target Q networks respectively, while r ( x t , a t ) is the obtained reward. In Table 2, we observe that the base DrQ algorithm is quite prone to failure for 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.

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 16. Results are for an expert policy level.

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 policy? InfoGating can be seen as a natural way to extract only information pertaining to 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 a 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 3 shows normalized success rates (each out of a 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 learnt 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

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 not come close to 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 as the pixel-level variant.

Adversarial InfoGating . We test the adversarial InfoGating algorithm using a multi-step inverse model loss, just as we did for cooperative InfoGating. Figure 4 shows how the adversarial masks tend to hide the agent body. Some additional image content is also erased, since a precise silhouette of the agent's pose would be highly predictive of the agent's pose. To test how useful this kind of masking is, we simultaneously train a separate encoder over the reverse of the adversarial masks. If 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 indeed see that the adversarial InfoGating model performs similarly to the cooperative version, thus verifying that it is an equally viable approach for learning parsimonious representations (see Table 5).

Original Observation

InfoGated

Observation

1

Effect 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 is masked and no useful information is captured.

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 Encoder . Since the masking network and the encoder process different kinds of input (unmasked and masked input, respectively), it is worth asking if the encoder can be better regularized by forcing it to be able to process both masked and unmasked images. We can also encourage the encoder to be robust to changes in the input/mask mapping learned by the masking network by occasionally swapping masks between images in a batch while training the encoder. We test both of these separately (Table 6) and observe that training the encoder on masked and unmasked inputs provides strong improvements over the base InfoGating model. This also has the added benefit of not using the masking network at evaluation time, and only using the encoder, 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.

50

100

150

200

= 1e-2

= 1e-1

= 5e-1

= 1.0

= 5e-2

mask scaling

performance

300

7 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-free 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. In 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 objectcentric representations without any explicit notion of objects forced into the model architecture or learning objective. For example, in settings where an agent's actions may affect multiple objects, and we want to 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

Wepresent 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 adversarial, while demonstrating that the gating can be learnt both at the input space or any intermediate feature space. We apply InfoGating to 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 of both out-of-distribution (visual D4RL distractor noise) and in-distribution (irrelevant Kitchen objects) generalization.

9 Acknowledgments

The authors would like to thank Charlie Snell, Colin Li, and Dibya Ghosh for providing feedback on an 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); a Canada CIFAR AI Chair, Amii; Compute Canada; Huawei; Mitacs; and NSERC.

References

Alessandro Achille and Stefano Soatto. Information dropout: Learning optimal representations through noisy computation. IEEE transactions on pattern analysis and machine intelligence , 40(12):2897-2905, 2018.

Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.00410 , 2016.

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.

Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. Advances in neural information processing systems , 32, 2019.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020.

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.

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 on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018 , volume 80 of Proceedings of Machine Learning Research , pages 1261-1269. PMLR, 2018. URL http://proceedings. mlr.press/v80/dietterich18a.html .

Yonathan Efroni, Dipendra Misra, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Provable rl with exogenous distractors via multistep inverse dynamics. arXiv preprint arXiv:2110.08847 , 2021.

Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing , 40(6):1662-1714, 2011.

Chelsea Finn, Ian J. Goodfellow, and Sergey Levine. Unsupervised learning for physical interaction through video prediction. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett, editors, Advances in Neural Information Processing Systems

29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain , pages 64-72, 2016. URL https://proceedings.neurips.cc/ paper/2016/hash/d9d4f495e875a2e075a1a4a6e1b9770f-Abstract.html .

Xiang Fu, Ge Yang, Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning , pages 3480-3491. PMLR, 2021.

Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. arXiv preprint arXiv:1910.11956 , 2019.

Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the thirteenth international conference on artificial intelligence and statistics , pages 297-304. JMLR Workshop and Conference Proceedings, 2010.

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603 , 2019.

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.

Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint arXiv:1903.12261 , 2019.

Riashat Islam, Manan Tomar, Alex Lamb, Yonathan Efroni, Hongyu Zang, Aniket Didolkar, Dipendra Misra, Xin Li, Harm van Seijen, Remi Tachet des Combes, et al. Agentcontroller representations: Principled offline rl with rich exogenous information. arXiv preprint arXiv:2211.00164 , 2022.

Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems , 34:21696-21707, 2021.

Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterization trick. Advances in neural information processing systems , 28, 2015.

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.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169 , 2021.

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.08229 , 2022.

Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. J. Mach. Learn. Res. , 17:39:1-39:40, 2016. URL http://jmlr.org/ papers/v17/15-522.html .

Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: differentiable architecture search. In International Conference on Learning Representations (ICLR) , 2019.

Cong Lu, Philip J Ball, Tim GJ Rudner, Jack Parker-Holder, Michael A Osborne, and Yee Whye Teh. Challenges and opportunities in offline reinforcement learning from visual observations. arXiv preprint arXiv:2206.04779 , 2022.

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.

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 , 33:3686-3698, 2020.

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-reinforcement learning. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019 . OpenReview.net, 2019. URL https: //openreview.net/forum?id=HyztsoC5Y7 .

Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A universal visual representation for robot manipulation. arXiv preprint arXiv:2203.12601 , 2022.

Andrew Y Ng. Feature selection, l 1 vs. l 2 regularization, and rotational invariance. In Proceedings of the twenty-first international conference on Machine learning , page 78, 2004.

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748 , 2018.

Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning , pages 27782787. PMLR, 2017.

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.

Amrith Setlur, Benjamin Eysenbach, Virginia Smith, and Sergey Levine. Adversarial unlearning: Reducing confidence along adversarial directions. arXiv preprint arXiv:2206.01367 , 2022.

Yuge Shi, N Siddharth, Philip Torr, and Adam R Kosiorek. Adversarial masking for selfsupervised learning. In International Conference on Machine Learning , pages 20026-20040. PMLR, 2022.

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-1958, 2014.

Alex Tamkin, Mike Wu, and Noah Goodman. Viewmaker networks: Learning views for unsupervised representation learning. arXiv preprint arXiv:2010.07432 , 2020.

Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 ieee information theory workshop (itw) , pages 1-5. IEEE, 2015.

Naftali Tishby, Fernando C Pereira, and William Bialek. The information bottleneck method. arXiv preprint physics/0004057 , 2000.

Manan Tomar, Utkarsh A Mishra, Amy Zhang, and Matthew E Taylor. Learning representations for pixel-based control: What matters and why? arXiv preprint arXiv:2111.07775 , 2021.

Pascal Vincent, Hugo Larochelle, Isabelle Lajoie, Yoshua Bengio, Pierre-Antoine Manzagol, and Léon Bottou. Stacked denoising autoencoders: Learning useful representations in a deep network with a local denoising criterion. Journal of machine learning research , 11(12), 2010.

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.

A InfoGating with SimSiam

Random masking has been used as an SSL objective to learn useful features for downstream object detection/classification. 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 [6] objective uses a cosine similarity loss between the representation of one view z 1 and a predicted output of the representation of the second view p = p ( z 2 ) , where the function p is called the predictor. This is in place of using the InfoNCE loss to define contrastive pairs for z 1 and z 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:

where D denotes the cosine similarity function:

Note that D is not a symmetric quantity as it uses a stop gradient operation on its second argument z .

Having described how SimSiam works, we can now go into how InfoGating is applied alongside it. Given two augmented views of the input image x 1 and x 2 , we mask both views to get x ig 1 and x ig 2 , then process them through the encoder to compute the SimSiam loss L ( z ig 1 , z ig 2 ) . Finally, we also add the original SimSiam loss, i.e. one over the unmasked inputs to the overall objective:

We test this version of InfoGating on CIFAR-10, while evaluating performance on Corrupted CIFAR10 [16]. Figure 6 shows that InfoGating leads to improved evaluation scores in comparison to the base SimSiam method, thus showcasing better robustness to test-time induced corruptions.

Speckle-Noise

Motion-Blur

Pixelate

Frost

Contrast

Brightness

0

20

40

60

Performance

SimSiam

SimSiam + InfoGating

B Visualizing InfoGating Masks

Learning info-gates directly from the downstream loss allows for adapting the mask based on the task at hand. On the other hand, a fixed random masking scheme offers limited benefits, largely pertaining to enhanced data augmentation. For instance, using an 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 7). 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 a 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 as in the inverse model case. Similarly, when the

Forward w/

Distractor

Inverse w/

Distractor

downstream objective is a self-supervised loss that focuses on all parts of the input (as opposed to only information useful to control an agent), we see that the learnt masks capture the edge boundaries of both background and foreground objects (see Figure 8). These different variations highlight the importance of choosing the right downstream loss as well as learning InfoGating masks in conjunction with the given loss.

C Extended Results

Throughout the paper, we use a 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 ig t and x ig t + k into corresponding encodings z ig t and z ig t + k . We then concatenate the two encodings along with the action to form a single encoding triplet. This is then passed through a 2-layer MLP to output logits for the energy value corresponding to the given ( z ig t , z ig t + k , a t ) triplet. The actions actually taken by the agent make up for a 'positive' triplet and the model is trained to output a low energy value. Similarly, a t is replaced by random actions ¯ a t to form a 'negative' triplet, for which the model is trained to output a high energy value.

We can consider a similar version that uses a forward contrastive loss instead. To that end, we simply encode the current action and observation pair ( x ig t , a t ) to ¯ z ig t ) while encoding the future/goal observation x ig t to a similar size vector z ig t + k . 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 a multi-step inverse model has. We test 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 t with the z ig t

representation, before projecting it to the same dimension as the goal representation z ig t + k . Finally, we penalize the info-gates for both current and goal observations as done previously to obtain the following objective:

where ¯ z ig t = g ( z ig t , a t ) is a projection function that maps the state and action embedding to a space of the same size as z ig t . Note that in the InfoNCE loss, the negatives now come from sampling different future observation representations from the batch.

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 8). In such a 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 to generate an embedding vector. Such a shared parameter setup by default ensures that the encoder has a chance to see both masked and unmasked images, which could help avoid overfitting to the distribution of masks output by the mask encoder.

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 3 is deployed as usual. The details for the masking network (which follows a UNet architecture) and the encoder are described in Table 9 and Table 10 respectively. Note that these network architectures correspond to the visual D4RL locomotion tasks. We add additional layers based on differences in input image sizes (84 x 84) for the Kitchen (256 x 256) and CIFAR/STL-10 experiments (32 x 32 / 96 x 96).

F Rebuttal

Table 10: Encoder Architecture .

Held-out

Distractor (Eval)

In-distribution

Distractor (Train)

Masked

Observation

Un-Masked

Observation

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n### 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.\n\nWEAKNESSES:\nI 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\nQUESTIONS:\n- (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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nWhile 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nMain 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*.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

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

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. 
* InfoGating is very general and could be applied to input or any intermediate representations, with good interpretability

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.
* 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.

QUESTIONS:
For the minimax objective, are there any nested loops when updating the two sets of parameters?
Is there code available?

LIMITATIONS:
The authors appropriately addresss ethical or social impacts of their research.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-26

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nThe proposed approach is not compared with other masking approaches for traditional augmentation and image classification settings.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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\nSTRENGTHS:\nS1. 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.\n\nWEAKNESSES:\nW1. 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\nQUESTIONS:\nQ1. 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\nLIMITATIONS:\nI found the limitations are well covered.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-28"}