Original: S PATIAL G ENERALIZATION OF V ISUAL I MITATION L EARNING WITH P OSITION -I NVARIANT R EGULARIZATION
Revised: S PATIAL G ENERALIZATION OF V ISUAL I MITATION L EARNING WITH P OSITION -I NVARIANT R EGULARIZATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: How the visual imitation learning models can generalize to novel unseen visual observations is a highly challenging problem.
Revised: How the visual imitation learning models can generalize to novel unseen visual observations is a highly challenging problem.

Original: Such a generalization ability is very crucial for their real-world applications.
Revised: Such a generalization ability is very crucial for their real-world applications.

Original: Since this generalization problem has many different aspects, we focus on one case called spatial generalization , which refers to generalization to unseen setup of object (entity) locations in a task, such as a novel setup of object locations in the robotic manipulation problem.
Revised: Since this generalization problem has many different aspects, we focus on one case called spatial generalization , which refers to generalization to unseen setup of object (entity) locations in a task, such as a novel setup of object locations in the robotic manipulation problem.

Original: 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.
Revised: 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.

Original: As a result, the models will perform poorly in novel object location setups.
Revised: As a result, the models will perform poorly in novel object location setups.

Original: Nevertheless, so far, it remains unclear how we can solve this problem effectively.
Revised: Nevertheless, so far, it remains unclear how we can solve this problem effectively.

Original: 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.
Revised: 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.

Original: To this end, we propose a novel, position-invariant regularizer for generalization.
Revised: To this end, we propose a novel, position-invariant regularizer for generalization.

Original: The proposed regularizer will penalize the imitation learning model when its features contain absolute, positional information of objects.
Revised: The proposed regularizer will penalize the imitation learning model when its features contain absolute, positional information of objects.

Original: We carry out experiments on the MAGICAL and ProcGen benchmark, as well as a real-world robot manipulation problem.
Revised: We carry out experiments on the MAGICAL and ProcGen benchmark, as well as a real-world robot manipulation problem.

Original: We find that our regularizer can effectively boost the spatial generalization performance of imitation learning models.
Revised: We find that our regularizer can effectively boost the spatial generalization performance of imitation learning models.

Original: Through both qualitative and quantitative analysis, we verify that our method does learn robust relational representations.
Revised: Through both qualitative and quantitative analysis, we verify that our method does learn robust relational representations.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Imitation learning is a class of algorithms that enable robots to acquire behaviors from human demonstrations (Hussein et al., 2017).
Revised: Imitation learning is a class of algorithms that enable robots to acquire behaviors from human demonstrations (Hussein et al., 2017).

Original: 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 et al., 2017).
Revised: 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 et al., 2017).

Original: In spite of its success, visual imitation learning methods still face many practical challenges.
Revised: In spite of its success, visual imitation learning methods still face many practical challenges.

Original: 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).
Revised: 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).

Original: In the literature, this generalization problem is also known as the robustness problem.
Revised: In the literature, this generalization problem is also known as the robustness problem.

Original: The problem covers many different aspects.
Revised: The problem covers many different aspects.

Original: For example, here we can identify two basic generalization capabilities: observational generalization and spatial generalization (Figure 1).
Revised: For example, here we can identify two basic generalization capabilities: observational generalization and spatial generalization (Figure 1).

Original: Observational generalization refers to the generalization to novel visual textures.
Revised: Observational generalization refers to the generalization to novel visual textures.

Original: The changes in background color, object texture, or ambient light in the robotic manipulation task are examples of observational generalization.
Revised: The changes in background color, object texture, or ambient light in the robotic manipulation task are examples of observational generalization.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Consider the task of letting a warehouse robot move a box to some target region.
Revised: Consider the task of letting a warehouse robot move a box to some target region.

Original: If we set the initial position of the box to a place that is not covered by the demonstration dataset, then the imitation learning methods must be able to perform spatial generalization so as to succeed.
Revised: If we set the initial position of the box to a place that is not covered by the demonstration dataset, then the imitation learning methods must be able to perform spatial generalization so as to succeed.

Original: In reality, the generalization challenge usually emerges as a
Revised: In reality, the generalization challenge usually emerges as a

Original: [Figure]
Revised: [Figure]

Original: Train
Revised: Train

Original: Test
Revised: Test

Original: Spatial Generalization
Revised: Spatial Generalization

Original: Remove!
Revised: Remove!

Original: Keep
Revised: Keep

Original: Absolute Info
Revised: Absolute Info

Original: Relation Info
Revised: Relation Info

Original: Examples of Visual Generalization
Revised: Examples of Visual Generalization

Original: What should be in the feature?
Revised: What should be in the feature?

Original: combination of different generalization capabilities.
Revised: combination of different generalization capabilities.

Original: In this paper, we focus on the study of spatial generalization .
Revised: In this paper, we focus on the study of spatial generalization .

Original: 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.
Revised: 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.

Original: Some work finds that vanilla deep visual imitation learning models strongly overfit to the absolute position of objects (Toyer et al., 2020), which suggests that they do not extract relational information of objects to make decisions like humans (Doumas et al., 2022).
Revised: Some work finds that vanilla deep visual imitation learning models strongly overfit to the absolute position of objects (Toyer et al., 2020), which suggests that they do not extract relational information of objects to make decisions like humans (Doumas et al., 2022).

Original: Since the representation learning methods can usually lead to good semantic representations (features) and lead to generalization, Chen et al.
Revised: Since the representation learning methods can usually lead to good semantic representations (features) and lead to generalization, Chen et al.

Original: (2021) investigate the use of self-supervised representation learning methods in visual imitation learning.
Revised: (2021) investigate the use of self-supervised representation learning methods in visual imitation learning.

Original: However, they find that these general-purpose representations fail to improve the generalization performance of vanilla visual imitation learning models effectively.
Revised: However, they find that these general-purpose representations fail to improve the generalization performance of vanilla visual imitation learning models effectively.

Original: Aside from these works, we also notice that some works propose variants of vision transformers (Dosovitskiy et al., 2021) to improve spatial generalization (Yuan et al., 2021), though these methods are not for imitation learning.
Revised: Aside from these works, we also notice that some works propose variants of vision transformers (Dosovitskiy et al., 2021) to improve spatial generalization (Yuan et al., 2021), though these methods are not for imitation learning.

Original: Moreover, they make additional assumptions, such as the availability of object information.
Revised: Moreover, they make additional assumptions, such as the availability of object information.

Original: So far, it remains unclear how to ensure spatial generalization in visual imitation learning.
Revised: So far, it remains unclear how to ensure spatial generalization in visual imitation learning.

Original: Based on these observations, our main insight into this problem is to explicitly remove the absolute, positional information from the learned features in the visual imitation learning models.
Revised: Based on these observations, our main insight into this problem is to explicitly remove the absolute, positional information from the learned features in the visual imitation learning models.

Original: Note that this does not mean that the decision-making process is not dependent on absolute information.
Revised: Note that this does not mean that the decision-making process is not dependent on absolute information.

Original: Rather, we expect that the model can extract the relational information (e.g., distance, direction) from the absolute information to make robust decisions.
Revised: Rather, we expect that the model can extract the relational information (e.g., distance, direction) from the absolute information to make robust decisions.

