Paper: DITTO: Offline Imitation Learning with World Models (ICLR-2024)

Reviewer comment (Novelty):
**Novelty**: the approach novelty is limited, as it combines the idea of performing imitation learning through some form of state matching/inverse reinforcement learning, which is not novel [1], with the idea of using world models to learn behaviour in imagination for high-dimensional inputs, which has become more and more popular in the last few years.

Author response:
It's true that our approach is related to the state divergence minimization literature, as we discuss in the paper. Our contribution is to provide a simple and performant instantiation of divergence minimization imitation learning, with the twist that we can do it through simple on-policy RL in a learned world model.

Reviewer comment (Presentation):
The presentation could be improved in several ways. There are some minor typos (see Questions section of the review), Figure 1 is not very informative (could show better how the world model and the expert trajectories are employed), Figure 3 has an unusual layout with two plots being much larger than the others, and the related work contains equations that are not completely relevant (e.g. Equation 4).

Author response:
Thank you for pointing out the writing issues, we will address these.

Reviewer comment (Evaluation):
In order for the work to be accepted, I think the evaluation should be improved and better reflect the motivating scenario. I suggest the authors complete the story with more experiments and baselines. I also recommend the authors look into more adequate metrics of comparison that better support the motivation of this work, e.g. that show that DITTO reduces covariance shift (also see comments in Weaknesses).

Author response:
RE: Appropriate metrics of *covariate* shift. This is why we included the latent distance plots, to point out how typical metrics such as action prediction accuracy fail to be predictive of the expected environment return, whereas our latent distance metric is predictive.

Reviewer comment (Writing):
There is a repetition in the abstract: `in the latent space of a learned world model` appears twice in two consecutive sentences.

Author response:
Thank you for pointing out the writing issues, we will address these.

Reviewer comment (Writing):
Typo in the introduction: `the both the strong`.

Author response:
Thank you for pointing out the writing issues, we will address these.

Reviewer comment (Writing):
In Figure 1 caption: `the intrinsic reward 8`.

Author response:
Thank you for pointing out the writing issues, we will address these.

Reviewer comment (Evaluation):
If the idea was to compare only against imitation learning approaches, different metrics to evaluate covariate shift should be adopted. I think the latent distance metrics is useful to provide additional insights, but cannot be used as a general tool to compare different approaches, as it is strongly dependent on the world model.

Author response:
RE: Appropriate metrics of *covariate* shift. This is why we included the latent distance plots, to point out how typical metrics such as action prediction accuracy fail to be predictive of the expected environment return, whereas our latent distance metric is predictive.

Reviewer comment (Theory):
(major) This paper does not have a clear and friendly structure for the readers. Many of the key formulas (e.g., ELBO objective) are placed in the Appendix, while the authors spend many spaces on related work part. The authors state that they theoretically show that their formulation induces a divergence bound between the expert and the learner. However, no formal theorem/lemma/proposition is presented in the paper. The core theoretical result in Eq 6 is based on the previous conclusion. Also, what is the purpose and the role of Appendix A? What conclusion can be derived based on corollary A.1?

Author response:
RE: "Many of the key formulas (e.g., ELBO objective) are placed in the Appendix, while the authors spend many spaces on related work part" We placed these details in the appendix because they are not novel, the world model is exactly that introduced in the Dreamerv2 paper.

RE: "However, no formal theorem/lemma/proposition is presented in the paper." This is incorrect. Appendix A gives a proof of the return bound.

Reviewer comment (Experiments):
(major) The evaluations are limited. The authors only conduct experiments on 5 Atari tasks and 1 continuous control task. I do not think these are sufficient to show the generality and the effectiveness of their method. More experiments on both the discrete control domain and the continuous control domain can definitely make this paper stronger.

Author response:
RE: "The evaluations are limited. The authors only conduct experiments on 5 Atari tasks and 1 continuous control task." Frankly, most imitation learning papers report results on 3-4 low-dimensional state-based environments. We are surprised by this pushback, we feel our evaluation is comprehensive and in more difficult environments than many related works.

Reviewer comment (Experiments):
(major) Potential issues with the baselines. The authors only compare their proposed DITTO method against some comparatively weak baselines like BC, D-GAIL, etc. As far as the reviewer can tell, there are some imitation learning methods that can achieve quite good performance on the Atari games, e.g., IQ-learn [1]. So, how does DITTO compete against it? Another recent offline imitation learning algorithm OTR [2] computes rewards using the optimal transport, so how does DITTO compete against OTR? Intuitively, we can choose not to use the world model and instead the optimal transport to calculate the rewards for the downstream RL algorithm. It is unclear whether DITTO can beat OTR in this scenario.

