Paper: Better Imitation Learning in Discounted Linear MDP (ICLR-2024)

Reviewer comment (Novelty):
The algorithmic designs and analysis techniques in this paper are not new. In terms of algorithmic designs, the main difference between ILARL, and existing IL algorithms OAL and OGAIL is the policy optimization step. However, the policy optimization algorithm in ILARL is highly similar to the one in [Sherman et al., 2023b]. For theoretical analysis, the key step to removing the persistent excitation assumption is the regret decomposition in Eq.(2), which reduces IL to online optimization with adversarial losses. However, such a regret decomposition has been presented in OAL. Furthermore, among the three types of errors, policy regret is the most difficult part to analyze. However, the analysis of the policy regret in Theorem 3 largely depends on existing techniques developed in OAL and [Sherman et al., 2023b].

Author response:
We agree that the techniques are similar to the existing works that you mentioned. In particular [Shani et al.,2021] and [Sherman et al., 2023b].  However, notice the following important differences.

1) The algorithms we present achieves a better regret bound than [Sherman et al., 2023b].We indeed exploit the fact that we have full information about the cost vector while [Sherman et al., 2023b] focuses on the bandit feedback. Moreover our regret bound also improves the result of [Zhong \& Zhang, 2023] for finite horizon adversarial MDPs.

2) We notice that the regularization in the policy improvement step allows to control the pseudo regret $ \sum^K_{k=1} < d^{\pi_k} - d^{\pi_E}, c_k >$ in the infinite horizon discounted setting.

3) In the finite horizon case, we obtain a further improvement from $\mathcal{O}(\epsilon^{-4})$ to $\mathcal{O}(\epsilon^{-2})$. This is possible only with the new insight that if we allow the cost player to update first, then, the $\pi$ player can know in advance the cost vector at the next round. We then, notice that LSVI-UCB originally designed for a fixed cost function still guarantees $\mathcal{O}(\sqrt{K})$ regret against a time changing cost vector sequence if the cost vector at the next round can be observed in advance.
This insight is missing in [Shani et al., 2021] that indeed updates the policy first and then the cost vector. This difference would make LSVI-UCB inapplicable. To fix the situation, the $\pi$ learner could use our Algorithm 1 but this would lead to a MDP trajectories bound of order $\mathcal{O}(\epsilon^{-4})$ which is clearly worst than the $\mathcal{O}(\epsilon^{-2})$ bound attained by our BRIG (Algorithm 4).

In addition, we highlight that the contribution of our submission could be very valuable for the member of the imitation learning community that might not know the power of the connection between imitation learning and online learning in adversarial MDPs.

Reviewer comment (Experiments):
The empirical evaluation is limited. This paper only considers a simple 2D environment. It is expected to verify the effectiveness of ILARL on more complicated tasks. Besides, this paper does not involve OGAIL for comparison.

Author response:
OGAIL is not implementable because it requires an integration oracle over the state space. See their definition of the features $\phi^h_k$ just after their equation 16 in this version https://proceedings.mlr.press/v162/liu22u/liu22u.pdf

Moreover, OGAIL is not designed for the discounted setting but rather for the finite horizon case.

*Regarding more complicated experiment*
Given that [Shani et al. 2021, Liu et al. 2022] and our current submission suggests the importance of exploration in imitation learning. Thus,  we are planning to apply exploration heuristics to the policy update in commonly used algorithms like OAL, GAIL and IQLearn and see if the performance is enhanced.
We think that the answer to this question is currently open. Indeed, in our understanding, the OAL authors provide several MuJoCo experiments but they do not try to implement approximately the optimistic bonuses that are needed to prove the theoretical guarantees.

At the same time, we think that this deep imitation learning extension is beyond the scope of the current submission.

Reviewer comment (Presentation):
Line 9 in Algorithm 3 is a little confusing. Algorithm 2 is a complete RL method that runs for K iterations while line 9 only corresponds to a one-iterate policy update.

Author response:
Thanks for catching this issue ! We clarified that the policy is updated performing only one iteration of Algorithm 2. We also provide a complete presentation in Algorithm 5.