Original: To this end, we propose a novel position-invariant regularizer called POINT.
Revised: To this end, we propose a novel position-invariant regularizer called POINT.

Original: This regularizer will penalize the imitation learning model when it finds that the learned feature highly correlates with absolute, positional information.
Revised: This regularizer will penalize the imitation learning model when it finds that the learned feature highly correlates with absolute, positional information.

Original: As a result, the imitation learning model has to discover more robust relational features.
Revised: As a result, the imitation learning model has to discover more robust relational features.

Original: To validate our idea, we test the proposed regularizer on the MAGICAL (Toyer et al., 2020) and ProcGen (Cobbe et al., 2020) benchmark, as well as a real-world robot manipulation problem.
Revised: To validate our idea, we test the proposed regularizer on the MAGICAL (Toyer et al., 2020) and ProcGen (Cobbe et al., 2020) benchmark, as well as a real-world robot manipulation problem.

Original: We find that our method can effectively improve spatial generalization performance.
Revised: We find that our method can effectively improve spatial generalization performance.

Original: Furthermore, we conduct qualitative and quantitative analysis and find that the imitation learning models can indeed learn relational features with our proposed regularizer.
Revised: Furthermore, we conduct qualitative and quantitative analysis and find that the imitation learning models can indeed learn relational features with our proposed regularizer.

Original: To summarize, our contributions in this paper are as follows.
Revised: To summarize, our contributions in this paper are as follows.

Original: • We define the spatial generalization problem of visual imitation learning models and propose a novel position-invariant regularizer called POINT to tackle this problem.
Revised: • We define the spatial generalization problem of visual imitation learning models and propose a novel position-invariant regularizer called POINT to tackle this problem.

Original: • We test our method on MAGICAL and ProcGen benchmarks, as well as a real-world robot manipulation problem.
Revised: • We test our method on MAGICAL and ProcGen benchmarks, as well as a real-world robot manipulation problem.

Original: We find that our proposed regularizer can improve the spatial generalization performance of previous imitation learning models effectively.
Revised: We find that our proposed regularizer can improve the spatial generalization performance of previous imitation learning models effectively.

Original: • Through qualitative and quantitative studies, we verify that our proposed regularizer does make the visual imitation learning models extract relational information.
Revised: • Through qualitative and quantitative studies, we verify that our proposed regularizer does make the visual imitation learning models extract relational information.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Imitation Learning Imitation learning (IL) is a classical method to solve sequential decisionmaking problems with expert demonstration data.
Revised: Imitation Learning Imitation learning (IL) is a classical method to solve sequential decisionmaking problems with expert demonstration data.

Original: Existing imitation learning methods consist of two classes of algorithms: behavioral cloning (Bain & Sammut, 1995) and inverse reinforcement learning (Ng et al., 2000).
Revised: Existing imitation learning methods consist of two classes of algorithms: behavioral cloning (Bain & Sammut, 1995) and inverse reinforcement learning (Ng et al., 2000).

Original: Behavioral cloning is a supervised learning method that directly fits experts’ actions.
Revised: Behavioral cloning is a supervised learning method that directly fits experts’ actions.

Original: Inverse reinforcement learning instead proposes to infer a reward function from experts’ demonstrations and use it to train an RL agent.
Revised: Inverse reinforcement learning instead proposes to infer a reward function from experts’ demonstrations and use it to train an RL agent.

Original: Both of these methods suffer from the generalization problem.
Revised: Both of these methods suffer from the generalization problem.

Original: 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).
Revised: 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).

Original: Generalizable Policy Learning The generalization problem of imitation learning has been a long-standing problem.
Revised: Generalizable Policy Learning The generalization problem of imitation learning has been a long-standing problem.

Original: Since the problem also exists in reinforcement learning and most of the solutions can be adapted to the case of imitation, we discuss the works from both fields here.
Revised: Since the problem also exists in reinforcement learning and most of the solutions can be adapted to the case of imitation, we discuss the works from both fields here.

Original: One existing branch of work is the domain randomization (Tobin et al., 2017).
Revised: One existing branch of work is the domain randomization (Tobin et al., 2017).

Original: The basic idea of domain randomization is to collect data from diverse setups, such as different backgrounds and textures.
Revised: The basic idea of domain randomization is to collect data from diverse setups, such as different backgrounds and textures.

Original: As a result, the trained model will be able to discover more robust features.
Revised: As a result, the trained model will be able to discover more robust features.

Original: Resembling this idea, another line of work tries to solve this problem with data augmentation (Yarats et al., 2020;
Revised: Resembling this idea, another line of work tries to solve this problem with data augmentation (Yarats et al., 2020;

Original: Hansen & Wang, 2021), which augments the input or the learned features according to some heuristics.
Revised: Hansen & Wang, 2021), which augments the input or the learned features according to some heuristics.

Original: For example, MixReg (Wang et al., 2020) proposes a mixture augmentation method and leads to better out-of-distribution generalization.
Revised: For example, MixReg (Wang et al., 2020) proposes a mixture augmentation method and leads to better out-of-distribution generalization.

Original: OREO (Park et al., 2021) proposes to mask out the objects in the feature map so that the model will not fit a specific object in the observation.
Revised: OREO (Park et al., 2021) proposes to mask out the objects in the feature map so that the model will not fit a specific object in the observation.

Original: CLOP (Bertoin &
Revised: CLOP (Bertoin &

Original: Rachelson, 2022) finds it important to shuffle the feature map locally to ensure robustness.
Revised: Rachelson, 2022) finds it important to shuffle the feature map locally to ensure robustness.

Original: Though not exactly a data augmentation approach, De Haan et al.
Revised: Though not exactly a data augmentation approach, De Haan et al.

Original: (2019) propose to randomly mask out the semantic components in the learned feature and check whether they are correlated with correct expert actions.
Revised: (2019) propose to randomly mask out the semantic components in the learned feature and check whether they are correlated with correct expert actions.

Original: Aside from works using diverse data or data augmentation, some works also approach this problem by incorporating proper inductive biases into the design of policy networks (Dasari & Gupta, 2020).
Revised: Aside from works using diverse data or data augmentation, some works also approach this problem by incorporating proper inductive biases into the design of policy networks (Dasari & Gupta, 2020).

Original: For instance, SORNet (Yuan et al., 2021) uses a variant of vision-transformer to solve robotics tasks that require relational reasoning.
Revised: For instance, SORNet (Yuan et al., 2021) uses a variant of vision-transformer to solve robotics tasks that require relational reasoning.

Original: However, SORNet is not aimed at imitation learning or reinforcement learning.
Revised: However, SORNet is not aimed at imitation learning or reinforcement learning.

Original: Zhou et al.
Revised: Zhou et al.

Original: (2022) propose to use an attentional architecture to solve a compositional, object-centric reinforcement learning problem.
Revised: (2022) propose to use an attentional architecture to solve a compositional, object-centric reinforcement learning problem.

Original: The idea of these methods is to leverage the relational inductive bias of the attention operations.
Revised: The idea of these methods is to leverage the relational inductive bias of the attention operations.

Original: Wen et al.
Revised: Wen et al.

