Policy Contrastive Imitation Learning
Authors: Jialei Huang
Zhao-Heng Yin
Yingdong Hu
Yang Gao
Venue: Submitted to ICLR 2023 (2023)
Keywords: adversarial imitation learning
contrastive learning

Abstract: Adversarial imitation learning (AIL) is a popular method that has recently achieved much success. However, the performance of AIL is still unsatisfactory on the more challenging tasks. We find that one of the major reasons is due to the low quality of AIL discriminator representation. 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. We propose a new method called Policy Contrastive Imitation Learning (PCIL) to resolve this issue. PCIL learns a contrastive representation space by anchoring on different policies and uses a smooth cosine-similarity-based reward to encourage imitation learning. 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. From a theoretical perspective, we show the validity of our method using the apprenticeship learning framework. Furthermore, our empirical evaluation on the DeepMind Control suite demonstrates that PCIL can achieve state-of-the-art performance. Finally, qualitative results suggest that PCIL builds a smoother and more meaningful representation space for imitation learning.

Under review as a conference paper at ICLR 2023
POLICY CONTRASTIVE IMITATION LEARNING
Anonymous authors
Paper under double-blind review
ABSTRACT
Adversarial imitation learning (AIL) is popular method that has recently achieved much success. However, the performance of AlL is still unsatisfactory on the more challenging tasks. We find that one of the major reasons is due the low quality of AIL discriminator representation. 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. We propose a new method called Policy Contrastive Imitation Learning (PCIL) to resolve this issue. PCIL learns contrastive representation space by anchor- ing on different policies and generates a smooth cosine-similarity-basss reward. Our proposed representation learning objective be viewed stronger ver- sion the AIL objective and provide a more meaningful comparison between the agent and the policy. From theoretical perspective, we show the validity of our method using the apprenticeship learning framework. Furthermore, our empirical evaluation on the DeepMind Control suite demonstrates that PCIL can achieve state-of-the-art performance. Finally, qualitative results suggest that PCIL builds a smoother and more meaningful representation space for imitation learning.
INTRODUCTION
Imitation is one of the fundamental capabilities of an intelligent agent (Hussein et al. 2017). Ani- mals and humans can acquire many skills by mimicking each other (Bymee) 2009). In engineering, imitation learning also enables many robotics applications. One mainstream class of imitation learn- ing algorithms the adversarial imitation learning (AIL) (Ho Ermon 2016). AIL converts the imitation task into a distribution matching problem and proposes to imitate it by training policy against an adversarial discriminator. All. has enjoyed great success on many imitation tasks: achieves superior performance (Ho & Ermon 2016 Kostrikov et al.) (018). and has been experi- mentally proven to alleviate some of the distributional drift issue, and can work even without expert actions (Torabi et al. Howeeer, 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. (02II. The fragility of the discriminator (Peng 20.11 8 not only leads to poor performance but also severely limits the applicability of AIL to a broader range of tasks.
Numerous techniques have been proposed to improve the performance of AlL., such as using reg- ularization and gradient penalties (Fu et al.) 2017] Kostrikov et al., 2018 Gulrajani et al., (017). Some works also propose to use different distribution metries (e.g.. KL divergence, Wasserstein dis- tance) (Xiao et al., 2019) for distribution matching and show some improvements. Though these methods show encouraging results, we notice that they ignore one crucial aspect of the problem: the representation of AIL's discriminator. To be specific, the discriminator in AlL is usually trained with binary classification loss that distinguishes expert transitions from agent transitions. This dis- criminator is then used to define rewards. However, since the only goal of the discriminator is to distinguish the expert from the agent, it does not necessarily learn good, smooth representation space that can provide reasonable comparison between the behavior of two Ideat represen- tations should be able to provide semantically meaningful signals to compare the expert policy and the agent policy.
In this paper, we propose a new algorithm called Policy Contrastive Imitation Learning (PCIL) to achieve this goal. Instead of training with binary- classifcatct objective, we propose to train a dis- criminator representation space with the contrastive learning loss. Our method differs from the prior representation learning approach in AIL in that we perform contrastive learning between different

---

Under review as a conference paper at ICLR 2023
AlL Representation Space
PCIL Representation Space
Decision Boundary
Expert
Sub-optimal
Random
Figure I: Comparison between the representation space of AIL method and our method. Since the AlL methods use binary classification objective to distinguish expert and non-expert transitions, the representation space is only required to separate two classes in two disjoint subspaces. So the embedding space is not required to be semantically meaningful enough, c.g. (Left) the distance between 2 expert data points may be even longer than the distance between expert data point and sub- optimal non-expert data point. (Right) We overcome this limitation by proposing PCIL. Our method enforces the compactness of the expert's representation. This ensures that the learned representation can capture common, robust features of the expert's transitions, which leads more meaningful representation space.
policies. More specifically, we push the expert's representation together and pull the agent pol- icy's representation away from them. We define the imitation learning reward via cosine similarity between the policy's 's and expert's transition.
As is shown in Figure [ the discriminator (binary classifier) might not have good representation space: the distance between two expert transitions can be even larger than the distance between an expert transition and agent transition. This implies that the discriminator may encode some common features of the expert's behavior and may use non-robust features to compare the behav- ior. The lack of proper structure in the representation space of traditional discriminators may yield low-quality AIL rewards. To alleviate this issue, we explicitly define a constraint on the represen- tation space, requiring that the distance between expert transitions be smaller than their distance to the agent's transitions. This is a stronger constraint on the discriminator's representation: it is easy to derive a binary classifier from our learned representation. However, the binary classifier's representation space does not necessarily satisfy objective.
From a theoretical perspective, we show the soundness of PCIL the apprenticeship learning frame- work. We also evaluate method empirically by benchmarking Control Suite (Tassa et al. 2018). Experimental results show that our method is able to achieve state- of-the- art results. Through ablation study and qualitative visualization, find that method is more effective than prior representation learning methods and able to provide a better representation space for imitation learning.
method the DeepMind
In summary, our contributions in this paper are as follows.
1. We point out new direction to improve the performance of the AIL methods, i.e., going beyond naive binary classification and leveraging more stable and meaningful representa- tion learning algorithms for imitation.
2. We propose algorithm called Policy Contrast Imitation Learning (PCIL) method that instantiates such an improvement and establishes its connection to apprenticeship learning from a theoretical perspective.
3. We evaluate method the DeepMind Control Suite and achieve state-of-the oftte-of-the-rtt per- formance. Through ablation studies, we highlight its essential difference from previous contrastive learning methods in AlL..
2

---

Under review as a conference paper at ICLR 2023
2 PRELIMINARIES
2.1 NOTATIONS
In this paper, we model the imitation leanring problem as a markov decision process MM (S, A. (s)s) (s' 8.a) . a.)) ,7). Here, S is the state space. A is the action space. Po(s) de- fines the initial state distribution, p(s' s. a) defines the transition dynamics. r(s, a, s') is the reward function. is the discount factor. The goal is to maximize the expected return of the learned policy n, which is defined by
8 rr (Sks Ok, 88411
( = E.
Σ
For the imitation learning problem, the algorithm does not have access to the reward function and the transition dynamics. Instead, it is provided with an expert demonstration dataset 'D sampled from an expert policy TE. which can perform well in M. Here, DD takes the form of {(sf af)). where (s. is sampled from Pag. the stationary state-action visiting distribution of gg. The imitation learning algorithm is then required to reproduce the expert's behavior with D.
2.2 ADVERSARIAL IMITATION LEARNING
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
The idea of GAIL is to minimize the divergence between Prg and Pr- uses a discriminator ((s to distinguish expert's transitions (s a 'D from the policy transitions (s, Px. which is trained by maximizing the objective
C (e (oee,.)))<nooo(( 0:))] + (offaa))>olllg D(s (eef,a))) ())]
To achieve imitation, the agent policy is then required to fool the discriminator, which can only be possible when the policy x resembles the expert Ty. Specifically, GAIL defines an adversarial reward r(s,, az) log(1 (st, az)). and trains to maximize the expected return with respect to this reward using policy RL algorithms.
3 POLICY CONTRASTIVE IMITAITON LEARNING
3.1 OVERVIEW
We propose novel represenaaioo--laaiing- approach called Policy Contrastive Imitation Learning to improve the AIL reward. The overview of PCIL is illustrated in Figure 2 Our key insight is to learn policy-contrastive representation space. 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. We will discuss the training of this representation in Section 3.2 and the reward design in Section 3.3 Then, will show the convergence of our algorithm in Section 3.4
3.2 CONTRASTIVE POLICY REPRESENTATION FOR IMITATION
The vanilla AIL algorithms are based on unconstrained representations and can be very non-rote non-robust. One possible approach to handle this problem is to learn a more meaningful representation by con- trastive learning. Researchers (He et al. 2020) have found that can usually learn semantically meaningful representation, leading to better performance
downstream tasks, such as classifica- tion. However, though is effective in the field of supervised learning, prior work (Chen et al.) 2021 has found that it does not greatly improve AlL much.
To understand the reason behind this, we first recall that the contrastive learning method learss by drawing the representation of one training sample towards similar positive sample Xp, and pulling away from a dissimilar, negative sample x,, Some heuristic rules determine the choice of positive and negative samples: the positive sample is usually defined as data augmentation of
3

---

Under review as a conference paper at ICLR 2023
Expert Trajectory
#1
Positive
Push together
Expert
Encoder
Pull away
Trajectory #2
Negative
Agent Trajectory
Representation Space
Figure 2: Illustration of our contrastive learning approach. We first select anchor state (the or- ange) from the expert trajectory. Then, we select positive state sample (the red) from another expert trajectory and a negative state sample (the green) from the agent trajectory. We map these selected states to the representation space. Finally, we push the representation of the anchor state and the positive state together and pull the representation of negative samples away from the representation of the anchor state.
Therefore, this augmentation decides what should be similar to the representation. Nevertheless, we argue that the definition of similarity in the previous works is not strong enough for AIL. This is because, in AIL, the representation should also be able to let discern good behavior from the bad behavior. However, the difference between good and bad behavior here can sometimes be very faint. For example, consider a result, fails to accomplish the task. The good and bad states right before this point may look very similar. More specifically, the difference can simply be minor particular physical measurement, like the distance. In this case, the representation should consider this as semantic component and be sensitive such difference to succeed. 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. As a result, the model may not distinguish between good and bad states effectively. Though combining the representation learning objective and the AIL objective may help combat this problem by enforcing a hard distinguishing constraint contrastive representation, in practice find this does not work well (Section
where robot misses the exact point execute certain action and, as
These observations motivate us to learn a representation that is semantically meaningful and able to distinguish between good and bad states. We find 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 negative samples. In the case of imitation learning, samples can be naturally divided into two categories, namely expert and non-expert samples. Then given the encoder &b S S that maps the state to a representation vector in nigh-dimensional sphere, we define its infoNCE representation loss function as follows:
n
L
E Xo (s.a xpp(pp.an ),DD x, (8, A:lmPP
[xxx)) log ((x)) (xxp)) Σ exp (
Here, Xp is some state-action pair from the expert transitions; x, is some transition from the expert data acting as positives; 8, some agent transition working as negatives. In other words, we require & to draw the expert samples towards each other and pull all the policy samples away from the expert samples.
Our proposed objective function is a strictly stronger constraint on the discriminator: The binary classification discriminator can find any hyper-plane that separates the two types of transitions, with constraint how the transitions are embedded. However, our objective enforces the pair- wise distance constraint between any triplets. One can derive a binary expert-police transition classifier from trained & by computing & (xxx)" (xx) L, where Xp is any expert transition, is the tran-

---

Under review as a conference paper at ICLR 2023
sition to be classified and is some threshold. However, on the other hand, the binary classification induced latent space might not satisfy our constraint. As illustrated in Figure [ the latent space for the binary classification discriminator might have some expert-expert pairs that are even further away than some expert-agent pairs.
Interestingly, 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. Here we also consider all the samples in the expert demonstrations as similar. However, unlike SCL, we do not require the samples from the agents to be similar to each other. This is because the agent transitions are generated from different policies during the training process.
3.3 SIMILARITY-BASER IMITATION REWARD
With representation that can capture the difference between good and bad states, we can then define a reward function encourage imitation leaming. Though using an AlL-style reward with this representation is still possible, we find that better choice is to use cosine similarity metric to define the reward. has several advantages: is bounded and appears relatively smooth in practice, leading to more stable learning. Concretely, define:
r(x) = (x)" EXEND (XE).
Nevertheless, in practice evaluating the latter expectation can be time-consuming since & is fre- quently updated. Therefore, we use random expert sample for the reward calculation. From this reward, we can see that a policy can only obtain high rewards when it frequently visits the expert's distribution. This naturally connects method to the distribution matching, and we provide the theoretical analysis of our algorithm in the following subsection.
3.4 THEORETICAL ANALYSIS
In this part, we show that PCIL be reduced Apprenticeship Learning (AL) (Abbeel & Ng. 2004). First, let us recall that an AL problem takes the following form (Ho & Ermon 2016):
min max ER
E (5,0)~D
[r(x)]
I
[r(x)],
x (S.a)~p
where R is a set of reward functions. AL plays min-max game between the policy x and the reward function r. Intuitively, in the inner loop we would like to find cost function such that the expert data's 's cummulative return is higher than that of the agent's, s, and their gap is maximized. Mcanwhile, the policy ı tries to minimize this gap.
Now, let consider how to reduce our objective to the AL formulation. For simplicity, we consider the case that we only have one negative sample. We notice that Equation is then
exp (xn) x
L =
E
log
(4)
exp (xx)) (e (xp) + exp (xx)) (x,))
As suggested by (Khosla et al. 2020), we can apply the Taylor expansion trick to approximate this loss function with the following form:
L E [x)) (x,) (xx) (13)
(5)
Note that we drop the constant terms and the scaling constant since they do not affect the optimiza-
tion objective. Moreover, since & embeds the data points to the sphere, we have [x)]]] = 1, Vx. As result, we can further expand each term above and have
L = E (xx)) ( (xx,)) ((xx) (xx)))((xx))))
(6)
Since the variables in this equation are independent from each other, we are minimizing
C= Ex, ~p ( (xx)) Ex, [[xx--(((xx))) (xx)))
(7)
Let the reward function (x) Exo~I ( (xo) ) xx) as we defined Equation 2 then mini-
mizing the Equation 7 is exactly doing the maximization of
Ex-D [x)( Ex-p, [ro(x)]
(
which is exactly the inner maximization loop of AL. Then optimizing the policy with respect to this ro is exactly the outer loop. Hence, our algorithm is reduced to AL..
5

---

Under review as a conference paper at ICLR 2023
4 EXPERIMENTS
In this section, we empirically evaluate PCIL on an extensive set of tasks from the DeepMind control suite (Tassa et al. 2018), widely used benchmark for continuous control. Our experiments designed to answer the following questions: (1) Can PCIL achieve expert performance, and how sample efficient is PCIL compared to state-of-the-an imitation learning algorithms?
(2) How does the representation space of PCIL differ from that of the AIL methods? (3) How does our method perform when we use different representation learning methods and reward design?
4.1 EXPERIMENTAL SETUP
Environments We experiment with 10 MuloCo Todorov et al., 2012) tasks provided by Deep- Mind Control Suite. The selected tasks cover various difficulty levels, ranging from simple control problems, such as the single degree of freedom cart pole, to complex nigh-dimensional tasks, such as the quadruped run. The episode length for all tasks is 1000 steps, where per-step ground truth environment reward is the unit interval [0, 1]. For each task, we train an expert policy using DrQ-v2 (Yarats et al. 2021) with the true environment reward function and use it to collect 10 demonstrations. We refer readers to Appendix A for the full task list and more details about the demonstrations.
Training Details To update the encoder, we randomly sample 128 expert transitions and 128 agent transitions from replay buffer. For arbitrary expert transition, any other expert transition is con- sidered positive sample, and all the agent transitions constitute the set of negative samples. We update the encoder by minimizing Equation with respect to these samples. We use DrQ-v2 (Yarats) (ct al. 2021) as the underlying RL algorithm to train the agent with the cosine similarity reward given in Equation 2
We use budget of 2M environment steps for all the experiments. Further implementation details can found in Appendix B
Baselines We compare PCIL to Behavioral Cloning (BC) and two major classes of imitation learn- ing algorithms:
1. Adversarial IRL:
We consider Discriminator Acttr-Critrriitii (DAC) (Kostrikov et al. 2018), state-of-the-aar AIL method that employs an unbiased AIL reward function and performs off- policy training to reduce environmental interactions.
Trajectory-matchin IRL: Primal Wasserstein Imitation Learning (PWIL) (Dadashi al. 2020) and Sinkhorn Imitation Learning (SIL) (Papagiannis & Li 2020) are two recently pro- posed trajectory- imitation learning methods. PWIL computes the reward based on an upper bound of Wasserstein distance. SIL computes the reward based on Sinkhorn distances (Cu- 2013).
To ensure a fair comparison, we implement all the baselines using the same RL algorithm. The implementation details of these algorithms are in the Appendix.
4.2 MAIN RESULTS
We show the performance curves of tasks in Figure which are averaged over three random seeds. 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. It achieves near-expert performance within our online sample budget in all considered tasks except Hopper Hop. In terms of sample effi- ciency, i.e., the number environment interactions required to solve a task, PCIL shows significant improvements over prior methods on five tasks: Cheetah Run, Finger Spin, Hopper Hop, Hopper Stand, and Quadruped Run. For the remaining tasks, PCIL achieves similar results compared with the state-of-te the-ahtr adversarial imitation learning method DAC. In particular, performance gain is larger on more difficult tasks (e.g.. Cheetah Run, Quadruped Run). On those easier tasks (e.g., Walker Stand, Walker Walk), the baselines also able to achieve strong results.
notice that PCIL's

---

Under review as a conference paper at ICLR 2023
Cheetah Run
Finger Spin
Hopper Hop
1000
Eptrode Rot am
820
200
130
0.3
1.3
0.0
0.5
1
2.0
6.5
1/3
Enstronmmnt Stoes (x1ef)
Environment Stopo (x111)
Environment Stops Ur
HopperStanne Stand
Quadruped Run
Walker Run
770
600
630
Eretode Det-rn
930
600
600
+40
200
Steps
Expert
BC
SIL
PWIL
DAC
PCIL (Ours)
Figure 3: Comparisons of algorithms on selected tasks. See Appendix D for more tasks. For every 20k environment steps, perform 10-episode rollouts of the policy without exploration noise and report average episode returns over the 10 episodes. We plot the mean performance over seeds together with the shaded regions, which represent 95% confidence intervals.
DAC Cheetah Run
DAC Hopper Hop
DAC Finger Spin
PCIL Cheetah Run
PCIL Hopper Hop
PCIL Finger Spin
Figure 4: t-SNE visualization results for DAC (top panels) and our PCIL method (bottom panels). The blue color indicates the expert's transition. The lighter (yellow) color indicates agent transitions with higher real reward while the darker (green) indicates lower real reward.
4.3 ANALYSIS OF REPRESENTATION SPACE
We visualize the representation space of PCIL and DAC using t-SNE (Van der Maaten & Hinton 2008 in Figure For DAC, since there no explicit representation learning in the discriminator, treat the last hidden layer of its discriminator as the representation. We randomly sample 128 expert transitions and 256 agent transitions for visualization. For fair comparison, all the transitions of

---

Under review as a conference paper at ICLR 2023
Methods
Finger Spin
Walker Run
Hopper Stand
Hopper Hop
PCL + Sim-reward (Ours)
964.2 224.2
778.4 #14.1
771.2 1142.1
243.2 200.1
TCN Sim-reward
975.2 +:4.1
180.5 +11.4
+1.2
6.8 =0.2
PCL + GAIL-reward
2.5 +2244
18.5 +6.5
1.2 +1.8
0.1 =0.1
TCN GAIL-reward
35.4 19.2
19.8 =4.2
1.5 #0.3
17.3 14.6
Ablation studies on the policy contrastive representation and similarity-baseer imitation reward. We report the average final returns on 4 selected tasks over 3 random seeds and standard deviations are given in error bars.
Table 1:
PCIL and DAC are selected from the same episode during the training process on the same task. We use color to indicate the real environment reward of the agent's 's transitions. The lighter (yellow) color indicates higher reward for agent transition while the darker (green) indicates a lower reward agent transition. The blue color indicates the expert's is transition.
We observe that in the representation space of PCIL, the expert transitions are concentrated cluster. Moreover, the distance between each agent transition to the cluster of expert transition is highly correlated with the real reward of that agent transition. This fact suggests that the contrastive objective of PCIL indeed induces meaningful representation space here. On the contrary, the representation space of DAC is much less structured. The expert transitions are scattered throughout the representation space of DAC. Moreover, identify that in the DAC's representation space, the agent's 's real ward does not correlate well with its distance to the expert transitions. These facts show that our method indeed learns better representation space.
4.4 ABLATION STUDIES
As described in our method, our method has two components: policy contrastive representation for imitation (Section (32) and similarity-base based imitation reward (Section 3.3). this part, we carry out ablation studies to analyze their effects. We first introduce design choices follows.
TCN Representation We replace proposed policy contrastive objective with other contrastive learning methods. For this purpose. we adopt the popular self- supervised representation learning method that leverages temporal information: T'meeCCntrastin Networks (TCN) (Sermanet et al.) 2018). In this case, the positive samples are selected within small window around the anchor sample, while the negative samples are selected from distant time steps in the same rollout trajectory. See Appendix(C for implementation details.
PCIL Representation
Similarity- -hased Reward v.s. GAlL-like Reward We also ablate the similarity- bbased reward in PCIL by replacing the similarity- imitation reward with a GAlL-like reward. Specifically, we train a linear binary classifier on the policy contrastive embedding space to distinguish expert or non-expert data. In this case, the embedding space is still trained by PCIL contrastive loss, and the GAIL reward's 's gradient is detached from the embedding network. We use the same reward predictor as other Gtll--tylle methods (Kostrikov al. 2018), i.e. log (x)) log(1 Dxx).
Analysis By comparing rows 1 and in Table we find that the approach with TCN encoder does not work in three out of four environments. This is because the optimizaiion goal of TCN is not to distinguish between expert and non-expert data. Thus the reward produced by comparing expert and non-expert data in the learned representation space is not necessarily meaningful. Note that the case in row 2 is no longer an adversarial IRL method. We also consider case (row in Table I] where we use TCN and GAlL-like reward predictor, but the performance of this method is poor. Moreover, we observe that in the absence of the similarity- based imitation reward (compare rows and 3 in Table [], method does not work. This is because space characteristics. As a result, we should use a distance- hissence-based metric to compute the reward. In conclusion, the two components of method necessary for achieving good performance.
representation space has metric-
8

---

Under review as a conference paper at ICLR 2023
5
RELATED WORK
5.1 IMITATION LEARNING
Imitation learning is class of algorithms that enables robot acquire behaviors from a demon- stration dataset. There are two classes of imitation learning algorithms: behavioral cloning (BC) and Inverse Reinforcement Learning (IRL) (Ng al. 2000). BC is simple supervised learning algo- rithm that directly fits the expert's s action. However, some work suggests that it has some drawbacks: it suffers from covariate shift problem (Ross al. 201T), and it hard to learn from demonstration dataset without expert actions (Torabi et al. 2018a). 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. In particular, recent branch of IRL is the AIL, which proposes to match agents' state-action distribution with experts via adversarial training. GAIL (Ho & Ermon 2016) proposed maximum entropy occupancy measure matching method which learns a discrim- inator to bypass the need to recover the expert's reward function. Later, several works proposed an improved version of the GAIL methods (Ghasemipour et al., 2020 Blondé al., 2022 Baram al., 2017 Kostrikov et al.) 2018, Fu et al. 2017). AIRL (Fu et al. 2017)) replaced the Shannon-Jensen divergence to measure similarity between state-action pair distributions. [Baram et al. (2017) bridges the GAIL framework to model-based reinforcement learning. DAC (Kostrikov et al., 2018) improved the sample efficiency by leveraging replay buffer without importance sampling and dealing with the absorbing state problem. In contrast to these works, we focus on the representation of AlL's discriminator and reformulating AIL in a contrastive embedding space.
divergence used in GAIL by Kullback-
5.2 REPRESENTATION LEARNING FOR POLICY LEARNING
In this work, we propose a representatioo-llariinn based approach to improve imitation learning. As imitation learning is major class of policy learning algorithms, we review works that self- supervised learning to improve policy learning in this part. Pioneer works (Mirowski et al.) [0016] Jaderberg et al., 2016 [Shelhamer al.) 2016 Lample & Chaplot 2017) explore using auxiliary ob- jectives (e.g.. predict some property of the environment) as a self- self-supprviiion signal. Recent works employ more general self-supervised objectives (Oord al.) 2018). In particular, [Srinivas et al. (2000) are based on contrastive representation learning. [Sermanet et al. (018) learn representation from multiview video using time contrastive learning. 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. (022). In imitation learning, Mandi et al. (2022) proposes to use contrastive learning for one-shot imitation learning in robotics.
Chen et al. 2021) investigate the use of representation for imitation learning. However, their result suggests that self-supervised representation learning only provides small improvement of imitation learning algorithms' performance. Our method differs from all these existing works by proposing to anchor different policies and learn discriminative self-supervised representation for imitation learning.
CONCLUSION
In this paper, we suggested a new approach to improve adversarial imitation learning algorithms: to learn a more meaningful, discriminative representation space for imitation. To this end, we proposed a new algorithm called PCIL. We conducted a theoretical analysis of our method and showed its connection to apprenticeship learning. We also conducted experiments on the DeepMind Control Suite and showed that PCIL could achieve state-of-the-art performance. Moreover, we used an ablation study to highlight its difference from the previous representation learning method. In the future, we will focus on further improving our loss function design. For example, can we anchor on the agent policies different training stages? 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.

---

Under review as a conference paper at ICLR 2023
7
REPRODUCIBILITY
We implement our algorithm according to parameters and details described in Appendix B and Section 41
We will release our code and data.
REFERENCES
Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the twenty-first international conference on Machine learning, pp. 1, 2004.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214-223. PMLR, 2017.
Nir Baram, Oron Anschel, Itai Caspi, and Shie Mannor. End-to-end differentiable adversarial imi- tation learning. International Conference on Machine Learning, pp. 390-399. PMLR, 2017.
Lionel Blondé, Pablo Strasser, and Alexandros Kalousis. Lipschitzness is all you need to tame off-policy generative adversarial imitation learning. Machine Learning, 11 (1):1143311)) 2022.
Richard W Byrne. Animal imitation. Current Biology, 19(3): ((3)):11-111 2009.
Xin Chen, Sam Toyer, Cody Wild, Scott Emmons, Ian Fischer, Kuang-Huei Lee, Neel Alex, Steven Wang. Ping Luo, Stuart Russell, al. An empirical investigation of representation learning for imitation. In Neural Information Processing Systems Datasets and Benchmarks Track (NeurIPS), 2021.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Robert Dadashi, Léonard Hussenot, Matthicu Geist, and Olivier Pietquin. Primal wasserstein imita- tion learning. arXiv preprint arXiv:2006.04678, 2020.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse rein- 2017.
forcement learning. arXiv preprint
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International conference on machine learning, pp. 1587-1596. PMLR, 2018.
Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu. divergence minimization perspective on imitation learning methods. In Conference on Robot Learning, pp. 1259-1277. PMLR, 2020.
Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Im- proved training of wasserstein gans. Advances in neural information processing systems, 30, 2017.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint 18Xiv 10122 2018.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9729- 2020.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
Ahmed Hussein, Mohamed Medhat Gaber, Eyad Blyan, and Chrisina Jayne. Imitation learning: A survey of learning methods. ACM Computing Surveys (CSUR), 50(2):1-35, 2017.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arxiv preprint arXiv: 161 005397, 2016.
10

