When Demonstrations Meet Generative World Models: A Maximum Likelihood Framework for Offline Inverse Reinforcement Learning

Summary

This paper focuses on learning from demonstration via offline inverse reinforcement learning (IRL). Offline IRL suffers a similar problem as offline reinforcement learning (RL) and imitation learning (IL), where the policy cannot generalize well on unseen states and actions---this problem is known as distribution shift. To address this problem, this paper proposes to first learn a dynamic model, and formulates a maximum likelihood (ML) objective to simultaneously recovers both the reward function and the policy. Notably, the policy is optimized using a maximum entropy objective along with pessimism based on the uncertainty of the learned dynamic model. This paper provides PAC-style bounds to quantify the amount of samples required to achieve $\varepsilon$-optimal solution to the MLE objective. The paper further provides an algorithm that obtains such a $\varepsilon$-optimal solution under specific assumptions. Finally, this paper provides empirical evaluation on the D4RL benchmarks. ## Contributions - A new MLE objective for recovering a policy that is close to the expert policy. - A theoretical analysis that describes the statistical guarantees of the objective - An algorithm that obtains to near-optimal solution under linear parameterization of the reward function - An empirical evaluation on standard benchmark, D4RL, that indicates statistically significant improvement over existing baselines in majority of the tasks.

Strengths

- The paper is well written and easy to follow in general---I particularly appreciate the presentation on providing formal statements followed by the high-level intuitions. - The paper proposes a novel formulation for offline inverse reinforcement learning. - The paper provides numerous theoretical justifications and an algorithm that is inspired by said analyses.

Weaknesses

- In practice, how do we ensure assumption 2 (ergodicity)? It seems like this assumption actually "hides" some part of the coverage requirement? - I am curious as to how the MLE objective connects to the policy error---I completely understand that if $\pi_\theta = \pi^E$ then the policy error is zero. However, it does not seem to me that achieving $\varepsilon$-error on the MLE (i.e. $L(\theta^*) - L(\hat \theta) \leq \varepsilon$) does not directly tell us the policy error. - I think the training description for BC is somewhat vague---on page 8, line 314: what exactly does it mean by "train the algorithm until convergence"? Do we have some form of validation checking for BC? [1, 2, 3] have results regarding how BC would perform based on specific validation. Secondly, was there any hyperparameter search on BC, ValueDICE, and CLARE? - Regarding the experiment on recovered rewards, what is the performance if we were to fix the reward to 0? Isn't it better if we were to consider the correlation between the true reward function and the obtained reward function? ## References [1]: Hussenot, L., Andrychowicz, M., Vincent, D., Dadashi, R., Raichuk, A., Ramos, S., ... & Pietquin, O. (2021, July). Hyperparameter selection for imitation learning. In International Conference on Machine Learning (pp. 4511-4522). PMLR. [2]: Mandlekar, A., Xu, D., Wong, J., Nasiriany, S., Wang, C., Kulkarni, R., ... & Martín-Martín, R. (2021). What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint arXiv:2108.03298. [3]: Ablett, T., Chan, B., & Kelly, J. (2023). Learning from Guided Play: Improving Exploration for Adversarial Imitation Learning with Simple Auxiliary Tasks. IEEE Robotics and Automation Letters.

Questions

- On page 3, equation 2a: $\theta$ parameterizes only the reward function, and $\pi_\theta$ corresponds to the policy obtained when $r$ is parameterized by $\theta$, correct? - On page 5, proposition 1: Is there any result regarding larger $\varepsilon$? It seems like we may want to sacrifice this approximation error. ## Possible Typos - Page 3, line 93: gamma \in (0, 1)? - Page 3, line 95: The $P$ should not be the same as the transition dynamics right? I feel we should be clear that we are either overloading notation or use another symbol. - Page 5, equation 11: What's $\epsilon$? Is it the same as $\varepsilon$? - Page 8, line 303: Missing space after "2)"

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

- The paper's proposed method requires uncertainty estimation which is still an active research problem, especially in the context of neural networks---as far as I understand the paper leverages existing work that lacks theoretical guarantee.

Summary

This paper addressed the issue of covariate shift in offline imitation learning. The authors extended the uncertainty-regularized model-based offline RL to the imitation learning setting. The key idea is to first learn transition dynamics from samples, and then solve an optimization problem that jointly seeks a policy such that it optimizes the learned transition dynamics accompanied by a reward model and maximizes the log-likelihood of actions in data. The authors provide theoretical guarantees for the maximization of action log-likelihood and empirical results for the learned policy.

