Latent Wasserstein Adversarial Imitation Learning

Thanks for your constructive advice. We answer questions next:

Q1. Notation.

Thank you for carefully reading our paper. We have fixed the notation in our revised paper.

1. clip threshold (description weakness 1). The correct understanding (Gaussian noise $\epsilon$ with variance $\sigma^2>0$ and clip threshold $c_0>0$) is added to the target action $a'$. We have modified the submission to clarify this.

2. There is no $c(s,s’)$ in Eq. 4 (description weakness 2). The $c(s,s’)$ is hidden in the constraint $\|\|f\|\|_L\leq 1$, as explained in our explanation to Eq. (1) (which is essentially $\frac{|f(s)-f(s')|}{c(s, s')}\leq 1$). To clarify this, we have updated Sec. 3.2.

3. The input for $f$ (description weakness 3; questions 3). All $f(\cdot, \cdot)$ in Sec. 3.3 and later use $f(\phi(s),\phi(s')-\phi(s))$ as input. Thus there is no inconsistency in our ablation. The reward function is using latent space. We have updated Sec. 3.3 to reflect this.

4. The summation index t in Eq. 5 and definition of $gamma$ (description weakness 4). In RL, $\gamma=1$ is usually used in finite-horizon MDPs with total reward being the sum over each timestep, while $\gamma<1$ is usually used in infinite-horizon MDPs with reward summing over infinite timesteps in the future using a decaying weight. We adopt the infinite horizon MDP framework. We have updated $\gamma$ to $[0, 1)$ in our preliminary section.

5. Explanation of $\alpha$ (description weakness 4). $\alpha$ is an important parameter in IQL. As IQL is an algorithm that penalizes over-optimistic evaluation for out-of-distribution data in offline RL, it uses an asymmetric critic objective where overestimation is punished harder than underestimation. Thus, we have $\alpha>0.5$, such that $| \alpha - \mathbb{I}(A < 0) |$ has greater values for advantage $A>0$.

6. What is Eq. 6 Minimizing (description weakness 4)? Thanks for pointing this out. The minimization is over the value function $V_\theta$ and we have updated Eq. (6) in our paper. The detailed objective is listed in Appendix A.

7. KR duality instead of Rubinstein duality. (questions 2) Thanks for pointing this out. We have updated Rubinstein duality in our paper to KR duality.

Q2. Why distinctive reward signals are better in Fig. 3b, as most rewards in latent space are close to 0 though, and why less tendency for riskier explorations is related to ICVF behavior? (experiment weakness 1 and 3)

The plotted reward corresponds to $\sigma(-f(\phi(s),\phi(s’)-\phi(s)))$. It is shown in Fig. 3b that there exist several trails of trajectories around the goal, which is an evidence that the latent space is more aware of the environmental dynamics recovered from the trajectories. While the advantage might not be that obvious for a simple 2D environment (note our ablation in Fig. 3c still shows that our method outperforms TD3 with ground truth reward), such awareness of trajectory dynamics can be much more informative for agents as illustrated in our updated Fig. 2. Without the embeddings, the chance that exploration of online inverse RL falls into low-reward areas increases, thus decreasing the value estimation of adjacent states which in turn harms exploration.

Q3. The reward in LWAIL is naturally denser than the original reward function in Maze2D. The paper should compare the performance of TD3 with sigmoid transformation of the spare reward (experiment weakness 2).

We adopt the sigmoid function to regulate the output of our neural networks for better stability. Use of this was shown in WDAIL [1]. Importantly, the sigmoid function is not related to a density, but rather a technique to stabilize training.

To further address the reviewer’s question, we ablated TD3 with a sigmoid function on the ground truth reward. The result is illustrated below. The result shows that a naive sigmoid mapping on the reward function does not help TD3.

Q4. the input of network and ablation study. (experimental weakness 4)

We use $\phi(s’) - \phi(s)$ rather than $\phi(s)$ to account for the state-transition effect of the action. As $s$ can reach $s’$ in the dataset within one step, the dynamic-aware embeddings $\phi(s)$ and $\phi(s’)$ are often close; by using $\phi(s’) - \phi(s)$, the network can focus on learning the difference between embeddings. Preliminary experiments showed that this improves training stability. We already included the ablation study of $f(s,s’-s)$ without ICVF embedding in Sec. 4.3 of our paper, illustrated in Fig. 6. We found that without ICVF, the method works generally worse. Also, for ablations to other contrastive learning-based embedding methods, see our updated Appendix C.2.

Q5. What does it mean by “policy’s understanding of state transitions” (questions 1)?

It means the policy is easier to learn from a state-transition-aware embedding space. If we view the embedding function as a part of the policy, then the policy is more “aware” of the state transitions dynamics when using the embedding.

Q6. How is the representation network $\phi$ being used in policy updates? (questions 4)

Training $\phi$ is a part of the offline training process. $\phi$ is frozen during online training. This is illustrated in Fig. 1.

Q7. I found it very hard to grasp the main idea in Section Overcoming Metric Limitations, both intuition-wise and methodology-wise. It would be great if this part can be improved (questions 5).

To improve understanding, we have updated the notation and explained the design of the rewards. We have also modified the illustrations and explanations of the latent space in our updated Fig. 2. Also, we included the details for ICVF in our Appendix A due to page limits.

Here we present the intuition of this section again briefly: The denominator of the 1-Lipschitz constraint forces existing Kantorovich-Rubinstein (KR) duality-based methods to use Euclidean distance in practice due to the gradient regularizer. However, the Euclidean distance is typically not environment dynamics-aware (see illustration in Fig. 1a). Hence, intuitively, we expected an embedding space in which the Euclidean distance is dynamics-aware to yield better performance. We find the ICVF embedding space to satisfy this property because the ICVF training process is inherently dynamics-aware. Indeed, our developed method to incorporate ICVF embeddings shows better performance.

References

[1] Z. Zhang et al. Wasserstein Distance Guided Adversarial Imitation Learning with Reward Shape Exploration. In DDCLS, 2020.

Dear reviewer Eovk,

Thanks again for your constructive review to improve our paper. As the author-reviewer discussion period is close to its end, we kindly invite you to further consider our response, which we believe has addressed all your concerns raised in the rebuttal. Thank you very much!

Dear reviewer Eovk,

Thanks again for your effort in providing constructive feedback on improving our paper! As the discussion period will come to an end in less than 24 hours, we would like to know if you have any remaining concern, so that we can address them. We are sincerely looking forward to hearing from you, and are always happy to further discuss with you.

Thanks for your constructive advice. We answer questions next:

Q1. How do ICVF embeddings contribute to accurate imitation from limited expert data?

We have discussed this issue in Sec. 3.2, line 203-215. We have modified the paper to highlight this more clearly. To summarize, ICVF embeddings contribute to accurate imitation by creating an embedded state space where the Euclidean distance is more aligned with dynamic differences between states, thus fixing the core problem of prior Wasserstein imitation learning with Kantorovich-Rubinstein (KR) duality: reliance on Euclidean distance between states.

To empirically validate the effectiveness of ICVF embeddings, we follow the challenging setting of prior work (e.g. PWIL) and test our method with limited expert trajectories. Results are shown in the table below. Our method can match the state distribution very well even when the expert data is scarce. Additionally, we have tested our method with multiple expert trajectories, and it has consistently shown excellent performance regardless of the number of expert trajectories.

Q2. No comparison between ICVF-embedding and other embeddings.

While there are embedding methods for RL/IL, many of them are not applicable to our case. For example, most empirical state embedding methods are for visual environments [1, 2] or for cross-domain dynamics matching [3, 4]. Among theoretical state embedding methods, low-rank MDPs are not applicable to the MuJoCo environment, and bi-simulation requires a reward signal which is unavailable in IL.

Nonetheless, we identify two contrastive learning-based baselines most suitable for our scenario: CURL [5] and PW-DICE [6]. Both methods use InfoNCE [7] as their contrastive loss for better state embeddings. Their difference: 1) CURL updates embeddings with an auxiliary loss during online training, while PW-DICE updates embeddings before all other training; 2) CURL compares the current state with different noises added as positive contrast examples, while PW-DICE uses the next states as positive contrast samples. The result is shown below in normalized reward (higher is better):

