Original: P OLICY C ONTRASTIVE I MITATION L EARNING
Revised: P OLICY C ONTRASTIVE I MITATION L EARNING

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: Adversarial imitation learning (AIL) is a popular method that has recently achieved much success.
Revised: Adversarial imitation learning (AIL) is a popular method that has recently achieved much success.

Original: However, the performance of AIL is still unsatisfactory on the more challenging tasks.
Revised: However, the performance of AIL is still unsatisfactory on the more challenging tasks.

Original: We find that one of the major reasons is due to the low quality of AIL discriminator representation.
Revised: We find that one of the major reasons is due to the low quality of AIL discriminator representation.

Original: Since the AIL discriminator is trained via binary classification that does not necessarily discriminate the policy from the expert in a meaningful way, the resulting reward might not be meaningful either.
Revised: Since the AIL discriminator is trained via binary classification that does not necessarily discriminate the policy from the expert in a meaningful way, the resulting reward might not be meaningful either.

Original: We propose a new method called Policy Contrastive Imitation Learning (PCIL) to resolve this issue.
Revised: We propose a new method called Policy Contrastive Imitation Learning (PCIL) to resolve this issue.

Original: PCIL learns a contrastive representation space by anchoring on different policies and uses a smooth cosine-similarity-based reward to encourage imitation learning.
Revised: PCIL learns a contrastive representation space by anchoring on different policies and generates a smooth cosine-similarity-based reward.

Original: Our proposed representation learning objective can be viewed as a stronger version of the AIL objective and provide a more meaningful comparison between the agent and the policy.
Revised: Our proposed representation learning objective can be viewed as a stronger version of the AIL objective and provide a more meaningful comparison between the agent and the policy.

Original: From a theoretical perspective, we show the validity of our method using the apprenticeship learning framework.
Revised: From a theoretical perspective, we show the validity of our method using the apprenticeship learning framework.

Original: Furthermore, our empirical evaluation on the DeepMind Control suite demonstrates that PCIL can achieve state-of-the-art performance.
Revised: Furthermore, our empirical evaluation on the DeepMind Control suite demonstrates that PCIL can achieve state-of-the-art performance.

Original: Finally, qualitative results suggest that PCIL builds a smoother and more meaningful representation space for imitation learning.
Revised: Finally, qualitative results suggest that PCIL builds a smoother and more meaningful representation space for imitation learning.

Original:  INTRODUCTION
Revised:  INTRODUCTION

Original: Imitation is one of the fundamental capabilities of an intelligent agent (Hussein et al., 2017).
Revised: Imitation is one of the fundamental capabilities of an intelligent agent (Hussein et al., 2017).

Original: Animals and humans can acquire many skills by mimicking each other (Byrne, 2009).
Revised: Animals and humans can acquire many skills by mimicking each other (Byrne, 2009).

Original: In engineering, imitation learning also enables many robotics applications.
Revised: In engineering, imitation learning also enables many robotics applications.

Original: One recent class of imitation learning algorithms is the adversarial imitation learning (AIL) (Ho & Ermon, 2016).
Revised: One mainstream class of imitation learning algorithms is the adversarial imitation learning (AIL) (Ho & Ermon, 2016).

Original: AIL converts the imitation task into a distribution matching problem and proposes to imitate it by training a policy against an adversarial discriminator.
Revised: AIL converts the imitation task into a distribution matching problem and proposes to imitate it by training a policy against an adversarial discriminator.

Original: AIL has enjoyed great success on many imitation tasks: it achieves superior performance, does not suffer from the compounding error issue (Ross et al., 2011), and can work even without expert actions. However, AIL is hard to train in practice, usually involving careful tuning of discriminator neural network sizes and learning rates, etc (Wang et al., 2017; Kim & 
Revised: AIL has enjoyed great success on many imitation tasks: it achieves superior performance (Ho & Ermon, 2016; Kostrikov et al., 2018), and has been experimentally proven to alleviate some of the distributional drift issue, and can work even without expert actions (Torabi et al., 2018b).However, AIL is hard to train in practice, usually involving careful tuning of discriminator neural network sizes and learning rates (Wang et al., 2017; Kim & Park, 2018; Orsini et al., 2021).

Original: Park, 2018; Orsini et al., 2021). The fragility of the discriminator  not only leads to poor performance but also severely limits the applicability of AIL to a broader range of tasks.
Revised: The fragility of the discriminator (Peng et al., 2018) not only leads to poor performance but also severely limits the applicability of AIL to a broader range of tasks.

Original: Numerous techniques have been proposed to improve the performance of AIL, such as using regularization and gradient penalties (Fu et al., 2017;
Revised: Numerous techniques have been proposed to improve the performance of AIL, such as using regularization and gradient penalties (Fu et al., 2017;

Original: Kostrikov et al., 2018; Gulrajani et al., 2017).
Revised: Kostrikov et al., 2018; Gulrajani et al., 2017).

Original: Some works also propose to use different distribution metrics (e.g., KL divergence, Wasserstein distance) (Xiao et al., 2019) for distribution matching and show some improvements.
Revised: Some works also propose to use different distribution metrics (e.g., KL divergence, Wasserstein distance) (Xiao et al., 2019) for distribution matching and show some improvements.

Original: Though these methods show encouraging results, we notice that they ignore one crucial aspect of the problem: the representation of AIL’s discriminator.
Revised: Though these methods show encouraging results, we notice that they ignore one crucial aspect of the problem: the representation of AIL’s discriminator.

Original: To be specific, the discriminator in AIL is usually trained with binary classification loss that distinguishes expert transitions from agent transitions.
Revised: To be specific, the discriminator in AIL is usually trained with binary classification loss that distinguishes expert transitions from agent transitions.

Original: This discriminator is then used to define rewards.
Revised: This discriminator is then used to define rewards.

Original: However, since the only goal of the discriminator is to distinguish the expert from the agent, it does not necessarily learn a good, robust representation space that can provide a reasonable comparison between the behavior of two agents.
Revised: However, since the only goal of the discriminator is to distinguish the expert from the agent, it does not necessarily learn a good, smooth representation space that can provide a reasonable comparison between the behavior of two agents.

Original: Ideal representations should be able to provide semantically meaningful signals to compare the expert policy and the agent policy.
Revised: Ideal representations should be able to provide semantically meaningful signals to compare the expert policy and the agent policy.

Original: In this paper, we propose a new algorithm called Policy Contrastive Imitation Learning (PCIL) to achieve this goal.
Revised: In this paper, we propose a new algorithm called Policy Contrastive Imitation Learning (PCIL) to achieve this goal.

Original: Instead of using a discriminator trained with a binary-classification objective, we propose to train a discriminator representation space with the contrastive learning loss.
Revised: Instead of training with a binary-classification objective, we propose to train a discriminator representation space with the contrastive learning loss.

Original: Our method differs from the prior representation learning approach in AIL in that we perform contrastive learning 
Revised: Our method differs from the prior representation learning approach in AIL in that we perform contrastive learning between different

Original: [Figure]
Revised: [Figure]

Original: between different policies.
Revised: policies.

Original: More specifically, we push the expert’s representation together and pull the agent policy’s representation away from them.
Revised: More specifically, we push the expert’s representation together and pull the agent policy’s representation away from them.

Original: We define the imitation learning reward via cosine similarity between the policy’s and expert’s transition.
Revised: We define the imitation learning reward via cosine similarity between the policy’s and expert’s transition.

Original: As is shown in Figure 1, the discriminator (binary classifier) might not have a good representation space: the distance between two expert transitions can be even larger than the distance between an expert transition and an agent transition.
Revised: As is shown in Figure 1, the discriminator (binary classifier) might not have a good representation space: the distance between two expert transitions can be even larger than the distance between an expert transition and an agent transition.

