Model-Enhanced Adversarial Inverse Reinforcement Learning with Model Estimation Reward Shaping in Stochastic Environments

We sincerely appreciate reviewer caVP's recognition for our contribution and comments. And below is our response.

Weakness 1:

Since the trained environment model differs from the real model, the generated trajectories are biased, which affects the learning of the reward model. Although the theoretical analysis in the paper discuss this, the algorithm design does not account for this issue.

Answer: In this work, the main contribution is to formulate transition model in the reward shaping to solve the MCE IRL problem under the stochastic setting. Improving transition model accuracy to mitigate distribution shift is not primary focus of our current work. Admittedly, the accuracy of transition model effect the performance reward learning, which we have extended an ablation study about this issue in Appendix E.2. We will explore different possibilities of transition model structures in future work.

Weakness 2:

• The theoretical upper bound in the paper is too loose (i.e., the theoretical description is too broad) and not tightly integrated with the designed algorithm. The theoretical conclusions do not adequately reflect the properties of the proposed algorithm
• The theoretical analysis of the reward model only considers the error of the learned environment model. However, the reward model is trained using both data generated by the environment model and data generated through interaction with the real environment. The reward model's error involves not only model errors but also data coverage, sampling errors, etc.
• The performance analysis only considers the reward model's error, but performance is also influenced by the RL algorithm itself. For example, the collected data is generated through interactions between the learned policy and the environment, and the learned policy is influenced by both the reward model and the RL algorithm.

Answer: Our main contribution is proposing a model-based reward shaping to framework solve MCE IRL problem under stochastic setting. Consequently, the theoretical bound presented in our paper concentrates on providing a theoretical framework to connect possible suboptimality of recovered policy with transition model learning error. Thus, under this theoretical framework, there could be other possible algorithms with refinement or extension building upon our bound by using more guaranteed accurate transition model or considering convergence and sampling errors from other policy optimization methods. And we have empirically investigated the effect of different policy optimization methods under our model-enhanced reward shaping in ablation study Appendix E.3.

Weakness 3:

Experiments: The experimental details are insufficient. Key parameters such as model rollout length, the proportion of model-generated data (and other parameters from Algorithm 1) are not provided. Additionally, details on the training of the environment model are lacking (e.g., is the environment model continuously updated during policy learning? On what data is each update based?).

Answer: We have provided more explicit implementation details in updated version. Please take a look at section 5.2 and Appendix F for details.

Reviewer caVP : if possible, can you respond to the rebuttal?

• Model-based RL is commonly employed in online settings, particularly with off-policy algorithms, as interactions with the real environment are crucial for ensuring the accuracy of the learned transition model. This accuracy, in turn, helps align the generated data samples with the true environmental dynamics [1,2]. While there are model-based approaches for offline RL, additional treatments are typically required to address the inaccuracies in transition learning and potential misalignment of generated samples [3]. Regardless of whether the setting is online or offline, model-based approaches generally do not assume prior knowledge about the transition dynamics [1], which is consistent with the assumptions made in our framework.
• As highlighted in the examples above, different reward mechanisms can result in similar behavior or visitation distributions for trajectories. Consequently, we do not believe that any fair normalization technique exists to establish a relationship between fundamentally different reward mechanisms. The purpose of recovering a reward function in IRL is not to precisely replicate the true reward but to enable the agent to reproduce the same visitation distribution as the expert policy[4]. This alignment is typically demonstrated empirically through comparable levels of performance. If the reviewer considers the robot examples provided earlier, it becomes clear that entirely different reward mechanisms can lead to similar agent behaviors. Given this, there is no inherent relationship that must exist between such divergent reward functions. Moreover, pursuing such a comparison is beyond the scope of IRL, where the goal is not true reward learning, as would be the case in world model learning, but rather behavior replication.

[1]Moerland, Thomas M., et al. "Model-based reinforcement learning: A survey." Foundations and Trends® in Machine Learning 16.1 (2023): 1-118.

[2]Nair, Ashvin V., et al. "Visual reinforcement learning with imagined goals." Advances in neural information processing systems 31 (2018).

[3]Yu, Tianhe, et al. "Mopo: Model-based offline policy optimization." Advances in Neural Information Processing Systems 33 (2020): 14129-14142.

[4]Arora, Saurabh, and Prashant Doshi. "A survey of inverse reinforcement learning: Challenges, methods and progress." Artificial Intelligence 297 (2021): 103500.