Reviewer comment (Theory):
The paper shows that the learned policy achieve the similar performance as the expert policy. I am wondering if there is any guarantee on the recovery of the true cost function.

Author response:
**The paper shows that the learned policy achieve the similar performance as the expert policy. I am wondering if there is any guarantee on the recovery of the true cost function.**

Not really, we think that recovering the true cost function from a single set of expert demonstration is unfortunately not possible without additional assumptions. There are some recent works along these lines [1,2,3] that we could add in the related works section if the reviewer feels that it would be helpful.

[1] **Reward identification in inverse reinforcement learning**, K Kim et al. - International Conference on Machine Learning, 2021

[2] **Identifiability in inverse reinforcement learning** H Cao et al. - NeurIPS, 2021

[3] **Identifiability and generalizability from multiple experts in Inverse Reinforcement Learning** Rolland et al. - NeurIPS, 2022

Reviewer comment (Theory):
The linear MDP assumption is restrictive. The contribution of the paper can be more significant if it can be extended to general MDPs.

Author response:
**The linear MDP assumption is restrictive. The contribution of the paper can be more significant if it can be extended to general MDPs.**

We have added a new section in the revised manuscript at page 14,15 and 16 where we discuss how the sample complexity of order $\mathcal{O}(\epsilon^{-2})$ for the finite horizon setting can be attained for bilinear classes which is a more general MDP class.

However, the resulting algorithm is not guaranteed to be computationally efficient.
The focus of our paper was to consider a smaller MDP class but provide a computationally efficient algorithm for this setting.

Reviewer comment (Theory):
Although the paper claims that it studies linear MDPs, Assumptions 1-3 are considering the finite state-action case.

Author response:
**Although the paper claims that it studies linear MDPs, Assumptions 1-3 are considering the finite state-action case.**

Thanks for this comment. That's true. We are still considering finite states and action spaces but we assume that the number of states is too large to be enumerated. The exact same proof would go through for infinite states but it would obfuscate the presentation slightly so we preferred to focus on the finite state case.
On the other end, the action set needs to be finite otherwise we would not have computationally efficient policy updates. For example, the update in Step 20 of Algorithm 1 would require to compute an integral over the action space.

Reviewer comment (Experiments):
The empirical study is performed on a articrafted MDP rather than a real reinforcement learning environment.

Author response:
**The empirical study is performed on a articrafted MDP rather than a real reinforcement learning environment.**

We think that the take away of this theoretical work is that exploration in the policy optimization phase can be beneficial for imitation learning.

A natural follow up on the applied imitation learning side is to to try some exploration heuristics for neural networks and insert them in the policy update of common deep imitation learning algorithms like GAIL, IQLearn and see if this would empirically enhanance their performance.

We think that this approach has the potential to lead to convincing empirical performance in more challenging MDPs but at the same times it looks to us beyond the scope of the current submission.

Reviewer comment (Theory):
Is it possible to extend the result to general MDPs rather than simple linear MDPs?

Author response:
**Is it possible to extend the result to general MDPs rather than simple linear MDPs?**

As mentioned previously, please chack page 14,15 and 16 of the update manuscript for an informal discussion of the extension to bilinear classes.

Reviewer comment (Writing):
The paper does not keep consistent notations: $\mathcal{A}$ is used for action space in MDP setting and an algorithm in definition 1; Cost function is utilised in MDP setting, but the numerical experiments adopt reward function setting.

Author response:
**Cost function is utilised in MDP setting, but the numerical experiments adopt reward function setting.**

Thanks, we will soon update a revision which will make clear that in the experiment section's plot we consider the cumulative reward which equals the opposite of the cumulative cost. Sorry if this created confusion. We made this choice because using costs in the regret analysis is common in the adversarial MDP literature while comparing algorithms with the cumulative rewards is common in imitation learning papers.

Reviewer comment (Theory):
The norm inequalities in assumptions 1-2 seem very technical, and it would be better if the authors could provide some insights about them.

Author response:
**The norm inequalities in assumptions 1-2 seem very technical, and it would be better if the authors could provide some insights about them.**