Original: (2022) propose
Revised: (2022) propose

Original: PrimeNet, which introduces a shortcut into the neural network to ensure that the model focus on more meaningful areas.
Revised: PrimeNet, which introduces a shortcut into the neural network to ensure that the model focus on more meaningful areas.

Original: Representation learning methods are another approach for generalization.
Revised: Representation learning methods are another approach for generalization.

Original: Some works use pretrained representation.
Revised: Some works use pretrained representation.

Original: For instance, Yen-Chen et al.
Revised: For instance, Yen-Chen et al.

Original: (2020) use a pretrained model on MS-COCO dataset (Lin et al., 2014) to provide representation for imitation learning in robotics manipulation.
Revised: (2020) use a pretrained model on MS-COCO dataset (Lin et al., 2014) to provide representation for imitation learning in robotics manipulation.

Original: The pretrained model over large datasets provides more robust object representations and can generalize the policy to unseen objects.
Revised: The pretrained model over large datasets provides more robust object representations and can generalize the policy to unseen objects.

Original: Besides pretrained representation, self-supervised representation learning (Mandi et al., 2022; Chen et al., 2021) on the demonstration dataset can also improve the generalization in some cases.
Revised: Besides pretrained representation, self-supervised representation learning (Mandi et al., 2022; Chen et al., 2021) on the demonstration dataset can also improve the generalization in some cases.

Original: However, Chen et al.
Revised: However, Chen et al.

Original: (2021) suggest that self-supervised representations do not improve the generalization performance in general.
Revised: (2021) suggest that self-supervised representations do not improve the generalization performance in general.

Original: We refer readers to (Kirk et al., 2021) for comprehensive knowledge of this field.
Revised: We refer readers to (Kirk et al., 2021) for comprehensive knowledge of this field.

Original: Our work differentiates from all these existing works by explicitly defining the spatial generalization problem and proposing to remove the nonrobust positional information from the representation.
Revised: Our work differentiates from all these existing works by explicitly defining the spatial generalization problem and proposing to remove the nonrobust positional information from the representation.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: Notations We model the sequential decision making problem as a Markov Decision Process
Revised: Notations We model the sequential decision making problem as a Markov Decision Process

Original: M = ( S , A , R , T ) .
Revised: M = ( S , A , R , T ) .

Original: S is the state space.
Revised: S is the state space.

Original: A is the action space.
Revised: A is the action space.

Original: R is the reward function.
Revised: R is the reward function.

Original: T is the transition dynamics.
Revised: T is the transition dynamics.

Original: The agent’s state at timestep t is s t ∈ S .
Revised: The agent’s state at timestep t is s t ∈ S .

Original: The agent takes action a t and receives reward r t =
Revised: The agent takes action a t and receives reward r t =

Original: R ( s t , a t ) .
Revised: R ( s t , a t ) .

Original: Its state at timestep t + 1 is then s t +1 ∼ T ( s t , a t ) .
Revised: Its state at timestep t + 1 is then s t +1 ∼ T ( s t , a t ) .

Original: The objective of the agent is to maximize the return (cid:80)
Revised: The objective of the agent is to maximize the return (cid:80)

Original: Tt =

Revised: Tt =

Original: γ t r t , where γ ∈ (0 , 1] is a discount factor.
Revised: γ t r t , where γ ∈ (0 , 1] is a discount factor.

Original: For the imitation learning problem studied here, the agent has no access to R and T , but it is provided with a fixed expert demonstration dataset D
Revised: For the imitation learning problem studied here, the agent has no access to R and T , but it is provided with a fixed expert demonstration dataset D

Original: = { τ i } .
Revised: = { τ i } .

Original: Here, each τ i =
Revised: Here, each τ i =

Original: ( s E 0 , a E 0 , s E 1 , a E 1 , ...
Revised: ( s E 0 , a E 0 , s E 1 , a E 1 , ...

Original: s ET , a ET ) is an expert trajectory that can achieve high performance (return) in M .
Revised: s ET , a ET ) is an expert trajectory that can achieve high performance (return) in M .

Original: Therefore, the agent should learn the behavior by leveraging the given demonstration dataset.
Revised: Therefore, the agent should learn the behavior by leveraging the given demonstration dataset.

Original: [Figure]
Revised: [Figure]

Original: Regularization
Revised: Regularization

Original: [Figure]
Revised: [Figure]

Original: NotPaired p = ( x , y ) i
Revised: NotPaired p = ( x , y ) i

Original: i i p
Revised: i i p

Original: = ( x , y )
Revised: = ( x , y )

Original: j j j
Revised: j j j

Original: Discriminator
Revised: Discriminator

Original: Gradient of fooling D
Revised: Gradient of fooling D

Original: [Figure]
Revised: [Figure]

Original: Behavioral Cloning One classical imitation learning algorithm is the Behavioral Cloning (BC).
Revised: Behavioral Cloning One classical imitation learning algorithm is the Behavioral Cloning (BC).

Original: BC turns the imitation learning problem into a supervised learning problem.
Revised: BC turns the imitation learning problem into a supervised learning problem.

Original: It fits the expert’s action a i given the observation
Revised: It fits the expert’s action a i given the observation

Original: s i .
Revised: s i .

Original: For the visual imitation learning problem, the BC model can be divided into two consecutive parts: a vision encoder f θ (which is usually a convolutional neural network), and a policy head π .
Revised: For the visual imitation learning problem, the BC model can be divided into two consecutive parts: a vision encoder f θ (which is usually a convolutional neural network), and a policy head π .

Original: The f θ first encodes s i to the feature f i
Revised: The f θ first encodes s i to the feature f i

Original: =
Revised: =

Original: f
Revised: f

Original: θ ( s i ) , and the π then uses it to predict the expert’s action.
Revised: θ ( s i ) , and the π then uses it to predict the expert’s action.

Original: The BC algorithm minimizes the following negative log-likelihood objective:
Revised: The BC algorithm minimizes the following negative log-likelihood objective:

Original: [Equation]
Revised: [Equation]

Original: Due to its simplicity, BC is widely used in visual imitation learning.
Revised: Due to its simplicity, BC is widely used in visual imitation learning.

Original: Therefore, we study the spatial generalization of BC in this paper.
Revised: Therefore, we study the spatial generalization of BC in this paper.

Original: GradCAM Gradient-weighted Class Activation Mapping (GradCAM) (Selvaraju et al., 2017;
Revised: GradCAM Gradient-weighted Class Activation Mapping (GradCAM) (Selvaraju et al., 2017;

Original: Chattopadhay et al., 2018) is a class of algorithms that can interpret the deep vision models by visualizing the important region for decision making.
Revised: Chattopadhay et al., 2018) is a class of algorithms that can interpret the deep vision models by visualizing the important region for decision making.

Original: Its basic idea is to calculate the gradient at each feature map and aggregate them into an importance map.
Revised: Its basic idea is to calculate the gradient at each feature map and aggregate them into an importance map.

Original: In this paper, we leverage GradCAM as a rough object detector in the proposed regularizer.
Revised: In this paper, we leverage GradCAM as a rough object detector in the proposed regularizer.

Original: We refer readers to the original paper for more technical details.
Revised: We refer readers to the original paper for more technical details.