Original: This implies that the discriminator may not encode some common features of the expert’s behavior and may use non-robust features to compare the behavior.
Revised: This implies that the discriminator may not encode some common features of the expert’s behavior and may use non-robust features to compare the behavior.

Original: The lack of proper structure in the representation space of traditional discriminators may yield low-quality AIL rewards.
Revised: The lack of proper structure in the representation space of traditional discriminators may yield low-quality AIL rewards.

Original: To alleviate this issue, we explicitly define a constraint on the representation space, requiring that the distance between expert transitions be smaller than their distance to the agent’s transitions.
Revised: To alleviate this issue, we explicitly define a constraint on the representation space, requiring that the distance between expert transitions be smaller than their distance to the agent’s transitions.

Original: This is a stronger constraint on the discriminator’s representation: it is easy to derive a binary classifier from our learned representation.
Revised: This is a stronger constraint on the discriminator’s representation: it is easy to derive a binary classifier from our learned representation.

Original: However, the binary classifier’s representation space does not necessarily satisfy our objective.
Revised: However, the binary classifier’s representation space does not necessarily satisfy our objective.

Original: From a theoretical perspective, we show the soundness of PCIL in the apprenticeship learning framework.
Revised: From a theoretical perspective, we show the soundness of PCIL in the apprenticeship learning framework.

Original: We also evaluate our method empirically by benchmarking our method on the DeepMind
Revised: We also evaluate our method empirically by benchmarking our method on the DeepMind

Original: Control Suite Tassa et al.
Revised: Control Suite (Tassa et al., 2018).

Original: Experimental results show that our method is able to achieve stateof-the-art results.
Revised: Experimental results show that our method is able to achieve stateof-the-art results.

Original: Through ablation study and qualitative visualization, we find that our method is more effective than prior representation learning methods and able to provide a better representation space for imitation learning.
Revised: Through ablation study and qualitative visualization, we find that our method is more effective than prior representation learning methods and able to provide a better representation space for imitation learning.

Original: In summary, our contributions in this paper are as follows.
Revised: In summary, our contributions in this paper are as follows.

Original: We point out a new direction to improve the performance of the AIL methods,
Revised: We point out a new direction to improve the performance of the AIL methods,

Original: i
Revised: i

Original: e., going beyond naive binary classification and leveraging more stable and meaningful representation learning algorithms for imitation.
Revised: e., going beyond naive binary classification and leveraging more stable and meaningful representation learning algorithms for imitation.

Original: We propose an algorithm called Policy Contrast Imitation Learning (PCIL) method that instantiates such an improvement and establishes its connection to apprenticeship learning from a theoretical perspective.
Revised: We propose an algorithm called Policy Contrast Imitation Learning (PCIL) method that instantiates such an improvement and establishes its connection to apprenticeship learning from a theoretical perspective.

Original: We evaluate our method on the DeepMind Control Suite and achieve state-of-the-art performance.
Revised: We evaluate our method on the DeepMind Control Suite and achieve state-of-the-art performance.

Original: Through ablation studies, we highlight its essential difference from previous contrastive learning methods in AIL.
Revised: Through ablation studies, we highlight its essential difference from previous contrastive learning methods in AIL.

Original:  P RELIMINARIES 2. N OTATIONS
Revised:  P RELIMINARIES 2. N OTATIONS

Original: In this paper, we model the imitation leanring problem as a markov decision process M =
Revised: In this paper, we model the imitation leanring problem as a markov decision process M =

Original: ( S , A , p 0 ( s ) , p ( s ′ | s, a ) , r ( s, a, s ′ ) , γ ) .
Revised: ( S , A , p 0 ( s ) , p ( s ′ | s, a ) , r ( s, a, s ′ ) , γ ) .

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

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

Original: p 0 ( s ) defines the initial state distribution.
Revised: p 0 ( s ) defines the initial state distribution.

Original: p ( s ′ | s, a ) defines the transition dynamics.
Revised: p ( s ′ | s, a ) defines the transition dynamics.

Original: r ( s, a, s ′ ) is the reward function.
Revised: r ( s, a, s ′ ) is the reward function.

Original: γ is the discount factor.
Revised: γ is the discount factor.

Original: The goal is to maximize the expected return of the learned policy π , which is defined by
Revised: The goal is to maximize the expected return of the learned policy π , which is defined by

Original: [Equation]
Revised: [Equation]

Original: For the imitation learning problem, the algorithm does not have access to the reward function and the transition dynamics.
Revised: For the imitation learning problem, the algorithm does not have access to the reward function and the transition dynamics.

Original: Instead, it is provided with an expert demonstration dataset D sampled from an expert policy π E , which can perform well in M .
Revised: Instead, it is provided with an expert demonstration dataset D sampled from an expert policy π E , which can perform well in M .

Original: Here, D takes the form of { ( s Ei , a Ei ) } , where ( s Ei , a Ei ) is sampled from ρ π E , the stationary state-action visiting distribution of π E .
Revised: Here, D takes the form of { ( s Ei , a Ei ) } , where ( s Ei , a Ei ) is sampled from ρ π E , the stationary state-action visiting distribution of π E .

Original: The imitation learning algorithm is then required to reproduce the expert’s behavior with D .
Revised: The imitation learning algorithm is then required to reproduce the expert’s behavior with D .

Original: 2. A DVERSARIAL I MITATION L EARNING
Revised: 2. A DVERSARIAL I MITATION L EARNING

Original: One popular class of the imitation learning algorithm is the adversarial imitation learning (AIL), whose vanilla version is Generative Adversarial Imitation Learning (GAIL) (Ho & Ermon, 2016).
Revised: One popular class of the imitation learning algorithm is the adversarial imitation learning (AIL), whose vanilla version is Generative Adversarial Imitation Learning (GAIL) (Ho & Ermon, 2016).

Original: The idea of GAIL is to minimize the divergence between ρ π E and ρ π .
Revised: The idea of GAIL is to minimize the divergence between ρ π E and ρ π .

Original: It uses a discriminator D ( s, a ) to distinguish expert’s transitions ( s Ei , a Ei ) ∼ D from the policy transitions ( s, a ) ∼ ρ π , which is trained by maximizing the objective
Revised: It uses a discriminator D ( s, a ) to distinguish expert’s transitions ( s Ei , a Ei ) ∼ D from the policy transitions ( s, a ) ∼ ρ π , which is trained by maximizing the objective

Original: [Equation]
Revised: [Equation]

Original: log( D ( s i , a i ))]
Revised: log( D ( s i , a i ))]

Original: + E ( s Ei ,a Ei ) ∼D
Revised: + E ( s Ei ,a Ei ) ∼D

Original: [log(1 − D ( s Ei , a Ei ))] .
Revised: [log(1 − D ( s Ei , a Ei ))] .

Original: To achieve imitation, the agent policy is then required to fool the discriminator, which can only be possible when the policy π resembles the expert π E .
Revised: To achieve imitation, the agent policy is then required to fool the discriminator, which can only be possible when the policy π resembles the expert π E .

Original: Specifically, GAIL defines an adversarial reward r ( s t , a t ) =
Revised: Specifically, GAIL defines an adversarial reward r ( s t , a t ) =

Original: − log(1 − D ( s t , a t )) , and trains π to maximize the expected return with respect to this reward using on-policy RL algorithms.
Revised: − log(1 − D ( s t , a t )) , and trains π to maximize the expected return with respect to this reward using on-policy RL algorithms.

Original:  P OLICY C ONTRASTIVE I MITAITON L EARNING 3. O VERVIEW
Revised:  P OLICY C ONTRASTIVE I MITAITON L EARNING 3. O VERVIEW

