Structural Generalization of Visual Imitation Learning with Position-Invariant Regularization
Authors: Zhao-Heng Yin
Yang Gao
Qifeng Chen
Venue: Submitted to ICLR 2023 (2023)

Abstract: How the visual imitation learning models can generalize to novel unseen visual observations is a highly challenging problem. Such a generalization ability is very crucial for their real-world applications. Since this generalization problem has many different aspects, we focus on one case called structural generalization, which refers to generalization to unseen task setup, such as a novel setup of object locations in the robotic manipulation problem. In this case, previous works observe that the visual imitation learning models will overfit to the absolute information (e.g., coordinates) rather than the relational information between objects, which is more important for decision making. As a result, the models will perform poorly in novel scenarios. Nevertheless, so far, it remains unclear how we can solve this problem effectively. Our insight into this problem is to explicitly remove the absolute information from the features learned by imitation learning models so that the models can use robust, relational information to make decisions. To this end, we propose a novel, position-invariant regularizer for generalization. The proposed regularizer will penalize the imitation learning model when its features contain absolute, positional information of objects. We carry out experiments on the MAGICAL and ProcGen benchmark, as well as a real-world robot manipulation problem. We find that our regularizer can effectively boost the structural generalization performance of imitation learning models. Through both qualitative and quantitative analysis, we verify that our method does learn robust relational representations.


Under review as a conference paper at ICLR 2023
SPATIAL GENERALIZATION OF VISUAL IMITATION LEARNING
WITH
REGULARIZATION
Anonymous authors Paper under double-blie blind review
ABSTRACT
How the visual imitation learning models can generalize to novel unseen visual observations is highly challenging problem. Such a generalization ability very crucial for their real-world applications. Since this generalization problem has many different aspects, we focus on one case called spatial generalization, which refers generalization to unseen setup of object (entity) locations in task, such as a novel setup of object locations in the robotic manipulation problem. In this case, previous works observe that the visual imitation learning models will overfit to the absolute information (e.g., coordinates) rather than the relational information between objects, which is more important for decision making. As result, the models will perform poorly in novel object location setups. Nevertheless, so far, it remains unclear how we can solve this problem effectively. Our insight into this problem is to explicitly remove the absolute information from the features learned by imitation learning models so that the models can use robust, relational information to make decisions. To this end, we propose novel, regularizer for generalization. The proposed regularizer will penalize the imitation learning model when its features contain absolute, positional information of objects. We carry out experiments on the MAGICAL and ProcGen benchmark, as well as a real-world robot manipulation problem. We find that our regularizer effectively boost the spatial generalization performance of imitation learning models. Through both qualitative and quantitative analysis, verify that method does learn robust relational representations.
INTRODUCTION
Imitation learning is class of algorithms that enable robots to acquire behaviors from human demon- strations (Hussein al. 2017j. The recent advance in deep learning has boosted the development of visual imitation learning and supported its applications like autonomous driving, robotic manipulation, and human-robot interaction (Hussein al. 2017).
In spite of its success, visual imitation learning methods still face many practical challenges. One major challenge is its ability to generalize to novel unseen visual observations, which is very common when we deploy the trained models (Toyer et al. 2020) Park et al.) 2021). In the literature, this generalization problem is also known as the robustness problem. The problem covers many different aspects. For example, here we can identify two basic generalization capabilities: observational generalization and spatial generalization (Figure D Observational generalization refers to the generalization to novel visual textures. The changes in background color. object texture, or ambient light in the robotic manipulation task are examples of observational generalization. Such kind of visual change does not affect the physics structure (e.g., the position of object and targets) and only requires the robot to reason about semantic meanings correctly. In contrast, spatial generalization refers to the generalization to unseen setup of objects' (entities) locations in one task, which instead requires physical common sense about space and object. Consider the task of letting a warehouse robot move box to target region. If set the initial position of the box to place that is not covered by the demonstration dataset, then the imitation learning methods must be able to perform spatial generalization
succeed. In reality, the generalization challenge usually emerges

---

Under review as a conference paper at ICLR 2023
Examples of Visual Generalization
What should in feature?
Train
Test
Train
Test
Absolute Info
Relation Info
Observational Generalization
Spatial Generalization
Remove!
Keep
Two kinds of visual generalization. The examples based the MAGICAL benchmark provided by Toyer et al. (2020), in which a robot is required to relocate a box target region. The left figure shows example of observational generalization, in which the only change during the testing phase the visual texture of objects. The middle figure shows an example of spatial generalization. The object setup in the testing phase is unseen. Right: To achieve spatial generalization, we suggest that absolute information should be removed from the feature while the relational information should be kept. We propose a novel, position-invarian regularizer for this purpose.
Figure 1:
Left and Middle:
combination of different generalization capabilities. In this paper, we focus on the study of spatial generalization.
For better spatial generalization, the visual imitation learning models should be able to obtain knowledge about objects and their spatial relations with proper inductive biases. Some work finds that vanilla deep visual imitation learning models strongly overfit to the absolute position objects (Toyer) al. 2020). which suggests that they do not extract relational information of objects to make decisions like humans (Doumas et al.] 2022). Since the representation learning methods can usually lead to good semantic representations (features) and lead generalization, Chen al. (2021) investigate the use of self-supervised representation learning methods in visual imitation learning. However, they find that these purppse representations fail improve the generalization performance of vanilla visual imitation learning models effectively. Aside from these works, we also notice that some works propose variants of vision transformers (Dosovitskiy et al.) 2021 to improve spatial generalization (uuan et al. 2021), though these methods are not for imitation learning. Moreover, they make additional assumptions, such as the availability of object information. So far, remains unclear how to ensure spatial generalization in visual imitation learning.
Based on these observations, our main insight into this problem is to explicitly remove absolute, positional information from the learned features in the visual imitation learning models. Note that this does not mean that the decision-maaing process not dependent on absolute information. Rather, expect that the model can extract the relational information (e.g.. distance, direction) from the absolute information to make robust decisions. To this end, we propose a novel position-invariane regularizer called POINT. This regularizer will penalize the imitation learning model when it finds that the learned feature highly correlates with absolute, positional information. As a result, the imitation learning model has discover more robust relational features. To validate test the proposed regularizer on the MAGICAL (Toyer et al. 2020) and ProcGen (Cobbe et al. 2020) benchmark, as well real-world robot manipulation problem. We find that our method can effectively improve spatial generalization performance. Furthermore, we conduct qualitative and quantitative analysis and find that the imitation learning models can indeed learn relational features with our proposed regularizer.
idea,
To summarize, our contributions in this paper are as follows.
We define the spatial generalization problem of visual imitation learning models and propose novel Dosition-invariant regularizer called POINT to tackle this problem.
We test our method on MAGICAL and ProcGen benchmarks, as well as a real-world robot manipulation problem. We find that our proposed regularizer can improve the spatial generalization performance of previous imitation learning models effectively.
Through qualitative and quantitative studies, we verify that our proposed regularizer does make the visual imitation learning models extract relational information.
2