We appreciate the comments provided by Reviewer Mu2o, and our responses are as follows.

Weakness 1:

The main technical contribution of the proposed method is to incorporate system transition estimation into reward shaping for inverse reinforcement learning. However, the training process for the transition model is not explained in the paper. The method used to train this transition model and the accuracy of the model itself likely have a substantial impact on the final results. I suggest the authors provide additional analysis or details on this aspect.

Answer: We have incorporated implementation details in section 5.2 and Appendix F of the updated version. We also extend an ablation study on how the transition model learning error affects performance which can be found in Appendix E.2.

Weakness 2:

It appears that the primary distinction between the proposed framework and AIRL is the incorporation of state transition probabilities within the reward function.

Answer: See General Response

Weakness 3:

Experimental results are presented for only three environments, which limits the strength of the evaluation. Expanding the experiments to additional environments would make the results more convincing. (Provide source code)

Answer: We have conducted more experiments on more environments. And we also provide the source code to reproduce our results.

Dear Reviewer Mu2o,

We understand this is a busy time of year, but we wanted to kindly remind you to review our rebuttal. Your feedback is invaluable for us to improve our paper. Please feel free to let us know if there are any specific points that require further clarification.

Thank you for your time and effort. We truly appreciate your contributions to this process!

Best, Authors

Reviewer Mu2o: if possible, can you respond to the rebuttal?

Thank you for your thoughtful response and questions! Below are our clarifications:

1. Major Contribution and Transition Model Design:
   The assumption of Gaussian distributions has been adopted in several works within the field [7,8]. While we acknowledge that real-world applications often involve various types of noise, extending our transition model's robustness to handle different forms of uncertainty is an exciting direction for future work. Our primary contribution lies in introducing a model-enhanced reward shaping framework to address the MCE IRL problem. While designing a transition model that better captures stochasticity is indeed important, it is not the primary focus of this work. Nevertheless, many well-established model structures, such as RNN-based RSSM in Dreamer [1], torchsde [2,3], and VAE-like structures [5], could be seamlessly integrated into our framework to address more complex or high-dimensional dynamics. Given the challenges of learning a perfect transition model under diverse stochastic environments, our approach prioritizes practicality and computational efficiency. For this study, we opted for a simple transition model, which is sufficient to demonstrate the efficacy of our method while balancing computational resources and the extensive range of experiments conducted.

2. Gaussian Noise and Normalization:
   Thank you for raising this point. In our current implementation, observations/environment objects are first normalized before Gaussian noise is applied. This ensures that the noise has a consistent effect across all dimensions, regardless of their original scale. Details of this implementation can be found in our anonymous code repository (linked in the General Response and Appendix F).

3. Transition Model Training and Clarity:
   We appreciate the reviewer pointing out potential clarity issues regarding the transition model training. To clarify:

   • Our approach differs from MBPO [4], which employs an ensemble of deterministic transition models with complex update rules. In contrast, we train a single transition model using a standard MLE loss for each update (as described in line 10 of our pseudocode).
   • Line 11 explains the varying ratio of samples drawn from different buffers, which is detailed further in Appendix F.
   • Regarding the pre-training phase, this is a standard procedure in off-policy algorithms to warm up the buffer [6]. Additionally, we update the transition model using the same MLE update loss in line 10 at each step during pre-training to mitigate potential divergence issues early in training. Detailed explanations can also be found in Appendix F.

We thanks again for reviewer's thorough feedback, and we hope for your continued consideration of our work.

[1]Hafner, Danijar, et al. "Dream to control: Learning behaviors by latent imagination." arXiv preprint arXiv:1912.01603 (2019). [2]Li, Xuechen, et al. "Scalable gradients for stochastic differential equations." International Conference on Artificial Intelligence and Statistics. PMLR, 2020. [3]Wang, Yixuan, et al. "Enforcing hard constraints with soft barriers: Safe reinforcement learning in unknown stochastic environments." International Conference on Machine Learning. PMLR, 2023. [4]Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in neural information processing systems, 32, 2019. [5]Girin, Laurent, et al. "Dynamical variational autoencoders: A comprehensive review." arXiv preprint arXiv:2008.12595 (2020). [6]Degris, Thomas, Martha White, and Richard S. Sutton. "Off-policy actor-critic." arXiv preprint arXiv:1205.4839 (2012). [7]Grunewalder, Steffen, et al. "Modelling transition dynamics in MDPs with RKHS embeddings." arXiv preprint arXiv:1206.4655 (2012). [8]Engel, Yaakov, Shie Mannor, and Ron Meir. "Reinforcement learning with Gaussian processes." Proceedings of the 22nd international conference on Machine learning. 2005.