Original:  M ETHOD
Revised:  M ETHOD

Original: In this section, we present our position-invariant regularizer for spatial generalization.
Revised: In this section, we present our position-invariant regularizer for spatial generalization.

Original: As is discussed before, our basic idea is to remove the absolute positional information from the features of the BC model.
Revised: As is discussed before, our basic idea is to remove the absolute positional information from the features of the BC model.

Original: In the following parts, we first present a formalized description of the idea and its practical challenges in Section 4.1.
Revised: In the following parts, we first present a formalized description of the idea and its practical challenges in Section 4.1.

Original: Then, we discuss how to handle these challenges with GradCAM and adversarial training in Section 4.2 and 4.3.
Revised: Then, we discuss how to handle these challenges with GradCAM and adversarial training in Section 4.2 and 4.3.

Original: Finally, we provide a summary of the overall algorithm in
Revised: Finally, we provide a summary of the overall algorithm in

Original: Section 4.4.
Revised: Section 4.4.

Original: We illustrate our method in Figure 2 and provide the pseudo code in Algorithm 1.
Revised: We illustrate our method in Figure 2 and provide the pseudo code in Algorithm 1.

Original: 4. F ORMULATION AND C HALLENGES
Revised: 4. F ORMULATION AND C HALLENGES

Original: For the tasks that involve spatial generalization, there usually exist multiple objects in the observed states, such as the agent, the target object, and the goal.
Revised: For the tasks that involve spatial generalization, there usually exist multiple objects in the observed states, such as the agent, the target object, and the goal.

Original: For the state s i , we denote each of these objects in s i as o ji , and their positions as ( x ji , y ji ) .
Revised: For the state s i , we denote each of these objects in s i as o ji , and their positions as ( x ji , y ji ) .

Original: Then, our idea can be formulated as the minimization problem of each I (( x j , y j ) , f ) , where I is the mutual information.
Revised: Then, our idea can be formulated as the minimization problem of each I (( x j , y j ) , f ) , where I is the mutual information.

Original: Note that we use the notation x j , y j , f to indicate the corresponding random variables of x ji , y ji , f i .
Revised: Note that we use the notation x j , y j , f to indicate the corresponding random variables of x ji , y ji , f i .

Original: However, this formulation leads to many practical challenges.
Revised: However, this formulation leads to many practical challenges.

Original: First, since each ( x ji , y ji ) is not provided directly by s i 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.
Revised: First, since each ( x ji , y ji ) is not provided directly by s i 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.

Original: However, both of these approaches can be difficult and tedious in practice.
Revised: However, both of these approaches can be difficult and tedious in practice.

Original: Second, even if we have ideal key-point detectors, we have to deal with a hard optimization problem in the summation form (cid:80)
Revised: Second, even if we have ideal key-point detectors, we have to deal with a hard optimization problem in the summation form (cid:80)

Original: j I (( x j , y j ) , f ) .
Revised: j I (( x j , y j ) , f ) .

Original: This can be intractable when there are many objects in the observed state.
Revised: This can be intractable when there are many objects in the observed state.

Original: Fortunately, we find that the previous works on the interpretation of deep learning models like
Revised: Fortunately, we find that the previous works on the interpretation of deep learning models like

Original: GradCAM provide useful tools to handle these challenges.
Revised: GradCAM provide useful tools to handle these challenges.

Original: It can reduce the problem to a much simpler form.
Revised: It can reduce the problem to a much simpler form.

Original: We discuss our observations as follows.
Revised: We discuss our observations as follows.

Original: 4. P ROBLEM R EDUCTION WITH G RAD CAM
Revised: 4. P ROBLEM R EDUCTION WITH G RAD CAM

Original: GradCAM is an interpretation method that can tell which part of the image is crucial in the decision process of a deep learning model.
Revised: GradCAM is an interpretation method that can tell which part of the image is crucial in the decision process of a deep learning model.

Original: Given a BC model ( f θ , π ) and input s , GradCAM outputs an importance heatmap of the same resolution as the input s .
Revised: Given a BC model ( f θ , π ) and input s , GradCAM outputs an importance heatmap of the same resolution as the input s .

Original: The heatmap indicates the importance of each pixel when we use this BC model for prediction.
Revised: The heatmap indicates the importance of each pixel when we use this BC model for prediction.

Original: One nice property of this generated heatmap is that it is smooth and usually coincides with the meaningful objects in the input s .
Revised: One nice property of this generated heatmap is that it is smooth and usually coincides with the meaningful objects in the input s .

Original: Therefore, we can consider the GradCAM as a rough object detector here.
Revised: Therefore, we can consider the GradCAM as a rough object detector here.

Original: We propose to sample p i
Revised: We propose to sample p i

Original: =
Revised: =

Original: ( x i , y i ) from the generated heatmap, and then minimize the I ( p , f ) .
Revised: ( x i , y i ) from the generated heatmap, and then minimize the I ( p , f ) .

Original: We find that this new objective can act as a proxy for the original objective in practice.
Revised: We find that this new objective can act as a proxy for the original objective in practice.

Original: Concretely, if p i is always far from a specific object like o k , then we know that o k is irrelevant to the decision process of the current model.
Revised: Concretely, if p i is always far from a specific object like o k , then we know that o k is irrelevant to the decision process of the current model.

Original: In this case, we conjecture that I (( x k , y k ) , f ) should be low enough to meet our requirement.
Revised: In this case, we conjecture that I (( x k , y k ) , f ) should be low enough to meet our requirement.

Original: On the contrary, if p i always coincides with a certain object like o l , then we actually minimize I ( p , f ) ≈
Revised: On the contrary, if p i always coincides with a certain object like o l , then we actually minimize I ( p , f ) ≈

Original: I
Revised: I

Original: ( x l , y l ) , f ) as we want.
Revised: ( x l , y l ) , f ) as we want.

Original: 4. L OSS F UNCTIONS
Revised: 4. L OSS F UNCTIONS

Original: Now, our remaining work is to reduce the mutual information I ( p , f ) .
Revised: Now, our remaining work is to reduce the mutual information I ( p , f ) .

Original: However, we find that jointly estimating and minimizing the mutual information in our vision-based tasks is hard in practice.
Revised: However, we find that jointly estimating and minimizing the mutual information in our vision-based tasks is hard in practice.

Original: Since our ultimate goal is to minimize the information of p in f , we instead propose an adversarial training framework to achieve this goal.
Revised: Since our ultimate goal is to minimize the information of p in f , we instead propose an adversarial training framework to achieve this goal.

Original: Specifically, we introduce a discriminator network D to play a two-player min-max game with the
Revised: Specifically, we introduce a discriminator network D to play a two-player min-max game with the

Original: BC model as follows.
Revised: BC model as follows.

Original: [Equation]
Revised: [Equation]

Original: In this min-max game, the discriminator D tries to tell the joint distribution of p and f , denoted as P p , f , from the product of their marginal distributions P p ⊗ f .
Revised: In this min-max game, the discriminator D tries to tell the joint distribution of p and f , denoted as P p , f , from the product of their marginal distributions P p ⊗ f .

