Distributionally Robust Imitation Learning

Abstract: We consider the imitation learning problem of learning a policy in a Markov Decision Process (MDP) setting where the reward function is not given, but demonstrations from experts are available. Although the goal of imitation learning is to learn a policy that produces behaviors nearly as good as the experts’ for a desired task, assumptions of consistent optimality for demonstrated behaviors are often violated in practice. Finding a policy that is distributionally robust against noisy demonstrations based on an adversarial construction potentially solves this problem by avoiding optimistic generalizations of the demonstrated data. This paper studies Distributionally Robust Imitation Learning (DRoIL) and establishes a close connection between DRoIL and Maximum Entropy Inverse Reinforcement Learning. We show that DRoIL can be seen as a framework that maximizes a generalized concept of entropy. We develop a novel approach to transform the objective function into a convex optimization problem over a polynomial number of variables for a class of loss functions that are additive over state and action spaces. Our approach lets us optimize both stationary and non-stationary policies and, unlike prevalent previous methods, it does not require repeatedly solving an inner reinforcement learning problem. We experimentally show the significant benefits of DRoIL’s new optimization method on synthetic data and a highway driving environment.


Reviews and author responses:

 Thank you for the clarifications. I stick to my initial assessment.

This paper studies Distributionally Robust Imitation Learning (DIRL) and reveals that Maximum Entropy (MaxEnt) can be a special case of DIRL. In the case of an "additive" loss function over state and actions, it also proposes a new algorithm to solve the adversarial training problem by using formulating the problem as a convex optimization problem over the space of marginal state-action distributions. Experiments show that the proposed algorithm can learn a good policy with a similar performance to the baseline, but much faster. Another set of experiments show that the proposed algorithm can find a better reward weight vector than the baseline algorithms, 
  The proposed algorithm is interesting but is limited to additive loss functions, discrete state/action spaces and requires a known transition model and a feature extractor. The connection between DIRL and MaxEnd seems to be quite simple and intuitive. The experiments seem good but I expect more baselines. The writing can be improved a lot. Notations are not consistent and sometimes not defined. 

Preliminaries: Do you assume $\phi$ is known and $\textbf{w}$ is unknown? 
Definition of $\mu(\pi)$: What is $k$? Is it $d$? 

Eq.1: What's the expectation over? The same question goes to $H(P)$ in L.136. 

Lemma.1: What are $\mathcal{T}, C, \tilde\Xi, \alpha_0, \alpha, p^*, q^*$? Does $q$ really mean the set of all mass probability functions? Are $\alpha_0, \alpha$ scalars or vectors? Do $\alpha$ or $\alpha_0$ depend on $C$ to guarantee that $p^* \in \Xi$ (or do we need it)? 

Theorem.1: What does "logarithmic loss" mean? Is it the cross-entropy loss like in Lemma.1? Please give an explicit form, especially if it's not too complicated. 
Theorem.1 proof: What's $\textbf{w}_0$? Is it a scalar or a vector? If it's a scalar, please use a consistent font style to avoid confusion. 

Lemma.2: Please put the assumptions into the lemma description instead of the proof. 

L.158: asa => as a
L.159: $\mathcal{L}(\hat \delta_i, \check \delta_j)$. 

Sec 4.2: What is an additive loss function? Is the logarithmic loss function additive? 

Eq.3, 4, 7, 8: Please use $\mathbf{\Omega}$ or $\Omega$ consistently. 

Sec 5: Is the transition matrix $\Gamma$ known? 

Experiments: I'm looking for more baselines. The most recent baseline used has been published for more than 10 years. I'm not an expert in this area so I don't know what's the recent progress and it might be possible that no directly comparable algorithms were developed in the last 10 years. 
1. As this paper uses an environment from "Levine et. al [16]", can the method in that paper also be compared? 
2. How does MaxEnt perform? 
3. How does Double Oracle perform? 
It's a little bit weird if we study the relationship between DRIL and MaxEnt but don't compare them in experiments. 

---

After reading the response from the authors, I increase my score from 4 to 5. But I think both the theoretical contribution is minor as the technical difficulty is not too much, and the experiments are limited due to strong assumptions. 

 I feel like the assumption that the transition model is known is very strong. I didn't see an obvious way to run the proposed algorithm in continuous state/action space, either. 

The paper tackles the problem of imitation learning in an environment with a discrete set of states and actions, known dynamics, and an (unknown) reward function that is linear with respect to known state-action features. The key idea is to formulate imitation as a two-player zero-sum game between the agent and demonstrator, and to efficiently compute the agent’s minimax strategy by solving a convex optimization problem. This approach, called Distributionally Robust Imitation Learning (DRIL), generalizes maximum-entropy inverse reinforcement learning to non-logarithmic losses. Simulation experiments in simple grid-world navigation and highway driving environments show that DRIL has a lower running time than comparable prior methods, can be used with a variety of different additive loss functions, and learns more effectively from noisy demonstrations than prior methods.  Overall, I enjoyed reading this paper. I look forward to the future work on extending DRIL to high-dimensional state and action spaces discussed in lines 388-391!