Original: We propose a novel representation-learning-based approach called Policy Contrastive Imitation
Revised: We propose a novel representation-learning-based approach called Policy Contrastive Imitation

Original: Learning to improve the AIL reward.
Revised: Learning to improve the AIL reward.

Original: The overview of PCIL is illustrated in Figure 2.
Revised: The overview of PCIL is illustrated in Figure 2.

Original: Our key insight is to learn a policy-contrastive representation space.
Revised: Our key insight is to learn a policy-contrastive representation space.

Original: Unlike the contrastive learning studied in previous AIL literature, the policy-contrastive representation here is obtained by anchoring on policy, leading to meaningful representation that can compare different policies.
Revised: Unlike the contrastive learning studied in previous AIL literature, the policy-contrastive representation here is obtained by anchoring on policy, leading to meaningful representation that can compare different policies.

Original: We will discuss the training of this representation in Section 3.2 and the reward design in Section 3.3.
Revised: We will discuss the training of this representation in Section 3.2 and the reward design in Section 3.3.

Original: Then, we will show the convergence of our algorithm in Section 3.4.
Revised: Then, we will show the convergence of our algorithm in Section 3.4.

Original: 3. C ONTRASTIVE P OLICY R EPRESENTATION FOR I MITATION
Revised: 3. C ONTRASTIVE P OLICY R EPRESENTATION FOR I MITATION

Original: The vanilla AIL algorithms are based on unconstrained representations and can be very non-robust.
Revised: The vanilla AIL algorithms are based on unconstrained representations and can be very non-robust.

Original: One possible approach to handle this problem is to learn a more meaningful representation by contrastive learning.
Revised: One possible approach to handle this problem is to learn a more meaningful representation by contrastive learning.

Original: Researchers  have found that it can usually learn semantically meaningful representation, leading to better performance on downstream tasks, such as classification.
Revised: Researchers (He et al., 2020) have found that it can usually learn semantically meaningful representation, leading to better performance on downstream tasks, such as classification.

Original: However, though it is effective in the field of supervised learning, people have found that it does not greatly improve AIL much (Chen et al., 2021).
Revised: However, though it is effective in the field of supervised learning, prior work (Chen et al., 2021) has found that it does not greatly improve AIL much.

Original: To understand the reason behind this, we first recall that the contrastive learning method learns by drawing the representation of one training sample x towards a similar positive sample x p , and pulling it away from a dissimilar, negative sample x n .
Revised: To understand the reason behind this, we first recall that the contrastive learning method learns by drawing the representation of one training sample x towards a similar positive sample x p , and pulling it away from a dissimilar, negative sample x n .

Original: Some heuristic rules determine the choice of positive and negative samples: the positive sample is usually defined as a data augmentation of x .
Revised: Some heuristic rules determine the choice of positive and negative samples: the positive sample is usually defined as a data augmentation of x .

Original: Therefore, this augmentation decides what should be similar to the representation.
Revised: Therefore, this augmentation decides what should be similar to the representation.

Original: Nevertheless, we argue that the definition of similarity in the previous works is not strong enough for AIL.
Revised: Nevertheless, we argue that the definition of similarity in the previous works is not strong enough for AIL.

Original: This is because, in AIL, the representation should also be able to let us discern good behavior from the bad behavior.
Revised: This is because, in AIL, the representation should also be able to let us discern good behavior from the bad behavior.

Original: However, the difference between good and bad behavior here can sometimes be very faint.
Revised: However, the difference between good and bad behavior here can sometimes be very faint.

Original: For example, consider a case where a robot misses the exact point to execute a certain action and, as a result, fails to accomplish the task.
Revised: For example, consider a case where a robot misses the exact point to execute a certain action and, as a result, fails to accomplish the task.

Original: The good and bad states right before this point may look very similar.
Revised: The good and bad states right before this point may look very similar.

Original: More specifically, the difference can simply be minor in a particular physical measurement, like the distance.
Revised: More specifically, the difference can simply be minor in a particular physical measurement, like the distance.

Original: In this case, the representation should consider this as a semantic component and be sensitive to such a difference to succeed.
Revised: In this case, the representation should consider this as a semantic component and be sensitive to such a difference to succeed.

Original: Unfortunately, the difference between the positive sample and the anchor sample in the previous representation learning methods is usually very large and overwhelms the difference between good and bad states.
Revised: Unfortunately, the difference between the positive sample and the anchor sample in the previous representation learning methods is usually very large and overwhelms the difference between good and bad states.

Original: As a result, the model may not distinguish between good and bad states effectively.
Revised: As a result, the model may not distinguish between good and bad states effectively.

Original: Though combining the representation learning objective and the AIL objective may help combat this problem by enforcing a hard distinguishing constraint over contrastive representation, in practice we find this does not work well (Section 4.4).
Revised: Though combining the representation learning objective and the AIL objective may help combat this problem by enforcing a hard distinguishing constraint over contrastive representation, in practice we find this does not work well (Section 4.4).

Original: These observations motivate us to learn a representation that is semantically meaningful and able to distinguish between good and bad states.
Revised: These observations motivate us to learn a representation that is semantically meaningful and able to distinguish between good and bad states.

Original: We find a surprisingly simple yet very effective approach: we can consider the samples drawing from the same policy as the positive samples and the samples from all the other policies as negative samples.
Revised: We find a surprisingly simple yet very effective approach: we can consider the samples drawing from the same policy as the positive samples and the samples from all the other policies as negative samples.

Original: In the case of imitation learning, our samples can be naturally divided into two categories, namely expert and non-expert samples.
Revised: In the case of imitation learning, our samples can be naturally divided into two categories, namely expert and non-expert samples.

Original: Then given the encoder Φ : S → S that maps the state to a representation vector in a high-dimensional sphere, we define its infoNCE representation loss function as follows:
Revised: Then given the encoder Φ : S → S that maps the state to a representation vector in a high-dimensional sphere, we define its infoNCE representation loss function as follows:

Original: [Equation]
Revised: [Equation]

Original: Here, x 0 is some state-action pair from the expert transitions; x p is some transition from the expert data acting as positives;
Revised: Here, x 0 is some state-action pair from the expert transitions; x p is some transition from the expert data acting as positives;

Original: ˜ x
Revised: ˜ x

Original: i is some agent transition working as negatives.
Revised: i is some agent transition working as negatives.

Original: In other words, we require Φ to draw the expert samples towards each other and pull all the policy samples away from the expert samples.
Revised: In other words, we require Φ to draw the expert samples towards each other and pull all the policy samples away from the expert samples.

Original: Our proposed objective function is a strictly stronger constraint on the discriminator.
Revised: Our proposed objective function is a strictly stronger constraint on the discriminator.

Original: The binary classification discriminator can find any hyper-plane that separates the two types of transitions, with no constraint on how the transitions are embedded.
Revised: The binary classification discriminator can find any hyper-plane that separates the two types of transitions, with no constraint on how the transitions are embedded.

Original: However, our objective enforces the pair-wise distance constraint between any triplets.
Revised: However, our objective enforces the pair-wise distance constraint between any triplets.

Original: One can derive a binary expert-policy transition classifier from a trained Φ by computing Φ( x 0 ) T Φ( x ) > t , where x 0 is any expert transition, x is the tran- sition to be classified and t is some threshold.
Revised: One can derive a binary expert-policy transition classifier from a trained Φ by computing Φ( x 0 ) T Φ( x ) > t , where x 0 is any expert transition, x is the tran- sition to be classified and t is some threshold.

Original: However, on the other hand, the binary classification induced latent space might not satisfy our constraint.
Revised: However, on the other hand, the binary classification induced latent space might not satisfy our constraint.