Original: Meanwhile, the BC model is trying to fool the discriminator by removing the information of p from f .
Revised: Meanwhile, the BC model is trying to fool the discriminator by removing the information of p from f .

Original: Applying the convergence theory of the generative adversarial network (GAN) (Goodfellow et al., 2020), we know that when f θ is a global minimizer of Equation 2, P p , f = P p ⊗ f , which implies that I ( p , f ) = 0 .
Revised: Applying the convergence theory of the generative adversarial network (GAN) (Goodfellow et al., 2020), we know that when f θ is a global minimizer of Equation 2, P p , f = P p ⊗ f , which implies that I ( p , f ) = 0 .

Original: Therefore this min-max game fulfills our requirement.
Revised: Therefore this min-max game fulfills our requirement.

Original: In practice, we train D to minimize the following binary classification loss function:
Revised: In practice, we train D to minimize the following binary classification loss function:

Original: [Equation]
Revised: [Equation]

Original: However, for the encoder f θ , we find that using −L D as the loss function for training will result in instabilities.
Revised: However, for the encoder f θ , we find that using −L D as the loss function for training will result in instabilities.

Original: We assume this is because the f i term is present in both of the two terms in Equation 2, which is different from that in the original GAN objective.
Revised: We assume this is because the f i term is present in both of the two terms in Equation 2, which is different from that in the original GAN objective.

Original: Therefore, we propose to use the following loss function for optimization, which we find works well empirically:
Revised: Therefore, we propose to use the following loss function for optimization, which we find works well empirically:

Original: [Equation]
Revised: [Equation]

Original: Combining the BC loss, the loss function to train the f θ and π is then
Revised: Combining the BC loss, the loss function to train the f θ and π is then

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: 4. S UMMARY OF THE A LGORITHM
Revised: 4. S UMMARY OF THE A LGORITHM

Original: Algorithm 1:
Revised: Algorithm 1:

Original: Initialize parameters of the BC model: f θ , π , and the discriminator D .
Revised: Initialize parameters of the BC model: f θ , π , and the discriminator D .

Original: Initialize target model ˜ f θ , ˜ π with f θ , π . while not converged do
Revised: Initialize target model ˜ f θ , ˜ π with f θ , π . while not converged do

Original: Sample a batch { ( s i , a i ) } from D .
Revised: Sample a batch { ( s i , a i ) } from D .

Original: Calculate each importance heatmap by H i =
Revised: Calculate each importance heatmap by H i =

Original: GradCAM( ˜ f θ , ˜ π, s i ) .
Revised: GradCAM( ˜ f θ , ˜ π, s i ) .

Original: Sample each p i from H i .
Revised: Sample each p i from H i .

Original: Optimize the BC model by Eqn. 5.
Revised: Optimize the BC model by Eqn. 5.

Original: Optimize the discriminator by Eqn. 3.
Revised: Optimize the discriminator by Eqn. 3.

Original: Periodically update ˜ f θ , ˜ π with f θ , π .
Revised: Periodically update ˜ f θ , ˜ π with f θ , π .

Original: Finally, we put everything together and summarize our algorithm in Algorithm 1.
Revised: Finally, we put everything together and summarize our algorithm in Algorithm 1.

Original: We use GradCAM++ (Chattopadhay et al., 2018), an improved version of GradCAM, to calculate the heatmap.
Revised: We use GradCAM++ (Chattopadhay et al., 2018), an improved version of GradCAM, to calculate the heatmap.

Original: We find that it can deal with multi-object observation better than GradCAM.
Revised: We find that it can deal with multi-object observation better than GradCAM.

Original: Similar to the target network in the deep reinforcement learning, we propose to use a target model in the GradCAM calculation.
Revised: Similar to the target network in the deep reinforcement learning, we propose to use a target model in the GradCAM calculation.

Original: This target model periodically copies the weights of the latest BC model.
Revised: This target model periodically copies the weights of the latest BC model.

Original: This ensures that the p i terms in Equation 2 will change slowly throughout the training so that the optimization of the min-max game between D and f θ becomes stable.
Revised: This ensures that the p i terms in Equation 2 will change slowly throughout the training so that the optimization of the min-max game between D and f θ becomes stable.

Original: Our python implementation of this algorithm is in the supplementary material and Appendix A.1.
Revised: Our python implementation of this algorithm is in the supplementary material and Appendix A.1.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In the experiments, we first test the performance of our method on two benchmarks that require spatial generalization: MAGICAL and ProcGen.
Revised: In the experiments, we first test the performance of our method on two benchmarks that require spatial generalization: MAGICAL and ProcGen.

Original: We study the generalization according to the IID protocol (Kirk et al., 2021).
Revised: We study the generalization according to the IID protocol (Kirk et al., 2021).

Original: This means that the training and testing task distributions are the same, though the test instance will be unseen.
Revised: This means that the training and testing task distributions are the same, though the test instance will be unseen.

Original: Then, we provide an analysis of our algorithm through both qualitative and quantitive studies.
Revised: Then, we provide an analysis of our algorithm through both qualitative and quantitive studies.

Original: Finally, we extend our method to a real robot manipulation problem.
Revised: Finally, we extend our method to a real robot manipulation problem.

Original: 5. T ASK S ETUP
Revised: 5. T ASK S ETUP

Original: MAGICAL The MAGICAL benchmark simulates a 2D robotic manipulation problem in a warehouse room.
Revised: MAGICAL The MAGICAL benchmark simulates a 2D robotic manipulation problem in a warehouse room.

Original: The tasks provided by the MAGICAL involve complex interactions between the agent and multiple objects, which require effective spatial generalization..
Revised: The tasks provided by the MAGICAL involve complex interactions between the agent and multiple objects, which require effective spatial generalization..

Original: In the experiments, we use a variant of its MatchRegion task.
Revised: In the experiments, we use a variant of its MatchRegion task.

Original: In this task, a robot is required to go across a square room to move some objects to a target region specified by a dashed rectangle.
Revised: In this task, a robot is required to go across a square room to move some objects to a target region specified by a dashed rectangle.

Original: We set up several task instances of the MatchRegion task: MatchRegion-Target-1, MatchRegion-Target-2,
Revised: We set up several task instances of the MatchRegion task: MatchRegion-Target-1, MatchRegion-Target-2,

Original: MatchRegion-Target-2-Distract, MatchRegion-Target-3, MatchRegion-Target-3-Distract.
Revised: MatchRegion-Target-2-Distract, MatchRegion-Target-3, MatchRegion-Target-3-Distract.

Original: We provide an illustration of these tasks in Figure 3.
Revised: We provide an illustration of these tasks in Figure 3.

Original: For each MatchRegion-Target-
Revised: For each MatchRegion-Target-

Original: X task (MR-T X ), there is no distractor object in the room, so the robot only needs to move all the X objects into the target location.
Revised: X task (MR-T X ), there is no distractor object in the room, so the robot only needs to move all the X objects into the target location.

Original: However, for the MatchRegion-Target- X -Distract task (MR-T X D), there is an additional distractor object in the room.
Revised: However, for the MatchRegion-Target- X -Distract task (MR-T X D), there is an additional distractor object in the room.

Original: This object is also randomly placed in the room during testing.
Revised: This object is also randomly placed in the room during testing.

