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

Authors: Simon Sinong Zhan, Qingyuan Wu, Philip Wang, Yixuan Wang, Ruochen Jiao, Chao Huang, Qi Zhu
Primary area: reinforcement learning
Keywords: Inverse Reinforcement Learning; Model-based Approach;

Abstract:
In this paper, we aim to tackle the limitation of the Adversarial Imitation Learning (AIL) method in stochastic environments where theoretical results cannot hold and performance is degraded. To address this issue, we propose a novel method which infuses the dynamics information into the reward shaping with the theoretical guarantee for the induced optimal policy in the stochastic environments. Incorporating our novel model-enhanced reward shaping, we present a novel Model-based IRL framework, which integrates transition model estimation directly into reward shaping. Furthermore, we provide a comprehensive theoretical analysis of the reward error bound and performance difference bound for our method. The experimental results in MuJoCo benchmarks show that our method can achieve superior performance in stochastic environments and competitive performance in deterministic environments, with significant improvement in sample efficiency, compared to existing baselines.

Decision: reject
Reviewer scores: 6, 5, 5, 5

Reviews and discussion:
<thread>
creation date:
1729435529883

thread_type:
official_review

<thread_object>
creator: reviewer_cavp

<content>
summary:
To address the issue of poor performance of traditional adversarial inverse reinforcement learning (AIRL) in stochastic environments, this paper incorporates environmental information into the reward model training and propose a new AIRL algorithm, which effectively improves performance in stochastic environments.

strengths:
1) Writing: The article is overall very well-written, with clear structure and logical flow, rigorous theoretical derivations, and sufficient experiments.

2) Algorithm: The approach is innovative. By integrating environmental information into the reward model learning process, the method effectively enhances the algorithm’s performance in stochastic environments.

3) Theory: The paper derives that minimizing the discriminator’s cross-entropy loss is equivalent to maximizing the log-likelihood under MCE IRL.

weaknesses:
1) Algorithm: 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.

2) Theory: 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. Specifically:

i) 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.

ii) 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.

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?).

questions:
1) The proposed algorithm uses the learned policy to continuously interact with the real environment (or a realistic simulated environment) to collect data $\mathcal{D}_{env}$. However, if the real environment is already known, why is there a need to train an environment model and generate data? From my perspective, having both seems contradictory. Typically, when the real environment is unknown, current research uses collected offline data to train the environment model, which is then used to augment the dataset in order to improve sample efficiency.

2) Since the paper primarily focuses on reward model learning, the author should compare and analyze the gap between the learned reward model and the real reward model (this can be done by comparing the results after normalizing both models).

3) The quality of the data generated by the learned environment model is significantly influenced by the rollout length $H$. The author should provide a brief discussion on how $H$ is selected, and explain how the environment model is updated during the policy learning process.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730380009184

thread_type:
official_review

<thread_object>
creator: reviewer_mu2o

<content>
summary:
This paper presents an enhanced version of Adversarial Inverse Reinforcement Learning (AIRL) by extending it to stochastic MDPs rather than deterministic ones. The authors revise the deterministic “next state” in AIRL to a probability distribution, which more accurately reflects real-world scenarios. The writing is clear and well-structured, making the paper easy to follow.

However, the key technique of transition modeling is not discussed, which could hinder reproducibility. Additionally, the experimental results are limited, making it challenging to fully assess the effectiveness of the proposed method.

strengths:
1. The motivation for this paper is clearly presented. While most previous work focuses on deterministic MDPs, this framework addresses stochastic settings, which are more representative of real-world scenarios.
2. The paper is grounded in strong theoretical foundations, and all key arguments are thoroughly supported by theoretical verification.
3. The writing is clear and well-structured, making the paper easy to follow.

weaknesses:
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.
2. It appears that the primary distinction between the proposed framework and AIRL is the incorporation of state transition probabilities within the reward function.
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.

questions:
Same as weakness. 

Since the transition model training is not well-explained in the main paper, the results are hard to reproduce. I suggest the authors to provide source code.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730631843011

thread_type:
official_review

<thread_object>
creator: reviewer_qe6h

<content>
summary:
The paper introduces a model-enhanced adversarial inverse reinforcement learning (IRL) approach that integrates transition model estimation into reward shaping, with theoretical guarantees and superior performance in stochastic environments.

strengths:
1. The paper presents a incremental framework over adversirial inverse RL that integrates model-based transition dynamics with adversarial IRL, potentially addressing performance degradation in stochastic environments, which is a good direction because considering the stochasticity in the environments is significant.
2. The authors show some theoretical guarantees of the optimality of the method and the transition model error.

weaknesses:
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.
2. The paper modifies the reward shaping term from $\phi(s_t)$ to $E_{\mathcal{T}}[\phi(s_{t+1})|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.
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.
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.
5. The experiments are too limited to thoroughly evaluate the effectiveness of the proposed method. Specifically, 
- The authors conduct experiments only on three relatively simple, low-dimensional environments and compare against just three baseline methods (see Tables 2 and 3). To strengthen the evaluation, the authors should consider testing on more challenging environments and expanding the comparisons to include additional algorithms, such as SAM, SQIL, and GAIfO, as mentioned in Table 1.
- Appendix D lists a few key hyperparameters, but several essential details are missing, which could significantly hinder reproducibility. Specifically, omitted hyperparameters include the batch size, discount factor, the number of discriminator training epochs per policy update, and exact 5 random seed values, etc. Additionally, while sharing code is not mandatory, its absence makes it even more challenging to accurately replicate the experimental setup and results.
- The authors should conduct more ablation studies. For example, with different noise scale in stochastic environments, to show the robustness of the proposed method.
- In Figure 4, both DAC and the proposed method show non-zero initial performance, which raises questions about the setup or initialization of these models. Additionally, it’s unclear why both methods terminate early in Figure 4, whereas they continue for the full duration in Figures 2 and 3.

questions:
Please see the above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730680997378

thread_type:
official_review

<thread_object>
creator: reviewer_pcdm

<content>
summary:
The paper proposes a model-enhanced reward-shaping method for adversarial inverse reinforcement learning in the stochastic environment where AIRL is likely to fail. By integrating the transition model to the account for the stochastic transition. Also, they provide an error bound for it.

strengths:
1, Propose a method that tackles the limitations of existing methods in the stochastic environment.

2, The method is simple and intuitive. 

3, The error bound in the paper is well connected with the paper.

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

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 )?

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

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

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? 

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?

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?