Original: As illustrated in Figure 1, the latent space for the binary classification discriminator might have some expert-expert pairs that are even further away than some expert-agent pairs.
Revised: As illustrated in Figure 1, the latent space for the binary classification discriminator might have some expert-expert pairs that are even further away than some expert-agent pairs.

Original: Interestingly, our approach echoes the supervised contrastive learning (SCL) (Khosla et al., 2020), which suggests that we consider all the samples in one class as similar, positive samples.
Revised: Interestingly, our approach echoes the supervised contrastive learning (SCL) (Khosla et al., 2020), which suggests that we consider all the samples in one class as similar, positive samples.

Original: Here we also consider all the samples in the expert demonstrations as similar.
Revised: Here we also consider all the samples in the expert demonstrations as similar.

Original: However, unlike SCL, we do not require the samples from the agents to be similar to each other.
Revised: However, unlike SCL, we do not require the samples from the agents to be similar to each other.

Original: This is because the agent transitions are generated from different policies during the training process.
Revised: This is because the agent transitions are generated from different policies during the training process.

Original: 3. S IMILARITY -B ASED I MITATION R EWARD
Revised: 3. S IMILARITY -B ASED I MITATION R EWARD

Original: With a representation that can capture the difference between good and bad states, we can then define a reward function to encourage imitation learning.
Revised: With a representation that can capture the difference between good and bad states, we can then define a reward function to encourage imitation learning.

Original: Though using an AIL-style reward with this representation is still possible, we find that a better choice is to use a cosine similarity metric to define the reward.
Revised: Though using an AIL-style reward with this representation is still possible, we find that a better choice is to use a cosine similarity metric to define the reward.

Original: It has several advantages: it is bounded and appears relatively smooth in practice, leading to more stable learning.
Revised: It has several advantages: it is bounded and appears relatively smooth in practice, leading to more stable learning.

Original: Concretely, we define:
Revised: Concretely, we define:

Original: [Equation]
Revised: [Equation]

Original: Nevertheless, in practice evaluating the latter expectation can be time-consuming since Φ is frequently updated.
Revised: Nevertheless, in practice evaluating the latter expectation can be time-consuming since Φ is frequently updated.

Original: Therefore, we use a random expert sample for the reward calculation.
Revised: Therefore, we use a random expert sample for the reward calculation.

Original: From this reward, we can see that a policy can only obtain high rewards when it frequently visits the expert’s distribution.
Revised: From this reward, we can see that a policy can only obtain high rewards when it frequently visits the expert’s distribution.

Original: This naturally connects our method to the distribution matching, and we provide the theoretical analysis of our algorithm in the following subsection.
Revised: This naturally connects our method to the distribution matching, and we provide the theoretical analysis of our algorithm in the following subsection.

Original: 3. T HEORETICAL A NALYSIS
Revised: 3. T HEORETICAL A NALYSIS

Original: In this part, wereformulate our method under the apprenticeship learning framework and show its theoretical validity.
Revised: In this part, we show that PCIL can be reduced to Apprenticeship Learning (AL) (Abbeel & Ng, 2004). First, let us recall that an AL problem takes the following form (Ho & Ermon, 2016): [Equation]

Original: Apprenticeship learning is a classical IRL framework, which states that matching the feature distribution of  the expert and the policy will lead to an optimal agent policy (if the true reward is a linear representation function).
Revised: where R is a set of reward functions. AL plays a min-max game between the policy π and the reward function r . Intuitively, in the inner loop we would like to find a cost function such that the expert data’s cummulative return is higher than that of the agent’s, and their gap is maximized.

Original: Here, we also show that our method is actually performing feature matching.
Revised: Meanwhile, the policy π tries to minimize this gap. Now, let us consider how to reduce our objective to the AL formulation.

Original: To see this, we use the case of one negative sample for the loss
Revised: For simplicity, we consider the case that we only have one negative sample.

Revised: We notice that Equation 1 is then

Original: Equation 1.
Revised: [Equation]

Original: Using Jensen’s inequality and the concavity of the log function, we can see that
Revised: As suggested by (Khosla et al., 2020), we can apply the Taylor expansion trick to approximate this loss function with the following form:

Original: [Equation]
Revised: [Equation]

Original: Therefore, optimizing the original loss function minimizes linear contrastive loss on the RHS.
Revised: Note that we drop the constant terms and the scaling constant since they do not affect the optimization objective.

Original: Since Φ maps the states toa compact space (i.e. the sphere S ), the RHS of Equation 3 is bounded.
Revised: Moreover, since Φ embeds the data points to the sphere, we have ∥ Φ( x ) ∥ 2 = 1 , ∀ x .

Original: When this minimum value is achieved, the optimal Φ ∗ satisfies:
Revised: As a result, we can further expand each term above and have

Original: [Equation]
Revised: [Equation]

Original: Intuitively, Φ tries to pull the negative policy samples in the opposite direction of expert samples. In apprenticeship learning, they also seek to maximize a similar objective in the form of w T [ E Φ( x p ) − E Φ(˜ x )] , and uses the r al ( x ) =
Revised: Since the variables in this equation are independent from each other, we are minimizing [Equation]

Original: Φ ( x ) T w as the reward function.
Revised: Let the reward function r θ ( x ) =

Original: However, the difference is that the representation Φ in apprenticeship learning is a fixed function defined by some prior knowledge, and the only parameter to optimize is the linear weight w . Here, we consider Φ as a learnable function.
Revised: E x 0 ∼D [Φ θ ( x 0 )]

Original: If we use Φ ∗ asthe representation in apprenticeship learning, we will find that the optimal w ∗ = E x E ∼ ρ πE
Revised: T Φ θ ( x ) as we defined in Equation 2, then minimizing the Equation 7 is exactly doing the maximization of

Original: Φ ∗ ( x E ) .
Revised: [Equation]

Original: In this case, the apprenticeship learning reward coincides with our reward: r al ( x ) = r ( x ) . In other words, our reward function (Equation 2) gives an improved form of apprenticeship learning IRL reward, in the sense that the representation is learned through contrastive learning.
Revised: which is exactly the inner maximization loop of AL.

Original: From this theoretical result, we observe that PCIL integrates representation learning and distribution matching in a unified loss function and frees us from jointly using representation loss and AIL loss like prior methods. As we will see later, our innovation leads to better performance than these methods.
Revised: Then optimizing the policy with respect to this r θ is exactly the outer loop. Hence, our algorithm is reduced to AL.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we empirically evaluate PCIL on an extensive set of tasks from the DeepMind control suite (Tassa et al., 2018), a widely used benchmark for continuous control.
Revised: In this section, we empirically evaluate PCIL on an extensive set of tasks from the DeepMind control suite (Tassa et al., 2018), a widely used benchmark for continuous control.

Original: Our experiments are designed to answer the following questions: (1) Can PCIL achieve expert performance, and how sample efficient is PCIL compared to state-of-the-art imitation learning algorithms?
Revised: Our experiments are designed to answer the following questions: (1) Can PCIL achieve expert performance, and how sample efficient is PCIL compared to state-of-the-art imitation learning algorithms?

Original: (2) How does the representation space of PCIL differ from that of the AIL methods?
Revised: (2) How does the representation space of PCIL differ from that of the AIL methods?

Original: (3) How does our method perform when we use different representation learning methods and reward design?
Revised: (3) How does our method perform when we use different representation learning methods and reward design?

Original: 4. E XPERIMENTAL S ETUP
Revised: 4. E XPERIMENTAL S ETUP

Original: Environments We experiment with 10 MuJoCo (Todorov et al., 2012) tasks provided by DeepMind Control Suite.
Revised: Environments We experiment with 10 MuJoCo (Todorov et al., 2012) tasks provided by DeepMind Control Suite.