---

Under review as a conference paper at ICLR 2023
Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. Advances in Neural Information Processing Systems, 33:18661-18673, 2020.
Kee-Eung Kim and Hyun Soo Park. Imitation learning via kernel mean embedding. In Thirty-Secone AAAI Conference on Artificial Intelligence, 2018.
Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tomp- son. Diseriminntto--ctoor--fii Addressing sample inefficiency and reward bias in adversarial imitation learning. arXiv preprint 1rXiv: 2018.
Iyy Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. 13649, 2020.
Guillaume Lample and Devendra Singh Chaplot. Playing fps games with deep reinforcement learn- ing. In Thirty-First AAAI Conference on Artificial Intelligence, 2017.
Timothy Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXivv1509.005 509.02979 2015.
Zhao Mandi, Fangchen Liu, Kimin Lee, and Pieter Abbeel. Towards more generalizable one-shot visual imitation learning. 2022 IEEE International Conference on Robotics and Automation (ICRA), 2022.
Piotr Mirowski, Razvan Pascanu, Fabio Viola, Hubert Soyer, Andrew Ballard, Andrea Banino, Misha Denil, Ross Goroshin, Laurent Sifre, Koray Kavukcuoglu, et al. Learning to navigate in complex environments. arXiv preprint arXiv: 2016.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A univer- sal visual representation for robot manipulation. arXiv preprint arXiv:2203. 12601, 2022.
Andrew Y Ng. Stuart Russell, et al. Algorithms for inverse reinforcement learning, In Icml, vol- ume 1, pp. 2, 2000.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint 0tt1::180.03777 2018
Manu Orsini, Anton Raichuk, Léonard Hussenot, Damien Vincent, Robert Dadashi, Sertan Girgin, Matthieu Geist, Olivier Bachem, Olivier Pietquin, and Marein Andrychowicz. What matters for adversarial imitation learning? Advances in Neural Information Processing Systems, 34:14656- 14668, 2021.
Georgios Papagiannis and Yunpeng Li. Imitation learning with sinkhorn distances. arXiv preprint MxXiv:2008.09167 2020.
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The unsurprising effectiveness pre-trained vision models for control. arXiv preprint arXiv:2003.03580 2022.
Xue Bin Peng, Angjoo Kanazawa, Sam Toyer, Pieter Abbeel, and Sergey Levine. Variational dis- criminator bottleneck: Improving imitation learning, inverse rl, and gans by constraining infor- mation flow. arXiv preprint 1.X.:1811000022. 2018.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. reduction of imitation learning and struc- tured prediction to no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pp. 627-635 JMLR Workshop and Conference Proceedings, 2011.
Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine, and Google Brain. Time-contrastivn networks: Self-supervised learning from video. In 2018 IEEE international conference on robotics and automation (ICRA), pp. 1134-1 141. IEEE, 2018.
11