One issue is that the experiments do not compare DRIL to recent imitation learning methods that aren't derived from MaxEnt IRL or adversarial games, but still attempt to perform state distribution matching:
 - RCE https://arxiv.org/abs/2103.12656
 - RED https://arxiv.org/abs/1905.06750
 - DRIL https://openreview.net/pdf?id=rkgbYyHtwB 
 - VDI https://arxiv.org/abs/2002.06473
 - IQ-Learn https://arxiv.org/abs/2106.12142 (this paper was released after the NeurIPS submission deadline, so I don't expect the authors to have somehow compared to it in advance, but I want to mention it anyway because it's relevant)

It would be nice to add comparisons to some subset of these methods -- not necessarily to show that the proposed DRIL method outperforms the prior methods, but primarily to show that the reason that the proposed DRIL method performs better than MMP, LEARCH, and the projection algorithm of Abbeel and Ng is, in fact, distributional robustness.

A few minor issues, typos, and questions:
 - The name “DRIL” conflicts with another DRIL from prior work: disagreement-regularized imitation learning https://openreview.net/pdf?id=rkgbYyHtwB
 - Line 173: do you mean exponential in the size of the state space?
 - Line 209: “construing” → constructing
 - Equation 5: should “p_0(s)” be p_0(s’)?
 - Line 241: why are the dual variables w also the reward weights w?
 - Line 294: “We run the both algorithms until convergences,” → “We run both algorithms until convergence.”
 - Why should the loss L ever be non-diagonal? It seems like you would only ever want to measure the difference between the agent’s occupancy of (s, a) and the demonstrator’s occupancy of the same state-action pair (s, a), rather than the agent’s occupancy of (s, a) and the demonstrator’s occupancy of some different state-action pair (s’, a’).

Update after reading author response
-----
Thank you for your response. I will keep my original score.
 Section 6 adequately addresses limitations and societal impacts. See the main review for suggestions on how to improve the writing.

 We thank you for your careful reading and detailed review, we address your questions and concerns in the following:

Preliminary: Yes, we assume $\phi$ is known and $w$ is unknown. To have a consistent notation, we will use $d$ as the feature dimension. Thanks for pointing this out. 

Lemma 1: $\mathcal{T}$ is a real-valued vector statistic, $C$ is represent a constant vector, $\tilde{\Xi}$ represents a set of distribution with empirical value of $\mathcal{T}$, $\alpha_0$ is a scalar and $\alpha$ is a vector. $\alpha$ and $\alpha_0$ depends on $C$. In the context of imitation learning, when the empirical feature vector is obtained from the demonstrations, there exist a $p^*$ with some $\alpha$ and $\alpha_0$.  For more information and better understating of Lemma 1 refer to [12]. 

As pointed out by another reviewer, q is a single distribution that is chosen from the set of all “possible” mass functions; we will clarify the definition of q. $p^*$ and $q^*$ represents the distributions at the equilibrium. 

In line 136. $p$ and $q$ are any (general) distribution, in our case, depending on the definition of a policy, they are distributions over policies.

Theorem 1: yes, logarithmic loss is cross entropy. $\mathbf{w}_0$ is a scalar, we will correct that. 

Lemma 2: Thanks for the suggestion, we agree that putting assumptions in the lemma will make the writing more clear. 

Additive loss function is defined over different states and actions. Consider the state-action distribution of a policy, an additive loss can be written as the sum of loss at each state and action.  Logarithmic loss is not additive. 

As you mentioned in your review, we assume the transition matrix $\Gamma$ is known. 

Experiment: Yes. We use "Levine et. al [16]" code and the environments that they provide. However, it will not be reasonable to compare our method with GPIRL, as our method similar to other baselines, considers a reward function that is linear in feature space.  GPIRL in "Levine et. al [16]" consider a non-linear relation and they aim to show the benefit of non-linear features over linear feature which is out of the scope of our work. 

The performance of MaxEnt compared to our proposed method depends on the choice of loss (the input of our algorithm). Any loss specific imitation learning problem will benefit from our method as take any additive loss function as an input. Our proposed method can potentially outperform MaxEnt with the right choice of loss function. Osa et al, in “An algorithmic perspective on imitation learning” discuss the role of different loss functions in imitation learning.

The policy performance of the Double Oracle method is similar to our proposed marginal method as they both solve the same objective function. There could be some small difference as double-oracle uses some level of approximation to represent an equilibrium using a sparse number of deterministic policies. However, as shown in the first experiment, the performance of obtained policies are almost identical. The main difference is in training time and practicality of our method which is shown in the experiments.


 We thank you for your careful reading and detailed review, we address your questions and concerns in the following: 