Original: The selected tasks cover various difficulty levels, ranging from simple control problems, such as the single degree of freedom cart pole, to complex high-dimensional tasks, such as the quadruped run.
Revised: The selected tasks cover various difficulty levels, ranging from simple control problems, such as the single degree of freedom cart pole, to complex high-dimensional tasks, such as the quadruped run.

Original: The episode length for all tasks is 1000 steps, where a per-step ground truth environment reward is in the unit interval [0 , 1] .
Revised: The episode length for all tasks is 1000 steps, where a per-step ground truth environment reward is in the unit interval [0 , 1] .

Original: For each task, we train an expert policy using
Revised: For each task, we train an expert policy using

Original: DrQ-v2 (Yarats et al., 2021) with the true environment reward function and use it to collectdemonstrations.
Revised: DrQ-v2 (Yarats et al., 2021) with the true environment reward function and use it to collectdemonstrations.

Original: We refer readers to Appendix A for the full task list and more details about the demonstrations.
Revised: We refer readers to Appendix A for the full task list and more details about the demonstrations.

Original: Training Details To update the encoder, we randomly sample 128 expert transitions and 128 agent transitions from a replay buffer.
Revised: Training Details To update the encoder, we randomly sample 128 expert transitions and 128 agent transitions from a replay buffer.

Original: For arbitrary expert transition, any other expert transition is considered a positive sample, and all the agent transitions constitute the set of negative samples.
Revised: For arbitrary expert transition, any other expert transition is considered a positive sample, and all the agent transitions constitute the set of negative samples.

Original: We update the encoder by minimizing Equation 1 with respect to these samples.
Revised: We update the encoder by minimizing Equation 1 with respect to these samples.

Original: We use DrQ-v2 (Yarats et al., 2021) as the underlying RL algorithm to train the agent with the cosine similarity reward given in Equation 2.
Revised: We use DrQ-v2 (Yarats et al., 2021) as the underlying RL algorithm to train the agent with the cosine similarity reward given in Equation 2.

Original: We use a budget of 2M environment steps for all the experiments.
Revised: We use a budget of 2M environment steps for all the experiments.

Original: Further implementation details can be found in Appendix B.
Revised: Further implementation details can be found in Appendix B.

Original: Baselines We compare PCIL to Behavioral Cloning (BC) and two major classes of imitation learning algorithms:
Revised: Baselines We compare PCIL to Behavioral Cloning (BC) and two major classes of imitation learning algorithms:

Original: Adversarial IRL: We consider Discriminator-Actor-Critic (DAC)
Revised: Adversarial IRL: We consider Discriminator-Actor-Critic (DAC)

Original: (Kostrikov et al., 2018), a state-of-the-art AIL method that employs an unbiased AIL reward function and performs offpolicy training to reduce environmental interactions.
Revised: (Kostrikov et al., 2018), a state-of-the-art AIL method that employs an unbiased AIL reward function and performs offpolicy training to reduce environmental interactions.

Original: Trajectory-matching IRL:
Revised: Trajectory-matching IRL:

Original: Primal Wasserstein Imitation Learning (PWIL) (Dadashi et al., 2020) and Sinkhorn Imitation Learning (SIL) (Papagiannis & Li, 2020) are two recently proposed trajectory-matching imitation learning methods.
Revised: Primal Wasserstein Imitation Learning (PWIL) (Dadashi et al., 2020) and Sinkhorn Imitation Learning (SIL) (Papagiannis & Li, 2020) are two recently proposed trajectory-matching imitation learning methods.

Original: PWIL computes the reward based on an upper bound of Wasserstein distance.
Revised: PWIL computes the reward based on an upper bound of Wasserstein distance.

Original: SIL computes the reward based on Sinkhorn distances (Cuturi, 2013).
Revised: SIL computes the reward based on Sinkhorn distances (Cuturi, 2013).

Original: To ensure a fair comparison, we implement all the baselines using the same RL algorithm.
Revised: To ensure a fair comparison, we implement all the baselines using the same RL algorithm.

Original: The implementation details of these algorithms are in the Appendix.
Revised: The implementation details of these algorithms are in the Appendix.

Original: 4. M AIN R ESULTS
Revised: 4. M AIN R ESULTS

Original: We show the performance curves of 6 tasks in Figure 3, which are averaged over three random seeds.
Revised: We show the performance curves of 6 tasks in Figure 3, which are averaged over three random seeds.

Original: More results on DeepMind control tasks are provided in Appendix D. We find that PCIL is able to outperform the existing methods on all of these tasks.
Revised: More results on DeepMind control tasks are provided in Appendix D. We find that PCIL is able to outperform the existing methods on all of these tasks.

Original: It achieves near-expert performance within our online sample budget in all considered tasks except Hopper Hop.
Revised: It achieves near-expert performance within our online sample budget in all considered tasks except Hopper Hop.

Original: In terms of sample efficiency, i.e., the number of environment interactions required to solve a task, PCIL shows significant improvements over prior methods on five tasks: Cheetah Run, Finger Spin, Hopper Hop, Hopper
Revised: In terms of sample efficiency, i.e., the number of environment interactions required to solve a task, PCIL shows significant improvements over prior methods on five tasks: Cheetah Run, Finger Spin, Hopper Hop, Hopper

Original: Stand, and Quadruped Run.
Revised: Stand, and Quadruped Run.

Original: For the remaining tasks, PCIL achieves similar results compared with the state-of-the-art adversarial imitation learning method DAC.
Revised: For the remaining tasks, PCIL achieves similar results compared with the state-of-the-art adversarial imitation learning method DAC.

Original: In particular, we notice that PCIL’s performance gain is larger on more difficult tasks (e.g., Cheetah Run, Quadruped Run).
Revised: In particular, we notice that PCIL’s performance gain is larger on more difficult tasks (e.g., Cheetah Run, Quadruped Run).

Original: On those easier tasks (e.g., Walker Stand, Walker Walk), the baselines are also able to achieve strong results.
Revised: On those easier tasks (e.g., Walker Stand, Walker Walk), the baselines are also able to achieve strong results.

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

Original: 4. A NALYSIS OF R EPRESENTATION S PACE
Revised: 4. A NALYSIS OF R EPRESENTATION S PACE

Original: We visualize the representation space of PCIL and DAC using t-SNE (Van der Maaten & Hinton, 2008) in Figure 4.
Revised: We visualize the representation space of PCIL and DAC using t-SNE (Van der Maaten & Hinton, 2008) in Figure 4.

Original: For DAC, since there is no explicit representation learning in the discriminator, we treat the last hidden layer of its discriminator as the representation.
Revised: For DAC, since there is no explicit representation learning in the discriminator, we treat the last hidden layer of its discriminator as the representation.

Original: We randomly sample 128 expert transitions and 256 agent transitions for visualization.
Revised: We randomly sample 128 expert transitions and 256 agent transitions for visualization.

Original: For a fair comparison, all the transitions of
Revised: For a fair comparison, all the transitions of

Original: [Table]
Revised: [Table]

Original: PCIL and DAC are selected from the same episode during the training process on the same task.
Revised: PCIL and DAC are selected from the same episode during the training process on the same task.

Original: We use color to indicate the real environment reward of the agent’s transitions.
Revised: We use color to indicate the real environment reward of the agent’s transitions.

Original: The lighter (yellow) color indicates a higher reward for agent transition while the darker (green) indicates a lower reward agent transition.
Revised: The lighter (yellow) color indicates a higher reward for agent transition while the darker (green) indicates a lower reward agent transition.

Original: The blue color indicates the expert’s transition.
Revised: The blue color indicates the expert’s transition.

Original: We observe that in the representation space of PCIL, the expert transitions are concentrated in a cluster.
Revised: We observe that in the representation space of PCIL, the expert transitions are concentrated in a cluster.