---

Under review as a conference paper at ICLR 2023
2
RELATED WORK
Imitation Learning Imitation learning (IL) is a classical method to solve sequential decision- making problems with expert demonstration data. Existing imitation learning methods consist of two classes of algorithms: behavioral cloning (Bain & Sammut 1995) and inverse reinforcement learning (Ng et al. (200). Behavioral cloning is a supervised learning method that directly fits experts' actions. Inverse reinforcement learning instead proposes to infer a reward function from experts' demonstrations and use it to train an RL agent. Both of these methods suffer from the generalization problem. In this paper, we mainly focus on behavioral cloning since it is simple yet effective in visual domains, and does not involve dangerous online interactions like IRL (Park et al. 2021
Generalizable Policy Learning The generalization problem of imitation learning has been a long-standing problem. Since the problem also exists in reinforcement learning and most of the solutions can be adapted to the case imitation, we discuss the works from both fields here. One existing branch of work is the domain randomization (Tobin et al. 2017). The basic idea of domain randomization is to collect data from diverse setups, such as different backgrounds and textures. As a result, the trained model will be able discover more robust features. Resembling this idea. another line of work tries to solve this problem with data augmentation (Yarats et al. 2020; Hansen & Wang 202T), which augments the input or the learned features according to some heuristics. For example, MixReg (Wang et al., 2020) proposes mixture augmentation method and leads to better dittrii-iitiitttion generalization. OREO (Park al., 2021) proposes to mask out the objects in the feature map so that the model will not fit specific object in the observation. CLOP (Bertoin & Rachelson 2022) finds it important to shuffle the feature map locally to ensure robustness. Though not exactly data augmentation approach, De Haan et al. (2019) propose to randomly mask out the semantic components in the learned feature and check whether they are correlated with correct expert actions. Aside from works using diverse data data augmentation, some works also approach this problem by incorporating proper inductive biases into the design of policy networks (Dasari to
& Gupta 2020). For instance, SORNet (Yuan et al., (021) uses variant of Dision-transfore solve robotics tasks that require relational reasoning. However, SORNet is not aimed at imitation learning or reinforcement learning. Zhou et al. (2022) propose to use an attentional architecture to solve a compositional, object-centric reinforcement learning problem. The idea of these methods to leverage the relational inductive bias of the attention operations. |Wen et al. (2022) propose PrimeNet, which introduces a shortcut into the neural network to ensure that the model focus on more meaningful areas. Representation learning methods are another approach for generalization. Some works use pretrained representation. For instance, Yen-Chen al. (2020) use pretrained model on MS-COCO dataset (Lin et al. 2014) to provide representation for imitation learning in robotics manipulation. The pretrained model over large datasets provides more robust object representations and can generalize the policy to unseen objects. Besides pretrained representation, self-supervised representation learning (Mandi et al., 2022: Chen al., 2202I) on the demonstration dataset can also improve the generalization in some cases. However, Chen et al. (202T) suggest that self-supervised representations do not improve the generalization performance in general.
We refer readers to Kirk et al. 2021 for comprehensive knowledge of this field. Our work differ- entiates from all these existing works by explicitly defining the spatial generalization problem and proposing to remove the nonrobust positional information from the representation.
3 PRELIMINARIES
Notations We model the sequential decision making problem a Markov Decision Process M = (S, A, R, T). S is the state space. A is the action space. R is the reward function. T is the transition dynamies. The agent's 's state timestep i is 8, E S. The agent takes action d, and receives reward re- = R(st,a,) at). Its state at timestep t + 1 is then St+1 ~ T T(s oz). The objective of the agent (0, 1] is discount factor.
to maximize the return Σ " where
For the imitation learning problem studied here, the agent has no access to R and T, but is provided with fixed expert demonstration dataset D {- Here, each T. (s a S. "f.aa....s......... s..... a is an expert trajectory that can achieve high performance (return) in M. Therefore, the agent should learn the behavior by leveraging the given demonstration dataset.
3

---

Under review as a conference paper at ICLR 2023
Gradient of fooling
Encoder fa
π
Action
Input s;
Feature fi
Periodic Copy
Model
Sample
Paired
fo
πt
Pr (xp))))
D
Sample
Target Model
Paired
Pj (x
Discriminator
Input $j
GradCAM**
Importance
Heatmap
Proposed Regularization
Figure 2: Overview of our method, The blue branch above is the common imitation learning (BC) pipeline. Our proposed regularizer is shown in the light pink box the bottom. The regularizer first uses the GradCAM+ algorithm to find out the important areas based on which the latest BC model makes decisions. Then it samples the coordinates from the discovered important areas and trains discriminator network D to calculate whether these sampled coordinates are paired with the feature fi. The BC model (encoder (a) is then trained fool the discriminator D. When the encoder so is able to fool D, the absolute positional information is removed from the feature as desired.
Behavioral Cloning One classical imitation learning algorithm is the Behavioral Cloning (BC). BC turns the imitation learning problem into a supervised learning problem. It fits the expert's 's action a given the observation s;. For the visual imitation learning problem, the BC model two consecutive parts: a vision encoder fa (which is usually convolutional neural network), and policy head æ. The fe first encodes s; the feature f; (s)). and the then uses it to predict the expert's 's action. The BC algorithm minimizes the following negative likelioeld objective:
be divided into
LBc (s..))) |- log ((a]]f(ss)))
(I)
Due to its simplicity, BC is widely used visual imitation learning. Therefore, we study the spatial generalization of BC in this paper.
GradCAM weighttd Class Activation Mapping (GradCAM) (Selvaraju et al. 2017
Chattopadhay et al. 2018) is a class of algorithms that can interpret the deep vision models by visualizing the important region for decision making, Its basic idea is to calculate the gradient at each feature map and aggregate them into an importance map. In this paper, we leverage GradCAM as rough object detector in the proposed regularizer. We refer readers to the original paper for more technical details.
METHOD
In this section, present our position-invariane regularizer for spatial generalization. As discussed before, our basic idea is to remove the absolute positional information from the features of the BC model. In the following parts, we first present formalized description of the idea and its practical challenges in Section 4.1 Then, we discuss how to handle these challenges with GradCAM and adversarial training in Section and
and 4.3] Finally, we provide a summary the overall algorithm in
Section
We illustrate our method in Figure and provide the pscudo code in Algorithm
4.1 FORMULATION AND CHALLENGES
For the tasks that involve spatial generalization, there usually exist multiple objects in the observed states, such the agent, the target and the goal. For the state 8,, denote each of these objects
4