---

Under review as a conference paper at ICLR 2023
Evan Shelhamer, Parsa Mahmoudich, Max Argus, and Trevor Darrell. Loss is its reward: Self- supervision for reinforcement learning. arXiv preprint arXiv:1612.07307, 2016.
David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference machine learning, pp. 387-395, PMLR, 2014.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representa- tions for reinforcement learning. arXiv preprint arXiv:2004.04136, 2020.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, al. Deepmind control suite. arXiv preprint arXiv: 100111011000000 00090, 2018.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pp. 5026-5033. IEEE, 2012.
Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. arXiv preprint 8a.
Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observation.
arXiv preprint arXiv: Tttv:1807.06155 2018b.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(1 1), 2008.
Ziyu Wang, Josh S Merel, Scott E Reed, Nando de Freitas, Gregory Wayne, and Nicolas Heess. Robust imitation diverse behaviors. Advances in Neural Information Processing Systems, 30, 2017.
Huang Xiao, Michael Herman, Joerg Wagner, Sebastian Ziesche, Jalal Etesami, and Thai Hong Linh. Wasserstein adversarial imitation learning. arXiv preprint arXiv:1906.08) 13, 2019.
Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. arXiv preprint arXiv:2203.06172 2022.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous con-
trol: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.0964 2021.
12