Original: Moreover, the distance between each agent transition to the cluster of expert transition is highly correlated with the real reward of that agent transition.
Revised: Moreover, the distance between each agent transition to the cluster of expert transition is highly correlated with the real reward of that agent transition.

Original: This fact suggests that the contrastive objective of PCIL indeed induces a meaningful representation space here.
Revised: This fact suggests that the contrastive objective of PCIL indeed induces a meaningful representation space here.

Original: On the contrary, the representation space of DAC is much less structured.
Revised: On the contrary, the representation space of DAC is much less structured.

Original: The expert transitions are scattered throughout the representation space of DAC.
Revised: The expert transitions are scattered throughout the representation space of DAC.

Original: Moreover, we identify that in the DAC’s representation space, the agent’s real ward does not correlate well with its distance to the expert transitions.
Revised: Moreover, we identify that in the DAC’s representation space, the agent’s real ward does not correlate well with its distance to the expert transitions.

Original: These facts show that our method indeed learns a better representation space.
Revised: These facts show that our method indeed learns a better representation space.

Original: 4. A BLATION S TUDIES
Revised: 4. A BLATION S TUDIES

Original: As described in our method, our method has two components: a policy contrastive representation for imitation (Section 3.2) and a similarity-based imitation reward (Section 3.3).
Revised: As described in our method, our method has two components: a policy contrastive representation for imitation (Section 3.2) and a similarity-based imitation reward (Section 3.3).

Original: In this part, we carry out ablation studies to analyze their effects.
Revised: In this part, we carry out ablation studies to analyze their effects.

Original: We first introduce our design choices as follows.
Revised: We first introduce our design choices as follows.

Original: PCIL Representation v.s. TCN Representation
Revised: PCIL Representation v.s. TCN Representation

Original: We replace our proposed policy contrastive objective with other contrastive learning methods.
Revised: We replace our proposed policy contrastive objective with other contrastive learning methods.

Original: For this purpose, we adopt the popular selfsupervised representation learning method that leverages temporal information: Time-Contrastive
Revised: For this purpose, we adopt the popular selfsupervised representation learning method that leverages temporal information: Time-Contrastive

Original: Networks (TCN) (Sermanet et al., 2018).
Revised: Networks (TCN) (Sermanet et al., 2018).

Original: In this case, the positive samples are selected within a small window around the anchor sample, while the negative samples are selected from distant time steps in the same rollout trajectory.
Revised: In this case, the positive samples are selected within a small window around the anchor sample, while the negative samples are selected from distant time steps in the same rollout trajectory.

Original: See Appendix C for implementation details.
Revised: See Appendix C for implementation details.

Original: Similarity-based Reward v.s. GAIL-like Reward
Revised: Similarity-based Reward v.s. GAIL-like Reward

Original: We also ablate the similarity-based reward in
Revised: We also ablate the similarity-based reward in

Original: PCIL by replacing the similarity-based imitation reward with a GAIL-like reward.
Revised: PCIL by replacing the similarity-based imitation reward with a GAIL-like reward.

Original: Specifically, we train a linear binary classifier on the policy contrastive embedding space to distinguish expert or non-expert data.
Revised: Specifically, we train a linear binary classifier on the policy contrastive embedding space to distinguish expert or non-expert data.

Original: In this case, the embedding space is still trained by PCIL contrastive loss, and the
Revised: In this case, the embedding space is still trained by PCIL contrastive loss, and the

Original: GAIL reward’s gradient is detached from the embedding network.
Revised: GAIL reward’s gradient is detached from the embedding network.

Original: We use the same reward predictor as other GAIL-style methods (Kostrikov et al., 2018), i.e. log( D ( x ))
Revised: We use the same reward predictor as other GAIL-style methods (Kostrikov et al., 2018), i.e. log( D ( x ))

Original: − log(1 − D ( x )) .
Revised: − log(1 − D ( x )) .

Original: Analysis By comparing rows 1 and 2 in Table 1, we find that the approach with TCN encoder does not work in three out of four environments.
Revised: Analysis By comparing rows 1 and 2 in Table 1, we find that the approach with TCN encoder does not work in three out of four environments.

Original: This is because the optimization goal of TCN is not to distinguish between expert and non-expert data.
Revised: This is because the optimization goal of TCN is not to distinguish between expert and non-expert data.

Original: Thus the reward produced by comparing expert and non-expert data in the learned representation space is not necessarily meaningful.
Revised: Thus the reward produced by comparing expert and non-expert data in the learned representation space is not necessarily meaningful.

Original: Note that the case in row 2 is no longer an adversarial IRL method.
Revised: Note that the case in row 2 is no longer an adversarial IRL method.

Original: We also consider a case (row 4 in Table 1) where we use TCN and GAIL-like reward predictor, but the performance of this method is poor.
Revised: We also consider a case (row 4 in Table 1) where we use TCN and GAIL-like reward predictor, but the performance of this method is poor.

Original: Moreover, we observe that in the absence of the similarity-based imitation reward (compare rowsand 3 in Table 1), our method does not work.
Revised: Moreover, we observe that in the absence of the similarity-based imitation reward (compare rowsand 3 in Table 1), our method does not work.

Original: This is because our representation space has metricspace characteristics.
Revised: This is because our representation space has metricspace characteristics.

Original: As a result, we should use a distance-based metric to compute the reward.
Revised: As a result, we should use a distance-based metric to compute the reward.

Original: In conclusion, the two components of our method are necessary for achieving good performance.
Revised: In conclusion, the two components of our method are necessary for achieving good performance.

Original:  RELATED WORK 5. I MITATION L EARNING
Revised:  RELATED WORK 5. I MITATION L EARNING

Original: Imitation learning is a class of algorithms that enables a robot to acquire behaviors from a demonstration dataset.
Revised: Imitation learning is a class of algorithms that enables a robot to acquire behaviors from a demonstration dataset.

Original: There are two classes of imitation learning algorithms: behavioral cloning (BC) and
Revised: There are two classes of imitation learning algorithms: behavioral cloning (BC) and

Original: Inverse Reinforcement Learning (IRL) (Ng et al., 2000).
Revised: Inverse Reinforcement Learning (IRL) (Ng et al., 2000).

Original: BC is a simple supervised learning algorithm that directly fits the expert’s action.
Revised: BC is a simple supervised learning algorithm that directly fits the expert’s action.

Original: However, some work suggests that it has some drawbacks: it suffers from covariate shift problem (Ross et al., 2011), and it is hard to learn from a demonstration dataset without expert actions (Torabi et al., 2018).
Revised: However, some work suggests that it has some drawbacks: it suffers from covariate shift problem (Ross et al., 2011), and it is hard to learn from a demonstration dataset without expert actions (Torabi et al., 2018a).

Original: Instead, IRL (Abbeel & Ng, 2004) proposes to recover the underlying policy by estimating the underlying reward function and then maximizing the overall return with this reward.
Revised: Instead, IRL (Abbeel & Ng, 2004) proposes to recover the underlying policy by estimating the underlying reward function and then maximizing the overall return with this reward.

Original: In particular, a recent branch of IRL is the AIL, which proposes to match agents’ state-action distribution with experts via adversarial training.
Revised: In particular, a recent branch of IRL is the AIL, which proposes to match agents’ state-action distribution with experts via adversarial training.

Original: GAIL (Ho & Ermon, 2016) proposed a maximum entropy occupancy measure matching method which learns a discriminator to bypass the need to recover the expert’s reward function.
Revised: GAIL (Ho & Ermon, 2016) proposed a maximum entropy occupancy measure matching method which learns a discriminator to bypass the need to recover the expert’s reward function.

