Non-Adversarial Inverse Reinforcement Learning via Successor Feature Matching

We thank the reviewer for their careful review and their thought-provoking questions. We respond to the questions below:

W1: "non-adversarial” property and base feature trained with adversarial learning by setting $\mathcal{L}_{feat}$ = Eq.(3)
In Eq 3, we show how the task of IRL which uses adversarial learning can be framed as minimizing the $L_2$ gap between expected features, leading to a framework that can use any mechanism to learn base features.

Regarding the non-adversarial property-- you are of course correct that our framework does not prohibit $\mathcal{L}_{feat}$ from being trained adversarially. However, our argument is a sort of converse-- the framework enables us to avoid adversarial feature learning, which leads to simpler algorithms, achieving greater performance than the adversarial baselines. That said, your hypothesis that adversarially-trained features can lead to even better performance is interesting.

Suppose one were to employ an adversarially-trained $\mathcal{L}_{feat}$-- as you suggest, one can do this via an adversarial IRL algorithm similar to MM. There are two possible outcomes that we might observe by training SFM with those features:

1. The performance is worse than our non-adversarial SFM results. In this case, we conclude that non-adversarial features should be preferred.
2. The performance is better than what we achieved in SFM. In this case, since SFM already generally outperforms baselines, we can use adversarial features to get better IRL agents, which again is a positive outcome.

So, in any case, we argue that SFM is a valuable tool for IRL. Continuing on this thread, we conduct experiments with adversarially (Adv) trained features. We observe in Fig 5 of the updated paper that adversarial features do well, but FDM still outperforms other base feature functions. (Note that we moved Fig 5 in the submitted paper to the Appendix due to space constraints, Fig 6 is now Fig 7 and Fig 5 presents the RLiable plot of base feature functions in the revision).

W2: Compared to Figs. 4 and 5, single-demonstration task performance in Fig 3 is only sometimes good
We have provided agents with single expert demonstrations in all our experiments. The major change between Fig 3 and Fig 5 is the architecture of the policy optimizer. In Figure 3, we compared the performance of SFM with state-only MM and GAIfO with a framework based on TD7 algorithm. Our implementation of baselines with TD7 framework significantly improved the performance when compared to using a TD3 optimizer (Fig 4 & 5) which explains the large gaps in performance for the adversarial baselines between Fig 3 and Fig 5. However, the performance of SFM is similar with both frameworks suggesting that SFM is more robust to policy optimizers.

In this work, we focus on the scenario of state-only demonstrations and use state-only base features. Without expert actions, it is challenging to estimate the state-action base features required to estimate the expert's SF. However, we would like to highlight that it is not a limitation and SFM can be easily extended to have state-action features when using demonstrations with expert action labels (as discussed between lines 282-288 in the paper).

W4: no supporting tables in the appendix
We have added tables in Appendix D showing the numerical results of the Figures.

Q1: How applying l2-loss in Eq.(4) can be understood for the policy $\pi_{\mu}$?
For the policy $\pi_{\mu}$, the l2-loss defined in objective will optimize the agent to match the expected features with the expert. As we observe in Eq 3, this would provide bounds of the performance gap between the expert and the agent, i.e. if the l2 gap is small, then the imitation gap between the expert and the learning agent should be small. To compute the gradients of the policy for this l2-based objective, we propose leveraging Proposition 1 & 2 that provides the gradients for any states sampled from the replay buffer.

Dear Reviewer ZRQG,

Thank you for your time and feedback, which helped improve the paper. During the rebuttal period, we added comparisons with other base feature functions (Fig 5): 1) Autoencoder over only states and not next states, actions, or goals (Q2); and 2) Adversarial where the base features are trained with Eq 3 (W1).

We believe our revision and responses have addressed most of your concerns. As the author-reviewer discussion period ends in 2 days, we would be happy to discuss and address any further questions during this time.

Best Regards,
Authors

Deal authors,

I appreciate your detailed explanation and additional experiments. In particular, I am impressed that the authors took W1 and Q2 into serious consideration in the revision. Currently, I have follow-up questions.