---

Under review as a conference paper at ICLR 2023
A
ENVIRONMENT'
use 10 continuous control tasks from the DeepMind control suite summary for each task is provided in Table
TTssa et al. 2018). The
Task
dim(S)
dim(A)
Finger Spin
6
2
Hopper Stand
14
4
Pendulum Swingup
2
1
Walker Stand
18
6
Walker Walk
18
6
Acrobot Swingup
4
1
Cheetah Run
18
6
Hopper Hop
14
4
Quadruped Run
56
12
Walker Run
18
6
Table 2: A detailed description of each tasks used in our experiments.
Demonstrations For each task, we train expert policies using DrQ-v2 (Yarats et al. 2021 on the actual environment rewards. We run seeds and pick the seed that achieves the highest return. Then use this expert policy to collect 10 demonstrations.
B ALGORITHM DETAILS
B.1 IMPLEMENTATION
RL agent We use DrQ-v2 as the underlying RL algorithm. DrQ-v2 is an off-policy actor-critic algorithm for continuous control. The core of DrQ-v2 is Deep Deterministic Policy Gradient (DDPG) (Lillicrap al., 2015) augmented with n-step returns. The critic is trained using clipped double Q-learning (Fujimoto et al.) 2018) to reduce the overestimation bias in the target value. The deterministic actor is trained using deterministic policy gradients (DPG) (Silver et al.) 2014). We also follows the setting of actor's 's and critic's 's neural network architectures in state-based DrQ-
v2 Yarats et al. (021))
Contrastive encoder The contrastive encoder is implemented as a 4 layer MLP with hidden size [256, 256, 256]. The output dimension is 64. Following the architecture in Yarats et al. (202T), the contrastive encoder, the critic and the actor share the same encoder backbone. This shared encoder trained with the gradient of the critic alone, which is also following the suggestion of Kostrikov ct al. (2200): Yarats et al. (2021). The input of this shared encoder is state s of transition.
Reward predictor Reward of the agent transition is computed according to Equation 2 Note that cosine similarity between expert data points is high due the optimization goal described in Equation ??. Thus, we randomly sample one expert transition from the expert replay buffer to compute agent reward. Empirically, we find that using the mean embedding of the expert data yields similar performance.
Gradient penalty In order to make the algorithm more stable, we use the gradient penalty tech- nique (Gulrajani et al., 2017) widely used in Wasserstein-GANY (Arjovsky et al. (017). make minor adjustments to accommodate strain the gradient norm of the discriminator' 's output with respect to its input. While for PCIL, the contrastive encoder's output needs more step. Specifically, the output embeddings are first used to calculate rewards following Equation 2 then we compute and penalize the gradient norm of the rewards. We 10 the weighting for the gradient penalty.
policy contrastive loss. GAlL-like methods usually con-
13

---

Under review as a conference paper at ICLR 2023
B.2 HYPERPARAMETERS
Table 3 lists the hyperparameters that are used for all baseline methods and our method. Expert data ratio in PCIL means the ratio between expert data and batch size. A ratio of 0.5 means that half of the batch expert data and the other half is the agent data. The contrastive learning usually needs temperature scaling after computing the cos-similarity, before computing the exponential. For simplicity, we ignored it in the main text. In the experiment, we follow prior contrastive learning work He al. 2020) and use a typical value of 0.07 for the temperature.
Methods
Parameter
Value
all methods
Replay buffer size
500k
Agent update frequency
2
Optimizer
Adam
Learning rate
le-4
Critic soft-update rate
0.01
Random seed
1,2,3
RL batch size
128
Discriminator training batch size
256
Hidden dim
256
PCIL
Expert data ratio
0.5
Contrastive temperature
0.07
Table The hyperparameters of baseline methods and our method.
C ABLATION STUDY IMPLEMENTATION DETAILS
C. TCN REPRESENTATION
TCN encoder shares the same network architecture as the PCL encoder. During the training process in TCN encoder, for a random sampled anchor, we use data point adjacent as positive pair and another random sampled data point negative pair. We also set contrastive temperature to 0.07 and batch size to 256, which the same as PCL encoder.
C.2 GAIL-LIKE REWARD
GAIL reward predictor can be seen as a simplified version GAIL discriminator which has only one linear classifier layer. The reward predictor is trained independently to distinguish whether the input is from expert data non-expert data with a binary classification loss.
D ADDITIONAL EXPERIMENTAL RESULTS
Figure shows the performance of PCIL on the other tasks from the DeepMind Control suite. We notice that the performance of some relatively easy tasks has saturated. All the baselines achieve expert performance on Pendulum Swingup. On Walker Stand and Walker Walk, PCIL competitive with DAC, which already demonstrates impressive sample efficiency.
14