6, If the transition model is unbounded, for example, when the transition model is hard to learn in more complex environments, can it still outperform AIRL? I believe more samples will be needed.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731949018439

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank all the reviewers for the effort they put in to improve our paper from various perspectives. This work aims to propose a novel model-based reward-shaping method to solve Maximum Causal Entropy IRL under stochastic settings. Below is our general response. First, we want to clarify that adversarial IRL/IL[1] is a category including all methods formulating the imitation learning problem into an adversarial training structure, and **AIRL[2] is a specific deterministic on-policy algorithm based on Maximum Entropy (ME) framework** falling into that category, so do plenty of others. To clarify, our approach is an **off-policy algorithm targeting stochastic setting** with an adversarial training framework as all the other adversarial IL/IRL works, but it is rooted in the **Maximum Causal Entropy (MCE) framework**, rather than the standard ME followed by AIRL, DAC, etc, leading to distinctive discriminator formulation and training scheme. The ME framework maximizes the entropy of a policy to encourage exploration without regard to environmental stochasticity, while the MCE framework specifically maximizes entropy over action sequences conditioned on past states and actions, ensuring that the policy accounts for the stochastic nature of the environment. We illustrate the difference between MCE and ME frameworks and resulting mathematical IRL problem formulations in lines 171-183. More details can be found in [3,4]. This nature together with **model-based off-policy** training fashion leads to our method's efficacy in performance and sample efficiency under stochastic settings, which we further validate this statement in our newly added experiments introduced in the next paragraph. We provide the theoretical insight in Proposition 5.1 to show the equivalence between our adversarial training objective and MCE objective in reward learning. 

**_We acknowledge that the original experiments and algorithmic details are not sufficient, thus we have updated our paper with more thorough experiments under stochastic settings, ablation studies, and algorithmic details including hyperparameters selection and source code (See newly updated PDF)._**  Detailed changes are in the following sections:
* [**Implementation Details and Source Code**] We have updated the algorithmic details of our method, including transition model learning and implementation specifics, which are now available in Section 5.2 (lines 306-310) and Appendix F (page 26). Additionally, we have included an anonymous link to our code repository in Appendix F, which contains the expert policies used for expert trajectory collections, as well as implementations of baseline methods and our proposed approach. (https://anonymous.4open.science/r/MBIRL-4C2F/README.md)
* [**More Environments**] We have added new experiments in the HalfCheetah-v4 and Walker2d-v4 environments under stochastic and deterministic conditions. Our experiments also support our claims of competitive performance under deterministic settings and superior performance under stochastic settings. These results are provided in lines 428-469 and lines 501-511.
* [**Sample Efficiency**] To examine the sample utilization efficiency of provided expert trajectories under stochastic settings, we conducted tests across all five environments using varying numbers of expert trajectories, from 5 to 1000. Our method demonstrated notable performance superiority across most environments, regardless of the quantity of expert trajectories provided. Results can be found in lines 471-485 and provided training curves of all environments can be found in Appendix D. The results indicate that our method consistently recovers expert-level performance even with limited expert trajectories, a capability that all baselines miss in such settings. Furthermore, as shown in Appendix D, our method also achieves expert-level performance with the fewest environmental steps in all environments.
* [**Ablation Studies**] We conducted three additional ablation studies, presented in Appendix E. The first study examines our method's robustness to different levels of stochasticity. The second study explores the effects of transition model learning errors on reward learning. The third study investigates under the same model-enhanced reward shaping formulation, how SAC and Model-based Policy Optimization methods affect the performance and sample efficiency differently. 


[1]Orsini, Manu, et al. "What matters for adversarial imitation learning?." Advances in Neural Information Processing Systems 34 (2021): 14656-14668.

[2]Fu, Justin, et al. "Learning robust rewards with adversarial inverse reinforcement learning." arXiv preprint arXiv:1710.11248 (2017).

[3]Brian D Ziebart, et al. Modeling interaction via the principle of maximum causal entropy. 2010.

[4]Gleave, Adam, et al. "A primer on maximum causal entropy inverse reinforcement learning." arXiv preprint arXiv:2203.11409 (2022).
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper addresses adversarial inverse reinforcement learning (IRL), with a view of supporting stochastic environments. To accomplish this, it uses reward shaping based on a learned transition model.

The main strengths of the paper are that (1) stochastic environments are important in practice and (2) the approach has thoretical justification.

However, important concerns remain. Specifically:
- the experiments are still limited, even after having been expanded in the rebuttal phase
- using performance profiles (https://agarwl.github.io/rliable/) would be nice for this type of paper
- the writing and clarity needs to improve by a lot

For these reasons, I recommend rejection.