We appreciate reviewer qE6H's comments. Following are our response.

Weakness 1:

The writing quality of the paper requires improvement. For instance, while the core idea is relatively straightforward, the title is overly long and complex, making it challenging to grasp before reading. The authors should also enhance the clarity of propositions and theorems by providing more detailed descriptions and insights, explicitly connecting them to the paper’s main ideas. Additionally, the numbering of theorems and their proofs in the appendix should be consistent and aligned for easier reference.

Answer: We have made changes accordingly including title in the updated paper. Changes are highlighted in cyan color.

Weakness 2:

The paper modifies the reward shaping term from $\phi(s_t)$ to $E_\tau [\phi(s_{t+1})\mid s_t,a_t]$, incorporating the transition model into the reward shaping to better handle stochastic environments. While the idea is conceptually straightforward, it introduces additional complexity in modeling the transition function, especially in real-world environments with inherent stochasticity in transitions.

Answer: In our setting, we only need to train a simple transition model as a two-layered low-dimensional MLP which is already sufficient to predict a few step horizon. Thus, we believe that comparing to the performance and sample efficiency improvement brought by this simple structure, which has been thoroughly investigated in our experiments, a little additional complexity is cost-effective.

Weakness 3:

In Section 5.3, the authors present a theoretical analysis of performance. However, the analysis appears to apply broadly to inverse RL settings with transition mismatches, offering limited theoretical insight into the specifics of the proposed method. In particular, Theorems 5.3 and 5.4 provide bounds on rewards and value functions across two MDPs, yet they do not address the optimality or performance bounds of the policy learned through the proposed approach. This generality detracts from a deeper understanding of how the method uniquely impacts policy performance.

Answer: Our main contribution is proposing a model-based reward shaping to framework solve MCE IRL problem under stochastic setting. Consequently, the theoretical bound presented in our paper concentrates on providing a theoretical framework to connect possible suboptimality of recovered policy with transition model learning error. Thus, under this theoretical framework, there could be other possible algorithms with refinement or extension building upon our bound by using more guaranteed accurate transition model or considering convergence and sampling errors from other policy optimization methods.

Weakness 4:

The authors claim that their method achieves improved sample efficiency, as demonstrated in Section 6 (line 439). However, it is unclear if the authors have accounted for the additional computational costs involved in modeling the transition functions. Moreover, it remains ambiguous whether the observed sample efficiency in policy learning is primarily due to the proposed method itself or the specific techniques introduced in Section 5.2 (line 316). Clarifying these aspects would strengthen the evaluation of the method’s true efficiency gains.

Answer: We have added an ablation study in Appendix E.3 investigating the performance difference between DAC, model-enhanced reward shaping with SAC, and our method. Our result shows that our model-enhance reward shaping technique significantly contributes to the sample efficiency and performance improvement. Model-based policy optimization techniques further improve the sample efficiency, however, it doesn't cause notable differences in terms of performance, which makes sense since both policy optimization methods should converge to the same optimality[1].

Dear Reviewer qE6H,

We understand this is a busy time of year, but we wanted to kindly remind you to review our rebuttal. Your feedback is invaluable for us to improve our paper. Please feel free to let us know if there are any specific points that require further clarification.

Thank you for your time and effort. We truly appreciate your contributions to this process!

Best, Authors

Thank you for taking the time to review our rebuttal and for acknowledging the improvements we made to address your concerns. We sincerely appreciate your effort in carefully considering our revisions and providing valuable feedback.

We would like to clarify that the additional experiments were conducted primarily to further validate the performance superiority, robustness, and sample efficiency of our novel off-policy model-enhanced reward shaping framework, which addresses the Maximum Causal Entropy IRL problem in stochastic settings. Importantly, the core methodology, theoretical contributions, and key statements regarding efficacy and efficiency remain unchanged from the initial submission. The additional experiments were specifically included to address the feedback received during the review process and provide stronger empirical support for our claims, not to alter the foundational aspects of the paper.

We believe we have thoroughly addressed all the concerns raised during the rebuttal period. However, we respect your overall stance and deeply value your perspective. If there are any specific aspects of the paper that remain unclear or insufficient, we would be happy to provide further clarifications or address those points directly.

Thank you again for your thoughtful feedback, and we hope for your continued consideration of our work.

We thank Reviewer PcDm for the comment. Following are our answers to the mentioned questions and concerns.

Weakness 1:

Experiments are only in Hopper and inverted pendulum, while the latter seems to be an easy setting for all the methods.

Answer: We have extended our experiments to additional environments (e.g., HalfCheetah-v4 and Walker2d-v4) with additional exploration on the effects of quantity of expert trajectories and a set of ablation studies. Please check the detailed illustration in General Response and our updated paper.

Weakness 2:

Some of the presentation is not clear. For example, the L11 in the algorithm is confused, Does it correspond to the L324-331? And why do you sample H-step trajectories (can you connect it with the sample efficiency somewhere)?

Answer: We have fixed a few presentation issues in the updated paper. Line 11 in the algorithm does correspond with Sample Efficiency section. Generating $H$ steps of trajectories from the learned transition model is a common approach for Model-based RL [1,2,3,4] to supplement real environmental samples for policy optimization.

Question 1:

Can you explain intuitively why not having the expectation in Eq(6) doesn't work?

Answer: The reasons here are two-fold. First, let's consider a stochastic MDP with continuous state $\mathcal{S}$ and action $\mathcal{A}$ space, which means $\mathcal{T}(s_{t+1}|s_t,a_t)\rightarrow[0,1]$, where $s_t,s_{t+1}\in\mathcal{S}, a_t\in\mathcal{A}$, and $\vert\mathcal{S}\vert\rightarrow\infty, \vert\mathcal{A}\vert\rightarrow\infty$. The intuitive explanation is that each $(s_t,a_t,s_{t+1})$ tuple can be regarded as a sample in a continuous distribution. Since state space and action space are continuous, there are numerous possible tuples requiring extensive sampling to recover the reward. Secondly, for different tuples above, each has a transition probability. There could be a scenario that we recover high reward to a tuple with extremely low transition probability, which in reality might be very unlikely to reach. This is called risk-seeking behavior[5], commonly caused by directly shifting ME-based algorithms to stochastic environments.

Question 2:

Why compare with a deterministic environment? Isn't it the same with other reward shaping method?

Answer: The reason for comparing under deterministic settings is to validate that our method has competitive performance with all the other baselines. And in the stochastic settings, our method can achieve much better performance compared to baselines.

Question 3:

Why the L11 in algorithm will mitigate the distribution shift? Can you give a little bit more background on why this induces distribution shift?

Answer: To clarify that Line 11 in the algorithm does not mitigate the distribution shift, we introduce the distribution shift mitigation method in lines 327 to 333 in the updated paper version. The reason for the distribution shift is that at the early stage of training, distribution of learned transition model $\hat{\mathcal{T}}$ is far from real $\mathcal{T}$, thus policy optimization based on samples from both transition dynamic might cause deviation without any treatment[1], since part of the samples might be unrealistic under real environment.

Question 4:

I am wondering can this idea be applied to the GAIL, which uses the transition model to generate a mixture of the D_{gen} and D_{env}? And also account for the stochastic of the env in the -logD for gail?

Answer: For GAIL, it is a direct imitation learning algorithm, thus it is not recovering reasonable rewards as IRL. Therefore, the learned transition model might not be helpful to GAIL. Additionally, GAIL is an on-policy algorithm, and model-based policy optimization approaches normally are off-policy, thus learned transition model might not be helpful on both end.

Question 5:

How do you vary the proportion of selecting synthetic data or environment data for policy optimization. And how do you select H? And together how they affect the sample efficiency?

Answer: At the beginning of the training, the proportion to select synthetic data is 0, since the learned transition model is not accurate at this point. As training progresses, the proportion increases linearly with the training environmental steps. For H, we set it to 2, but from past experience anything below 8 should work. Details of these hyper-parameters selection schemes can be found in Appendix F. Generally, increased proportion and H will result in more trajectories generated. However, depending on the learned transition model capacity, as H increases, the generated trajectories might become less realistic, resulting in suboptimal policy performance.

Dear Reviewer PcDm,

We understand this is a busy time of year, but we wanted to kindly remind you to review our rebuttal. Your feedback is invaluable for us to improve our paper. Please feel free to let us know if there are any specific points that require further clarification.

Thank you for your time and effort. We truly appreciate your contributions to this process!

Best, Authors

Reviewer PcDm: if possible, can you comment on the rebuttal?