---

Under review as a conference paper at ICLR 2023
Acrobot Swingup
Pendulum Swingup
1000
800
200
Episade Beturn
250
800
Eptsade
400
50
1.3
0.3
Environment Steps (teps(x10")
Environment Steps (x200)
Walker Stand
Walker Walk
1000
Episoda Return
800
800
600
400
200
-200
Environment Steps (xtof)
Environment Steps (x310)
Expert
BC
SIL
PWIL
DAC
PCIL (Ours)
Figure Comparisons of algorithms on the other tasks.
15


Reviews:
[{"review_id": "aR8mejNh8m", "reviewer": "Reviewer_dovc", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The main idea in this paper is to do adversarial imitation learning (AIL) with a discriminator trained via the infoNCE loss, rather than the binary cross-entropy loss. The paper argues that this choice of loss function will result in \"smoother\" and more \"semantically meaningful\" reward signals for training the policy. Experimentally, the paper shows that the proposed method improves performance (often significantly) compared with DAC and other imitation learning baselines on the challenging image-based DM Control benchmark.", "summary_of_the_review": "Overall, the empirical results from the paper look strong, and I like the simplicity of the proposed method. I am a bit concerned that the math behind the proposed method doesn't check out. If this math is clarified, and if the writing clarity is improved (see detailed comments above), I will advocate for accepting the paper.", "strength_and_weaknesses": "Strengths\n* The proposed method is simple, and would be easy to apply on top of many different AIL algorithms.\n* The experimental results are quite strong.\n* Presentation. The text does a good job of clearly explaining the proposed method. The figures are generally easy to read.\n\nWeaknesses\n* I am a bit unsure about some of the math (see details below)\n* Some of the writing is imprecise; some has grammatical errors (see below).\n\n**Concern about the math**: My main concern about the math is that it's not entirely clear what the learned reward function (Eq. 2) actually corresponds to. Here's my attempt to figuring this out. Let $p(x, y) $ be the distribution over positive examples, and $p(x)p(y)$ be the corresponding marginal distributions. In our setting, we have $x = (s, a)$ and $p(x, y) = \\rho_E(s, a)\\rho_E(s', a')$. For simplicity, let's assume that we just use 1 negative example. Then the marginal distributions are $p(x) = \\rho_E(s, a)$ and $p(y) = \\frac{1}{2}\\rho_E(s', a') + \\frac{1}{2} \\rho_\\pi(s', a')$. The optimal critic $f(x, y)$ satisfies $e^{f(x, y)} = \\frac{p(x, y)}{p(x)p(y)c(x)}$, where $c(x)$ is an arbitrary positive function [1, 2]. In our setting, this means we have $e^{f(s, a, s', a')} = \\frac{\\rho_E(s, a)\\rho_E(s', a')}{\\rho_E(s, a) \\frac{1}{2}(\\rho_E(s', a') + \\rho_\\pi(s', a'))c(s, a)}$. The $\\rho_E(s, a)$ terms in the numerator and denominator cancel, so we have $e^{f(s, a, s', a')} = \\frac{\\rho_E(s', a')}{\\frac{1}{2}(\\rho_E(s', a') + \\rho_\\pi(s', a'))c(s, a)}$. Written in terms of the representations, we have $\\Phi(s, a)^T \\Phi(s', a') = \\log \\rho_E(s', a') - \\log \\left(\\frac{1}{2}(\\rho_E(s', a') + \\rho_\\pi(s', a')) \\right) - \\log c(s, a)$. The only term that depends on $s, a$ is the $\\log c$ term, which is entirely arbitrary. If this is correct, it suggests that the proposed method maximizes a completely arbitrary reward function, not minimizing any sort of divergence measure. I suspect there's a bug somewhere in my math above, but hopefully it will be helpful in analyzing what the learned reward function actually is.\n\n**Other questions and concerns**:\n* \"fragility of the discriminator leads to ...\" -- How do we know that the problem is with the discriminator, and not with some other part of the learning algorithm? Consider adding a citation/evidence.\n* Fig 1 -- I found this figure a bit misleading. The claim that \"the embedding space is not required to be semantically meaningful\" would seem to suggest that classification doesn't lead to good representation spaces. Not only is this likely untrue (e.g., ImageNet does learn semantically meaningful features), but also the paper later claims that the infoNCE loss (an instance of classification) _does_ yield good representations. The claim that \"our method enforces compactness\" could be easily imposed on existing methods by normalizing their representations, without requiring any changes to the objectives.\n* \"As a result, the model may not distinguish between good and bad states effectively.\" -- I didn't understand this point. It seems like existing AIL methods should be able to distinguish this, if given enough data and enough model capacity. Why do they fail?\n* Sec 3.4 -- I found this section a bit confusing, partially because it was unclear exactly what the aims were. It might be helpful to exactly write out the apprenticeship learning IRL reward (in the general case), and then show that the proposed method corresponds to a special case. An ungenerous reading of this section is that it just says that the proposed method is maximizing a linear reward function, which is trivially true for any AIL method (the reward is a linear function of the last-layer features).\n* Sec 4.2 -- One potential confounder in the experiments is the scale of the learned reward functions. It might be good to study this. E.g., if the scale of the DAC rewards is forced to be similar to the PCIL rewards, does DAC perform better?\n* Sec. 4.3 -- I'm not sure the evidence provided in Fig 4 supports the claims in this section. To me, the DAC representations look like they *better* separate expert vs non-expert data, and that the distance from the expert data is *more* correlated with the rewards (as compared with the PCIL representations).\n\n\n**Minor writing comments**\n* \"to encourage imitation learning\" -- Unclear what this means. Taken literally, it seems like an imitation learning method wouldn't need extra \"encouragement\" to do imitation learning.\n* \"One recent class\" -- 2016 isn't particularly recent.\n* \"error issue (Ross et al\" -- This makes it seem like a 2011 paper proved a result about a 2016 paper, which would be impossible.\n* \"superior performance\" -- Cite.\n* \"without expert actions\" -- Cite.\n* \"etc\" -- Avoiding using \"etc\" in technical writing.\n* \"semantically meaningful signals\" -- Unclear what this means\n* \"Instead of using a discriminator ... we propose to train a discriminator representation\" -- This seems like a type error, comparing a discriminator to a representation.\n* \"the binary classifier's representation space does not necessarily satisfy our objective\" -- I didn't understand this sentence.\n* \"Suite Tassa\" -- Use \\citep instead of \\citet here.\n* \"very non-robust\" -- Robust w.r.t. what criterion? Add a citation. Stylistically, I'd remove the \"very,\" too.\n* \"Researchers have found\" -- Cite.\n* \"people have found\" --> \"prior work~\\citep{...} has found\"\n* Fig 2 -- This figure is great for explaining the method!\n* \"definition of similarity in the previous works\" -- What is this definition? It'd be helpful to include the definition here.\n* \"Though combining ...\" -- I didn't understand the strawman construction here.\n* \"pair-wise distance constraint\" -- What is this constraint?\n* \"$\\Phi(x_0)^T\\Phi(x) > t$\" -- How was this derived?\n* \"the binary classification induced latent space might not satisfy our constraint.\" -- I didn't understand this point.\n* \"Theoretical Analysis\" -- I'd recommend renaming this \"Connections with Feature Matching\" because the section doesn't include any formal theorems/lemmas.\n* Eq. 3 -- Perhaps cite [3] or similar prior work that uses this objective\n* \"we will find that the\" -- How is this proven?\n* \"three random seeds\" -- I'd highly recommend running at least 2 additional seeds.\n* Fig 3 -- increase the size of the xlabel and ylabel fonts.\n* Fig 4 -- Add a takeaway to the caption, so the reader knows what they are supposed to learn from this plot.\n* \"distant time steps in the same trajectory\" -- I'd recommend highlighting this, using it to motivate why sampling positive examples from different trajectories is better than sampling positive examples from the same trajectory.\n* Table 1 -- In column 1, both the first and second rows should be highlighted. The caption should specify the number of random seeds and indicate what the error bars correspond to. Table captions should appear before the table.\n* \"representation space has metric-space characteristics\" -- Every representation space forms a metric space.\n* I'd highly recommend running a grammar checker on the paper (e.g., copy+paste into a Google Doc and run the grammar checker)\n* Conclusion -- This is well written.\n* \"Russel, et al.\" -- Why are the rest of the authors not listed?\n* Table 2 -- Are these image-based tasks or state-based tasks? Captions should appear above the table.\n* Gradient penalty -- I'd recommend showing an equation or pseudocode to help clarify this section.\n* \"we ignored it in the main text\" -- Please add this detail back to the main text.\n\n[1] https://arxiv.org/pdf/1809.01812.pdf\n\n[2] http://proceedings.mlr.press/v97/poole19a/poole19a.pdf\n\n[3] https://arxiv.org/abs/2106.04156", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "XACbcjhDvx", "reviewer": "Reviewer_1ZsP", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a new adversarial imitation learning method, where a discriminator is trained to learn a representation such that the distance between expert trajectories are shorter than the distance between agent-expert trajectories. This representation is then used to calculate an imitation reward for the agent to encourage it to generate trajectories that are more expert-like. The paper shows that this representation amounts to feature matching between the expert and the agent. The results on MuJoCo tasks show that the proposed method significantly outperforms the state-of-the-art adversarial imitation learning methods. ", "summary_of_the_review": "This paper proposes an interesting and novel idea with a nice theoretical backup. The results look solid with a minor caveat that I pointed out above. Thus, I believe that this paper is interesting enough to be presented at ICLR. ", "strength_and_weaknesses": "[Strength]\n\n* The idea of using contrastive learning and similarity-based reward is novel and interesting. \n* The theoretical result is also novel and nicely justifies the proposed method.\n* The performance on various MuJoCo tasks looks strong.\n* The analysis of the learned representation makes the method more convincing.\n* The paper is well-written. \n\n[Weakness]\n\n* Lack of robustness analysis. In the introduction, the paper motivates the problem by pointing out how fragile the discriminator is in existing adversarial imitation learning algorithms. However, there is no experiment demonstrating the robustness of the proposed method. It would be even more convincing if the paper showed that the proposed method is much more robust to the choice of the hyperparameters and architectures of the proposed discriminator. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "STAKpAe85c", "reviewer": "Reviewer_igbW", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this paper, to address known issues of GAIL such as instabilities and non-robust representations, the authors propose an approach that additionally learns a representation in a contrastive way, by rewarding expert states to be nearby in the latent space, and expert-nonexpert states to be far apart. They contain some theoretical considerations as well as experiments.", "summary_of_the_review": "The paper seeks to address an important problem, and contains some interesting ideas, but is too premature in terms of writing and analysis of the method.", "strength_and_weaknesses": "**Strengths:**\n\nI fully buy the motivation: GAIL is difficult to stabilize and representations may be unrobust. Also they may semantically not be meaningful.\n\nTherefore, the idea to have an additional loss to improve representations makes sense, and contrastive learning may be a sensible way.\n\n**Weaknesses**:\n\nI think the paper is too premature in terms of both, writing and method.\n\nSec. 3.4 claims to deliver a theoretical analysis, but it stays pretty vague. Eq 3 seems central, but the derivation to get there from Eq. 2 is just informally described instead of adding at least 1-2 formal lines, especially since there are some many different $x$ version etc. where one can get lost in the expression.\n\nOverall, intuitively I cannot follow the method's idea. \n* I feel such a representation is not robust to distributional changes etc. I think it often happens that in one instance (say, for one initialization), a state-action pair is good (taken by the expert), while in another instance it is bad (taken by non-expert). I.e., whether a state is good or bad can heavily depend on initial state, random fluctuations, etc. E.g., from one starting position in a maze, the expert may never want to visit a certain state, while from another one, the state would lie on the shortest path. It's hard for me to make this intuitive counterargument specific, and maybe GAIL suffers from similar problems, but it does not seem to be a robust approach.\n* It seems it does something similar as GAIL, but in a less principled way. But then why do we need to add it in the first place?\n* Also I'm confused by relating the tuning of the representation, to tuning the policy. I think these are different things and I have a hard time seeing that Sec. 3.4 justifies this work from IRL/AL principles.\n\nSome (minor) points re. writing:\n* The overall method, i.e., GAIL + the new representation mechanism, is described too little; I feel it is just Fig. 2.\n* \"leanring\" p3\n* \"markov\" p3\n* x is only informally introduced as state-action-pair. I strongly recommend to introduce all important notation in a formal, central way.\n* what does \"anchoring on policy\" mean (p4)?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject