Noise-conditioned Energy-based Annealed Rewards (NEAR): A Generative Framework for Imitation Learning from Observation

Dear reviewer, thank you for your feedback. We have revised the paper. Please find clarifications below. We hope these sufficiently establish the contribution of the paper and its place within the IL from the observation field. Note: some lines in the text are coloured to indicate suggested modifications to other reviewers. These will be turned back to normal after the rebuttals.

Comparison with DiffAIL [1]: [1] also applies inverse reinforcement learning using diffusion models. Could the authors explain NEAR's major advantages compared with other works that integrate diffusion models into IRL?

• Thank you for pointing out the DiffAIL [1] paper. This is indeed an interesting paper. There are two main differences between NEAR and DiffAIL. First, DiffAIL is a combination of diffusion and adversarial techniques. It relies on the same objective as older adversarial methods like GAIL (please refer to METHOD: DIFFUSION ADVERSARIAL IMITATION LEARNING in [1]) and adds the diffusion loss within this term (eq. 17 in [1]). Hence, it might still be prone to the challenges of adversarial optimisation that we try to eliminate in our paper.
• Further, the use of the diffusion mechanism in DiffAIL is different to the use of score-based models in NEAR. In NEAR, we use score-based models to learn an energy function prior to learning the policy. We then use this fixed energy function to train the policy. Instead, DiffAIL trains the diffusion models simultaneously in hopes that it provides a more accurate reward than the simple adversarial classifier. Hence, while both algorithms use score-based models, their purpose in our opinion is quite different.
• Finally, DiffAIL has significantly higher computational costs because of the added load of learning the diffusion-based classifier (Table 4 in [1]). In comparison NEAR only has a fraction of the computation load (Table 7 in our paper).
• Having said that, we agree that it is important to mention this paper in ours. We have now added DiffAIL to our related work (line 72). We also plan to add this as a baseline in our future work on NEAR.

For instance, as the authors mention that GAIL can work well with single-clip data, how does NEAR compare in such settings, especially regarding data efficiency and performance? Further justification or empirical evaluation would be appreciated.

• Thank you for pointing out the single-clip performance of NEAR. We also trained NEAR in data-limited single-clip settings. These results are provided in Table 1 and Figure 4 in the paper (mummy-style walking and spin-kick are single-clip tasks). From our experiments, we found that NEAR performs poorly in single-clip settings while AMP is less affected by data unavailability. We believe that this is because of the added challenge of accurately learning an energy function via score-based modelling in limited data.
• We have now also provided additional information on the training resources (Table 5) and computational efficiency (Table 7) of NEAR. We find that in most cases (including single-clip settings), NEAR has better RL computational efficiency than AMP. Although, overall, it still requires slightly more training time.
• We hope that this explanation and additional results clarify NEAR’s performance under limited data.

Comparison with SMILING [2]: [2] introduces a non-adversarial framework using diffusion models for imitation from observation and provides theoretical analysis. Could the authors elaborate on the differences and advantages of NEAR compared with [2]? Including further empirical results comparing NEAR with [2] would strengthen the paper and clarify NEAR's contributions relative to existing methods.

• While SMILING [2] is indeed very interesting research, it is a recently proposed algorithm (submitted on ArXiV on 17 Oct 2024 which is after the ICLR 2025 submission date) and is currently also under peer review at ICLR. Hence, unfortunately, we could not include this in our literature search or empirical analysis.

A significant concern is the absence of direct comparison with other works that apply diffusion models to imitation learning from observation.

• To the best of our knowledge, NEAR is novel in its use of score-based models and energy-based rewards for imitation learning from observation (excluding SMILING [2]).
• When considering other adversarial techniques, we chose AMP as a baseline because it is an improved formulation of previous state-only methods like GAiFO [4] and is the strongest state-of-the-art baseline in this field (it is still a prominent algorithm used for learning such complex policies [5]).