• In revised Fig 5, it seems the new Adv version performs in second place. Does this fit well with your original claims, considering the adv version might run slightly more efficiently without acquiring next-states? I still have a minor concern that the "Non-adversarial IRL" title might be misleading.
• I can agree that this method might be suitable for certain RL envs with physics simulations. What are your current thoughts on whether SFM could work well for other environments, such as high-dim vision-based and real-world linguistic domains?

Dear Reviewer ZRQG,

We thank you for the response and address the questions below:

Does this fit well with your original claims
While adversarial features performed reasonably well, we would like to emphasize that this does not go against our claims. In particular, non-adversarial features achieve equally good (and often better) performance, and as we discuss next, they do so without sacrificing efficiency.

Adv version might run slightly more efficiently without acquiring next-states?
Firstly, training adversarial features is quite challenging in general. For our experiments, we had to tune multiple parameters like the gradient penalty coefficient, learning rate decay, and updating the features quite less frequently as compared to the policy. Secondly, for training the adversarial features, the agent collects on-policy samples by generating trajectories with the current policy (inspired from adversarial reward learning approaches like GAIL [1] and HyPE [2]). The states visited in this on-policy dataset are used for updating the adversarial features (as described in equation 13). However, for training non-adversarial features, the agent could just learn it from transitions in the replay buffer without additional samples which makes it more efficient at learning. We tried learning adversarial features by sampling states from the replay buffer, but it was not performing well in our experiments. Therefore, non-adversarial feature learning is more efficient.

Whether SFM could work well for other environments?
This is of course an interesting question. Classically, IRL methods have struggled on such domains, especially in the state-only setting. However, by reducing the IRL problem to RL, we hope that we can achieve strong performance in such domains by leveraging recent representation learning techniques in RL. Given the promising results of RL in such domains, we believe this is an intriguing avenue for future work.

We would be happy to answer any further questions and hope you consider increasing the score.

Best Regards,
Authors

References

[1] Ho et al., Generative Adversarial Imitation Learning, NeurIPS 2016.
[2] Ren et al., Hybrid Inverse Reinforcement Learning, ICML 2024.

We thank the reviewer for their thoughtful review and their enthusiasm in our work. With regard to both of your questions, these are valid points, and we have elaborated on extending SFM to stochastic policies as well as exploration in our general response.

We hope we addressed most of your concerns and would be happy to answer any further questions.

We thank the reviewer for their thorough assessment of our work and their insightful comments. We have provided an extension to stochastic policies and discussed about exploration in the general response. We respond to the remaining questions below:

W3: State only base features
In this work, we train SFM with state-only base features to deal with state-only demonstrations; if we have action labels in expert demonstrations, SFM can easily support action-dependent base features. Indeed, the only area where state-only base features are required is where we estimate the expert's SFs: if we have expert actions, we can estimate SFs for action-dependent base features. This is a fairly natural constraint---without access to expert actions, it is unlikely that learned action-conditioned base features will lead to improved imitation.

Q3: Comparison with offline IRL approaches:
While SFM is an online IRL algorithm, it makes it hard to compare with offline methods. Typically, interactive methods are proven theoretically to have better performance as they deal with covariate shift by allowing agents to recover from its own mistakes. Moreover, offline methods assume presence of optimal or sub-optimal dataset whereas this work provides agents access only to a single offline trajectory. In this regard, we skipped comparison with offline IRL methods in this work. We agree an interesting direction of future work is an offline version of SFM by borrowing tricks from the offline RL literature. Lastly, we compared with Behavior Cloning (BC) as an offline imitation learning baseline where the agent was not able to learn with a single trajectory with the expert action information in demonstrations.

Q4: Computational requirements vs adversarial methods:
In terms of computational complexity, SFM and adversarial methods are fairly similar---per training iteration, SFM updates its SFs and its policy, while adversarial methods would update a discriminator/reward function and its policy. However, since the SF and policy objectives in SFM do not form a min-max game, the two gradient updates can occur simultaneously---we do not need to consider complicated training dynamics like in adversarial methods, which induce more hyperparameters that can be tricky to tune.