Assuming bounded features is common in the Linear MDP setting see for example the paper by Jin et al., 2019.

The intuition is that since the quantities that we are assuming to lie in the column span of the matrix $\Phi$ are bounded (the cost and the transition probabilities are bounded by one). Therefore, it would not make sense to use potentially unbounded features to represent the cost and the transition dynamics as an inner product.

Reviewer comment (Theory):
The mathematical formulation for state value function in finite time horizon is pretty strange. I suppose the summation should take from 1 to $h$?

Author response:
**The mathematical formulation for state value function in finite time horizon is pretty strange. I suppose the summation should take from 1 to h?**

There is a small typo, the sum should be from $h$ to $H$ as standard. See for example equation 4.2.1 in Puterman's book.

Reviewer comment (Theory):
The infinite horizon trajectories, according to the description in section 2, have random length sampled from the geometric distribution. Why geometric distribution is adopted here? The sampled number is still finite, so the cost in the time horizons greater than the sampled number is set to zero?

Author response:
**The infinite horizon trajectories, according to the description in section 2, have random length sampled from the geometric distribution. Why geometric distribution is adopted here?**

The geometric distribution is the standard choice in Discounted Infinite Horizon MDP (see Chapter 6 in Puterman's book).

Reviewer comment (Theory):
In algorithm 3, line 6, the proposed algorithm project $w^{k+1}$ to the unit ball. How to ensure that the projected $w$ still constitute an adversarial costs in $[0, 1]$, as assumed in the MDP setting in section 2? Similar question happens to algorithm 4, line 7.

Author response:
**In algorithm 3, line 6, the proposed algorithm project $w^{k+1}$ to the unit ball. How to ensure that the projected $w$ still constitute an adversarial costs in [0, 1], as assumed in the MDP setting in section 2? Similar question happens to algorithm 4, line 7.**

This is ensured by Assumptions 1 and 2 respectively. In particular, we assume that each adversarial cost $c^k$ can be written as linear combination of the column span of $\Phi$ with weights $w^k$ such that $||w^k||_2 \leq 1$. In the imitation learning contest, projecting in the unit ball ensures this last bound on the weights.

Reviewer comment (Theory):
It seems that the proposed algorithms have never updated matrix $M$, in assumption 1 and 2. Does this mean that the true transition kernel is not estimated or involved in the algorithms?

Author response:
**It seems that the proposed algorithms have never updated matrix $M$, in assumption 1 and 2. Does this mean that the true transition kernel is not estimated or involved in the algorithms?**

Interesting question! Indeed, the matrix $M$ is never estimated by the Algorithm. The problem in doing this is that the matrix $M$ has dimensions $d \times S $. Therefore, estimating it would unavoidably lead to an algorithm which requires memory linear in $S$. To avoid such an issue our algorithm only tries to estimate the quantity $M V^k$ where $V^k$ are the value functions sequence produced by the Algorithm.

The advantage of this approach is that $M V^k$ is just a $d$ dimensional vector so our algorithm enjoys memory and running time which is independent of $S$.

Reviewer comment (Theory):
The matrix $\Phi$ is already known, according to assumption 1 or 2. But assumption 3 claims that the learner has access to $\Phi$. What is the difference between the matrix $\Phi$ in assumption 3 and $\Phi$ in assumption 1 and 2?

Author response:
**The matrix  Φ is already known, according to assumption 1 or 2. But assumption 3 claims that the learner has access to Φ. What is the difference between the matrix Φ in assumption 3 and Φ in assumption 1 and 2?**

The difference is in which cost functions we require to be linear in the features. In Assumption 1 and 2 we require the sequence of changing costs to be linear in the features. Assumption 3, which is needed in the imitation learning setting, we additionaly assume that the true unknown cost function is linear in the features.

Reviewer comment (Theory):
As stated in remark 1, the results in theorem 1 and 2 hold with high probability. So theorem 1 and 2 actually state that the trajectories numbers are independent of $\delta$?

Author response:
**As stated in remark 1, the results in theorem 1 and 2 hold with high probability. So theorem 1 and 2 actually state that the trajectories numbers are independent of δ ?**

The dependence on $\delta$ does not appear in Theorem 1 and 2 because we are stating the result in expectation. As highlighted in Remark 1, we prove this result passing first via a high probability result (Theorems 5 and 6) and then as shown in Lemma 6 we set $\delta = \epsilon$ to obtain the result given in Theorem 1 and 2.

Check Theorem 5 and 6 to appreciate how the number of required trajectory depend on $\delta$.

Reviewer comment (Presentation):
What is the y-axis in figure 1 and figure 2?

Author response:
**What is the y-axis in figure 1 and figure 2?**

It is the normalized cumulative return. We normalize it in a way that the expert performance averaged over $10$ seeds equals 1 and the one of the uniform policy equals 0.

Reviewer comment (Novelty):
The paper improves the prior state of the art in the best known sample complexity for IL in the discounted and finite horizon settings, and the technical novelty is moderate. The analysis is largely to improve the best known results for linear MDPs with adversarial costs in the full information setting. This is novel and may be of independent interest, but largely borrows insights from previous work, Viano et al (2022), the analysis of UCB in Jin et al and analysis of no-regret algorithms (MWE). The contribution feels somewhat like an A+B(+C) type result.

Author response:
We would like to start the rebuttal with 2 clarification. First, we would like to emphasize the fact that the improvement is not only on the dependence on the dimension or the accuracy $\epsilon$ but it is on avoiding the limiting persistent excitation assumption.
This is very important as we explain in our common answer. 

Second, the connection to adversarial MDP was not pointed out in Viano et al. 2022. It appeared in (Shani et al., 2021) for the finite horizon tabular setting and then it went somehow forgotten in the imitation learning community. In this work, we adapt it to finite and infinite horizon Linear MDP. New insights are required when compared to (Shani et al., 2021) to prove the improved bound in the finite horizon setting.


**The paper improves the prior state of the art in the best known sample complexity for IL in the discounted and finite horizon settings, and the technical novelty is moderate. The analysis is largely to improve the best known results for linear MDPs with adversarial costs in the full information setting. This is novel and may be of independent interest, but largely borrows insights from previous work, Viano et al (2022), the analysis of UCB in Jin et al and analysis of no-regret algorithms (MWE). I think it's still a nice contribution, but feels somewhat like an A+B(+C) type result.**

Please also appreciate the fact that there are some key differences when compared with the ideas used in the UCB analysis of Jin et al. . In particular, their analysis crucial relies on the fact that they can prove that for every $k$, and for every state action pair $s,a$ and for every stage $h$, it holds that $Q_h^k(s,a) > Q_h^\star(s,a)$ (upper bound because in Jin's paper they consider rewards instaed of costs) with high probability. In the infinite horizon, unfortunately we can not establish this property so we have to exploit a different form of optimism which is that the $Q$ values are lower bounds to the ideal value iterations updates. That is, $Q^k(s,a) \leq c^k(s,a) + \gamma P V^k(s,a)$ for every state action pair $s,a$ and round $k$ with high probability.

Reviewer comment (Presentation):
Overall, the writing of the paper is ok, but there is room for improvement in terms of presentation. The related work section can be organized in a much better way. This is important to put into context the results in the paper. There are several lines of work related to this one, which makes it all the more important to structure the related section more effectively.

Author response:
**Overall, the writing of the paper is ok, there is room for improvement in terms of the presentation. The related work section can be organized in a much better way. This is important to put into context the results in the paper. There are several lines of work related to this one, and so it's all the more important to structure the related section in a better way.**

We are happy to work to improve the writing further! Which lines of work do you feel we are missing ?

Reviewer comment (Experiments):
The experimental evaluation in the paper is very limited, and appears to be conducted on a very simple environment. While the theory in the paper is the major contribution, it would have been helpful to see a more comprehensive evaluation. Typically on harder environments, it is quite difficult to achieve the expert's performance, but this is not the case in any of the experiments in the paper. Although this point is not reducing the reviewer’s score, if the paper gets rejected, the author is encouraged to run more comprehensive experiments.

Author response:
**The experimental eval in the paper is very limited, and what seem to be on a very simple environment. While the theory in the paper is the major contribution, it would have been helpful to see a more comprehensive evaluation. I am not reducing my score for the paper because of this point, but if the paper gets rejected, I encourage the author to run more comprehensive experiments. Typically on harder environments, it is quite difficult to achieve the expert's performance, but this is not the case in any of the experiments in the paper.**

Thanks for the suggestion. We think that the take away of this theoretical work is that exploration in the policy optimization phase can be beneficial for imitation learning.

A natural follow up that we are planning is to to try some exploration heuristics for neural networks and insert them in the policy update of common deep imitation learning algorithms like GAIL, IQLearn and see if this would empirically enhanance their performance.

Reviewer comment (Theory):
Minor: "Therefore, the policy suboptimality scale as $H^4 \log |\Pi| / \epsilon^2$". Isn’t the policy suboptimality precisely $\epsilon$?

Author response:
**Therefore, the policy suboptimality scale as ...**

This is a typo, we meant that the required number of expert demonstrations scale as $\frac{H^4 \log \Pi}{\epsilon^2}$. We corrected this in our revision. Thanks for catching it !

Reviewer comment (Theory):
The standard BC reduction, for the finite horizon setting, argues that BC achieves a $O(H^2)$ suboptimality scaling in finite-horizon settings. This is in contrast to the discussion in the discounted setting on page 1 for BC. The reviewer is not aware of a work or analysis which states that BC requires $\widetilde{O} (1/(1-\gamma)^4)$ demonstrations in the discounted setting, and suggests citing a paper here.

Author response:
**The standard BC reduction, for the finite horizon setting argues that BC achieves a $O(H^2)$ suboptimality scaling in finite-horizon settings. This is in contrast to the discussion in the discounted setting on page 1 for BC. I am not aware of a work or analysis which states that BC requires $\mathcal{O}((1 - \gamma)^{-4})$ demonstrations in the discounted setting. It would be helpful to cite a paper here.**

Good question. This is a potentially confusing point. The common bound for BC says that the policy suboptimality error for the policy output denoted as $\pi_B$ is bounded as

$$
V^{\pi_{B}} - V^{\pi_E} \leq (1 - \gamma)^{-2} \mathbb{E}_{s \sim d^{\pi_E}} \sum_a | \pi_E(a|s) - \pi_B (a|s)|
$$

so here the scaling is $(1 - \gamma)^{-2}$ as you mentioned. However, for general experts, the dependence become quartic when we ask how many expert trajectories are necessary to achieve $V^{\pi_{B}} - V^{\pi_E} \leq \epsilon$.
In this case, we can use that by Hoeffding we have that with high probability 
$$
\mathbb{E}_{s \sim d^{\pi_E}} \sum_a | \pi_E(a|s) - \pi_B(a|s) |  \leq \sqrt{\frac{\log (\Pi/\delta)}{N_E}}
$$ 
where $N_E$ is the number of expert trajectories.
Therefore 

$$
V^{\pi_{B}} - V^{\pi_E} \leq (1 - \gamma)^{-2} \sqrt{\frac{\log (\Pi/\delta)}{N_E}}
$$

So, finally to have that the right hand side equal to $\epsilon$ we need $N_E \geq \frac{\log (\Pi/\delta)}{(1 - \gamma)^4\epsilon^2}$.
In this last bound the horizon dependence shows up with quartic dependence. The same result has been proven in (Li \& Zhang, 2022).

Reviewer comment (Theory):
It is unclear whether the results in this setting hold beyond linear MDPs, for example for bilinear classes or Bellman rank bounded MDPs. It would have been nice to include a discussion about this point.

Author response:
**Do the results in this setting hold beyond linear MDPs, say for bilinear classes or Bellman rank bounded MDPs? It would have been nice to include a discussion about this point.**

Yes, we think that we can extended the result to finite horizon bilinear classes ! However, the resulting algorithm can not be ensured to be computationally efficient. We added an informal discussion about this extension at page 14,15 and 16 of our revised draft.

The main focus of our paper is to consider the smaller class of linear MDP but give an efficient algorithm for it.