Dear reviewer, thank you for your feedback, we appreciate the positive comments and the clearly explained questions/limitations. We have now uploaded a revised paper and supplementary material. Please find clarifications and paper changes below (revisions based on your feedback are coloured red). We hope these sufficiently establish the contribution of the paper.

It is unclear how Figure 1 has been generated. It is used for illustration purposes, however, some more information about the energy function and the adversarial reward are necessary. I also do not understand the choice of using different scales for the two rewards, and why the energy is high around the agent's trajectory while the adversarial reward is low. I would ask the authors to detail how the energy function and the adversarial reward have been learnt (also in the supplementary material if it does not fit the main text).

• Figure 1 Clarity:
  • Thank you for pointing out the clarity issues with Figure 1. This was also pointed out by another reviewer and we have taken several steps to ensure that the figures are clear to the reader.
  • We have now added an additional explanation of the maze domain in the Figure 1 caption. We have also added another appendix section (Appendix B.3 MAZE DOMAIN DETAILS) detailing the procedures to generate Figure 1.
  • We explain these briefly here. The energy-based reward function in Figure 1 was learnt using the energy-based NCSN model described in this paper. We trained the model with $\sigma=20.0$ on expert trajectories in the maze domain. The adversarial reward function was obtained by training AMP in the maze domain and saving the checkpoints at a comparable point in training.
  • We request the reviewer to kindly provide feedback on the revised figure clarity.
• Scaling Differences:
  • The scale difference between the energy-based and adversarial rewards is a result of the way in which the energy-based model and discriminator are defined.
  • In AMP, the discriminator is a least-squares regression classifier (please refer to Section 5.2 in [1]) and returns a value between [-1,1]. We chose to keep the same scale in the illustration to maintain correctness with the original AMP paper.
  • The energy function in score-based models is typically un-normalised. However, we chose to illustrate it on a scale of [0,1] to avoid unnecessary complications in the illustration.
  • We agree that this results in a somewhat confusing comparison between the two reward functions. We are happy to remove the scaling details altogether as they are not integral to the figure’s meaning.
• Energy/Discriminator Value Clarifications:
  • Typically, low energies indicate closeness to the data distribution. However, for simpler downstream RL, we chose to invert the sign of the learnt energy function such that maximising the energy reward means closeness to the expert (please refer to 4.1 LEARNING ENERGY FUNCTIONS)
  • The adversarial reward around the data distribution is low because of the fluctuating nature of the discriminator's decision boundary. In this case, $p_G$ has a high intersection with $p_D$. The decision boundary is shown such that the discriminator prediction is low (close to -1) near $p_G$ and gradually changes to be higher towards the remaining (non-intersecting) portion of $p_D$. In our experiments, we recorded highly fluctuating adversarial rewards for AMP. We have now also included a new appendix section (Appendix E) providing additional empirical evidence of the instability and non-smoothness in AMP. We hope this inclusion clarifies the ideas behind Figure 1.

It would be important to explain why other non-adversarial imitation learning methods are excluded from the evaluation and their differences and similarities with NEAR.

• To the best of our knowledge, NEAR is novel in its use of score-based models and energy-based rewards for imitation learning from observation. Other non-adversarial techniques for imitation learning (like [4] and its variants) primarily operate in a behaviour-cloning fashion. They hence require the actions executed by the expert. In contrast, NEAR operates in a state-only fashion.

While the evaluation metric "spectral arc length" seems to favor NEAR, motions generated with AMP are generally closer to the ground truth.

• We disagree with the claim that motions generated by AMP are closer to the ground truth. We believe that NEAR is closer to the ground truth in most cases. The provided supplementary videos show qualitatively that our approach generates closer-to-ground truth motions. We have now also included additional videos of the ground truth motions. Please, indicate to us if the supplementary videos helped.

I apologize if I sound pedantic, but I still have a few doubts after your last answer.

First of all, thanks for updating the figure, I find that the new one nicely shows an important property of your algorithm and a benefit of removing the adversarial component from the training process. However, this was not really what I meant with my comment, but I am still happy it lead to an improvement.