The result shows that 1) state embeddings generally aid learning; and 2) our proposed method works best.

Q3. What is the difference between WDAIL and LWAIL without ICVF embedding?

1. Note, the ICVF embedding is a core part of our LWAIL method, differentiating our work from prior works using Wasserstein imitation learning with KR duality. By adopting an embedding space where the Euclidean distance is aligned with the dynamic differences between the states, we address the core problem of prior works, i.e., the use of an inadequate distance.

2. This being said, without ICVF embedding, our LWAIL method still differs from WDAIL:

a) WDAIL needs expert actions, but LWAIL can learn from action free expert demonstrations and still outperforms WDAIL as shown in Tab. 1 of the paper. Note, learning from action free expert demonstrations addresses a much harder task due to uncertain environment dynamics [8], and is widely applicable e.g., for learning from video demonstrations (where the expert action is unavailable), or learning from a different embodiment (where the expert action is not applicable).

b) Our LWAIL uses TD3 as the downstream RL algorithm rather than PPO with entropy regularizer as adopted by WDAIL. We ablate the downstream RL algorithm in our LWAIL and show that TD3 outperforms PPO. This is intuitive, as off-policy algorithms are more robust to the change of rewards in the adversarial process.

c) Our method without ICVF embedding differs from WDAIL in technical details. For instance, we propose a normalization in the Wasserstein discriminator’s input $f(\phi(s), \phi(s’)-\phi(s))$ which stabilizes the algorithm.