Strengths

1. The issue of covariate shift is indeed important for offline IRL. 2. The efficacy of the algorithm is partially supported by empirical results. 3. Analysis is provided for the model-based part of this algorithm.

Weaknesses

1. The paper is somewhat hard to follow. See questions below. 2. The effect of overcoming the distributional shift is not emphasized in the experiment section. None of the experiments was carried out on small datasets where the coverage of state–action space is limited. In fact, the medium datasets in D4RL contain 1M transitions, and the medium-expert versions contain 2M transitions. The datasets for results in Figure 2 contain 5000 expert demonstrations, which might correspond to 5M transitions if each expert trajectory contains 1000 transitions. 3. The paper does not have an informative conclusion part.

Questions

1. Why do you include discounting in the log-likelihood term? 2. Why do you include the entropy term in eq. 3a? 3. The theoretical analysis relies on the relation between optimal policy and optimal soft Q-function (eq.4 and eq. 14). But as mentioned in line 226--230, practical implementations utilize an actor-critic framework to approximate the optimal policy. How does this approximation affect the analysis? 4. Eq. 15 shows a clear resemblance with the reward maximization part in max-entropy IRL (eq. (1)). It suggests that the proposed algorithm seems to replace the online interaction of max-entropy IRL with sample generation using the learned dynamics. Then, what is the motivation to maximize the log-likelihood of expert actions? A suggestion. One concern of this approach is that to estimate the transition dynamics we need a sufficient amount of data, but the covariate shift issue occurs when we only have a small amount of data. I would suggest the authors include results on a few expert demonstrations (e.g. 10~100) to verify their claim.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

No. there is no discussion on limitations.

Summary

Offline inverse reinforcement learning (IRL) is a method for finding an unknown reward function optimized by an agent from demonstrations using only a finite dataset. The most common framework is maximum entropy (MaxEnt) IRL, which attempts to find a reward function that induces a policy which achieves the same expected reward as the trajectories from the expert demonstrations while maximizing its entropy. Prior work has shown that this is equivalent to finding a policy which maximizes the likelihood of the demonstrations under the constraint that this policy comes from solving a MaxEnt RL problem. This formulation as a bi-level optimization problem reduces the computational burden that results from alternating between finding the policy and updating the reward. However, it requires access to the true dynamics of the environment, which is incompatible with the offline IRL setup. Instead, this work proposes to learn the dynamics model in an uncertainty-aware fashion and incorporate a measure of this uncertainty in the learned reward function. This results in a two-stage procedure: 1) fitting a dynamics model from transition samples in the dataset and 2) recovering the reward function using the maximum likelihood (ML) formulation of the IRL problem. To perform the second step, the authors propose a novel decomposition of the upper-level objective, which consists of a surrogate objective that is more computationally tractable to optimize. The authors provide statistical guarantees about the optimality of the recovered reward function in terms of dataset coverage, a concept common in offline RL. Importantly, their bounds depend on dataset coverage on expert-visited state-action pairs, not the full joint space.

Strengths

- Offline inverse RL is an important area for tackling challenging sequential decision making problems in potentially safety-critical applications. - The paper is well organized and clearly written. It does a good job explaining the novelty and results and provides enough information to support its claims. - The paper presents an extensive experimental evaluation on several benchmarks, comparing to both model-based and model-free offline IRL algorithms and existing imitation learning approaches. Their algorithm outperforms these baselines in most cases. - The authors provide a nice analysis of surrogate objective and its relation to the true upper-level objective. This provides a nice motivation for optimizing the surrogate instead, which is more computationally tractable. - The authors present a nice optimality guarantee of the stationary point of their algorithm in terms of the surrogate objective in the case where the reward function is linear in a feature vector of states and actions. They also relate this stationary point to the true optimal solution of the original problem. - The additional reward transfer experiments indicate that the learned reward function may transfer to dynamics models trained on different state-action distributions. This appears to hold even if the state-action coverage used to train the reward function is close to expert-visited states.

Weaknesses