Q5: Sensitivity to SF architecture:
In all of our experiments, we simply used the standard architecture used for training the TD3 critic (with a modified head to predict multivariate features). We did not rigorously evaluate the sensitivity in this regard, however our experiments show that the most natural architecture choice tends to work very well.

We would be happy to answer any further questions. We hope we addressed most of your concerns and you consider increasing the score.

Dear Reviewer w2Zs,

Thank you for your time and feedback, which helped improve the paper. During the rebuttal period, we introduced a variant of SFM incorporating stochastic policies (addressing Q2 and W2) and expanded our discussion on extending SFM to address exploration challenges in reinforcement learning (as part of our response to Q1 and W1).

We believe our revision and responses have addressed most of your concerns. As the author-reviewer discussion period ends soon, we would be happy to address any further questions during this time and hope you consider increasing the score.

Best Regards,
Authors

We thank the reviewer for their thoughtul review and their enthusiasm about our work. We respond to the questions below:

W1: Performance of IQ-Learn
As discussed in Jena et al. [3], some imitation learning methods can hack the rewards using the terminal states and survival instincts. However, we conduct experiments on DMControl suite where each task is infinite horizon tasks with fixed episode length-- a scenario where agents cannot hack through survival biases. We believe this might be true for IQ-Learn due to which the method does not perform well on infinite horizon tasks. Note that recent works [4] have demonstrated similar performance of IQ-Learn with the DMControl tasks.

W2: The choice of baselines could be expanded
We appreciate your suggestions about further baselines to include. We do note, that we actually significantly strengthened the baselines in our work by incorporating the latest techniques in deep RL / IRL, massively improving performance over their original versions. We have added OPOLO [2] as another baseline in our revision, and while it performed competitively with other baselines, it struggled on the quadruped tasks---ultimately, SFM still outperforms all baselines. Our hypothesis is that OPOLO heavily depends on a good Inverse Dynamics Module (IDM) which is hard to learn in the quadruped domain (something we saw in our experiments in Fig 7 of updated draft). I2L [1] is an interesting method that deals with expert demonstrations collected in a different MDP than the learner's MDP to tackle challenges with changing dynamics. However, this is not the setting we tackle in our paper; indeed, none of the baseline methods are designed to perform well in this setting, and it is not the setting of our experiments. That said, we believe similar tricks from I2L can be leveraged to learn base features that are agnostic to transition dynamics to extend SFM to this setting, and for now we leave it as a very compelling direction for future work.

Q: What are the hyperparameter search space for SFM and baselines respectively?
For adversarial baselines, we observed that the choice of policy optimizer (between TD7 and TD3) significantly affected the performance. Other hyperparameters for tuning the discriminator (reward model) include the coefficient of gradient penalty loss, learning rate decay and frequency of updates during training. For SFM, the SF network uses the standard architecture used for training the TD3 critic (with a modified head to predict multivariate features). For base feature functions, we directly use the architecture and implementation provided in Hilp~[5]. The hyperparameters for SFM include the policy optimizer, base feature function and the dimension of base features. We found that SFM works ``out of the box'' with the most natural choice of hyparparameters (namely, architecture of the base feature function and policy optimizer). As such, we did not extensively assess the sensitivity to these parameters, because we found empirically that they do not need to be tuned. Finally, we stress that SFM is fairly robust to hyperparameters and policy optimizers across environments.

We believe our revision has addressed all of your comments, and if so, we would be grateful if you would consider increasing your score. Otherwise, we are eager to discuss if you have any further questions or comments.

References

[1] Gangwani et al., State-only imitation with transition dynamics mismatch, ICLR 2020.
[2] Zhu et al., 'Off-policy imitation learning from observations.', NeurIPS 2020.
[3] Jena et al., Addressing reward bias in adversarial imitation learning with neutral reward functions, arXiv:2009.09467.
[4] Pirotta et al., Fast Imitation via Behavior Foundation Models, ICLR 2023.
[5] Park et al., Foundation Policies with Hilbert Representations, ICML 2024