Original: Later, several works proposed an improved version of the GAIL methods (Ghasemipour et al., 2020; Blond´e et al., 2022; Baram et al., 2017; Kostrikov et al., 2018; Fu et al., 2017).
Revised: Later, several works proposed an improved version of the GAIL methods (Ghasemipour et al., 2020; Blond´e et al., 2022; Baram et al., 2017; Kostrikov et al., 2018; Fu et al., 2017).

Original: AIRL (Fu et al., 2017) replaced the Shannon-Jensen divergence used in GAIL by Kullback-Leibler divergence to measure similarity between state-action pair distributions.
Revised: AIRL (Fu et al., 2017) replaced the Shannon-Jensen divergence used in GAIL by Kullback-Leibler divergence to measure similarity between state-action pair distributions.

Original: Baram et al.
Revised: Baram et al.

Original: (2017) bridges the GAIL framework to model-based reinforcement learning.
Revised: (2017) bridges the GAIL framework to model-based reinforcement learning.

Original: DAC (Kostrikov et al., 2018) improved the sample efficiency by leveraging a replay buffer without importance sampling and dealing with the absorbing state problem.
Revised: DAC (Kostrikov et al., 2018) improved the sample efficiency by leveraging a replay buffer without importance sampling and dealing with the absorbing state problem.

Original: In contrast to these works, we focus on the representation of AIL’s discriminator and reformulating AIL in a contrastive embedding space.
Revised: In contrast to these works, we focus on the representation of AIL’s discriminator and reformulating AIL in a contrastive embedding space.

Original: 5. R EPRESENTATION L EARNING F OR P OLICY L EARNING
Revised: 5. R EPRESENTATION L EARNING F OR P OLICY L EARNING

Original: In this work, we propose a representation-learning-based approach to improve imitation learning.
Revised: In this work, we propose a representation-learning-based approach to improve imitation learning.

Original: As imitation learning is a major class of policy learning algorithms, we review works that use selfsupervised learning to improve policy learning in this part.
Revised: As imitation learning is a major class of policy learning algorithms, we review works that use selfsupervised learning to improve policy learning in this part.

Original: Pioneer works (Mirowski et al., 2016;
Revised: Pioneer works (Mirowski et al., 2016;

Original: Jaderberg et al., 2016; Shelhamer et al., 2016; Lample & Chaplot, 2017) explore using auxiliary objectives (e.g., predict some property of the environment) as a self-supervision signal.
Revised: Jaderberg et al., 2016; Shelhamer et al., 2016; Lample & Chaplot, 2017) explore using auxiliary objectives (e.g., predict some property of the environment) as a self-supervision signal.

Original: Recent works employ more general self-supervised objectives (Oord et al., 2018).
Revised: Recent works employ more general self-supervised objectives (Oord et al., 2018).

Original: In particular, Srinivas et al.
Revised: In particular, Srinivas et al.

Original: (2020) are based on contrastive representation learning.
Revised: (2020) are based on contrastive representation learning.

Original: Sermanet et al.
Revised: Sermanet et al.

Original: (2018) learn representation from multiview video using time contrastive learning.
Revised: (2018) learn representation from multiview video using time contrastive learning.

Original: Some other methods also explore the use of self-supervised representation pretrained on environment data (Ha & Schmidhuber, 2018) or from real-world images (Xiao et al., 2022; Parisi et al., 2022; Nair et al., 2022).
Revised: Some other methods also explore the use of self-supervised representation pretrained on environment data (Ha & Schmidhuber, 2018) or from real-world images (Xiao et al., 2022; Parisi et al., 2022; Nair et al., 2022).

Original: In imitation learning,
Revised: In imitation learning,

Original: Mandi et al.
Revised: Mandi et al.

Original: (2022) proposes to use contrastive learning for one-shot imitation learning in robotics.
Revised: (2022) proposes to use contrastive learning for one-shot imitation learning in robotics.

Original: Chen et al.
Revised: Chen et al.

Original: (2021) investigate the use of representation for imitation learning.
Revised: (2021) investigate the use of representation for imitation learning.

Original: However, their result suggests that self-supervised representation learning only provides a small improvement of imitation learning algorithms’ performance.
Revised: However, their result suggests that self-supervised representation learning only provides a small improvement of imitation learning algorithms’ performance.

Original: Our method differs from all these existing works by proposing to anchor on different policies and learn a discriminative self-supervised representation for imitation learning.
Revised: Our method differs from all these existing works by proposing to anchor on different policies and learn a discriminative self-supervised representation for imitation learning.

Original:  CONCLUSION
Revised:  CONCLUSION

Original: In this paper, we suggested a new approach to improve adversarial imitation learning algorithms: to learn a more meaningful, discriminative representation space for imitation.
Revised: In this paper, we suggested a new approach to improve adversarial imitation learning algorithms: to learn a more meaningful, discriminative representation space for imitation.

Original: To this end, we proposed a new algorithm called PCIL.
Revised: To this end, we proposed a new algorithm called PCIL.

Original: We conducted a theoretical analysis of our method and showed its connection to apprenticeship learning.
Revised: We conducted a theoretical analysis of our method and showed its connection to apprenticeship learning.

Original: We also conducted experiments on the DeepMind Control
Revised: We also conducted experiments on the DeepMind Control

Original: Suite and showed that PCIL could achieve state-of-the-art performance.
Revised: Suite and showed that PCIL could achieve state-of-the-art performance.

Original: Moreover, we used an ablation study to highlight its difference from the previous representation learning method.
Revised: Moreover, we used an ablation study to highlight its difference from the previous representation learning method.

Original: In the future, we will focus on further improving our loss function design.
Revised: In the future, we will focus on further improving our loss function design.

Original: For example, can we anchor on the agent policies at different training stages?
Revised: For example, can we anchor on the agent policies at different training stages?

Original: It will also be interesting to extend the proposed representation learning method in the relaxed setting of IL, like the scene where we can access both the reward and demonstration.
Revised: It will also be interesting to extend the proposed representation learning method in the relaxed setting of IL, like the scene where we can access both the reward and demonstration.

Original:  R EPRODUCIBILITY
Revised:  R EPRODUCIBILITY

Original: We implement our algorithm according to parameters and details described in Appendix B and
Revised: We implement our algorithm according to parameters and details described in Appendix B and

Original: Section 4.1.
Revised: Section 4.1.

Original: We will release our code and data.
Revised: We will release our code and data.

Original: A E NVIRONMENTS
Revised: A E NVIRONMENTS

Original: We use 10 continuous control tasks from the DeepMind control suite (Tassa et al., 2018).
Revised: We use 10 continuous control tasks from the DeepMind control suite (Tassa et al., 2018).

Original: The summary for each task is provided in Table 2.
Revised: The summary for each task is provided in Table 2.

Original: [Table]
Revised: [Table]

Original: Demonstrations For each task, we train expert policies using DrQ-v2 (Yarats et al., 2021) on the actual environment rewards.
Revised: Demonstrations For each task, we train expert policies using DrQ-v2 (Yarats et al., 2021) on the actual environment rewards.

Original: We run 3 seeds and pick the seed that achieves the highest return.
Revised: We run 3 seeds and pick the seed that achieves the highest return.

Original: Then we use this expert policy to collect 10 demonstrations.
Revised: Then we use this expert policy to collect 10 demonstrations.

Original: B A LGORITHM D ETAILS B. I MPLEMENTATION
Revised: B A LGORITHM D ETAILS B. I MPLEMENTATION

Original: RL agent We use DrQ-v2 as the underlying RL algorithm.
Revised: RL agent We use DrQ-v2 as the underlying RL algorithm.

Original: DrQ-v2 is an off-policy actor-critic algorithm for continuous control.
Revised: DrQ-v2 is an off-policy actor-critic algorithm for continuous control.