My misunderstanding was not about r(s, s') changing throughout the training, I understand that the function is learnt from a static dataset and not updated. My question is simpler than that. I don't know the details about the state, but I believe it includes the position of the agent in the 2D maze (or some representation of it). So I believe that to plot the energy function you fixed s and plotted r(s,s') where in s' you modify the 2D position in the maze. Please correct me if I am wrong.

Now, s depends on the current position of the agent in the maze. In the previous figure, you plotted r(s, s') for the energy function, but you used a different s for the adversarial reward (the agent was in a different location). That's why I said it would be more useful if you plotted r: s' -> r(s,s') keeping the same s, but using first r=energy and then r=adversarial reward. And then I proposed that you could take s within a trajectory and show how the "heatmap" of r: s' -> r(s,s') changes as a function of s (you could stack frames one after the other). I assume it would show some interpretable pattern, e.g., that the states behind the agent become lower reward states as it moves, as the reward should promote the agent to go forward and not backward. I believe this would immediately convince the reader that the energy of a transition is a good reward function.

Dear reviewer, thank you for your feedback. We have now submitted a revised paper. Please find clarifications and paper changes below. The text in some parts is coloured to indicate which revisions are associated with which reviewer (this will be reverted after rebuttals). We hope these clarifications help explain the contribution of our paper and request you to reconsider your score.

NEAR requires training a noise-conditioned energy model, which can be computationally intensive. A detailed comparison of training costs relative to other methods, particularly in terms of time and resources, would be beneficial.

• Thank you for highlighting this. In their typical use case (very high dimensional image/video generation) score-based generative models are indeed computationally intensive. However, we believe that the typical use case of score-based models is not comparable to their use in NEAR (relatively low dimensional state-transition pairs). We train the energy-based reward function for less than 150k training iterations in all cases. This is more than sufficient to learn the energy-based reward functions that are used in NEAR.
• We have now included two additional tables providing information on the training costs. Table 5 in Appendix B.2.4 provides more information on the training iterations and environment interactions of both NEAR and AMP in our experiments. Both algorithms' RL policies are trained for the same number of environment interactions. Table 7 in Appendix C compares the wall-clock training times of NEAR and AMP. We find that overall NEAR requires slightly more training time than AMP. Across all tasks, NCSN contributes to less than 30% of the total computational time of NEAR, with RL accounting for the majority of the computation load and wall-clock time. Interestingly, AMP always has a larger reinforcement learning computational load. This is expected as AMP is learning both the policy and the reward function simultaneously.

Could you provide a comparison of NEAR’s computational requirements (e.g., training time, GPU hours) relative to other baselines like AMP?

• Please see the previous answer and the newly included Table 5 and Table 7.

NEAR is only compared against one baseline - AMP and it doesn't seem to always be the winner (but has higher variance in most cases) despite the additional complexity of learning an energy network.

• To the best of our knowledge, AMP is the only comparable baseline for imitation learning from observation in such high-dimensional physics-based problems. It is an improved formulation of previous state-only methods like GAiFO [1] and is the strongest state-of-the-art baseline in this field (it is still a prominent algorithm used for learning such complex policies [2]).
• We agree that NEAR does have a slightly higher overall computational load than AMP (please also see the answer above and Tables 5 and 7). However, we would also like to point out that our research mainly aims to provide a better and more stable alternative to adversarial imitation learning. To that end, we believe that NEAR has a better mechanism for learning reward functions that is not prone to instability, non-stationarity and learning non-smooth rewards. While NEAR does have more variance in some tasks, this is mainly due to the limitations of annealing. We believe that the underlying reward learning mechanism of NEAR is still superior to that of AMP.
• To better highlight these limitations of AMP, we have now added another appendix section with additional explanations and empirical analysis (Appendix E) on the AMP discriminator. We hope that this additional information helps situate NEAR as a better alternative for reward learning in imitation learning from observation problems.

NEAR’s effectiveness is primarily evaluated in humanoid tasks, which are continuous and physics-driven. The framework’s applicability in other types of imitation learning tasks, especially those with discrete actions or diverse goal-oriented, is not fully explored.

• We believe that the scope of our paper aligns well with the experimental analyses from related works like [1,3,4,5,6,7,8]. All these prior works in the field of generative imitation learning also only conduct experiments in continuous domains. We believe that NEAR is situated amongst these works and our experiments are comparable to these related works.

Dear reviewer, thank you for your feedback. We have revised the paper accordingly. Please find clarifications and paper changes below (text revisions based on your feedback are coloured blue). We hope these sufficiently establish the contribution and presentation of the paper.

It is unclear to me what “until horizon” means in line 230 and 231

• Thank you for pointing this out. The horizon here refers to the rollout horizon used during policy training. The policy is rolled out for a horizon of 16 steps, then the collected transitions and rewards are added to the replay buffer to update the policy network weights. We have also explained this in detail in Appendix B.2.2 REINFORCEMENT LEARNING.
• We realized the term “until horizon” might be unclear in the algorithm. To improve clarity, we have now explicitly mentioned this in the algorithm’s setup (line 219) and have pointed to the appendix for additional details.

I wish there was a future work section

• Thank you for this suggestion. We agree that NEAR can be extended via several interesting ideas. We have now added a short future work paragraph on lines 498-511.

I feel like as well the sentences from 71-77 are pretty vague and I would appreciate you explicitly stating the challenges clearly and then comparing your work. The related work also seems to be wrapped in the intro which is fine but I would appreciate a more explicit comparison of other methods

• We realize that lines 71-77 (in the old manuscript) could indeed be improved.
• We have rewritten the introduction section (71 - 85) to clearly establish our contributions. We would appreciate it if the reviewer informs us if the revised one is clearer.
• Related works relevant to NEAR rely on adversarial reward learning. This mechanism has been briefly explained in lines 62-71. We have rewritten lines 79-85 to clearly establish our contribution and differentiate it from the mechanism used in prior works. We would appreciate it if the reviewer could provide feedback on the updated lines.

The significance of this work seems good but it would be nice if the contribution was more explicitly stated. My interpretation was that the contribution is the use of energy-based diffusion to train imitation learning but it would be better not to leave it up to the reader.

• Yes, the reviewer is correct in understanding that our contribution is to use energy-based models as reward functions to train imitation learning policies. Our approach allows us to learn this energy model offline and use it as a reward in the RL part in contrast with adversarial approaches that require an online update of the discriminator during the RL phase and lead to unstable min-max problems. We have modified lines 84-85 to state this explicitly.

I wish the tables were bigger sized, I’m pretty sure the page limit is 10 so you should be able to simply increase the size.

• Yes, we agree that the results tables could be enlarged. In light of other reviews, we might have to add a few more details to the paper. For the moment we have not increased the table size but are happy to do so if there is space available after adding other suggested modifications. We hope the reviewer understands our prioritisation of clarifying content in the revised paper.

The statistical rigor is ok. You say each is trained 5 times and it is the results are averaged across 20 trials. 5 seems low to me here. Especially since the confidence intervals seem to be sort of wide and overlapping. Doing maybe 20 or 30 would be much better unless this is prohibitively expensive.

• Yes, we understand that more trials would be beneficial. While running additional trials is not an unrealistic expectation, we chose to do 5 trials for the following reasons.
• In our preliminary experiments on a subset of imitation tasks, we did not find a significant difference in results when using more trials.
• While the computational requirements of a single trial are not too expensive, they grow to be significantly expensive for 20+ trials (across all 10 tasks in the paper). Unfortunately, being on a university-level computational budget, we are limited in terms of the available compute hours and GPU resources.

Generally, I think more intuition could have been used in the paragraphs that starts at 167, 189 and 154.

• We have now included additional intuition behind score-based models and NCSN on lines 169, 182, 190. We hope this also sufficiently clarifies the idea of the DIST() function mentioned on line 208.