Original: 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.
Revised: 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.

Original: As we will discuss later, even the existence of one
Revised: As we will discuss later, even the existence of one

Original: [Table]
Revised: [Table]

Original: distractor object can lead to a significant increase of generalization difficulty.
Revised: distractor object can lead to a significant increase of generalization difficulty.

Original: The study of more distractors is carried out in the analysis part.
Revised: The study of more distractors is carried out in the analysis part.

Original: For each of the tasks above
Revised: For each of the tasks above

Original: [Table]
Revised: [Table]

Original: we collect its human demonstration dataset by ourselves.
Revised: we collect its human demonstration dataset by ourselves.

Original: For each demonstration trajectory, we randomly set up the initial position of the objects, target region, and the robot.
Revised: For each demonstration trajectory, we randomly set up the initial position of the objects, target region, and the robot.

Original: Note that the demonstration dataset provided by the original MAGICAL benchmark uses a fixed initial position setup.
Revised: Note that the demonstration dataset provided by the original MAGICAL benchmark uses a fixed initial position setup.

Original: However, we find that this setup is too strict for spatial generalization if no other dataset or pretraining tasks are available to provide prior knowledge.
Revised: However, we find that this setup is too strict for spatial generalization if no other dataset or pretraining tasks are available to provide prior knowledge.

Original: For MR-T1, we collect 50 trajectories.
Revised: For MR-T1, we collect 50 trajectories.

Original: For each of the other tasks, we collect 100 trajectories.
Revised: For each of the other tasks, we collect 100 trajectories.

Original: The collection of all these trajectories takes 2 hours.
Revised: The collection of all these trajectories takes 2 hours.

Original: We also study the outcome of using a different number of trajectories in the later analysis part.
Revised: We also study the outcome of using a different number of trajectories in the later analysis part.

Original: ProcGen The ProcGen benchmark is an arcade-game-based benchmark.
Revised: ProcGen The ProcGen benchmark is an arcade-game-based benchmark.

Original: These arcade games involve interactions between many semantic objects, including the player, monsters (or traps), bonuses, and tools.
Revised: These arcade games involve interactions between many semantic objects, including the player, monsters (or traps), bonuses, and tools.

Original: Though it is designed to test the generalization ability of RL agents, we use it to test the spatial generalization of imitation learning here.
Revised: Though it is designed to test the generalization ability of RL agents, we use it to test the spatial generalization of imitation learning here.

Original: We use 8 environments: Coinrun,
Revised: We use 8 environments: Coinrun,

Original: Jumper, Ninja, Leaper, Miner, Starpilot, Fruitbot, and Bigfish.
Revised: Jumper, Ninja, Leaper, Miner, Starpilot, Fruitbot, and Bigfish.

Original: We follow the ‘easy’ generalization evaluation protocol used in previous works.
Revised: We follow the ‘easy’ generalization evaluation protocol used in previous works.

Original: We first train a PPO agent (Schulman et al., 2017) to collect demonstrations in the first 200 levels of ProcGen and use these demonstrations to train the imitation learning model.
Revised: We first train a PPO agent (Schulman et al., 2017) to collect demonstrations in the first 200 levels of ProcGen and use these demonstrations to train the imitation learning model.

Original: We use 64 expert trajectories in the experiment.
Revised: We use 64 expert trajectories in the experiment.

Original: Then for the evaluation, we test the imitation learning models at all levels.
Revised: Then for the evaluation, we test the imitation learning models at all levels.

Original: We remove the background of the game throughout training and testing since this requires observational generalization.
Revised: We remove the background of the game throughout training and testing since this requires observational generalization.

Original: 5. B ASELINES
Revised: 5. B ASELINES

Original: For the vanilla BC policy, we train an IMPALA (Espeholt et al., 2018) policy, whose encoder is a residual convolutional neural network.
Revised: For the vanilla BC policy, we train an IMPALA (Espeholt et al., 2018) policy, whose encoder is a residual convolutional neural network.

Original: We also try vision-transformer (Dosovitskiy et al., 2021) 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.
Revised: We also try vision-transformer (Dosovitskiy et al., 2021) 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.

Original: Then, we implement the following baselines for comparison.
Revised: Then, we implement the following baselines for comparison.

Original: (1) Dropout .
Revised: (1) Dropout .

Original: Dropout (Srivastava et al., 2014) is a widely-used regularization method in deep learning.
Revised: Dropout (Srivastava et al., 2014) is a widely-used regularization method in deep learning.

Original: Many works also find it useful for improving the performance of imitation learning.
Revised: Many works also find it useful for improving the performance of imitation learning.

Original: (2) Crop Crop is a popular data augmentation method used in the state-of-the-art visual RL algorithms like DrQ (Yarats et al., 2020) and RAD (Laskin et al., 2020), it encourages translation
Revised: (2) Crop Crop is a popular data augmentation method used in the state-of-the-art visual RL algorithms like DrQ (Yarats et al., 2020) and RAD (Laskin et al., 2020), it encourages translation

Original: Ours
Revised: Ours

Original: [Figure]
Revised: [Figure]

Original: Object Dropout Ours 3.
Revised: Object Dropout Ours 3.

Original: Estimated 𝐼((𝑥 𝑖 ,𝑦 𝑖 ),𝐟) invariance to the perspective.
Revised: Estimated 𝐼((𝑥 𝑖 ,𝑦 𝑖 ),𝐟) invariance to the perspective.

Original: Cutout Cutout (DeVries & Taylor, 2017) is a data augmentation method.
Revised: Cutout Cutout (DeVries & Taylor, 2017) is a data augmentation method.

Original: During training, it randomly erases part of the input image.
Revised: During training, it randomly erases part of the input image.

Original: MixReg MixReg (Wang et al., 2020) is a data augmentation method.
Revised: MixReg MixReg (Wang et al., 2020) is a data augmentation method.

Original: It generates interpolated, synthetic training data based on the original training dataset.
Revised: It generates interpolated, synthetic training data based on the original training dataset.

Original: OREO OREO
Revised: OREO OREO

Original: (Park et al., 2021) is an object-aware data augmentation method.
Revised: (Park et al., 2021) is an object-aware data augmentation method.

Original: It randomly drops the features of certain objects.
Revised: It randomly drops the features of certain objects.

Original: (6) CLOP (Bertoin &
Revised: (6) CLOP (Bertoin &

Original: Rachelson, 2022)
Revised: Rachelson, 2022)

Original: CLOP is a locality-aware data augmentation method.
Revised: CLOP is a locality-aware data augmentation method.

Original: It randomly shuffles the feature map of the policy network in a small, local region.
Revised: It randomly shuffles the feature map of the policy network in a small, local region.

Original: 5. R ESULTS
Revised: 5. R ESULTS

Original: MAGICAL The result on MAGICAL is shown in Table 1.
Revised: MAGICAL The result on MAGICAL is shown in Table 1.

Original: The performance is defined by the success rate of the trained policy, which is the number of target objects that are successfully transferred to the target region, divided by the total number of target objects.
Revised: The performance is defined by the success rate of the trained policy, which is the number of target objects that are successfully transferred to the target region, divided by the total number of target objects.