Q4. Is the observation that ICVF-trained embedding provides a much more dynamic-aware metric than Euclidean distance consistent across environments?

Yes. To verify this, we have updated Fig. 2, to include all MuJoCo environments evaluated in our experiment section and by plotting more steps of a trajectory in each environment. We also plot ground truth reward obtained on each state with different colors (brighter is higher) for a better understanding of the dynamic-aware property. We observe that the high-reward areas are clustered in the embedded latent space, while in the original state space they are scattered. The result shows that we provide a more informative embedding for the agent.

References

[1] L. Meng et al. Unsupervised State Representation Learning in Partially Observable Atari Games. In CAIP, 2023.

[2] P. Sermanet et al. Time-Contrastive Networks: Self-Supervised Learning from Video. ArXiv, 2017.

[3] Y. Duan et al. One-Shot Imitation Learning. In NIPS, 2017.

[4] T. Franzmeyer et al. Learn what matters: cross-domain imitation learning with task-relevant embeddings. In NeurIPS, 2022.

[5] A. Srinivas et al. CURL: Contrastive Unsupervised Representations for Reinforcement Learning. In ICML, 2020.

[6] K. Yan et al. Offline Imitation from Observation via Primal Wasserstein State Occupancy Matching. In ICML, 2024.

[7] A. Oord et al. Representation Learning with Contrastive Predictive Coding. ArXiv, 2018.

[8] Z. Zhu et al. Off-Policy Imitation Learning from Observations. In NeurIPS, 2020.

Dear reviewer ySZG,

Thanks again for your effort in providing constructive feedback on improving our paper! As the discussion period will come to an end in less than 24 hours, we would like to know if you have any remaining concern, so that we can address them. We are sincerely looking forward to hearing from you, and are always happy to further discuss with you.

Thanks for your constructive advice. We answer questions next:

Q1. ICVF and WGAIL are existing methods, and state-only imitation learning is not a new topic.

1. While we agree that ICVF and WGAIL are existing methods, this does not imply a lack of novelty for our method. Our developed Wasserstein imitation learning method with ICVF is based on an important, novel insight, overlooked in prior works which employ the Kantorovich-Rubinstein (KR) dual. The insight is that ICVF embeddings create an embedded state space where the Euclidean distance is more aligned with dynamic differences between states, thus fixing the core problem of prior Wasserstein imitation learning works with KR duality: reliance on the Euclidean distance between states. To our best knowledge, nobody has tried to address this problem.

2. It is true that state-only imitation learning is not a new topic, but this again does not imply a lack of novelty. In contrast, prior work in this field shows that we are focusing on a popular and important setting in imitation learning.

3. We have compared our method to a variety of baselines, and showed that our method outperforms prior works. The results show that our proposed solution successfully benefits from the distance metric which others have not identified. We think this is a valuable contribution for our community.

Q2. The more robust policy has not been validated with specific results.

To validate robustness of our policy, we provide results with subsampled expert trajectories, a widely-adopted scenario in many prior works such as PWIL and IQlearn. Only a small portion of the complete expert trajectories are present. Our subsample ratio is 10, i.e., we take 1 expert state pair out of adjacent 10 pairs. We use 10 expert trajectories as demonstrations.

We observe that our method can deal with highly incomplete trajectories, underlining its robustness. We also conduct experiments showing the robustness of our method when the dynamics between expert demonstrations and the actual environment don’t match. Results are presented in Q9 of our response.

Q3. What’s the motivation for performing imitation learning in the latent space?

Our method is an imitation learning method which uses state distribution matching to bring the agent and the expert close. However, as shown in Fig. 2, states close according to the Euclidean metric are not always close in the actual state space. This leads to suboptimal matching results. Worse still, the Euclidean metric is an inherent part of all KR duality-based methods (explained in the preliminary section), which is taken for granted in prior works.