Weight vector: For more clarification, we will add some description of weight vector before theorem 1.
 
Marginal Distribution: Yes, the constraints (projection step) can be seen, loosely, as the inner RL optimization. However, the proposed projection is quadratic with an analytical O(n log n) solution. We can further discuss this on the paper. 

Additional Environment: Correct, we learn the policy on the source domain and apply it on additional random environments. 

 We thank you for your careful reading and detailed review, we address your questions and concerns in the following: 


Experiments: Thanks for suggesting the relevant works to compare our method with. We will take them into consideration to compare these methods with DRIL. Since our work is one of the pioneer works in using distributional robustness in imitation learning, we will try to include as many comparisons with other baselines as possible in this work.

DRIL: Thanks for pointing out the name conflict, we will attempt to find an alternative abbreviation. 

Dual variables: As explained in line 241-244, dual variables can be seen as reward weight vector. 

Non-diagonal loss: non-diagonal elements of loss function force two different states and actions of the learner and demonstrator to have relatively similar occupancy which can be used in designing a particular loss of interest.  


 We thank you for your careful reading and detailed review, we address your questions and concerns in the following: 

Discrete State-action Space: For more clarity, we will emphasize that our work considers discrete state - action spaces. 

Expectation: We consider a general definition that potentially can deal with: 1. Distribution over policies 2. The stochastic “behavior” of a policy (e.g. state distribution of a stochastic policy). 

Demonstrator’s margin on matching behavior: Thanks for pointing out an interesting direction, we implemented this notion using a $l_2$ regularizer term on the demonstrator. For some values of the regularizer coefficient, a better policy can be obtained, however the improvement is not significant. However, we believe further research is needed for better understanding. 

Lemma 1:  Yes, q is a single distribution that is chosen from the set of all “possible” mass functions; we will clarify the definition of q. 

Loss function: As an example, Ratliff et al [30] use different heuristics to obtain a better policy. Osa et al, in “An algorithmic perspective on imitation learning” discuss the role of different loss functions in imitation learning. Our main objective was to show how the choice of different loss functions would affect the quality of a produced policy. However, we also believe adding more baselines will benefit our work and result in a better comparison with other methods. 

Further discussion on experimental results: we agree that more discussion will benefit our work: One simple reason can be if the ratio of suboptimal decisions varies in each training data batch, the performance of learned policy also varies. In our particular case, we sample a new set of data for each size of training data. Since we sample demonstrations from a stochastic policy, the variance in the quality of demonstrations might vary and relying too much on the demonstration data (as opposed to our robust method) might result in poor performance. Also as the number of demonstrations increases, the variance in obtained features reduces. So the better performance of  the methods when the number of samples is smaller could be from the inaccuracy of estimated future expectation, and as the feature expectation gets closer to the true feature expectation of the demonstrated policy, their performance declines. 


The paper extends on a previous method [5] for distributional robust inverse reinforcement learning. Both works learn a reward function, the agent's policy, and an adversarial policy that aims to incur maximum loss (w.r.t. a given loss function) on the agent while being constrained to matching given features in expectation.

The contributions of the submission are
* showing that MaxEnt-IRL is a special case of that formulation (using the cross-entropy between adversary and agent as loss)
* presenting a more efficient optimization procedure compared to the previous ("double oracle"-approach)  Originality
---------------
I think that both above-listed contributions are novel and significant enough to justify publication.

Quality
-----------
The derivations seem correct, but I did not check them thoroughly. Using distributional robustness in an IRL setting makes a lot of sense and the approach seems sound.

Clarity
----------
The presentation is mostly clear. 

* Reading the reference [5] helped me a lot, because it made it a bit more clear that the weight vector corresponds to the Lagrangian of the feature matching constraint (as in MaxEnt). It would be helpful to add this somewhere before Theorem 1
 
* It was not immediately clear to me that the expectation of the loss was computed w.r.t. p, e.g. for the log-loss in line 124. Its easy to guess, but it would be better to indicate it in the equations.

* The benefit of changing to marginal distribution also needs to be made clearer. The paper mainly talks about how linear and nice everything becomes, but from what I understand the constraint that the marginal distributions must be realizable does not make the problem easier. For example, the optimization for Eq. 4 still requires the inner reinforcement learning. The need for inner RL only seems to be removed in Section 5, when by expressing the weight in terms of the distributions, the order of optimizing the policy and the reward get switched. I think at this point the paper should draw the connection to GAIL which also achieves faster imitation learning (without explicitly learning the rewards) by moving from $\text{IRL} \circ \text{RL}$ to $\text{RL} \circ \text{IRL}$.

* Line 336 states "We also evaluated each method on the original environment and on 4 additional random environments" on a transfer setting. I assume you evaluated the methods by evaluating the _policies_ that have been learned on the source domain on additional random environments? The paper should be more explicit here.