Original: We observe that our method is able to achieve state-of-the-art results and outperform the baselines by a large margin.
Revised: We observe that our method is able to achieve state-of-the-art results and outperform the baselines by a large margin.

Original: Concretely, it improves the success rate by about 30%.
Revised: Concretely, it improves the success rate by about 30%.

Original: Besides, we 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.
Revised: Besides, we 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.

Original: This shows that they may solve some common issues in the generalization problem.
Revised: This shows that they may solve some common issues in the generalization problem.

Original: However, their performance gap from our method suggests that we tackle a different issue here, which is overfitting to absolute positions.
Revised: However, their performance gap from our method suggests that we tackle a different issue here, which is overfitting to absolute positions.

Original: ProcGen
Revised: ProcGen

Original: The result on ProcGen is shown in Table 1.
Revised: The result on ProcGen is shown in Table 1.

Original: The performance is calculated as the normalized agent’s reward in the test environment (with respect to the random agent’s and expert’s performance).
Revised: The performance is calculated as the normalized agent’s reward in the test environment (with respect to the random agent’s and expert’s performance).

Original: Though we observe some improvement, the performance gain is smaller in this case compared to that of MAGICAL.
Revised: Though we observe some improvement, the performance gain is smaller in this case compared to that of MAGICAL.

Original: We hypothesize that this is because the decision-making process of
Revised: We hypothesize that this is because the decision-making process of

Original: ProcGen is highly dependent on the local image patches, as suggested by CLOP.
Revised: ProcGen is highly dependent on the local image patches, as suggested by CLOP.

Original: In this case, the BC model does not need to perform long-range reasoning across objects frequently, and suffers less from the risk of overfitting to absolute positional information.
Revised: In this case, the BC model does not need to perform long-range reasoning across objects frequently, and suffers less from the risk of overfitting to absolute positional information.

Original: Our result here also suggests that we need task-specific inductive biases so as to achieve good generalization performance.
Revised: Our result here also suggests that we need task-specific inductive biases so as to achieve good generalization performance.

Original: 5. A NALYSIS
Revised: 5. A NALYSIS

Original: Qualitative Results To understand whether our method is robust, we use GradCAM++ to visualize the importance heatmap of the learned models.
Revised: Qualitative Results To understand whether our method is robust, we use GradCAM++ to visualize the importance heatmap of the learned models.

Original: For simplicity, we show the result on the
Revised: For simplicity, we show the result on the

Original: MatchReigion-Target-1-Distract task.
Revised: MatchReigion-Target-1-Distract task.

Original: We compare the result of our model to the model trained with dropout here (Figure 4).
Revised: We compare the result of our model to the model trained with dropout here (Figure 4).

Original: We notice that the dropout model tends to focus on the red distractor object rather than the correct target object.
Revised: We notice that the dropout model tends to focus on the red distractor object rather than the correct target object.

Original: In contrast, our model is able to focus on the correct objects.
Revised: In contrast, our model is able to focus on the correct objects.

Original: Even when the distance between the agent and the object is large, it can attend to the agent and the object simultaneously.
Revised: Even when the distance between the agent and the object is large, it can attend to the agent and the object simultaneously.

Original: The visualization results suggest that our regularizer indeed leads to relational features even when the vision network IMPALA does not have an explicit relational inductive bias.
Revised: The visualization results suggest that our regularizer indeed leads to relational features even when the vision network IMPALA does not have an explicit relational inductive bias.

Original: To understand this, we further use MINE estimator (Belghazi et al., 2018) to measure the MI between
Revised: To understand this, we further use MINE estimator (Belghazi et al., 2018) to measure the MI between

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: the learned representation and the position of different objects.
Revised: the learned representation and the position of different objects.

Original: Our representation indeed contains much fewer bits about the absolute position compared with dropout as expected.
Revised: Our representation indeed contains much fewer bits about the absolute position compared with dropout as expected.

Original: Unseen Number of Distractors A robust model should base its decision on robust relational information.
Revised: Unseen Number of Distractors A robust model should base its decision on robust relational information.

Original: As a result, for the MAGICAL tasks, it should be able to ignore the distractor and generalize to an unseen number of distractors.
Revised: As a result, for the MAGICAL tasks, it should be able to ignore the distractor and generalize to an unseen number of distractors.

Original: Therefore, we test whether our model trained on
Revised: Therefore, we test whether our model trained on

Original: MR-T1D (where only one distractor presents) can generalize to MR-T1D with the unseen number of distractors (e.g., 0, 2, 3).
Revised: MR-T1D (where only one distractor presents) can generalize to MR-T1D with the unseen number of distractors (e.g., 0, 2, 3).

Original: We also compare the results with the previous models.
Revised: We also compare the results with the previous models.

Original: The result is shown in Figure 5.
Revised: The result is shown in Figure 5.

Original: We find that our model is able to generalize to the case of 0, 2, 3, though the performance is lower than the case of 1 (training scenario).
Revised: We find that our model is able to generalize to the case of 0, 2, 3, though the performance is lower than the case of 1 (training scenario).

Original: In contrast, the prior model, such as the dropout model, fails in these unseen cases totally.
Revised: In contrast, the prior model, such as the dropout model, fails in these unseen cases totally.

Original: This echoes our qualitative analysis results.
Revised: This echoes our qualitative analysis results.

Original: Number of Demonstrations We also study whether the proposed method works when the amount of expert demonstrations is limited.
Revised: Number of Demonstrations We also study whether the proposed method works when the amount of expert demonstrations is limited.

Original: For this purpose, we test our method on the MAGICAL with 25% , 50% , 75% of expert demonstrations.
Revised: For this purpose, we test our method on the MAGICAL with 25% , 50% , 75% of expert demonstrations.

Original: We show the averaged performance in Figure 6.
Revised: We show the averaged performance in Figure 6.

Original: We find that our method can achieve consistent improvement, though the performance decreases as the dataset becomes smaller.
Revised: We find that our method can achieve consistent improvement, though the performance decreases as the dataset becomes smaller.

Original: This result suggests that we still require a certain amount of diverse data to achieve spatial generalization.
Revised: This result suggests that we still require a certain amount of diverse data to achieve spatial generalization.

Original: 5. R EAL -W ORLD E XPERIMENTS
Revised: 5. R EAL -W ORLD E XPERIMENTS

Original: [Table]
Revised: [Table]

Original: Finally, we test whether our method scales to the real-world  manipulation problem.
Revised: Finally, we test whether our method scales to the real-world pickand-place manipulation problem.

Original: We extend the MR-T1D to a UR10 robot arm with a Robotiq parallel-jaw gripper (Figure 7).
Revised: We extend the MR-T1D to a UR10 robot arm with a Robotiq parallel-jaw gripper (Figure 7).

Original: As suggested by Hsu et al.
Revised: As suggested by Hsu et al.

Original: (2022), we use a gripper camera and a workspace camera to provide observation.
Revised: (2022), we use a gripper camera and a workspace camera to provide observation.

Original: 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.
Revised: 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.

Original: We use the proposed regularizer to regularize the workspace branch.
Revised: We use the proposed regularizer to regularize the workspace branch.

Original: We collect 75 human demonstrations for training.
Revised: We collect 75 human demonstrations for training.