Author response:
RE: IQ-learn and OTR. Thanks for pointing these out. It appears IQ-learn uses online interactions in the Atari environments? They are only fully offline in the state-based environments, if we've understood their paper correctly? It's difficult to directly compare since the experts we imitate have different performance profiles across the environments, but their results indeed look comparable, except for the online interaction part. The OTR paper is also interesting, but it appears they don't test in visual environments, only state-based. Indeed, the difficulty is getting a useful measure of state in visual domains to do some kind of divergence minimization or optimal transport approach. It would be interesting to run an OT approach on the world-model latent space! It seems people think visual imitation learning is trivial, but if you look at many of these papers, they *do not* test in pixel-based environments.

Reviewer comment (Experiments):
(major) It is unclear how different choices of the base algorithm affect the learning efficiency and the final performance. The authors only adopt simple actor-critic architecture with the actor being updated using the REINFORCE algorithm and the critic being updated with the $\lambda$-return. A valuable and necessary study is, how different base algorithms affect the performance of DITTO.

Author response:
RE: Lambda return and inner-RL details: These are exactly the Dreamer policy learning procedure, there's nothing exotic or particular here. Indeed we imagine that using more modern approaches such as PPO for the inner RL would probably result in a stronger agent. But we are focusing here on the conceptual problem of doing imitation using world models, not the RL details in the inner loop. The fact that DITTO is this strong without an optimized inner RL loop is probably a good sign.

Reviewer comment (Experiments):
(minor) The author state in Section 2.1 that  
> These methods are effective when low-dimensional proprioceptive states and actions are available but have not yet demonstrated strong performance in high-dimensional observation environments  
I do not think it is difficult to extend these algorithms into the pixel-based variants, do you think these methods can have poor performance in the pixel-based tasks, given that we introduce an encoder to these algorithms?

Author response:
RE: "I do not think it is difficult to extend these algorithms into the pixel-based variants, do you think these methods can have poor performance in the pixel-based tasks, given that we introduce an encoder to these algorithms?" Well, yes. If you look at SQIL and DRIL, you'll see that adapting standard algorithms like GAIL to the visual setting results in terrible performance. Pixel-based imitation learning is understudied. We hope our work helps to highlight this area of research.

Reviewer comment (Experiments):
If possible, please include some stronger and more recent baselines.

Author response:
RE: IQ-learn and OTR. Thanks for pointing these out. It appears IQ-learn uses online interactions in the Atari environments? They are only fully offline in the state-based environments, if we've understood their paper correctly? It's difficult to directly compare since the experts we imitate have different performance profiles across the environments, but their results indeed look comparable, except for the online interaction part. The OTR paper is also interesting, but it appears they don't test in visual environments, only state-based. Indeed, the difficulty is getting a useful measure of state in visual domains to do some kind of divergence minimization or optimal transport approach. It would be interesting to run an OT approach on the world-model latent space! It seems people think visual imitation learning is trivial, but if you look at many of these papers, they *do not* test in pixel-based environments.

Reviewer comment (Novelty):
The literature review is inadequate. For example, there lacks a discussion about works within offline imitation learning domain, e.g., [2,3].

Author response:
Regarding the literature review being inadequate: your mentioned [2] and [3] evaluate on 3, and 4 low-dimensional state-based environments, respectively. We evaluate on 6 pixel-based environments. There are indeed many related works in the imitation learning literature, but to maintain focus on the hard aspects of partial observability which we would like to tackle, we do not consider works which have so far only demonstrated success in state-based environments. Instead, we adapt mature methods like BC and GAIL to the world model setting. If you compare our adapted baselines to recent methods, you'll find they perform competitively.

Reviewer comment (Presentation):
The paper is not well-organized. Some important details are missed in the main body, e.g., the training objective of the world model.

Author response:
RE: "The paper is not well-organized. Some important details are missed in the main body, e.g., the training objective of the world model." Since these details are exactly identical to prior works (adapted straight from Dreamer), we do not consider them significant contributions and placed them in the Appendix.

Reviewer comment (Experiments):
Some offline model-based RL baselines are missed [4]. Besides, some IL baselines that handles high dimensional state should be considered, such as AIRL [5] and DICE [6].