Significance
-----------------
I think that the new insights might be valuable for some researchers, in particular, because I think that distributional robustness is a logical and sound approach to deal with the non-optimality of the expert and the ill-posedness of IRL. The current work is somewhat limited (which is also mentioned in the paper) because it is restricted to simple problems, such as Gridworlds. Correspondingly, the comparisons with LEARCH and MMP seem a bit outdated. I do not think that this a major problems, since the current contribution can still proof useful in the future, but it diminishes the significance nevertheless. Limitations and societal impact was sufficiently addressed.

For discrete low-dimensional Markov decision processes, the paper proposed a distributionally robust imitation learning framework to learn a policy with respect to the worst-case demonstrator (subject to feature matching constraints from data). The maximum generalized entropy imitation learning framework naturally includes the predominant maximum entropy IRL. Compared to previous double oracle approach, the contribution of this paper is then transforming the original minimax problem over policies (in Definition 1) to optimization over occupancy measures (state-action distributions), such that only a polynomial number of variables need to be solved instead of exponential. The policies are later obtained based on the link between occupancy measures and stochastic policies. Besides, the flexibility of generalized entropy allows us to incorporate more domain knowledge into the design of loss function, although this point can be better demonstrated with experiments.  The paper is trying to address an important problem in imitation learning, that expert demonstrations may not be consistently optimal (suboptimal demonstrations). The proposed approach is based on finding a policy that matches the worst-case demonstrator, thus avoiding optimistic generalization. However, since the worst-case demonstrator is still under the feature matching constraint, it seems if the expert demonstrations is not near optimal, the solution will still be quite suboptimal, even with the distributionally robust framework. Some recent works attempts to address this problem by assuming access to extra information such as ranking/scoring associated with the demos, which seems to be more plausible to address this challenging scenario.

That being said, I think incorporating the notion of generalized entropy from Grünwald P D & Dawid A P.'s paper to formulate the robust imitation learning problem is interesting. It is natural that the maximum entropy IRL framework is a special case of DRIL since maximum entropy is a special case of generalized entropy (e.g. in Grünwald P D & Dawid A P.'s paper). From this perspective, I also think the traditional maximum generalized entropy notion (Grünwald P D & Dawid A P.'s paper) should be discussed in more details as preliminaries to provide background and allow readers understand the contributions properly.

Below are some detailed comments/questions:
- When introducing state, action spaces in preliminaries, please explicitly say these are (low-dimensional) discrete spaces. discrete MDP is unclear (might be discrete/continuous time).
- In definition 1, eq 1, what distributions is the expectation taken over? Just based on the definition, I don't see any other random variables and the loss already takes both learner and demonstrator as input.
- In distributionally robust optimization, we usually find an adversary within some lp-ball around some data, hence tolerating some small error/differences around demos. Do you think this notion is necessary/can be implemented: in the minimax problem, the adversary (worst-case) demonstrator need to match the feature statistics of the data up to some small error (rather than perfect feature matching).
- I think Eq 2 is true only for deterministic dynamics.
- In lemma 1, q is not "the set of all mass probability functions", but a single mass function?
- In line 135, need to add ref to the generalized entropy function.
- Section 6.2 is trying to demonstrate that by smartly choosing loss functions, we can incorporate proper domain knowledge to the definition of generalized entropy, which is also an important question to think about when using generalized entropy in practice. However in the provided example, the best euclidean distance seems to be already the default loss choices, which is not a surprising example. Any other examples where using a smart choice of loss function leads to better performance than an obvious/default choice? Also, the baselines such as random or zero-one seem to be too weak, and the popular logarithm loss in MaxEnt IRL is missing?
- Some experimental details may be better discussed perhaps in the Appendix. For example, I'm not clear why in the setting of Fig 3, more examples lead to worse performance for the baselines. The negative societal impact is properly discussed. For comments on technical contents, see main review section.

Meta-review: This paper is somewhat borderline. The paper shows an interesting connection between distributional robustness and MaxEnt IRL, showing that MaxEnt IRL is a special case of distributionally robust IL. Using DRO in IRL is interesting & underexplored. The paper also has a minor contribution of making distributionally robust IL more efficient. The main downsides of the paper are as follows. First, the approach is limited to discrete state and action spaces, severely limiting its practical relevance. As a result, the experimental comparisons are also outdated. Second, the connection between DRO and MaxEnt IRL has limited technical difficulty since it simply invokes a lemma about exponential families (though simplicity can also be viewed as a good thing).

Overall, this is a scenario where the paper is correct and generally well-executed but has potentially limited expected impact and limited technical difficulty. Given that impact is subjective and at least half of the reviewers found the ideas to be interesting, it seems like others in the community would also get value out of the paper.

Decision: accept