Original: We compare our method to dropout with different numbers of distract objects.
Revised: We compare our method to dropout with different numbers of distract objects.

Original: The result is shown in Table 2.
Revised: The result is shown in Table 2.

Original: Our method also achieves a large improvement in this problem.
Revised: Our method also achieves a large improvement in this problem.

Original: We provide detailed setup and qualitative results in the Appendix A.4 and A.3.
Revised: We provide detailed setup and qualitative results in the Appendix A.4 and A.3.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this paper, we studied the spatial generalization problem of imitation learning.
Revised: In this paper, we studied the spatial generalization problem of imitation learning.

Original: To solve this problem, we proposed a novel position-invariant regularizer to remove the absolute positional information from the features.
Revised: To solve this problem, we proposed a novel position-invariant regularizer to remove the absolute positional information from the features.

Original: Through experiments on the MAGICAL and ProcGen benchmarks, as well as a robot manipulation system, we confirmed that previous methods do overfit to the absolute position and showed that our proposed approach can effectively solve this problem.
Revised: Through experiments on the MAGICAL and ProcGen benchmarks, as well as a robot manipulation system, we confirmed that previous methods do overfit to the absolute position and showed that our proposed approach can effectively solve this problem.

Original: We hope that the proposed method can inspire future generalization research.
Revised: We hope that the proposed method can inspire future generalization research.

Original:  R EPRODUCIBILITY S TATEMENT
Revised:  R EPRODUCIBILITY S TATEMENT

Original: We provide the python source code of the main algorithm and the full model in the supplementary material.
Revised: We provide the python source code of the main algorithm and the full model in the supplementary material.

Original: We will make our project publicly available.
Revised: We will make our project publicly available.

Original: A A PPENDIX A. I MPLEMENTATION OF THE A LGORITHM
Revised: A A PPENDIX A. I MPLEMENTATION OF THE A LGORITHM

Original: We provide the core training code of our project in the supplementary material, which comes from our whole project.
Revised: We provide the core training code of our project in the supplementary material, which comes from our whole project.

Original: The code shows all the training details and the model implementation.
Revised: The code shows all the training details and the model implementation.

Original: It is divided into three files.
Revised: It is divided into three files.

Original: • bc.py
Revised: • bc.py

Original: This python file provides the main training routine.
Revised: This python file provides the main training routine.

Original: The main class is the BehaviorCloning .
Revised: The main class is the BehaviorCloning .

Original: The training routine is its train method.
Revised: The training routine is its train method.

Original: • model.py
Revised: • model.py

Original: This python file provides the implementation of BC policy model ( Model ) and the discriminator ( DiscrimModel ).
Revised: This python file provides the implementation of BC policy model ( Model ) and the discriminator ( DiscrimModel ).

Original: network.py
Revised: network.py

Original: This python file provides the implementation of the IMPALA encoder ( ResNetBase ) used by the BC policy model.
Revised: This python file provides the implementation of the IMPALA encoder ( ResNetBase ) used by the BC policy model.

Original: Note that these files are for illustration purpose.
Revised: Note that these files are for illustration purpose.

Original: The whole runnable project will be made publicly available.
Revised: The whole runnable project will be made publicly available.

Original: A. H YPERPARAMETERS
Revised: A. H YPERPARAMETERS

Original: The size of image observation is (64 , 64) in all the experiments.
Revised: The size of image observation is (64 , 64) in all the experiments.

Original: We use the Adam optimizer (Kingma & Ba, 2015) for the optimization of both the discriminator and the BC model.
Revised: We use the Adam optimizer (Kingma & Ba, 2015) for the optimization of both the discriminator and the BC model.

Original: The learning rate of the optimizers is 0.0003.
Revised: The learning rate of the optimizers is 0.0003.

Original: We also add a weight decay of 0.0001 to the BC model.
Revised: We also add a weight decay of 0.0001 to the BC model.

Original: We use a batchsize of 128 throughout the experiments ( batch size ).
Revised: We use a batchsize of 128 throughout the experiments ( batch size ).

Original: We will sample 4 points for each batch ( sample n ).
Revised: We will sample 4 points for each batch ( sample n ).

Original: We repeat the training of discriminator for 4 steps ( repeat ).
Revised: We repeat the training of discriminator for 4 steps ( repeat ).

Original: For the discriminator ( DiscrimModel ), its embedding size is set to 64 ( e dim ).
Revised: For the discriminator ( DiscrimModel ), its embedding size is set to 64 ( e dim ).

Original: For the choice of λ , we use λ = 1 . 0 for the MAGICAL, and 0 .
Revised: For the choice of λ , we use λ = 1 . 0 for the MAGICAL, and 0 .

Original: 1 otherwise.
Revised: 1 otherwise.

Original: All the other details and hyperparameters can be found in the source code in the supplementary material.
Revised: All the other details and hyperparameters can be found in the source code in the supplementary material.

Original: A. Q UALITATIVE R ESULTS OF THE M ANIPULATION P ROBLEM
Revised: A. Q UALITATIVE R ESULTS OF THE M ANIPULATION P ROBLEM

Original: In this section, we provide some qualitative results of the real-world manipulation problem.
Revised: In this section, we provide some qualitative results of the real-world manipulation problem.

Original: Recall that in this task, the robot is required to move a red cube to a target location specified by a green area.
Revised: Recall that in this task, the robot is required to move a red cube to a target location specified by a green area.

Original: We show the importance heatmap of the dropout model (Figure 8) and our model (Figure 9).
Revised: We show the importance heatmap of the dropout model (Figure 8) and our model (Figure 9).

Original: As is shown in the figures, we find that dropout model tends to attend more to the round distractor object compared with our model.
Revised: As is shown in the figures, we find that dropout model tends to attend more to the round distractor object compared with our model.

Original: However, due to the visual complexity, we find that our model sometimes may attend the shadow in the background.
Revised: However, due to the visual complexity, we find that our model sometimes may attend the shadow in the background.

Original: A. S ETUP OF R EAL R OBOT
Revised: A. S ETUP OF R EAL R OBOT

Original: 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.
Revised: 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.

Original: We use a discretized action space, so it will become easier for human to collect the demonstration via the keyboard.
Revised: We use a discretized action space, so it will become easier for human to collect the demonstration via the keyboard.

Original: There areactions in total, [ x, y, z ] move [backward/forward] by 2cm, and gripper [open/close] by 1cm.
Revised: There areactions in total, [ x, y, z ] move [backward/forward] by 2cm, and gripper [open/close] by 1cm.

Original: We only consider to regularize the workspace camera.
Revised: We only consider to regularize the workspace camera.

Original: This is because the gripper-camera can only observe a very limited area on the workspace in our case and the gripper is always at the center.
Revised: This is because the gripper-camera can only observe a very limited area on the workspace in our case and the gripper is always at the center.

Original: As a result, the spatial generalization issue is minor for the gripper-camera
Revised: As a result, the spatial generalization issue is minor for the gripper-camera

Original: and we find it do not help much to regularize it.
Revised: and we find it do not help much to regularize it.

Original: Our attends less to the
Revised: Our attends less to the

Original: However, due to that sometimes may attend
Revised: However, due to that sometimes may attend