SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $(\mathbf{x}_t^{i g},\mathbf{a}_t)$  to  $\bar{\mathbf{z}}_t^{i g})$  while encoding the future/goal observation  $\mathbf{x}_t^{i g}$  to a similar size vector  $\mathbf{z}_{t + k}^{i g}$ . 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  $\mathbf{a}_t$  with the  $\mathbf{z}_t^{i g}$
representation, before projecting it to the same dimension as the goal representation  $\mathbf{z}_{t + k}^{ig}$ . Finally, we penalize the info-gates for both current and goal observations as done previously to obtain the following objective:
$$
\mathcal {L} = \operatorname {I n f o N C E} \left(\bar {\mathbf {z}} _ {t} ^ {i g}, \mathbf {z} _ {t + k} ^ {i g}\right) + \lambda \left(\| i g (\mathbf {x} _ {t}) \| _ {1} + \| i g (\mathbf {x} _ {t + k}) \| _ {1}\right), \tag {9}
$$
where  $\bar{\mathbf{z}}_t^{ig} = g(\mathbf{z}_t^{ig},\mathbf{a}_t)$  is a projection function that maps the state and action embedding to a space of the same size as  $\mathbf{z}_t^{ig}$ . 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 "fwd + IG" adds InfoGating. For reference, training a forward dynamics model without distractors has a return of  $173.6 \pm 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
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.
Table 8: Sharing Masking Network and Encoder Parameters. Experiment setup follows Table 16. Default InfoGating version is marked in gray.
case	easy	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 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).
Table 9: UNet Architecture.
Down Sampling (InfoGating Network)
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 64, stride 1, pad 1, GroupNorm, ReLU
3 x 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)
3 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
1 x 1 conv2d 128, stride 1, pad 0, ReLU
3 x 3 conv2d 128, stride 1, pad 0, ReLU
4 x 4 conv2d 256, stride 2, pad 0, ReLU
FC 256, LayerNorm, ReLU
F Rebuttal
Figure 9: InfoGating masks for held-out distractors (during evaluation time), and for in-distribution distractors (during training). The IG masks generalize quite well to unseen distractors.
Table 11: Visual D4RL Locomotion Training Details.
batch size	128
λ	0.1
IG warm-up	5k steps
cropping padding	4
frame_stack	3
actionrepeat	2
buffer_size	100000
learning_rate	1e-4
eval Episodes	10
Table 12: Kitchen Training Details.
batch size	32
λ	schedule(0.1, 3.0, 2k steps)
IG warm-up	1k steps
cropping padding	none
frame_stack	1
num_demos	5 trajectories
learning_rate	1e-3
IG learning_rate	1e-4
actionrepeat	1
eval_trajecories	50
Table 13: CIFAR/STL-10 Training Details.
batch size	32
momentum	0.9
weight_Decay	1e-6
λ	0.06
IG warm-up	100 steps
cropping scale	(0.2, 1.0)
learning_rate	0.3
IG learning_rate	0.3
num_epochs	60
Table 14: Noisy Evaluation. Comparing InfoGating performance with multi-step inverse dynamics when training on noisy distractors and evaluating on held-out / out-of-distribution distractors (instead of noise-free evaluation as in Table 1 of the main paper). We also evaluate performance on unmasked observations (right-last column).
case	IQL	DRIML	Inv w/ Dropout	Inv w/ VIB	Inv w/ RCAD	Inv w/ IG	Inv w/ IG (unmasked eval)
expert
easy	33.8 ± 6.6	14.0 ± 5.8	68.4 ± 12.2	56.9 ± 1.3	41.9 ± 11.9	92.3 ± 19.4	206.4 ± 5.8
medium	12.6 ± 4.3	12.0 ± 8.0	28.4 ± 5.2	29.9 ± 6.6	13.6 ± 8.7	48.9 ± 7.0	60.2 ± 6.9
hard	17.4 ± 2.9	10.6 ± 2.1	17.1 ± 3.9	17.3 ± 3.9	9.9 ± 3.8	17.6 ± 4.5	12.5 ± 6.4
medium
easy	47.5 ± 12.9	26.1 ± 0.4	54.4 ± 4.4	84.2 ± 10.0	25.1 ± 12.2	53.1 ± 7.8	78.0 ± 27.5
medium	13.0 ± 4.8	16.1 ± 3.7	22.3 ± 2.1	18.8 ± 6.4	5.6 ± 1.1	48.1 ± 8.3	130.0 ± 21.2
hard	5.5 ± 4.8	15.3 ± 1.6	11.5 ± 2.8	8.6 ± 1.7	4.2 ± 0.5	5.0 ± 1.0	56.4 ± 14.7
Table 15: Walker Results. Comparing InfoGating performance with multi-step inverse dynamics when training on noisy distractors and evaluating on a held-out distractors (instead of noise-free evaluation as in Table 1).
case	IQL	DRIML	Inv w/ Dropout	Inv w/ VIB	Inv w/ RCAD	Inv w/ IG
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  $\lambda$ . Comparing InfoGating performance when using three different  $\lambda$  networks vs one  $\lambda$  network. All experimental settings are the same as in Table 1 of the main paper. For the multiple  $\lambda$  version, we do not train on mix of unmasked and masked observations.
case	Inv w/ IG (λ = 0.1, 0.2, 0.3)	Inv w/ IG (λ = 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