Original: The core of DrQ-v2 is Deep Deterministic Policy Gradient (DDPG) (Lillicrap et al., 2015) augmented with n -step returns.
Revised: The core of DrQ-v2 is Deep Deterministic Policy Gradient (DDPG) (Lillicrap et al., 2015) augmented with n -step returns.

Original: The critic is trained using clipped double Q-learning (Fujimoto et al., 2018) to reduce the overestimation bias in the target value.
Revised: The critic is trained using clipped double Q-learning (Fujimoto et al., 2018) to reduce the overestimation bias in the target value.

Original: The deterministic actor is trained using deterministic policy gradients (DPG) (Silver et al., 2014).
Revised: The deterministic actor is trained using deterministic policy gradients (DPG) (Silver et al., 2014).

Original: We also follows the setting of actor’s and critic’s neural network architectures in state-based DrQv2 Yarats et al.
Revised: We also follows the setting of actor’s and critic’s neural network architectures in state-based DrQv2 Yarats et al.

Original: Contrastive encoder
Revised: Contrastive encoder

Original: The contrastive encoder is implemented as a 4 layer MLP with hidden size [256, 256, 256].
Revised: The contrastive encoder is implemented as a 4 layer MLP with hidden size [256, 256, 256].

Original: The output dimension is 64.
Revised: The output dimension is 64.

Original: Following the architecture in Yarats et al.
Revised: Following the architecture in Yarats et al.

Original: (2021), the contrastive encoder, the critic and the actor share the same encoder backbone.
Revised: (2021), the contrastive encoder, the critic and the actor share the same encoder backbone.

Original: This shared encoder is trained with the gradient of the critic alone, which is also following the suggestion of Kostrikov et al.
Revised: This shared encoder is trained with the gradient of the critic alone, which is also following the suggestion of Kostrikov et al.

Original: ; Yarats et al.
Revised: ; Yarats et al.

Original: The input of this shared encoder is state s of a transition.
Revised: The input of this shared encoder is state s of a transition.

Original: Reward predictor Reward of the agent transition is computed according to Equation 2.
Revised: Reward predictor Reward of the agent transition is computed according to Equation 2.

Original: Note that cosine similarity between expert data points is high due to the optimization goal described in Equation 3.
Revised: Note that cosine similarity between expert data points is high due to the optimization goal described in Equation ?? .

Original: Thus, we randomly sample one expert transition from the expert replay buffer to compute agent reward.
Revised: Thus, we randomly sample one expert transition from the expert replay buffer to compute agent reward.

Original: Empirically, we find that using the mean embedding of the expert data yields similar performance.
Revised: Empirically, we find that using the mean embedding of the expert data yields similar performance.

Original: Gradient penalty In order to make the algorithm more stable, we use the gradient penalty technique (Gulrajani et al., 2017) widely used in Wasserstein-GANs (Arjovsky et al., 2017).
Revised: Gradient penalty In order to make the algorithm more stable, we use the gradient penalty technique (Gulrajani et al., 2017) widely used in Wasserstein-GANs (Arjovsky et al., 2017).

Original: We make minor adjustments to accommodate our policy contrastive loss.
Revised: We make minor adjustments to accommodate our policy contrastive loss.

Original: GAIL-like methods usually constrain the gradient norm of the discriminator’s output with respect to its input.
Revised: GAIL-like methods usually constrain the gradient norm of the discriminator’s output with respect to its input.

Original: While for PCIL, the contrastive encoder’s output needs one more step.
Revised: While for PCIL, the contrastive encoder’s output needs one more step.

Original: Specifically, the output embeddings are first used to calculate rewards following Equation 2, then we compute and penalize the gradient norm of the rewards.
Revised: Specifically, the output embeddings are first used to calculate rewards following Equation 2, then we compute and penalize the gradient norm of the rewards.

Original: We use 10 as the weighting for the gradient penalty.
Revised: We use 10 as the weighting for the gradient penalty.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: YPERPARAMETERS
Revised: YPERPARAMETERS

Original: Table 3 lists the hyperparameters that are used for all baseline methods and our method.
Revised: Table 3 lists the hyperparameters that are used for all baseline methods and our method.

Original: Expert data ratio in PCIL means the ratio between expert data and batch size.
Revised: Expert data ratio in PCIL means the ratio between expert data and batch size.

Original: A ratio of 0.5 means that half of the batch is expert data and the other half is the agent data.
Revised: A ratio of 0.5 means that half of the batch is expert data and the other half is the agent data.

Original: The contrastive learning usually needs a temperature scaling after computing the cos-similarity, before computing the exponential.
Revised: The contrastive learning usually needs a temperature scaling after computing the cos-similarity, before computing the exponential.

Original: For simplicity, we ignored it in the main text.
Revised: For simplicity, we ignored it in the main text.

Original: In the experiment, we follow prior contrastive learning work He et al.
Revised: In the experiment, we follow prior contrastive learning work He et al.

Original: and use a typical value of 0 .
Revised: and use a typical value of 0 .

Original: 07 for the temperature.
Revised: 07 for the temperature.

Original: [Table]
Revised: [Table]

Original: C A BLATION STUDY IMPLEMENTATION DETAILS C. TCN REPRESENTATION
Revised: C A BLATION STUDY IMPLEMENTATION DETAILS C. TCN REPRESENTATION

Original: TCN encoder shares the same network architecture as the PCL encoder.
Revised: TCN encoder shares the same network architecture as the PCL encoder.

Original: During the training process in TCN encoder, for a random sampled anchor, we use data point adjacent to it as positive pair and another random sampled data point as negative pair.
Revised: During the training process in TCN encoder, for a random sampled anchor, we use data point adjacent to it as positive pair and another random sampled data point as negative pair.

Original: We also set contrastive temperature to 0 .
Revised: We also set contrastive temperature to 0 .

Original: 07 and batch size to 256 , which is the same as PCL encoder.
Revised: 07 and batch size to 256 , which is the same as PCL encoder.

Original: C. GAIL- LIKE REWARD
Revised: C. GAIL- LIKE REWARD

Original: GAIL reward predictor can be seen as a simplified version of GAIL discriminator which has only one linear classifier layer.
Revised: GAIL reward predictor can be seen as a simplified version of GAIL discriminator which has only one linear classifier layer.

Original: The reward predictor is trained independently to distinguish whether the input is from a expert data or non-expert data with a binary classification loss.
Revised: The reward predictor is trained independently to distinguish whether the input is from a expert data or non-expert data with a binary classification loss.

Original: D A DDITIONAL E XPERIMENTAL R ESULTS
Revised: D A DDITIONAL E XPERIMENTAL R ESULTS

Original: Figure 5 shows the performance of PCIL on the other 4 tasks from the DeepMind Control suite.
Revised: Figure 5 shows the performance of PCIL on the other 4 tasks from the DeepMind Control suite.

Original: We notice that the performance of some relatively easy tasks has saturated.
Revised: We notice that the performance of some relatively easy tasks has saturated.

Original: All the baselines achieve expert performance on Pendulum Swingup .
Revised: All the baselines achieve expert performance on Pendulum Swingup .

Original: On Walker Stand and Walker Walk , PCIL is competitive with DAC, which already demonstrates impressive sample efficiency.
Revised: On Walker Stand and Walker Walk , PCIL is competitive with DAC, which already demonstrates impressive sample efficiency.

Original: [Figure]
Revised: [Figure]

Original: Expert
Revised: Expert

Original: BC
Revised: BC

Original: SIL
Revised: SIL

Original: PWIL
Revised: PWIL

Original: DAC
Revised: DAC

Original: PCIL (Ours)
Revised: PCIL (Ours)