Author response:
RE: "Some offline model-based RL baselines are missed" We are not doing offline RL. There are no rewards in our trajectories, so these methods simply do not apply.  

The DICE [1] paper evaluates on 4 state-based environments. Again, we evaluate on 6 pixel-based environments, and do not compare with that method since it hasn't been demonstrated in this setting. AIRL is essentially a variant of GAIL, and evaluates again only on low-dimensional state-based environments.

Reviewer comment (Experiments):
The paper does not provide any ablation studies or experiments to evaluate the impact of different components or design choices of DITTO. For example, how does changing the reward function, the distance measure, or the horizon length affect the results? How does DITTO compare with other model-based reinforcement learning methods that do not use imitation learning? How does DITTO perform on different types of environments, such as continuous control or navigation tasks?

Author response:
RE: "How does DITTO compare with other model-based reinforcement learning methods that do not use imitation learning?" DITTO does not address the same problem as MBRL methods.

RE: "How does DITTO perform on different types of environments, such as continuous control or navigation tasks?" We included results on a continuous control task in the appendix - bipedal walked from pixels alone. Please take a look!

Reviewer comment (Novelty):
The novelty of the method seems limited, which looks like a simple combination of Dreamer and GAIL / BC.

Author response:
RE: "The novelty of the method seems limited, which looks like a simple combination of Dreamer and GAIL / BC." How so? GAIL uses a learned, adversarial setup for policy learning. BC simply clones actions. We discuss why both of these methods fail to produce robust imitation policies, and introduce a simple, novel intrinsic reward which can be optimized with off-the-shelf on-policy RL methods. We demonstrate that out method outperforms world-model augmented versions of GAIL and BC.

Reviewer comment (Experiments):
The experiment setup is not convincing, it is too simple and all results are only evaluated on a small set of simple tasks.

Author response:
RE: "The experiment setup is not convincing, it is too simple and all results are only evaluated on a small set of simple tasks" These are standard environments for visual imitation learning. Please see e.g. [1][2], and compare their baseline results with ours. Most imitation learning papers evaluate on much simpler, state-based tasks.

Reviewer comment (Theory):
It is a little confusing that you prove a bound for the reward function given in Eq. (7) but actually use Eq. (8). The reason stated in the paper is that Eq. (7) is computationally expensive, but there is no clear relation between the formulations of Eq. (7) and Eq. (8). Could you please explain how you design the distance function in Eq. (8)?

Author response:
RE: "there is no clear relation between the formulations of Eq. (7) and Eq. (8)" As we stated in the paper, the result holds for any distance function. We chose the particular form of equation 8 based on its empirical performance. The theoretical result does not depend on the details of the distance function. Please see the Appendix and the cited related work for details.

Reviewer comment (Evaluation):
In Fig. 2, the latent distance is defined by $1-r_{int}$, where $r_{int}$ is also defined by yourself. It might be unfair to compare this indicator. What are the results of more common distance measures? Or could you give reasons why this measure is better than other measures?

Author response:
RE: "In Fig. 2, the latent distance is defined by 1−rint, where rint is also defined by yourself. It might be unfair to compare this indicator. What are the results of more common distance measures? Or could you give reasons why this measure is better than other measures?" The point here was to demonstrate that the latent distance measure we introduce is more predictive of final performance, compared to more common metrics of imitation performance, e.g. action prediction accuracy. We expect the on-policy latent distance to be more predictive of real world performance for reasons laid out in the paper, and we indeed find that to be the case.

Reviewer comment (Evaluation):
The submission does not support the claim that it addresses the problem of covariate shift with sufficient evidence.

Author response:
Equation 5 establishes the relationship between the expert and learner return gap, and bounds it by their state distribution divergence. Our method bounds this state distribution divergence in the latent space of a learned world model. We give a proof that this bounds the return difference. The learner is penalized by a novel intrinsic reward which we introduce to minimize the covariate shift during training episodes - this is the central idea of the algorithm.

Reviewer comment (Experiments):
How is the covariate shift problem embodied in the offline imitation learning problem? More empirical studies may be needed to verify the conclusion of the submission.

Author response:
Equation 5 establishes the relationship between the expert and learner return gap, and bounds it by their state distribution divergence. Our method bounds this state distribution divergence in the latent space of a learned world model. We give a proof that this bounds the return difference. The learner is penalized by a novel intrinsic reward which we introduce to minimize the covariate shift during training episodes - this is the central idea of the algorithm.

