SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\mathbf{x}_t$  and the goal observation  $\mathbf{x}_{t + k}$  using the masks from  $ig(\mathbf{x}_t)$  and  $ig(\mathbf{x}_{t + k})$  respectively. We then process both masked observations through the encoder to compute the corresponding representations  $\mathbf{z}_t^{ig} = f(\mathbf{x}_t^{ig})$  and  $\mathbf{z}_{t + k}^{ig} = f(\mathbf{x}_{t + k}^{ig})$ . These are then optimized through a loss similar to Eq. 3:
$$
\mathcal {L} = \operatorname {I n f o N C E} \left(\left(\mathbf {z} _ {t} ^ {i g}, \mathbf {z} _ {t + k} ^ {i g}\right), \mathbf {a} _ {t}\right) + \lambda \left(\left| \left| i g \left(\mathbf {x} _ {t}\right) \right| \right| _ {1} + \left| \left| i g \left(\mathbf {x} _ {t + k}\right) \right| \right| _ {1}\right) \tag {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 a 2-layer MLP over the frozen representations. We have tried replacing BC-based 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:
$$
\mathcal {L} _ {i g, \theta} = \left(Q _ {\theta} \left(\mathbf {x} _ {t} ^ {i g}, \mathbf {a} _ {t}\right) - \mathbf {r} \left(\mathbf {x} _ {t}, \mathbf {a} _ {t}\right) - \gamma Q _ {\theta} ^ {\prime} \left(\mathbf {x} _ {t + 1}, \mathbf {a} _ {t + 1}\right)\right) ^ {2} + \lambda | | i g (\mathbf {x} _ {t}) | | _ {1}, \tag {7}
$$
where  $Q_{\theta}$  and  $Q_{\theta}'$  denote the current and target Q networks respectively, while  $\mathbf{r}(\mathbf{x}_t, \mathbf{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.
case	easy	medium	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 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.
Table 3: Finetuning with Behavior Cloning on Kitchen-v3. We compare fine-tuning performance of BC with and 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.
env	ImageNet	CLIP
BC w/ IG	BC	BC w/ IG	BC
knob1-on	17.6 ± 3.1	7.3 ± 4.2	13.3 ± 1.49	11.3 ± 0.94
ldoor-open	9 ± 5.3	5 ± 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 3: Visualizing Masks Produced by InfoGating. Masks improve during training on the Kitchen-sdoor-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 cheetah-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 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.
Table 4: InfoGating in the Feature Space. We compare the Variational Information Bottleneck with InfoGating in feature space. The experimental setup is the same as in Table 16. Results are for an expert policy level.
case	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 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
Figure 4: Left. Original distractor-based image. Right. Learnt adversarial mask over the original image.
InfoGated Observation
Effect of Mask Sparsity. We test what range of  $\lambda$  values leads to improved performance (Figure 5). As is expected, when  $\lambda$  is too high, the entire input is masked and no useful information is captured.
Table 5: Adversarial InfoGating. We compare differences in the representations learnt with Cooperative vs Adversarial InfoGating. The experimental setup is the same as in Table 16. Results are for an expert policy level.
case	easy	medium	hard	overall
cooperative	176.2 ± 9.1	97.0 ± 5.7	44.8 ± 18.4	105.9
adversarial	202.3 ± 2.64	84.8 ± 41.2	4.6 ± 1.7	96.8
Similarly, when  $\lambda$  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.
Figure 5: Scaling  $\lambda$  leads to different masks which reveal different amounts of information. For small  $\lambda$ , the masked observation still contains distracting/non-salient information, thus hurting performance. Similarly, when masking is too aggressive, too much information is lost and performance goes down.
Table 6: Mixing Masking and Unmasked Inputs. Experiment setup follows Table 16. Default InfoGating version is marked in gray. Results are for an expert policy level.
case	easy	medium	hard	overall
w/ mix input	176.2 ± 9.1	97.0 ± 5.7	44.8 ± 18.4	106
shuffle 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
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 object-centric 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.
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 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.
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.
[1] 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.
[2] Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.00410, 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 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.
[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 on Machine Learning, ICML 2018, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 1261-1269. PMLR, 2018. URL http://proceedings.mlrpress/v80/dietterich18a.html.
[8] 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.
[9] Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing, 40(6):1662-1714, 2011.
[10] 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
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  $\mathbf{z}_1$  and a predicted output of the representation of the second view  $\mathbf{p} = p(\mathbf{z}_2)$ , where the function  $p$  is called the predictor. This is in place of using the InfoNCE loss to define contrastive pairs for  $\mathbf{z}_1$  and  $\mathbf{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:
$$
\mathcal {L} _ {\text {S i m S i a m}} \left(\mathbf {z} _ {1}, \mathbf {z} _ {2}\right) = \mathcal {D} \left(\mathbf {p} _ {1}, \mathbf {z} _ {2}\right) / 2 + \mathcal {D} \left(\mathbf {p} _ {2}, \mathbf {z} _ {1}\right) / 2,
$$
where  $\mathcal{D}$  denotes the cosine similarity function:
$$
\mathcal {D} (\mathbf {p}, \mathbf {z}) = - \frac {\mathbf {p}}{\| \mathbf {p} \| _ {2}} \cdot \frac {\mathbf {z}}{\| \mathbf {z} \| _ {2}},
$$
Note that  $\mathcal{D}$  is not a symmetric quantity as it uses a stop gradient operation on its second argument  $\mathbf{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  $\mathbf{x}_1$  and  $\mathbf{x}_2$ , we mask both views to get  $\mathbf{x}_1^{ig}$  and  $\mathbf{x}_2^{ig}$ , then process them through the encoder to compute the SimSiam loss  $\mathcal{L}(\mathbf{z}_1^{ig}, \mathbf{z}_2^{ig})$ . Finally, we also add the original SimSiam loss, i.e. one over the unmasked inputs to the overall objective:
$$
\begin{array}{l} \mathcal {L} = \mathcal {L} _ {\text {S i m S i a m}} \left(\mathbf {z} _ {1} ^ {i g}, \mathbf {z} _ {2} ^ {i g}\right) + \mathcal {L} _ {\text {S i m S i a m}} \left(\mathbf {z} _ {1}, \mathbf {z} _ {2}\right) \\ + \lambda \left(\left| \left| i g \left(\mathbf {x} _ {1}\right) \right| \right| _ {1} + \left| \left| i g \left(\mathbf {x} _ {2}\right) \right| \right| _ {1}\right) \tag {8} \\ \end{array}
$$
We test this version of InfoGating on CIFAR-10, while evaluating performance on Corrupted CIFAR-10 [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.
Figure 6: Generalizing to 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 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
InfoNCE w/o Distractor
InfoNCE w/ Distractor
Forward w/Distractor
Inverse w/Distractor
Figure 7: Effect of Losses on InfoGating Masks. InfoNCE w/o Distractor (mean score: 208.0) learns almost perfect masks, while InfoNCE w/ Distractor (mean score: 10.4) learns masks that are spread across the 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 as well.
Figure 8: Visualizing Masks Learned by InfoGating. STL-10 images alongside the masked image. The masking network gradually learns to trace the edge boundaries for both foreground and background objects.
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)  $\mathbf{x}_t^{ig}$  and  $\mathbf{x}_{t + k}^{ig}$  into corresponding encodings  $\mathbf{z}_t^{ig}$  and  $\mathbf{z}_{t + k}^{ig}$ . 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  $(\mathbf{z}_t^{ig},\mathbf{z}_{t + k}^{ig},\mathbf{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,  $\mathbf{a}_t$  is replaced by random actions  $\bar{\mathbf{a}}_t$  to form a 'negative' triplet, for which the model is trained to output a high energy value.