- The alternating optimization scheme discussed in this work appears to be identical to that presented in [21]. If true, that is fine, as the main contribution of the work lies in modeling the conservative MDP and providing novel bounds in the offline setting. However, it should be made explicit in the paper and mentioned in the contributions. - Theorem 2 seems very similar to Theorem 5.4 in [21], except that the Q function approximation error is considered explicitly. If this is true, it should be discussed that this is the novelty in the text. - Section 6 should discuss the differences in the three dataset types used, as the current text does not explain what they entail. This makes it difficult to understand the performance of the proposed algorithm in each setting without carefully looking at the Appendix. It should also talk about the purpose of using these different datasets. From the Appendix, it appears that they are only used to train the dynamics model. Thus, they are evaluating the effect of dataset coverage around the expert on performance. This should be explicitly discussed in the paper. I know space is limited, but these are important details that should be in the main text. - A minor comment is that it would be nice for the main paper to end with a conclusion section rather than ending abruptly. And this conclusion should mention limitations of the current method.

Questions

- Is there a difference in the alternating optimization scheme presented here and the one in [21]? - How does Theorem 2 in this text relate to Theorem 5.4 in [21]? - Am I correct in assuming that the three different datasets are used to evaluate how the coverage of the dataset around expert state-action pairs affects the algorithm's performance?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

There is no discussion of limitations in the paper. The paper would be made a lot stronger if this was discussed in a conclusion section.

Summary

In this paper the authors present a two level maximum likelihood based framework for offline inverse reinforcement learning, where both a world model and a reward model are learnt from expert demonstrations. In this two level algorithm the outer level or loop involves estimating the reward function, while the inner loop estimates the optimal policy for the chosen reward function in a conservative MDP setting, where a penalty is added which is loosely proportional to the uncertainty in the learnt model. The authors provide theoretical guarantees for the performance of their algorithm under fairly standard technical assumptions. They also show the numerical performance of their algorithm on 3 MuJoCo environments, comparing them with other state of the art offline RL algorithms.

Strengths

1. The paper is novel, clearly written and is easy to comprehend. 2. The authors have stated their results formally in the form of Lemmas and Theorems and have proved them in the supplementary material. This analysis proves the validity and utility of their proposed approach. 3. While model based offline inverse RL has been studied, I think the theoretical guarantees from this paper are novel and important.

Weaknesses

1. The authors have demonstrated performance on only 3 environments, in which in one of the cases, their proposed algorithm is not the best.

Questions

1. How is the state action visitation measure defined in Eq (5)? Specifically, I did not understand the term $P(S_t = s|s_0 \sim \eta)$. Does this imply that the transitions are action independent? 2. In the proof of Lemmma 1, how is the equation below the one labelled (iii) obtained? 3. How is the last step in Eq(30) in Proof of Lemma 1 obtained? 4. [Typo] Line 806 "bonded" -> "bounded". 5. I have understood most of the proofs as they are well written with adequate justification, but since I am not sure about the definition of the state-action visitation measure, I was unable to verify some proofs in the supplementary material.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

The authors address the limitations of this work and also some suggestions to overcome some of these limitations.

Summary

The paper "Understanding Expertise through Demonstrations: A Maximum Likelihood Framework for Offline Inverse Reinforcement Learning" proposes an innovative approach of offline inverse reinforcement learning. After a deep theoretical analysis of the inter-dependence between errors arising from dynamics modeling from limited offline data and performance gaps of resulting policies, authors propose an efficient algorithm to practically exploit conclusions for reward/policy learning. A small experimental section validates the approach.

Strengths

- Very well written with every notation well defined and every choice well justified - Very interesting problem and strong theoretical analysis - A practical algorithm that looks easy to reproduce - Good results

Weaknesses

- My main concern is that there is very few discussion about model uncertainty U in the paper, and particularly in section 3. I am surprised to not see it involved in the derivations and bounds, with no assumptions about it (except that it is bounded). No real meaning is given to it and it seems that it could be removed without changing anything in the theoretical conclusions. Is it true ? If yes, why introducing it in that part ? Also its impact is therefore not well understood from the theoretical analysis, which is a little be limitative to me (as it looks to have importance). - Still on U, I feel that experimental results on the choice of U would have been very useful.

Questions

- Could authors give more insights about U, both theoretically if possible, and by some experimental results that show its impact ?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

.

Thanks

Many thanks for these insightful answers, that well confortate the score I assigned to this submission.