Reviewer comment (Evaluation):
In Figure 3, why does DITTO exceed average expert performance in 3/5 environments, and is this desirable in offline imitation learning?

Author response:
We give theoretical results for the exact distribution matching form of the DITTO algorithm, but for computational cost reasons approximate it to match the nearest latent from the initial expert trajectory, not from any expert trajectory. This biases the algorithm towards the mode of the distribution. We speculate that this is beneficial in certain environments (i.e. that the mode expert behavior outperforms the full expert behavior), as discussed in [1].  

[1]: F. Sasaki and R. Yamashina. Behavioral cloning from noisy demonstrations. In ICLR, 2021.

Reviewer comment (Theory):
I am in wonder why the proposed method can handle the distribution shift issue in imitation learning. Please refer to Q1 below.

Author response:
We are not making any claims about the bias or generalization capabilities of the learned model itself. The point is that the return difference between expert and learner is bounded by the divergence of their induced distributions. Our method is about getting an estimate of the learner's distribution in the model, and explicitly penalizing the divergence from the expert over multiple time-steps, to address compounding error problems. We show that the world model latent space provides an excellent emasure of state distribution divergence. We show that minimizing this divergence using on-policy RL results in state-of-the-art performance in the offline setting (competitive with online methods)

Reviewer comment (Theory):
The theoretical discussions in Sec. 3.1 are confusing to me. Please refer to Q2 below.

Author response:
Equations 5 and 6 are related in the following way: equation 5 gives a bound on the return difference in terms of the state distribution divergence. Equation 6 shows how the true state distribution divergence is bounded by the divergence in the model. Therefore, bounding the divergence in the model bounds the true state divergence, thereby bounding the return difference.

Reviewer comment (Experiments):
The experimental results are not quite sufficient. There are several Atari games chosen in the experiments, but the reason for choosing them remain unclear. Why not report the results on more games? Furthermore, it is difficult to understand what kind of internal feature space that has been learned by the model. In my view, if learning the world model can address the distribution shift issue, one plausible reason is that the learned model feature space has strong generalization ability. While this remains unclear to me due to the missing of the analysis on the learned feature space.

Author response:
RE: "There are several Atari games chosen in the experiments, but the reason for choosing them remain unclear." As we stated in the paper, these are the benchmarks used in [1], a well-known SOTA method in the community, as well as other top-performing imitation learning methods, e.g. [2]. In addition to those Atari experiments, we also gave results in a pixel-based continuous control environment. Please see the appendix for those results.

RE: "it is difficult to understand what kind of internal feature space that has been learned by the model [...] one plausible reason is that the learned model feature space has strong generalization ability." Indeed - our results support this view.

[1]: Kiante Brantley, Wen Sun, and Mikael Henaff. Disagreement-regularized imitation learning. In ICLR, 2020.

[2]: Siddharth Reddy, Anca D. Dragan, and Sergey Levine. Sqil: Imitation learning via reinforcement learning with sparse rewards. In ICLR, 2020.

Reviewer comment (Theory):
As described in the paper, one of the central challenges for imitation learning is that the learner can face significant different data distribution to expert's, since her policy can be very different from which the expert uses. While if we use expert demonstrations only to learn the environment model, why can't this model still be biased? If we don't allow the learner to interact with the environment, how can we build accurate on situations that the expert seldom faces?

Author response:
We are not making any claims about the bias or generalization capabilities of the learned model itself. The point is that the return difference between expert and learner is bounded by the divergence of their induced distributions. Our method is about getting an estimate of the learner's distribution in the model, and explicitly penalizing the divergence from the expert over multiple time-steps, to address compounding error problems. We show that the world model latent space provides an excellent emasure of state distribution divergence. We show that minimizing this divergence using on-policy RL results in state-of-the-art performance in the offline setting (competitive with online methods)

Reviewer comment (Theory):
I don't understand why Eq. 5 and Eq.6 are correlated. Furthermore, the "adaptation error" term in Eq. 6 captures the error of model learning, which is crucial for understanding the proposed approach. While further analysis is lacking in the paper.

Author response:
Equations 5 and 6 are related in the following way: equation 5 gives a bound on the return difference in terms of the state distribution divergence. Equation 6 shows how the true state distribution divergence is bounded by the divergence in the model. Therefore, bounding the divergence in the model bounds the true state divergence, thereby bounding the return difference.