We propose to address this issue by introducing a latent space, where states cluster in a dynamic-aware manner. This makes the notion of “being close” more informative. We have provided an ablation study with and without ICVF embedding in our paper (Fig. 6 in Sec. 4.3). We reproduce the results in the tables below (the metric is normalized reward; higher is better). See our updated Appendix C.2 for ablations to other contrastive learning-based state embedding methods, where our method also improves.

Q4. Code is missing. Thanks for pointing this out. We will provide our code upon acceptance. We have updated the reproducibility statement accordingly.

Q5. discussion on the limitations Thanks for pointing this out. We have updated the submission to include this discussion. We post the limitations here for convenience:

Limitations. Similar to other prior adversarial imitation learning methods such as WDAIL, our pipeline requires an iterative update of the actor-critic agent and the discriminator during online training. The update frequency needs to be balanced during training. Also, testing our method on more complicated environments, such as image-based ones, is an interesting avenue for future research.

Q6. More complex visual environments.

We select vector-based environments for our evaluation as the quality of the embeddings in visual environments largely depends on visual features extracted from frames (which leads to success for methods like DrQ [1] and RAD [2]) instead of a dynamic-aware property. We think this entangles contributions from dynamics-aware embeddings and visual features, making it harder to clearly assess the contribution.

It is also worth noting that visual environments are not a standard in prior works. For example, BCO, WDAIL, OPOLO and DACfO all do not assess performance on visual environments.

This being said, it is valuable future work to scale to visual environments.

Q7. Comparison of leveraging ICVF with existing GAIL-based methods.

It is worth noting that the classic GAIL method aims to minimize the $f$-divergence between the generated state occupancy and the expert occupancy. This ignores the underlying distance between the states. Therefore, these methods do not fit our motivation, which is to employ an embedding where Euclidean distance aligns with dynamic-aware differences.

To show that our proposed solution outperforms existing GAIL-based methods with ICVF embedding, both Wasserstein-based (IQlearn, WDAIL) and $f$-divergence based, we have conducted an experiment and listed the result below (in average reward; higher is better). We found that 1) our method outperforms existing methods with ICVF embedding, and 2) ICVF does not necessarily improve the performance of existing methods, probably due to other components of our method (e.g., downstream RL algorithm as TD3 is more robust in an adversarial framework; see our updated Appendix C.4 for ablations).

Q8. Have the authors explored applying LWAIL in environments with multi-modal state distributions?

Yes, this is shown in our updated Fig. 2. For instance, in the halfcheetah and ant environment we observe that there exists clusters of states in the embedding space, and the agent transits quickly between the clusters. This implies a multimodal state distribution.

Q9. How does LWAIL handle situations where the ICVF-pretrained latent space does not align well with the environment’s true dynamics?

It is worth noting that the very motivation of LWAIL is to find a latent space which aligns well with the environment’s true dynamics. Despite this, we agree that there might be cases where the latent space employed in LWAIL does not align with the true dynamics due to inaccurate data, e.g., mismatched dynamics between expert demonstrations and the actual environment. To test such cases, we use the halfcheetah mismatched experts scenario analyzed in SMODICE [3]: for expert demonstration, the torso of the cheetah agent is halved in length, thus causing inaccurate alignment. We compared our methods with the results reported in the SMODICE paper. Below are the final average normalized rewards (higher is better):

The result shows that our method is robust to mismatched dynamics.

References

[1] D. Yarats et al. Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels. In ICLR, 2021.

[2] M. Laskin et al. Reinforcement Learning with Augmented Data. In NeurIPS, 2020.

[3] Y. J. Ma et al. Smodice: Versatile Offline Imitation Learning via State Occupancy Matching. In ICML, 2022.

[4] Liu. F et al. State Alignment-based Imitation Learning. In ICLR, 2020.

[5] K. Zolna et al. Offline Learning from Demonstrations and Unlabeled Experience. In Offline RL Workshop @ NeurIPS, 2020.

Dear reviewer hrKU,

Thanks again for your constructive review to improve our paper. As the author-reviewer discussion period is close to its end, we kindly invite you to further consider our response, which we believe has addressed all your concerns raised in the rebuttal. Thank you very much!

Dear reviewer hrKU,

Thanks again for your effort in providing constructive feedback on improving our paper! As the discussion period will come to an end in less than 24 hours, we would like to know if you have any remaining concern, so that we can address them. We are sincerely looking forward to hearing from you, and are always happy to further discuss with you.