---

Under review as a conference paper at ICLR 2023
in 8, as of, and their positions (z), v.) Then, our idea can be formulated the minimization problem of each T((x', y ).rr), f), where 7 is the mutual information. Note that we use the notation x², y to indicate the corresponding random variables of x?, vl ft. However, this formulation leads to many practical challenges. First, since each (i/ v²) not provided directly by s; and should be inferred, we have to either train some object key-point detectors to detect the underlying objects in the training set, or annotate the objects by ourselves. However, both of these approaches can be difficult and tedious in practice. Second, even if we have ideal key-point detectors, we have to deal with hard optimization problem the summation form Σ; y"), f). This can intractable
when there are many objects in the observed state.
Fortunately, we find that the previous works on the interpretation of deep learning models like GradCAM provide useful tools to handle these challenges. It can reduce the problem to much simpler form. We discuss our observations follows.
4.2 PROBLEM REDUCTION WITH GRADCAM
GradCAM is an interpretation method that can tell which part of the image is crucial in the decision process of deep learning model. Given BC model (fo: xt) and input s, GradCAM outputs an importance heatmap of the same resolution as the input 8. The heatmap indicates the importance of each pixel when that it is smooth and usually coincides with the meaningful objects in the input s. Therefore, we can consider the GradCAM as rough object detector here.
this BC model for prediction. One nice property of this generated heatmap is
We propose to sample = (Ii, Ve) from the generated heatmap, and then minimize the T(p, f). We find that this new objective can act as a proxy for the original objective in practice. Concretely, if ps is always far from a specific object like o², then we know that o is irrelevant to the decision process of the current model. In this case, we conjecture that I((x", )) f) should be low enough to meet our requirement. On the contrary, if always coincides with a certain object like o', then we actually minimize (p, f) [(xx, y"), f) as we want.
4.3
Loss FUNCTIONS
Now, our remaining work is to reduce the mutual information fp, However, we find that jointly estimating and minimizing the mutual information in our vision-based tasks is hard in practice. Since ultimate goal is to minimize the information of p in f, we instead propose adversarial training framework to achieve this goal.
Specifically, we introduce discriminator network D to play a two-player min-max game with the BC model as follows.
min max E(s;,a;) [log D(p,, fi) + log(1 1 D(p,, f())
(2)
In this min-max game, the discriminator D tries to tell the joint distribution of p and f. denoted as Pp.. from the product of their marginal distributions IP Peor. Meanwhile, the BC model is trying to fool the discriminator by removing the information of p from f. Applying the convergence theory of the generative adversarial network (GAN) (Goodfellow et al., (2000), we know that when fe is global minimizer of Equation 2 Pp.f ppor, which implies that 7 f) 0. Therefore this min-max game fulfills our requirement.
In practice, we train D to minimize the following binary classification loss function:
LD - -. [log D(ps, ft) + log(1 - D(p,, f())]
However, for the encoder fe. find that using LD as the loss function for training will result in instabilities. We assume this is because the fi term is present in both of the two terms in Equation 2 which is different from that in the original GAN objective. Therefore, we propose use the following loss function for optimization, which we find works well empirically:
Lreg (e [log D(ps, S+)]
(4)
Combining the BC loss, the loss function train the fo and is then
f. LBC Xhreg Be |- log क Ta (o:))) + A log D(ps, ft)
(5)
5

---

Under review as a conference paper at ICLR 2023
Distractor
MR-TI
MR-T1D
MR-T2
MR-T2D
MR-T3
Figure 3: The MAGICAL tasks used in our experiments. The grey robot is required to move the target objects (we mark them with red dots) to the target region. The red curve shows a possible plan solve the task (the interaction details like releasing box omitted). The long horizontal nature of this task brings additional challenges aside from the spatial generalization problem.
4.4 SUMMARY OF THE ALGORITHM
Finally, we put everything together and summa- rize our algorithm in Algorithm T We use Grad- CAM++ (Chattopadhay et al.) 2018). an improved version of GradCAM, calculate the heatmap. We find that it can deal with multi-object observation better than GradCAM. Similar the target network in the deep reinforcement learning. we propose to use a target model in the GradCAM calculation. This target model periodically copies the weights of the latest BC model. This ensures that the p: terms in Equation will change slowly throughout training so that the optimization of the min-max game between D and fa becomes stable. Our python implementation this algorithm in the supple- mentary material and Appendix A.
Algorithm 1: POINT
Initialize parameters of the BC model: fo. т, and the discriminator D. Initialize target model fo. म with fo, while not converged do Sample a batch {(s) 0:)} from D. Calculate each importance heatmap
by H; GradCAM s.).
Sample each p; from H4.
Optimize the BC model by Equ. 5
Optimize the discriminator by Equ.
Periodically update fo, & with fo, π.
EXPERIMENTS
In the experiments, we first test the performance of our method on two benchmarks that require spatial generalization: MAGICAL and ProcGen. We study the generalization according to the IID protocol (Kirk et al. 202T). This means that the training and testing task distributions are the same, though the test instance will be unseen. Then, we provide an analysis our algorithm through both qualitative and quantitive studies. Finally, extend our method to real robot manipulation problem.
5.1 TASK SETUP
MAGICAL The MAGICAL benchmark simulates a 2D robotic manipulation problem in a ware- house room. The tasks provided by the MAGICAL involve complex interactions between the agent and multiple objects, which require effective spatial generalization..
In the experiments, we use a variant of its MatchRegion task. this task, a robot is required to go across a square room to move some objects to target region specified by dashed rectangle. We set up several task instances of the MatchRegion task: MaachhRegion-TTrrget.... MatchRegion.Targett TarchReiion-Targete an illustration of these tasks in Figure 3 For each MatchRegion-Targe Target-3-Distract We provide task (MR-TX), there is no distractor object in the room, the robot only needs to move all the X objects into the target location. However, for the MatchRegion Target-X-Distract -Distract task (MR-TXD) there is an additional distractor object in the room. This object is also randomly placed in the room during testing. The existence of this distractor object not only increases the risk of learning spurious features but also adds to the difficulty of learning secure motions. As will discuss later, even the existence of one
Tarsett--

---

Under review as a conference paper at ICLR 2023
Table 1: Evaluation result on the MAGICAL and ProcGen benchmark. We show the average score on three random seeds. Our method can achieve state-of-the-arr results compared with the baselines.
Method
Vanilla
Dropout
Crop
Cutout
MixReg
OREO
CLOP
Ours
MR-TI
0.09
0.28
0.42
0.19
0.26
0.21
0.16
0.63
20.05
MR-TID
0.19
0.32
0.44
0.27
0.41
0.27
0.21
0.60
10.06
10.06
MR-T2
0.25
0.48
0.46
0.43
0.44
0.37
0.32
0.75
10.03
20.03
20.04
20.05
10.05
10.05
20.07
10.07
MR-T2D
0.27
0.35
0.38
0.32
0.33
0.27
0.23
0.70
10.06
10.03
20.01
20.03
10.03
10.03
20.04
20.04
MR-T3
0.23
0,51
0.47
0,32
0,48
0,42
0.35
0.66
10.02
10.03
10
10.04
10.05
10.04
10.07
20.08
Coinrun
0.75
0.81
0.80
0.76
0.79
0.69
0.82
0.87
20.05
+0.03
20.02
20.03
#0.03
20.03
20.05
#0.02
Jumper
0,63
0.61
0,64
0.66
0,67
0,63
0.68
0.72
10.07
10.06
10.03
#.004
#0.03
10.08
+0.02
10.04
Ninja
0.65
0.69
0.71
0.72
0,68
0.67
0.67
0.79
+00006
+0.02
40.02
+0.02
+0.02
#0.01
+0.02
20.02
Leaper
0.48
0.52
0.47
0.46
0.59
+0.07
20.05
Miner
0.16
0.28
0.25
0.21
0,23
0.22
0.19
0.35
+0.07
+0.05
+0.02
#0,04
Starpilot
0.33
0.53
0.61
0.39
0.55
0.35
0.42
0.71
10.05
10.14
20.08
20.06
10.10
10.03
20.04
20.12
Fruitbot
0.01
0.17
0.48
0.00
0,21
0.09
0.36
0.56
10.04
10.11
20.01
20.06
10.07
10.05
10.06
20.02
Bigfish 0.05
0.08
0.16
0,10
0.09
0.09
0.12
0.16
10.04
10.04
20.03
10.05
10.01
10.02
10.02
20.04
distractor object can lead to a significant increase of generalization difficulty. The study of more distractors is carried the analysis part.
For each of the tasks above, we collect its human demonstration dataset by ourselves. For each demonstration trajectory, we randomly set up the initial position of the objects, target region, and the robot. Note that the demonstration dataset provided by the original MAGICAL benchmark uses fixed initial position setup. However, we find that this setup too strict for spatial generalization if no other dataset or pretraining tasks are available to provide prior knowledge. For MR-T1, we collect 50 trajectories. For each of the other tasks, we collect 100 trajectories. The collection of all these trajectories takes 2 hours. We also study the outcome of using different number of trajectories in later analysis part.
ProcGen The ProcGen benchmark is an arcade-gamt game- based benchmark. These arcade games involve interactions between many semantic objects, including the player, monsters (or traps), bonuses, and tools. Though it is designed to test the generalization ability of RL agents, we use test the spatial generalization of imitation learning here. We use environments: Coinrun, Jumper, Ninja, Leaper, Miner, Starpilot, Fruitbot, and Bigfish. We follow the 'casy" generalization evaluation protocol used in previous works. We first train a PPO agent (Schulman et al. 2017) to collect demonstrations in the first 200 levels of ProcGen and use these demonstrations train the imitation learning model. We use 64 expert trajectories in the experiment. Then for the evaluation, test the imitation learning models all levels. We remove the background of the game throughout training and testing since this requires observational generalization.
5.2 BASELINES
For the vanilla BC policy, we train an IMPALA (Espeholt al., 2018) policy, whose encoder is residual convolutional neural network. We also try transtorrnt (Dosovitskiy et al. [202T) and relational network (Santoro et al. 2017) that have relational biases, but we find that they perform worse than IMPALA and do not report their results here. Then, implement the following baselines for comparison, (1) Dropout. Dropout (Srivastava et al. 2014) is a widely-used regularization method in deep learning. Many works also find it useful for improving the performance of imitation learning. (2) Crop Crop is popular data augmentation method used in the state-of the-th-rrttt visual RL algorithms like DrQ (Yarats et al. 2020) and RAD (Laskin et al. 2020 it encourages translation

---

Under review as a conference paper at ICLR 2023
Estimated 1((x),y))
Object Dropout
Ours
314
1.26
3.02
1.09
2.55
1.12
Dropout
Ours
Figure 4: Left: The GradCAM+ importance heatmap of the dropout model and our model on the MR-TID task. The red region indicates the most important region, while the dark blue indicates the least important region. The results suggest that the dropout model attends to the red distractor and is not robust. contrast, our model is able to attend to correct objects. Right: MI between feature and object's position. Our model contains less bits about absolute position information.
invariance to the perspective. (3). Cutout Cutout (De & Taylor 2017) is a data augmentation method. During training, it randomly erases part the input image. (4). MixReg MixReg (Wang et al. 2020) is data augmentation method. It generates interpolated. synthetic training data based on the original training dataset. (5). OREO OREO (Park et al. 2021) an object-aware data
augmentation method. It randomly drops the features of certain objects. (6) CLOP (Bertoin & Rachelson 2022) CLOP is a locality-awe data augmentation method. It randomly shuffles the feature map of the policy network in small, local region.
5.3 RESULTS
MAGICAL The result on MAGICAL is shown in Table 1 The performance is defined by the success rate of the trained policy, which the number target objects that are successfully transferred to the target region, divided by the total number of target objects. We observe that method is able to achieve hate-offthe--an results and outperform the baselines by a large margin. Concretely, improves the success rate by about 30%. Besides, find that most of the previous regularization methods do increase the success rate of the vanilla version and their results are similar to each other. This shows that they may solve some common issues in the generalization problem. However, their performance gap from our method suggests that we tackle a different issue here, which is overfitting absolute positions.
ProcGen The result ProcGen is shown in Table The performance is calculated the normalized agent's 's reward in the test environment (with respect to the random agent's 's and expert's performance). Though we observe some improvement, the performance gain is smaller in this case compared to that of MAGICAL. We hypothesize that this is because the decision-making process of ProcGen highly dependent on the local image patches, as suggested by CLOP. In this case, the BC model does not need to perform long-range range reasoning across objects frequently, and suffers less from the risk of overfitting to absolute positional information. Our result here also suggests that we need task-specific inductive biases so to achieve good generalization performance.
5.4 ANALYSIS
Qualitative Results To understand whether our method is robust, we use GradCAM+ to visu- alize the importance heatmap the learned models. For simplicity, we show the result on the MatchReigion 1-Distract task. We compare the result of our model to the model trained with dropout here (Figure 4). We notice that the dropout model tends to focus on the red distractor object rather than the correct target object. In contrast, our model is able to focus Even when the distance between the agent and the object is large, it can attend to the agent and the object simultaneously. The visualization results suggest that regularizer indeed leads relational features even when the vision network IMPALA does not have an explicit relational inductive bias. To understand this, we further use MINE estimator (Belghazi al. 2018) measure the MI between
the correct objects.
8

---

Under review as a conference paper at ICLR 2023
0.8
Succes Rate
Success Rate
Gripper Camera
Workspace
0.4
0.4
0
25%
50%
75%
100%
MixRog
CutDut
# of Distractors
Dropout
MixReg
CutOut
Dataset Size
CLOP
DRED
Durs
CLOP
OREO
Ours
Figure 5: The generalization performance to different num- ber of distractors on MR-TID.
Figure 6: The variation of per- formance on MAGICAL using the datasets of different sizes.
Figure 7: The setup of real- world robot manipulation exper- iments.
the learned representation and the position of different objects. Our representation indeed contains much fewer bits about the absolute position compared with dropout as expected.
Unseen Number of Distractors information. As a result, for the MAGICAL tasks, it should be able to ignore the distractor and generalize an unseen number of distractors. Therefore, MR-TID (where only one distractor presents) can generalize to MR-TID with the unseen number of distractors (e.g., 0, 2, 3). We also compare the results with the previous models. The result is shown in Figure 5] We find that our model is able to generalize to the case of 0, 2, 3, though the performance lower than the case of (training scenario). contrast, the prior model, such as the dropout model, fails in these unseen cases totally. This echoes our qualitative analysis results.
robust model should base its decision robust relational
whether our model trained on
Number of Demonstrations We also study whether the proposed method works when the amount of expert demonstrations is limited. For this purpose, we test our method on the MAGICAL with 25%, 50%, 75% of expert demonstrations. We show the averaged performance in Figure 6 We find that our method can achieve consistent improvement, though the performance decreases as the dataset becomes smaller. This result suggests that we still require a certain amount of diverse data to achieve spatial generalization.
5.5 REAL-W WORLD EXPERIMENTS
Finally, we test whether our method scales to the real-world pick- and-place manipulation problem. We extend the MR-TID a UR10 robot arm with a Robotiq parallcl-jaw gripper (Figure (7). As suggested by Hsu al. 2022), we use a gripper camera and a workspace camera to provide observation. For the BC model, we use two separate IMPALA encoders to process each camera image, concatenate their output features along with the z-coordinate of gripper, and feed them into an MLP. We use the proposed regularizer to regularize the workspace branch. We collect 75 human demonstrations for training. We compare our method to dropout with different numbers of distract objects. The result is shown in Table 2 Our method also achieves a large improvement in this problem. We provide detailed setup and qualitative results in the Appendix A.4 and A.3
Table 2:
The success rate of the real- world experiments. Our method is also effective here. Each test consists of 20 trials.
Method
Dropout
Ours
0 Dis. Obj
35%
55%
Dis. Obj 35%
60%
2 Dis. Obj
20%
50%
3 Dis. Obj
10%
45%
CONCLUSION
In this paper, we studied the spatial generalization problem imitation learning. To solve this problem, proposed novel position-invarian regularizer to remove the absolute positional information from the features. Through experiments on the MAGICAL and ProcGen benchmarks, as well as a robot manipulation system, confirmed that previous methods do overfit to the absolute position and showed that our proposed approach can effectively solve this problem. We hope that the proposed method can inspire future generalization research.

---

Under review as a conference paper at ICLR 2023
7 REPRODUCIBILITY STATEMENT
We provide the python source code of the main algorithm and the full model in the supplementary material. We will make our project publicly available.
REFERENCES
Michael Bain and Claude Sammut. A framework for behavioural cloning. In Machine Intelligence 15, pp. 103-129, 1995.
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In International Conference on Machine Learning (ICML), 2018.
David Bertoin and Emmanuel Rachelson. Local feature swapping for generalization reinforcement learning. In International Conference Learning Representations (ICLR), 2022.
Aditya Chattopadhay, Anirban Sarkar, Prantik Howlader, and Vineeth N Balasubramanian. Grad- cam++: Generalized gradient-based visual explanations for deep convolutional networks. In IEEE Winter Conference on Applications of Computer Vision (WACV), 2018.
Xin Chen, Sam Toyer, Cody Wild, Scott Emmons, Ian Fischer, Kuang-Huu Lee, Neel Alex, Steven H Wang, Ping Luo, Stuart Russell, et al. An empirical investigation of representation learning for imitation. In Neural Information Processing Systems Datasets and Benchmarks Track (NeurIPS), 2021.
Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International Conference Machine Learning (ICML), 2020.
Sudeep Dasari and Abhinav Gupta. Transformers for one-shot visual imitation. In Conference Robot Learning (CoRL), 2020.
Pim De Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. In Neural Information Processing Systems (NeurIPS), volume 32, 2019.
Terrance DeVeiie and Graham W Taylor. Improved regularization convolutional neural networks with cutout. arXiv preprint arXiv: arXiv: 170500452, 2017.
Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. International Conference on Learning Representation (ICLR), 2021.
Leonidas AA Doumas, Guillermo Puebla, Andrea Martin, and John Hummel. A theory of relation learning and cross- cross-domain generalization. Psychological review, 2022.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan. Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International Conference Machine Learning (ICML), 2018.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11):139-144, 2020.
Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmen- tation. In IEEE International Conference on Robotics and Automation (ICRA), 2021.
Kyle Hsu, Moo Jin Kim, Rafael Rafailov, Jiajun Wu, and Chelsea Finn. Vision-based manipulators need to also see from their hands. In International Conference on Learning Representation (ICLR), 2022.
10

---

Under review as a conference paper at ICLR 2023
Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning: A survey of learning methods. ACM Computing Surveys (CSUR), 50(2):1-35, 2017.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. survey of generalisation in deep reinforcement learning. arXiv preprint arXiv:2111.09794 2021.
Misha I Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Rein- forcement learning with augmented data. In Neural Information Processing Systems (NeurIPS), 2020.
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In European Conference Computer Vision (ECCV), 2014.
Zhao Mandi, Fangchen Liu, Kimin Lee, and Pieter Abbeel. Towards more generalizable one-shot visual imitation learning. 2022 IEEE International Conference on Robotics and Automation (ICRA), 2022.
Andrew Ng. Stuart Russell, et Algorithms for inverse reinforcement learning. In International Conference Machine Leurning (ICML), volume 1, pp. 2, 2000.
Jongjin Park, Younggyo Seo, Chang Liu, Li Zhao, Tao Qin, Jinwoo Shin, and Yan Liu. Object- aware regularization for addressing causal confusion in imitation learning. In Neural Information Processing Systems (NeurIPS), 2021.
Adam Santoro, David Raposo, David G Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. A simple neural network module for relational reasoning. In Neural Information Processing Systems (NeurIPS), 2017.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1007.066777 2017.
Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based localiza- tion. In International Conference on Computer Vision (ICCV), pp. 118-626,770 2017.
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 (JMLR), 19(1):1929-1958 2014.
Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), 2017.
Sam Toyer, Rohin Shah, Andrew Critch, and Stuart Russell. The magical benchmark for robust imitation. In Neural Information Processing Systems (NeurIPS), 2020.
Kaixin Wang, Bingyi Kang, Jie Shao, and Jiashi Feng. Improving generalization in reinforcement learning with mixture regularization. In Neural Information Processing Systems (NeurIPS), 2020.
Chuan Wen, Jianing Qian, Jierui Lin, Jiaye Teng, Dinesh Jayaraman, and Yang Gao. Fighting fire with fire: Avoiding dnn shortcuts through priming. In International Conference on Machine Learning (ICML), 2022.
Denis Yarats, lya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representa- tions (ICLR), 2020.
Lin Yen-Chen Andy Zeng, Shuran Song, Phillip Isola, and Tsung- Yi Lin. Learning to see before learning to act: Visual pre-training for manipulation. In IEEE International Conference on Robotics and Automation (ICRA), 2020.
11

---

Under review as a conference paper at ICLR 2023
Wentao Yuan, Chris Paxton, Karthik Desingh, and Dieter Fox. Someet: Spatial object-centric representations for sequential manipulation. In Conference on Robot Learning (CoRL), 2021.
Allan Zhou, Vikash Kumar, Chelsea Finn, and Aravind Rajeswaran. Policy architectures for compo- sitional generalization in control. arXiv preprint urXiv:2203.05966 2022.
12

---

Under review as a conference paper at ICLR 2023
A APPENDIX
A.1 IMPLEMENTATION OF THE ALGORITHM
We provide the core training code of our project in the supplementary material, which comes from our whole project. The code shows all the training details and the model implementation. It is divided into three files.
bc. py This python file provides the main training routine. The main class is the BehaviorCloning, The training routine is its train method.
model. This python file provides the implementation of BC policy model (Model) and the discriminator (Discrimkodel).
network. py This python file provides the implementation of the IMPALA en- coder (ResNet used by the BC policy model.
Note that these files are for illustration purpose. The whole runnable project will be made publicly available.
A.2 HYPERPARAMETERS
The size of image observation is (64, 64) in all the experiments. We the Adam optimizer (Kingma & Ba 2015) for the optimization of both the discriminator and the BC model. The learning rate of the optimizers is 0.0003. We also add weight decay of 0.0001 to the BC model. We batchsize of 128 throughout the experiments (batch.size), We will sample points for each batch (sample.n). We repeat the training of discriminator for steps (repeat), For the discrim- inator IDiscrimMode ), its embedding size is set to 64 (e_dim). For the choice of X, we use = 1.0 for the MAGICAL, and 0.1 otherwise. All the other details and hyperparameters can be found in the source code in the supplementary material.
A.3 QUALITATIVE RESULTS OF THE MANIPULATION PROBLEM
In this section, we provide some qualitative results of the real-works manipulation problem. Recall that in this task, the robot is required to move red cube to target location specified by green area. We show the importance heatmap of the dropout model (Figure 8 and our model (Figure As is shown in the figures, we find that dropout model tends to attend more the round distractor object compared with our model. However, due to the visual complexity, we find that our model sometimes may attend the shadow in the background.
A.4 SETUP REAL ROBOT
The observation of the robot include two resized 64 64 RGB images coming from the gripper camera and the workspace camera, and the z coordinate of the gripper. We use a discretized action space, it will become easier for human collect the demonstration via the keyboard. There are 8 actions in total, {x, y, 2] move [backward/forward] by 2cm, and gripper [open/close] by lcm. We only only observe
consider regularize the workspace camera. This is because the gripper-camera very limited area on the workspace in our case and the gripper is always at the center. As a result, the spatial generalization issue minor for the gripper-camera and we find do not help much to regularize it.
13

---

Under review as a conference paper at ICLR 2023
Figure 8: The GradCAM++ importance heatmap of dropout model in the real-worke manipulation
problem. The dropout model tends to attend the round distractor object.
Figure 9: The GradCAM+ importance heatmap of our model in the real-world manipulation problem. Our model attends less to the round distractor object. However, due to the visual complexity, we find that our model sometimes may attend the shadow in the background.
14


Reviews:
[{"review_id": "h3Ruv_rq3HR", "reviewer": "Reviewer_khtv", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces a position-invariant regularizer to remove the absolute information from features while preserving the robust relational information. The effectiveness of the proposed method in improving generalization is evaluated on MAGICAL, Procgen benchmark and a real-world robot manipulation problem.", "summary_of_the_review": "Overall, I think this paper is ok but requires some improvements to meet the acceptance bar.", "strength_and_weaknesses": "Strength:\n* The paper is in general well written.\n* The idea of leveraging GradCAM to handle the challenges (Sec. 4.2) is interesting and clever. \n\nWeaknesses and Comments:\n* The experiments on Procgen are a little bit limited. It would be better to evaluate the proposed method on all 16 games.\n* Apart from BC, does the position-invariant regularizer also work for other (maybe a SOTA one) imitation learning methods?\n* Compared to the baselines, how much extra cost (computation, time) does the adversarial training bring?\n* I suggest the authors using the metrics recommended in [1] to compare different methods.\n\n[1] Agarwal, Rishabh, et al. \"Deep reinforcement learning at the edge of the statistical precipice.\" Advances in neural information processing systems 34 (2021): 29304-29320.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "NnaVfXmczj", "reviewer": "Reviewer_efvU", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a method for learning 2D visual representations which explicitly discourage the model from memorizing absolute positions of objects on the scene, while still enabling the model to reason about their relative arrangement in order to derive action policies.", "summary_of_the_review": "I am a bit torn. While I can't find any major flaw in the paper, it remains the case that it's the scientific equivalent of using a nuclear bomb to hammer in a nail, and only succeeding in inserting it halfway. There is scientific relevance to such study, but little that would generalize beyond the narrow setting of the problem. My confidence level in this assessment is relatively low.\n[edit-post rebuttal: some concerns well-addressed. 5->6]", "strength_and_weaknesses": "Strengths:\n- interesting approach to solving the problem, and solid demonstration that solving this particular 'structural generalization' problem matters to improving BC policies.\n- paper is thorough, with a very good description of the method, and a sensible set of experimental results that provide convincing evidence.\n- code is provided\n- real-world experiments are provided, which greatly enhances the value of the experimental validation for robotics problems.\n\nWeaknesses:\n- The language used in the paper is sometimes getting in the way of clarity. Calling relative position invariance 'structural generalization' and expressing everything under that lens makes things unnecessarily difficult to read at first glance.\n[edit post rebuttal: addressed]\n- The level of complexity thrown at the problem in order to (partially) solve it is a definite concern: first you need to provide a saliency model, which are notoriously brittle to all but the simplest real-world images, and then you need to learn a complex discriminator model on top of it. This likely makes the approach prohibitive in practice, and means its usefulness is limited to being one purely academic step in addressing the problem.\n[edit post-rebuttal: Acknowledging that http://gradcam.cloudcv.org/ was shown to be robust to real-world scenarios in prior works. Validation of its use in real-world downstream tasks is still missing (not a huge concern) and excessive complexity-to-outcomes ratio argument still stands]\n- A number of works (https://transporternets.github.io/ comes to mind) bypass the issue of learning position-invariant representations by decoupling the policy representation from the visual representation learning, and having the policy only reason about the 'delta' between visual representations, while still enabling end-to-end differentiability of the architecture. I would have loved to see more discussion of such approaches.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "D5NVDGDZEF", "reviewer": "Reviewer_xoNn", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes an approach for inducing invariance to absolute object positions in visual imitation learning policies by determining the policy’s regions of attention in the input image and using a discriminator to ensure the learned visual representations do not encode information about their position. In experiments on abstract block pushing tasks, ProcGen and a real robot block pushing task the introduced regularization leads to better policy performance than a number of previous regularizers.", "summary_of_the_review": "I have a number of concerns with regards to the proposed method and it’s experimental evaluation. In particular, I believe the paper is too broadly scoped for the proposed method, its not clearly explained how the experiments test the desired generalization capabilities and vital baselines are missing. Thus, I cannot recommend acceptance of the paper.", "strength_and_weaknesses": "# Strengths\n- the approach is well-explained and easy to follow\n- it is demonstrated across three different environments, including a real robot\n- Fig 2 makes it quite easy to understand the approach\n- the discussion of related work is comprehensive\n\n# Weaknesses\n\nThere are several major concerns regarding the problem formulation, approach and experimental evaluation:\n\n(A) **limited definition of structural generalization**: The paper defines structural generalization as generalization to scenes with different absolute positions of objects (and thus argues that in order to allow policies to perform structural generalization, we need to remove absolute position information form the learned representations). However, this is a very limited definition of structural generalization: there are many other aspects of such generalization, e.g. generalization to new instances of objects, new distractors, new scene layouts (that also involve relative location changes), more instances of the objects (eg pick up two cups after being trained on picking up one cup). As a result the scoping of the paper seems overly general for the actual invariances that the model introduces. Alternative approaches mentioned in the related work, e.g. the object centric approach by Zhou et al. 2022, can achieve a broader range of the aforementioned structural policy generalization.\n\n(B) **objective can encourage information-less representations**: The introduced objective uses a discriminator to decrease mutual information between the regions of attention of the policy (computed via GradCAM) and the learned representation. There are two ways to reduce this mutual information: by removing position information from the representation or by uniformly attending to all parts of the image (so that the GradCAM attention visualization is uninformative). The latter solution is undesirable and does not align with the intuition of the proposed method, but it seems that the objective does not explicitly encourage the first over the second solution.\n\n(C) **experiments don’t describe how structural generalization is tested**: The description of the experimental evaluation does not explicitly outline how structural generalization is tested. What is the structural difference between the training demonstrations and the test tasks for the different environments? This information is crucial and without it understanding the experimental results is hard. (For example in the MAGICAL environment experiments the demonstrations already use randomized object positions, so it is hard to understand what unseen structural generalization is required at test time)\n\n(D) **missing baselines: random crop data augmentation**: The proposed method is an approach for inducing a certain invariance in the policy (absolute position invariance). As mentioned in the related work, an alternative for inducing invariances is data augmentation. While the experimental evaluation compares to multiple augmentation methods, it surprisingly omits random image cropping. This data augmentation has been shown to be the most effective for visual policy learning in prior work (DrQ, RAD) and also aims to induce positional invariance. At the same time it is substantially easier than the proposed approach. Thus, a comparison to this baseline and a qualitative comparison of the learned representation is crucial.\n\n\nThere is also a number of smaller points of critique:\n\n(E) qualitative analysis does not investigate positional information in representation: the presented qualitative analysis (Fig 4) does not actually investigate the learned representation or the core property the paper tries to induce in it, positional invariance. It merely shows that the policy learns to attend to the relevant objects, but I would expect any functioning policy to show such patterns. Instead, some measure of the mutual information between the learned representation and the absolute object position should be reported and compared to baselines.\n\n(F) dropout model attends to distractor, but not relevant objects: following the previous point, it is unclear to me why the dropout policy in Fig 4 would learn to attend to the distractor. I would expect any policy that properly learns to solve the task to attend to the relevant objects, not the distractor. The main qualitative difference between the baselines and the proposed approach should be in the degree to which the learned representation contains absolute position information. This could suggest that the baselines had some training issues unrelated to positional information.\n\n(G) no ablation studies: the paper does not perform ablation studies of the proposed approach — for example it could be good to ablate the simplified vs full loss function (eq. 3 vs 4) and the introduced target network.\n\n(H) not enough details about the robot experiment: the description of the robot experiment is too short, not enough details are provided eg about the action space of the agent, about why only the workspace camera encoder is regularized or how structural generalization is tested between training and test (see point (C) above).\n\n# Questions\n\n- How would the method perform for tasks with moving cameras where the 2D GradCAM projection might not accurately reflect the 3D position of the attended objects?\n\n- Why does the proposed regularization, which aims to remove absolute object position, make the policy more robust to distractors? This is not intuitively clear to me.\n\n- How would the method perform in tasks where absolute object position is important?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject