DITTO: Offline Imitation Learning with World Models
Authors: Branton DeMoss
Paul Duckworth
Nick Hawes
Ingmar Posner
Venue: Submitted to ICLR 2023 (2023)
Keywords: world models
imitation learning
reinforcement learning

Abstract: We propose DITTO, a fully offline approach to imitation learning which addresses the problem of covariate shift without access to an oracle or any additional online interactions. By unrolling agent policies in the latent space of a learned world model and penalizing drift from expert demonstrations, we can use online reinforcement learning algorithms to learn policies which solve the imitation objective, without access to the underlying environment or reward function. Decoupling policy and world model learning lets us leverage datasets of any quality to learn latent representations which provide a natural reward signal for imitation learning, avoiding the need for complex adversarial or sparse imitation-inducing rewards. Compared to competitive baselines, our method achieves state-of-the-art performance in a variety of challenging environments from pixel observations alone.

Under review as a conference paper at ICLR 2023
DITTO: OFFLINE WORLD MODELS
IMITATION
LEARNING
WITH
Anonymous authors Paper under double-blind review
ABSTRACT
We propose DITTO, an offline imitation learning algorithm which uses world models and on-policy reinforcement learning addresses the problem of covari- ate shift, without access to an oracle or any additional online interactions. We discuss how world models enable offlin,, on-policy imitation learning, and pro- pose simple intrinsic reward defined in the world model latent space that induces imitation learning by reinforcement learning. Theoretically, we show that our for- mulation induces divergence bound between expert and learner, in turn bounding the difference in reward. We test our method on difficult Atari environments from pixels alone, and achieve state-of-the-an performance in the offline setting.
INTRODUCTION
Generating agents which can capably act in complex environments is challenging. In the most difficult environments, hand-designed controllers are often insufficient and learning-based methods must be used to achieve good performance. When we can exactly specify the goals and constraints of the problem using reward function, reinforcement learning (RL) offers approach which has been extremely successful at solving a range of complex tasks, such as the strategy games of Go (Silver et al. 2016). Starcraft (Vinyals al. 2019), and poker (Brown et al. (020), and difficult real world control problems like quadrupedal locomotion (Lee et al.) 2020),. datacenter cooling (Lazic al. 2018). and chip placement (Mirhoseimi et al.) 2020). RL lifts the human designer's 's work from explicitly designing good policy, to designing good reward function. However, this optimization often results in policies that maximize reward undesirable ways that their designers did not intend [Lehman et al. 2018)- the so-called reward hacking phenomenon (Krakovna et al., 2020). To combat this, practitioners spend substantial effort observing agent failure modes and tuning reward functions with extra regularization terms and weighting hyperparameters to counteract undesirable behaviors (Wang et al. 2022) Peng et al. 2017).
Imitation learning (IL) offers alternative approach to policy learning which bypasses reward spec- ification by directly mimicking the behavior of an expert demonstrator. The simplest kind of IL, behavior cloning (BC), trains agent to predict an expert's actions from observations, then acts these predictions at test time. This approach fails to account for the sequential nature of decision problems, since decisions at the current step affect which states are seen later. The distribution of states seen at test time will differ from those seen during training unless the expert training data cov- ers the entire state space, and the agent makes no mistakes. This distribution mismatch, covariate shift, leads to compounding error problem: initially small prediction errors lead to small changes in state distribution, which lead larger errors, and eventual departure from the training distribution altogether (Pomerleau 1989). Intuitively, the agent has not learned how to act under its own induced distribution. This was formalized in the seminal work of Ross Bagnell (2010), who gave a tight regret bound on the difference in return achieved by expert and learner, which is quadratic in the episode length for BC.
Follow-up work in Ross al. 2011 showed that linear bound regret can be achieved if the agent learns online in an interactive setting with the expert: Since the agent is trained under its own distribution with expert corrections, there is no distribution mismatch at test-time. This works well when online learning is safe and expert supervision can be scaled, but is untenable in many real-world use-cases such robotics, where online be unsafe, time-consuming. otherwise

---

Under review as a conference paper at ICLR 2023
infeasible. On the one hand, we want to generate data on on-policy avoid covariate shift, but on the other hand, we may not be able to afford to learn online due to safety or other concerns.
HHa & Schmidhuber (2018) propose two-stage approach to policy learning, where agents first learn to predict the environment dynamics with a recurrent neural network called a "world model" (WM), and then learn the policy inside the WM alone. This approach is desirable since it enables on-policy learning offline, given the existence of the world model. Similar model-based learning methods have recently achieved success in standard online RL settings (Hafner et al. 2021). ), and impressive zero-shot transfer of policies trained solely in the WM to physical robots (Wu et al. 2022).
In this paper, we propose an imitation learning algorithm called Dream Imitation (DITTO), which addresses the tension between offline and on-policy learning, by training an agent using on-policy RL inside a learned world model. Specifically, we define reward which measures divergence be- tween the agent and expert demonstrations in the latent space of the world model, and show that optimizing this reward with RL induces imitation on the expert. We discuss the relationship be- tween our method and the imitation learning as divergence minimization framework (Ghasemipour et al. 2019), and show that our method optimizes a similar bound without requiring adversarial train- ing. We compare our method against behavior cloning and generative adversarial imitation learning (GAIL, Ho & Ermon (2016)), which we adapt to the world model setting, and show that we achieve better performance and sample efficiency in challenging Atari environments from pixels alone.
Our main contributions are summarized as follows:
We discuss how world models relieve the tension between offline and on-policy learning methods, which mitigates covariate shift from offline learning.
We demonstrate the first fully offline model-based imitation learning method that achieves strong performance on Atari from pixels, and show that our agent outperforms competitive baselines adapted to the offline setting.
We show how imitation learning can naturally be cast as a reinforcement learning problem in the latent space of learned world models, and propose a latent-matching intrinsic reward which compares favorably against commonly used adversarial and sparse formulations.
2 RELATED WORK
2.1 IMITATION BY REINFORCHMENT LEARNING
Imitation learning algorithms can be classified according to the set of resources needed to produce a good policy. Ross al. (2011) give strong theoretical and empirical results in the online interactive setting, which assumes that we can both learn while acting online the real environment, and that interactively query an expert policy to e.g. provide the learner with the optimal action in the current state. Follow-up up works have progressively relaxed the resource assumptions needed to produce good policies. [Sasaki & Yamashina (2021) show that the optimal policy can be recovered with a modified form of BC when learning from imperfect demonstrations, given a constraint on the expert sub-optimality bound. BBrntley et al.] (2020) study covariate shift in the online, non- interactive setting, and demonstrate an approximately linear regret bound by jointly optimizing the BC objective with novel policy ensemble uncertainty cost, which encourages the learner return to and stay in the distribution of expert support. They achieve this by augmenting the BC objective with the following uncertainty cost term:
E
1 E E
Var (T(a)s)) = E
(a(s))))
Σ (als)))
(I)
This term measures the total variance of a policy ensemble 11g = [ JE} trained on disjoint subsets of the expert They optimize the combined BC plus uncertainty objective using standard online RL. algorithms, and show that this mitigates covariate shift.
Inverse reinforcement learning (IRL) can achieve improved performance over BC by first learning reward from the expert demonstrations for which the expert is optimal, then optimizing that reward with on-policy reinforcement learning. This two-step process, which includes on-policy RL in the
2

---

Under review as a conference paper at ICLR 2023
second step, helps IRL methods mitigate covariate shift due to train and test distribution mismatches. However, the learned reward function can fail to generalize outside of the distribution of expert states which form its support.
recent line of work treats IRL as divergence minimization: instead directly copying the expert actions, they minimize divergence measure between expert and learner state distributions
min TD (p". p
where 0² (s,a) a) = (1 Et P 5, a) is the discounted state-action distribution induced by n. and D is divergence measure between probability distributions. The popular GAIL algorithm (Ho & Ermon 2016) constructs minimax game in the style of GANs (Goodfellow al., 2014) between the learner policy æ. and a discriminator Dv which learns to distinguish between expert and learner state distributions
max min (s.))=p [- lolog Dw(s, a)] + (x, |- log (1 (s(s,a)) DR
This formulation minimizes the Sensen-Shannee divergence between the expert and learner policies, and bounds the expected return difference between agent and expert. However, [ang et al. (2019) point out that adversarial reward learning is inherently unstable since the diseriminator is always trained to penalize the learner state sate-action distribution, even if the learner has converged to the expert policy. This finding is consistent with earlier work (Brock et al. (2019) which observed discriminator overfitting, necessitating early stopping to prevent training collapse. Multiple works have reported difficulty getting GAIL work reliably in pixel-based observation environments (Brantley al., 2020) (Reddy et al., 20020).
To combat problems with adversarial training. Wang al. 2019)) and Reddy et (020) consider reducing IL to RL on an intrinsic reward
r(s,a)= -(0 if (s, DE
00 otherwise
where D is the expert dataset. While this sparse formulation is impractical e.g. in continuous action settings, they show that a generalization of the intrinsic reward using support estimation by random network distillation (Burda al. 2019) results in stable learning that matches the performance of GAIL without the need for adversarial training. Ciosek (2022) showed that this formulation equivalent to divergence minimization under the total variation distance, and produced bound on the difference in extrinsic reward achieved between the expert and a learner trained with this approach.
2.2 OFFLINE LEARNING
Kumar 2019)) showed how naive application of Bellman backups in off-policy RL results in incorrect, overly optimistic value estimation. These off-policy bootstrapping errors are further compounded in the model-based setting, since states where either the learned dynamics or reward functions generalize poorly will be found and incorrectly backed out by naïve Q-learner. This is acceptable in the online setting, since areas of state space which the agent overly-optimistic about will tend to be visited, resulting in natural corrections to inaccuracies; but the offline setting does not share this property. To counteract these problems, offline RL methods either constrain the learner policy to stay on the expert distribution directly (Wu al. [2019). or use pessimistic methods to discourage value estimates from becoming too large out distribution (Kumar et al. 2020 (Rigter
et al. 2022).
2.3 WORLD MODELS
World models have recently emerged as promising approach to model-based learning. Ha & Schmidhuber (2018) defined the prototypical two-part model: a variational autoencoder (VAE) is trained to reconstruct observations from individual frames, while recurrent state-space model (RSSM) is trained to predict the VAE encoding of the next observation, given the current latent state and action. World models can be used train agents entirely inside the learned latent space,
3

---

Under review as a conference paper at ICLR 2023
without the need for expensive decoding back to the observation space. Hafner et al. (2020) intro- duced Dreamer, an RL agent which is trained purely in the latent space of the WM. and successfully transfers to the true environment at test-time. [Wu et al. (202)) showed that the same approach can be used to simultaneously learn a model and agent policy to control a physical quadrupedal robot online, without the control errors usually associated with transferring policies trained only in simu- lation to physical system (Hwangbo et al. 2019).
In this work, we propose the use of world models to address a number of common problems in imita- tion learning. Intrinsic rewards which induce imitation learning, like those introduced in Reddy et al.) (2020) and Wang al. (2019). can pose challenging online learning problems, since the rewards are sparse or require tricky additional training procedures to work in nigh-dimensional observation spaces. Similarly, approaches like GAIL (Ho & Ermon 2016) and AIRL (Fu et al., 2018) require adversarial on-policy training that is difficult to make work in practice. In contrast, our approach remedies both the online learning and reward specification problems by performing safe offline pol- icy learning solely inside the compact latent space of the world model, and uses natural divergence measure as reward: distance between learner and expert in the world model latent space, This pro- vides conceptually simple and dense reward signal for imitation by reinforcement learning, which we find outperforms competitive approaches in data efficiency and asymptotic performance.
3 DREAM IMITATION
We study imitation learning in partially observable Markov decision process (POMDP) with dis- crete time-steps and actions, and high dimensional observations generated by an unknown environ- ment. The POMDP M is composed of the tuple M = (S,A, X. R. T T,,, 7). where s E S is the state space,
A is the action space, z X' is the observation space, y is the discount factor, and r = R(s, a) is the reward function. The transition dynamics are Markovian, and given by 8141 T(- St, az). The agent does not have access to the underlying states, and only receives observations represented by x, ~ [[(- s). The goal is to maximize the discounted sum of extrinsic (environment) rewards B[E,7'r.), which the agent does not have access to.
Training proceeds in two parts: we first lear a world model from recorded sequences of observa- tions, then train an actor-critic agent imitate the expert in the world model. The latent dynamics of the world model define a fully observable Markov decision process (MDP), since the model states &, are Markovian. Model-based rollouts always begin from an observation drawn from the expert demonstrations, and continue for fixed set of time steps H, the agent training horizon. The agent rewarded for matching the latent trajectory of the expert.
3.1 PRELIMINARIES
We show that bounding the leamner-xpert state distribution divergence in the world model also bounds their return difference in the actual environment, and connect our method to the IL as di- vergence mininization framework (Ghasemipour et al., 2019). Rafailov et al. (202T) showed that for learned dynamics model 7 whose total variation from the true transitions is bounded such that Div Drv(T(s, T(s (s.o)) a)) S a ((s E x A and Rmux max(),a) R(s, a) then
Rmax
Rmax DAv (pM
F. -(T
(5)
learning error
adoptation error
where J(T, is the expected return of policy π in MDP M, and M is the "imagination MDP" induced by the world model. This implies the difference between the expert return and the leamer return the true environment is bounded by two terms, 1) term proportional to the model ap- proximation error a, which could in principle be reduced with more data, and 2) a model domain adaptation error term, which captures the generalization error of a model trained under data from one policy, and deployed under another. Rafailov et al.] (2021) also show that bounding the diver- gence between latent distributions upper bounds the true state distribution divergence. Formally, given latent representation of the transition history 2 q(E a<t) and a belief distribution P(st Igg, ) = P(s,) 2)), then if the policy conditions only on the latent representation & such that the belief distribution is independent of the current action P(s: 2t, 4z) = P(s, =t),
4

---

Under review as a conference paper at ICLR 2023
E
21
Latent distance
Espert teajectryy
Learner Injectories
Figure 1:
The learner begins from random expert latent states during training, and generates policy latent trajectories in the world model. The intrinsic reward encourraes the learner to recover from its mistakes over multiple time steps match the expert trajectory.
then the divergence between the latent state distribution of the expert and learner upper bounds the divergence between their true state distribution:
(s( a) ll p a)) S Dy(oā((s
as) s ()(e)((zaa) || P (z, (z,a)) a))
(6)
Where D is a generic f-divergence, e.g. KL or TV. This result, along with equation 5 suggests that minimizing divergence in the model latent space is sufficient to bound the expected expert-learner return difference.
Reward To bound expert-leamer state distribution divergences, prior approaches have focused on sparse indicator function rewards (Ciosek, 2022), or adversarial reward learning (Ghasemipour al. 2019). We propose new formulation, which rewards the agent for matching the expert latent state-action pairs over an episode. In particular, for arbitrary distance function d, agent state-action latent z?", and a set expert state-action latents DE:
int rt (z =
min d(z), a eeEDg
This function rewards matching the agent's state-action pairs the expert's, as studied in Ciosek (002).. The major difference is that we also smooth in the latent space, meaning exact match isn t required for a reward. We show in appendix how to make this relaxed reward compatible with the theoretical results from Ciosek (2022). J. In particular, we prove that maximizing this reward bounds the total variation in latent-state distributions between the expert and learner, as well as bounding their extrinsic reward difference.
Intuitively, matching latent states between the learner and expert is easier than matching observa- tions, since the representations learned from generative world model training should provide much richer signal of state similarity. In practice, the minimization over De can be computationally ex- pensive, so modify the objective 7 to exactly match learner latent states expert latents from the same time-step, as shown in Figure I In particular, we randomly sample consecutive expert latents zi++ from Dg and unroll the agent from the same starting state the world model, yielding sequence of agent latents H- Finally, we compute reward at each step t as follows:
ri ( =1 d(=
2
max
2
This formulation changes our method from distribution matching to mode seeking, since states fre- quently visited by the expert will receive greater reward in expectation. We found that this modified dot product reward empirically outperformed L2 and cosine- cosine-similarity metrics.
3.2 WORLD MODEL
Dataset World model training can be performed using datasets generated by policies of any quality, since the model only predicts transition dynamics. The transition dataset is composed of N episodes e, of sequences of observations It, actions a,: D = [ at)
||ewll E N).
Model architecture We adapt the architecture proposed by Hafner al. (021), which is com- posed of an image encoder, a recurrent state-space model (RSSM) which learns the transition dy- namics, and decoder which reconstructs observations from the compact latent states. The encoder
5

---

Under review as a conference paper at ICLR 2023
uses convolutional neural network (CNN) to produce representations, while the decoder is trans- posed CNN. The RSSM predicts a sequence of length T deterministic recurrent states (h₂) 1hx)?=o- each of which are used to parameterize two distributions posterior state 2y is a function the current observation Zi and recurrent state h,, while the stochas- tic prior state is trained match the posterior without access to the current observation. The current observation is reconstructed from the full model state, which is the concatenation of the deterministic and stochastic states š (t), The RSSM components are:
stochastic hidden states. The stochastic
Model state:
8. (he, 21)
Recurrent state:
h = f a1-1)
Prior predictor:
fs Pol h44)
Posterior predictor:
Zt ~ 2 he,
Image reconstruction:
2, ( 86)
All components are implemented as neural networks, with combined parameter vector Φ. Since the prior model predicts the current model state using only the previous action and recurrent state, without using the current observation, we can use to learn behaviors without access to observa- tions or decoding back into observation space. The prior and posterior models predict categorical distributions which are optimized with straight-throut gradient estimation (Bengio et al. 2013). All components of the model are trained jointly with modified ELBO objective:
min (R
Σ
log Po (t) + 3DKu ( |l Po he)
(10)
where DKL. (q || p) denotes KL balancing [Hafner et al. 2021). which is used to control the regu- larization of prior and posterior towards each other with parameter 6,
DKLB(L || p) 5 Dg sg())) +(1 8) (sg(s( P)
(11)
posterior regularizer
prior regularizer
and sg(-) is the stop gradient operator. The idea behind KL balancing is that the prior and poste- rior should not be regularized at the same rate: the prior should update more quickly towards the posterior, which encodes strictly more information.
3.3 AGENT
Agent architecture The agent is composed of stochastic actor which samples actions from a learned policy with parameter vector 8, and deterministic critic which predicts the expected discounted sum of future rewards the actor will achieve from the current state with parameter vector db. Both the actor and critic condition only on the current model state &, which is Markovian:
Actor: at ~ TT (a; Be)
(12)
Critic: (Be) R Exo vs D
We train the critic to regress the A-target (Sutton & Barto 2005
V = ri + ((1 ()ux() + \V22), VÅ ug
(13)
which lets control the temporal-difference (TD) learning horizon with the hyperparameter A. Setting À = 0 recovers 1-step TD learning, while = I recovers unbiased Monte Carlo returns, and intermediate values represent an exponentially weighted sum of n-step returns. In practice use 0.95. To train the critic, we regress the X-target directly with the objective:
min E TH PP
Σ (vv( s3(()))
(14)
There is loss the last time step since the target equals the critic there. We follow Mnih et al. (013),, who suggest using a copy of the critic which updates its weights slowly, called the target network, to provide the value bootstrap targets.

---

Under review as a conference paper at ICLR 2023
The actor is trained to maximize the discounted sum of rewards predicted by the critic. We train the actor to maximize the same X-target as the critic, and add an entropy regularization term to encour- age exploration and prevent policy collapse. We optimize the actor using REINFORCE gradients (Williams) 2004) and subtract the critic value predictions from the X-targets for variance reduction. The full actor loss function is:
[()) Eu .P
V
log To(a; 8e)sg(V, ub(()))
( (8+)
(15)
reinforze
catropy regularizer
Algorithm Learning proceeds in two phases: First, train the WM on all available demonstra- tion data using the ELBO objective [O] Next, we encode expert demonstrations into the world model latent space, and the on-policy actor critic algorithm descrihed above optimize the intrinsic reward 8 which measures the divergence between agent and expert over time latent space. In principle, any on-policy RL algorithm could be used in place of actor-critic. We describe the full procedure in Algorithm
Algorithm Dream Imitation (DITTO)
1: Require demonstration data D (Ir, Ce 2111) n MENN
|leal
2: Initialize world model parameters é
3: while not converged do
World model learning
4: Draw Bum transition sequences {(x,, ae, 1
5: Compute all sequential RSSM components according to eqn
6: Update o with ELBO loss TO
7:
end while
8:
Initialize actor and critic parameters 0,
9:
while not converged do
b Agent training
10:
Draw Buc expert latent state sequences ($ D
11:
Generate trajectories (&5, a, I+H with at To( 8T)
T=1
12:
Compute rewards r- 87, 8 and values Vig (37)
13:
Compute X-returns V.A Tt. ((1 A)o(87+1) (V
VA =((8
Update critic on A-targets: D (655 =g(22
15:
Update actor with eqn 15
16:
end while
EXPERIMENTS
To the best of our knowledge, behavior cloning in paxel-based observation environments. Prior works generally focus on improv- ing behavior cloning (Sasaki Yamashina 2021), or study mixed setting with some online inter- actions allowed (Rafailov et al. 2021) (Kidambi et al. (021). To demonstrate the effectiveness of world models for imitation learning, we train without any interaction with the environment.
the first to study completely offline imitation learning without
4.1 AGENTS
To test the performance of our algorithm, we compare DITTO to standard baseline method, be- havior cloning, and to two methods which we introduce in the world model setting.
Behavior cloning We train BC model end-to-end from pixels, using convolutional neural network architecture. Compared to prior works which study behavior cloning from pixels Atari games (Hester al. 2017) stronger results, even in games where is trained with lower-scoring data.
al.
et al. 2020). our implementation achieves
Dream agents We adapt GAIL (Ho & Ermon 2016) and BC to the world model setting, which we dub D-GAIL and D-BC respectively. D-GAIL and D-BC both receive world model latent states instead of pixel observations. The D-BC agent is trained with maximum-likelihon estimation

---

Under review as a conference paper at ICLR 2023
BeamRider
Breakout
8002
DITTO
400
700C
BBC
360
6002
BC
300
Export
250
5000
Retur
Ratum
200
4000
150
3000
100
2000
50
1002
0
200
000
1000
200
400
600
800
1000
Expert trajectoriss
Papert trajechntion
MsPacman
Qbert
1200
2003
20000
1000
1000
Rm1m
8C0
10000
500
5000
250
Expert Ingedoorss
Expert Inojedtoiis
Experi Ingeciorica
Figure 2: Results on five Atari environments from pixels, with fixed horizon H 15. The plots show mean final performance across runs with 20 validation simulations, and shaded regions show +1 standard error. The experts are strong pre-trained PPO agents from the RL Baselines3 Zoo.
the expert demonstrations in latent space, with additional entropy regularization term which found stabilized learning:
LBC = RE (B,0)=De [- log (T( HJC (T(())
(16)
The D-GAIL agent is trained on-policy in the world model using the adversarial objective from Equation 3 The D-GAIL agent optimizes its learned adversarial reward with the same actor-critic formulation used by DITTO, described in Section 3.3 We train both DITTO and D-GAIL with
fixed horizon H = 15. At test-time, the model mosel-based agent policies are composed with the world model encoder and RSSM convert high nigh-dimensional observations into latent representations.
All model- based policies in our experiments use an identical multi-layer perceptron (MLP) archi- tecture for fair comparison in terms of the policies' representation capacity, while the BC agent is parameterized by a stacked CNN and MLP architecture which mirrors the world model encoder plus agent policy. We found that D-GAIL was far more stable than expected, since prior works (Reddy) et al. 2020) (Brantley et al. 2020) reported negative results training GAIL on Atari games from pixels in the easier online setting. This suggests that world models may be beneficial for representa- tion learning even in the online case, and that other online algorithms could be improved by moving them to the world model setting.
We evaluate our algorithm and basclines on 5 Atari environments, using strong PPO agents (Schul- man al.] 2017 from the RL Baselines3 Zoo (Raffin 2020) as expert demonstrators, using N - [4, 8, 30, 60, 125, 250, 500, 1000) expert episodes to train the agent policies in the world model. To train the world models, we generate 1000 episodes from a pre-trained policy, either PPO advantage actor-critic (A2C) (Mnih et al. 2016), which achieves substantially lower reward compared to PPO. Surprisingly, we found that the A2C and PPO-trained world models performed similarly, and that only the quality of the imitation episodes affected final performance. We hy- pothesize that this is because the A2C and PPO-generated datasets provide similar coverage of the environment. It appears that the world model can learn environment dynamies from broad classes
8

---

Under review as a conference paper at ICLR 2023
of datasets long they cover the state distribution well. The data-genenn gedtarenneatiig policy's quality is relevant for imitation learning, but appears not to be for dynamics learning, apart from coverage.
4.2 RESULTS
Figure plotts the performance of DITTO against our proposed world model haselines and standard BC. In MsPacman and Qbert, most methods recover expert performance with the least amount of data we tested, and are tightly clustered, suggesting these environments are easier to learn good policies from low amounts of data. D-GAIL exhibited adversarial collapse twice in MsPacman, an improvement over standard GAIL, which exhibits adversarial collapse much more frequently in prior works which study imitation learning from pixels in Atari (Reddy et al.) (020) et al.) 2020). In contrast, DITTO always recovers or exceeds average expert performance in all tested environments, and matches or outperforms the baselines in terms of both sample efficiency and asymptotic performance.
CONCLUSION
Addressing covariate shift in imitation learning is a long-standing problem. this work we pro- posed DITTO, method which addresses this problem in the challenging offline setting where environment interactions are allowed while learning. The offline setting exacerbates covariate shift, since the agent learns from and acts under unrelated distributions, and cannot easily estimate its own induced state distribution to perform off-policy corrections. Recent strong II. methods such as DRIL (Brantley et al. 2020) and V-MAIL (Rafailov al. 2021 achieve excellent sample effi- ciency in terms of expert demonstrations, but still require environment interaction to perform on- policy learning. Our agent consistently recovers expert performance without online interaction in the tested Atari environments from pixel observations, and matches or exceeds the performance of strong baselines which we implement in the world model setting.
DITTO learns by decomposing imitation learning into two parts: First, world model is learned to approximate the underlying environment dynamics from all available demonstration data, regardless of the quality of the policy which generated it. Next, expert demonstrations are encoded into tra- jectories the world model latent space. Finally, DITTO produces on-policy rollouts in the world model, and optimizes an intrinsic reward which measures the agent's drift from expert trajectories. By optimizing this reward with standard on on-policy reinforcement learning algorithms, DITTO learns to recover from its own mistakes across multiple time-steps, and match the expert state distribution, which we prove bounds the agent-expert return difference in the true environment.
Decoupling dynamics and policy learning lets us learn the world model from any historical demon- stration data, and provides us with a surprisingly rich signal for imitation learning: the latent repre- sentations learned via generative world modeling. Our latent space reward extends the recent work on imitation by reinforcement learning (Ciosek) 2022) to difficult pixel-based observation environ- ments, and contrasts with recent IL methods which need to employ adversarial (Fu al., 2018) or sparse rewards Wang et al. (2019) to induce imitation, which can be difficult to train. Furthermore, other methods from reinforcement and imitation learning are orthogonal to and compatible with ours, and could be combined with DITTO to achieve greater performance and sample efficiency, e.g. by adding behavioral cloning rewards or uncertainty costs DITTO's objective.
6 REPRODUCIBILITY STATEMENT
In order to ensure reproducibility of the results, all hyperparameters used in our experiments are included in Appendix A. The code for the experiments will be made available for the reviewing process and we intend to release the code for the camera-ready version of the paper.
REFERENCES
Yoshua Bengio, Nicholas Léonard, and Aaron C. Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. ArXiv, abs/1 308.3432 2013. 6

---

Under review as a conference paper at ICLR 2023
Kianté Brantley, Wen Sun, and Mikael Henaff. Disagreement-regulaiizeed imitation learning. In ICLR, 2020.
9
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. In International Conference Learning Representations, 2019. URL
https:://ppenre//
net htt//oooummm// wow.mniddBBLssr1D.
Noam Brown, Anton Bakhtin, Adam Lerer, and Qucheng Gong. Combining deep reinforcement learning and search for imperfect-information games. ArXiv, abs/2007.13544, 2020. |
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov, Exploration by random network distillation. ArXiv, abs/1810. 12844, 2019.
Kamil Ciosek. Imitation learning by reinforcement learning,
ICLR. 2022. URL https://
wttp://eeeoorrmmmddddaaeeeeeaeep 3
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse rein-
forcement learning. ArXiv, 1bss/710.11124, 2018.
Seyed Kamyar Seyed Ghasemipour, Richard S. Zemel, and Shixiang Shane Gu. divergence min- imization perspective imitation learning methods. In CoRL, 2019. 5
lan J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farles Farde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative adversarial nets. In NIPS, 2014. 3
David R Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. ArXiv, abs/ 0tbs/1809.01// 809.01999, 2018. 2
Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learn- ing behaviors by latent imagination. ArXiv, abs/1912.01603, 2020.
Danijar Hafner, Timothy P. Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. ArXiv, abs/2010.02193, 2021.
Todd Hester, Matej Vecerík, Olivier Pietquin, Mare Lanctot, Tom Schaul, Bilal Piot, Andrew Sendonaris, Gabriel Dulac-Amold, Ian Osband, John P. Agapiou, Joel Z. Leibo, and Au- drunas Gruslys. Learning from demonstrations for real world reinforcement learning. ArXiv, abs/1704.03732, 2017. 7
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In NIPS, 2016.
Jemin Hwangbo, Joonho Lee, Alexey Dosovitskiy, Dario Bellicoso, Vassilios Tsounis, Vladlen Koltun, and Marco Hutter. Learning agile and dynamic motor skills for legged robots. Sci- ence Robotics, 4(26))eaa55872, 2019. doi: 10.1116/sciroboticc..aiii aau5872. aauutirr URL https://s///
sscience. wilieneee.or//doi/abb.t.
1126/ www.ssotiiiiccccoooooooooo aau5872
4
Anssi Kanervisto, Joonas Pussinen, and Ville Hautamäki. Benchmarking end-to-end behavioural
cloning on video games. 2020 IEEE Conference on Games (CoG), pp. 558-565, 2020.
Rahul Kidambi, Jonathan Chang, and Wen Sun. Mobile: Model-hased imitation learning from observation alone. In NeurlPS, 2021. 7
Victoria Krakovna, Jonathan Uesato, Vladimir Mikulik, Matthew Rahtz, Tom Everitt, Ramana Kumar, 7ac Kenton, Jan Leike, and Shane Legg. Specification gam- ing: the flip side of ai ingenuity, 2020. https: //wp:: deepmind. com/blog/ specification gaming- the- flip- side of ai- iii-ingenuity. 2022-9-12. II
Last accessed on
Aviral Kumar, Justin Fu, G. Tucker, and Sergey Levine. Stabilizing off-policy q-learning via boot- strapping error reduction. In NeurIPS, 2019.
Aviral Kumar, Aurick Zhou, G. Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. ArXiv, abs/2006.04779, 2020.
10

---

Under review as a conference paper at ICLR 2023
Nevena Lazic, Craig Boutilier, Tyler Lu, Echern Wong, Binz Roy, M. K. Ryu, and Gregory P. Imwalle. Data center cooling using model- predditive control. In NeurIPS, 2018. T
Joonho Lee, Jemin Hwangbo, Lorenz Wellhausen, Vladlen Koltun, and Marco Hutter. Learning quadrupedal locomotion over challenging terrain. Science Robotics, 5, 2020.
Joel Lehman, Jeff Clune, Dusan Misevic, Christoph Adami, Lee Altenberg, Julie Beaulieu, Pe- ter John Bentley, Samuel Bernard, Guillaume Beslon, David M. Bryson, Patryk Chrabaszez, Nick Cheney, Antoine Cully, Stéphane Doncieux, Fred C. Dyer, Kai Olav Ellefsen, Robert Feldt, Stephan Fischer, Stephanie Forrest, Antoine Frénoy, Christian Gagné, Léni K. Le Goff, Laura M. Grabowski, Babak Hodjat, Frank Hutter, Laurent Keller, Carole Knibbe, Peter Kreah, Richard E. Lenski, Hod Lipson, Robert MacCurdy, Carlos Maestre, Risto Mikkkulainen, Sara Mitri, David E. Moriarty, Jean-Baptiste Mouret, Anh M Nguyen, Charles Ofria, Mare Parizeau, David P. Par- sons, Robert T. Pennock, William F. Punch, Thomas Ray, Mare Schoenauer, Eric Shulte, Karl Sims, Kenneth O. Stanley, François Taddei, Danesh Tarapore, Simon Thibault, Westley Weimer, Richard A. Watson, and Jason Yosinksi. The surprising creativity of digital evolution: A col- lection of anecdotes from the evolutionary computation and artificial life research communities. Artificial Life, 26:274-306, 2018.
Azalia Mirhoseini, Anna Goldie, Mustafa Yazgan, Joe Wenjie Jiang, Ebrahim M. Songhori, Shen Wang, Young-Joon Lee, Eric Johnson, Omkar Pathak, Sungmin Bae, Azade Nazi, Jiwoo Pak, Andy Tong, Kavya Srinivasa, Will Hang. Emre Tuncer, Ananda Babu, Quoc V. Le, James Laudon, Richard Ho, Roger Carpenter, and Jeff Dean. Chip placement with deep reinforcement learning. ArXiv, abs/2004. 10746, 2020. T
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen, Charlic Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning- Nature, 553,, 2015. 6
Volodymyr Mnih, Adrià Puigdomènech Badia, Mchdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. ICML, 2016. X
Xue Bin Peng, Glen Berseth, KangKang Yin, and Michiel van de Panne. Deeploco: dynamic locomotion skills using hierarchical deep reinforcement learning. ACM Trans. Graph., 36:41:1- 41:13, 2017.
Dean A. Pomerleau. Alvinn: An autonomous land vehicle in neural network. In NeurlPS, 1989.
Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Visual adversarial imitation
learning using variational models. ArXiv, abs/2107.08829, 2021. 4
Antonin Raffin.
RI
basclines3
200.
https: com/ RM/
rl baselines3 200, 2020.
Siddharth Reddy, Anea D. Dragan, and Sergey Levine. Sqil: Imitation learning via reinforcement
learning with sparse rewards. In International Conference on Learning Representations, 2020.
URL httt:t::///peeeeeeeet net forum? id 81xKd24twB
Mare Rigter, Bruno Lacerda, and Nick Hawes. Rambo-rl: Robust adversarial model-based offline reinforcement learning. ArXiv, abs/2204.12581 2022.
Stéphane Ross and J. Andrew Bagnell. Efficient reductions for imitation learning. In AISTATS, 2010. !I
Stéphane Ross, Geoffrey J. Gordon, and J. Andrew Bagnell. A reduction of imitation learning and structured prediction no-regret online learning. In AISTATS, 2011. -
Fumihiro Sasaki and Ryota Yamashina. Behavioral cloning from noisy demonstrations. In ICLR,
2021. 2
11

---

Under review as a conference paper at ICLR 2023
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. ArXiv, abs/1707.0634477 2017. 8
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, L. Sifre, George van den Driess- che, Julian Schrittwieser, loannis Antonoglou, Vedavyas Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy P. Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of go with deep neural networks and tree search. Nature, 529:484-445 489, 2016. [
Richard S. Sutton and Andrew G. Barto. Reinforcement learning: An introduction. IEEE Transac- lions on Neural Networks, 16:285-286, 2005.
Oriol Vinyals, Igor Babuschkin, Wojeiech M. Czarnecki, Michaël Mathieu, Andrew Dudzik, Juny- oung Chung, David H. Choi, Richard Powell, Timo Ewalds, Petko Georgiev, Junhyuk Oh, Dan Horgan, Manuel Kroiss, Ivo Danihelka, Aja Huang. L. Sifre, Trevor Cai, John P. Agapiou, Max Jaderberg, Alexander Sasha Vezhnevets, Rémi Leblond, Tobias Pohlen, Valentin Dalibard, David Budden, Yury Sulsky, James Molloy, Tom Le Paine, Caglar Gulcehre, Ziyun Wang. Tobias Pfaff, Yuhuai Wu, Roman Ring, Dani Yogatama, Dario Wünsch, Katrina McKinney, Oliver Smith, Tom Schaul, Timothy P. Lillicrap, Koray Kavukcuoglu, Demis Hassabis, Chris Apps, and David Silver. Grandmaster level in starcraft using multi-agent reinforcement learning. Nature, pp. 1-5, 2019.
Ruohan Wang, Carlo Ciliberto, Pierluigi Vito Amadori, and Y. Demiris. Random expert distillation: Imitation learning via expert policy support estimation. In ICMI,, 2019.
Tsun-Hsuan Wang, Alexander Amini, Wilko Schwarting, Igor Gilitschenski, Sertaç Karaman, and Daniela Rus. Learning interactive driving policies via data-driven simulation. 2022 International Conference on Robotics and Automation (ICRA), pp. 7745-7752, 2022.
Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcmeet learning. Machine Learning, 8:229-256, 2004. 7
Philipp Wu, Alejandro Escontrela, Danijar Hafner, Ken Goldberg, and Abbeel. Daydreamer:
World models for physical robot learning. ArXiv, abs/2206. 14176, 2022. 4
Yifan Wu, G. Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. ArXiv, abs/1911.11361,, 2019. 3
Ruohan Zhang, Zhuode Liu, L. Guan, Luxin Zhang, Mary M. Hayhoe, and Dana H. Ballard. Atari- head: Atari human eye-tracking and demonstration dataset. Proceedings of the ... AAAI Confer- ence on Artificial Intelligence AAAI Conference on Artificial Intelligence, 34 4:681 1-8200, 2020.
A PROOF OF DIVERGENCE REWARD BOUND
We prove a corollary of proposition I from Ciosek 2022). Ciosek (002)) uses many intermediate re- sults and definitions, so we encourage the reader to reference their work while reading understand this proof.
Corollary A.1. Suppose we also have another imitation learner, which uses the same data-set of size N, and still satisfies Assumption 3, but instead trains some other intrinsic reward, Ry which satisfies (for some > 0):
Ryyn (s, a) = 1, (s) E
< (i < c, otherwiss
Let p² be the limiting state state-action distribution of this imitation learner. Then:
|lp² p 11l
1 S ?
E R 2 E,e [R] 7
12

---

Under review as a conference paper at ICLR 2023
Proof. Lemma 5 trivially still holds with Ri instead of Rint Ri 2 Riat always, vp, pp R
2 [p[Rm]] Hence the bound holding true for Epd [mmd] implies holds for
int
2 [ ind too.
Lemma holds with replaced by the result is E.
[E (1-
mimix
We do this by considering their proof in Appendix D. The properties of the intrinsic reward are utilised in just one paragraph, after equation 25. This is done in stating that Σe T + B [and] E [Rm]. This is not true for Rm Let be the limiting chance of the ex- Pa-
LMc
and B+1 pert agreeing with the ni imitation agent. Almost by definition, Ee LMre T
Pa and B+1
Note that E,' R agree, and at most 1 - e if we disagree. Hence, using K = Ept [ml we have K S Pa (1 Pa)(1
Pa (1 Pa)(1 c); yield reward of 1 every time
1 + Pae, hence pa 21- A.
So, taking limits as done in the original proof, we have:
p [R] 2 PulE [R] (1 Pa) 4mix
= 2 Pa Pa(EE (R) + 4Tmix) 4Tmix
K. (EE p= [R] + 4Tmix) 4Tmix
2 (1
Now, combining these lemmas is exactly in section 4.4 in Ciosek (2022). 7 as required.
The factor of L carries
forward, yielding Ep" [R] 2 E,
B HYPERPARAMETERS
Table 1: Experimental hyperparameters
Description
Symbol
Value
Number of world model training episodes
1000
Number of expert training episodes
NE
{4,8,
10, 60, 125, 250, 500, 1000)
World model training batch size
Bum
50
World model training sequence length
L
50
Agent training batch size
Bac
512
Agent training horizon
15
Discount factor
7
0.95
DDA)) parameter
X
0.95
KL-Ba AL-Balancing weight
3
0.1
KL-Balancing trade-off parameter
0.8
Actor-critic entropy weight
7
5 10
Behavior cloning entropy weight
0.1
Optimizer
Adam
All learning rates
3 10-4
Actor-critic target network update rate
100 steps
C
TIME HORIZON ABLATION
13

---

Under review as a conference paper at ICLR 2023
Breakout
400
350
300
250
Return
200
150
100
DITTO
D-GAIL
50
Expert
D-BC
4
8
10
12 Train horizon
14
16
18
20
Figure 3: We ablate the agent learning horizon H fixed data, for N 125 episodes. We find that the environment seems to have a characteristic planning horizon, which for Breakout appears at This result suggests that learning act multiple steps is beneficial in some environments. D-BC is always trained to predict the next expert action, and has no planning horizon.
14


Reviews:
[{"review_id": "st9uss75Dq", "reviewer": "Reviewer_ZmXh", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper introduces DITTO, a method for fully offline imitation learning in POMDPS high dimensional observation spaces. The idea progresses by first learning a dynamics model or 'world model'. Then, an RL agent is trained in this latent space to minimize the divergence to the learned state-action density to the expert state-action density. ", "summary_of_the_review": "The paper proposes a new method for explicitly model-based offline imitation learning. The method is exciting and shows promise. However, the experimental results are weak, not evaluating on the state-of-the-art modern offline imitation learning baselines. This is understandable since the authors were not aware of these methods. However, I don't think that a paper which lacks comparison to current benchmarks should be recommended for acceptance. I would be very happy to reconsider if the authors could provide comparisons to the methods mentioned above, and demonstrate how DITTO compares.\n\nUpdate after rebuttal:\n\nThanks to the authors for the replies and additional baselines. Based on these, I have increased my recommendation to 6. ", "strength_and_weaknesses": "Strengths\n+ The proposed method is very neat, and explained well. The improvement obtained from an explicit model-based approach is great to see, and encouraging for model-based RL in general.\n+ The theory is nice to have, as is the discussion of the fundamental challenges facing imitation learning. \n+ The experimental results seem promising to start with. \n\nWeaknesses\n+ The authors write 'to the best of our knowledge we are the first to study completely offline imitation learning without behavior cloning in high-dimensional observation environments'. In fact, there are several prior works which deal with this setting. They include ValueDice [1], IQ-Learn [2], EDM [3] AvRIL [4], all of which achieve decent performance. There are almost certainly more, of which you can probably identify most by looking at the papers that cite [3] .\n\n  [1] Kostrikov et al, Imitation Learning via Off-Policy Distribution Matching, ICLR 2019\n\n  [2] Garg et al, IQ-Learn: Inverse soft-Q Learning for Imitation, Neurips 2021\n\n  [3] Jarrett et al, Strictly Batch Imitation Learning, Neurips 2020\n\n  [4] Chan et al, Scalable Bayesian Inverse Reinforcement Learning, ICLR 2021\n\n  Several of these methods are non-adversarial and so do not suffer from the (very reasonable) issues raised by the authors with regards to GAIL. \n\n+ I think I may be misunderstanding, but is there a reason that RL is required to learn the actor in the world model? It seems that if everything is controlled and latent, we can just differentiate through the policy directly (using a Gumbel-softmax) for issues with discrete distributions.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "l5Gr0wMMR-", "reviewer": "Reviewer_Nmum", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces an offline imitation learning algorithm. The proposed algorithm consists of two parts: (1) to train a world model using demonstration of any quality and (2) to train a policy using $\\lambda$-TD where the reward function is computed as Equation (7) (i.e. intrinsic reward). Finally, the authors empirically show that the proposed method outperforms baselines in offline setting.", "summary_of_the_review": "Solving offline imitation learning is an important, but the proposed algorithm is too empirical.\nThis paper uses only 5 Atari domains and there is no theoretical analysis.", "strength_and_weaknesses": "## **Strength**\n\nDITTO shows good empirical results.\n\n## **Weakness**\n**Algorithm**\n\n- DITTO uses an intrinsic reward function. However, there are no justification to use this reward function. At least there should be some consideration of the ideal case (e.g. in theory, DITTO is equivalent to a distribution matching between current policy and expert policy)\n\n**Experiments**\n\n- In Figure 3. D-BC always achieves near-expert performance when the number of expert trajectories is large enough. Compared to D-BC, DITTO shows slightly better performance. Therefore, it seems that more experiments are needed to show the superiority of DITTO.\n- I am not sure if it is a good thing to outperform the expert in (offline) imitation learning. Why DITTO exceeds average expert performance?\n\n**Related Works**\n\n- One of the simplest ways to solve offline imitation learning with learned world model is to use HIDIL [1], which aims to solve an offline imitation learning problem given a misspecified simulator.\n- In addition, existing (model-free) offline IL algorithms [2-4] can also be used for policy learning on learned world model.\n\n[1] Jiang, Shengyi, Jingcheng Pang, and Yang Yu. \"Offline imitation learning with a misspecified simulator.\" NeurIPS. 2020.\n\n[2] Kim, Geon-Hyeong, et al. \"DemoDICE: Offline imitation learning with supplementary imperfect demonstrations.\" ICLR. 2022.\n\n[3] Xu, Haoran, et al. \"Discriminator-weighted offline imitation learning from suboptimal demonstrations.\" ICML. 2022.\n\n[4] Ma, Yecheng Jason, et al. \"SMODICE: Versatile Offline Imitation Learning via State Occupancy Matching.\" ICML. 2022.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "XKWr6J14BtT", "reviewer": "Reviewer_pKRJ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes an offline imitation learning (IL) method supported by a word model. The world model is obtained by training on the transitions collected in a dataset, without any direct access to the actual environment. A intrinsic reward is also designed that measures the divergence between expert and agent in the trained word model. The approach is well motivated. Empirical results shows that the proposed method outperforms baselines adapted to the world-model, namely GAIL and BC.", "summary_of_the_review": "The paper proposes a fully offline IL method by training a world model and leverage standard RL for imitation learning. The method is well motivated and shows improved performance over baselines. Additional ablation and experiments are needed to better understand the advantages of the proposed method.", "strength_and_weaknesses": "Strength:\n1. The paper is well written and easy to follow. All elements of the proposed model appear well-motivated. The proposed method is sound, and appears a novel combination of various techniques.\n\n2. Empirically, the evaluation shows the proposed method outperforming the baselines appropriately adapted to the world-model setting.\n\nWeakness:\n1. The empirical evaluation, while carefully designed, is limited in scope and does not explicitly highlight/explain why the proposed method outperforms the baselines. For instance, my understanding is that the only difference between GAIL and DITTO is that GAIL uses adversarial reward. However, from this singleton comparison it is difficult to assess if adversarial reward is worse than intrinsic reward. Related to this, there are various recent adversarial formulation that outperforms GAIL and they should also be considered in the evaluation to more broadly evaluate the impact of the reward formulation. As an example, PWIL [1] is fairly similar to DITTO in terms of how reward is computed.\n\n2. The quality of the world-model and reward functions are not independently analyzed. In the current experiments, the size of the dataset directly affects both the quality of the world model and of the reward. However, it is important to separate these two effects, as the authors also argued that learning the world model can be done using transitions of any quality. Potential ablation may include separate the datasets used for world model and reward computation to isolate the two effects.\n\n3. The method appears directly applicable to Mujoco environments. It would be nice to see if world model is still beneficial in a vector-based observation/state setting.\n\nOverall, any additional experiments to improve our understanding of the proposed method and why it outperforms existing ones would be beneficial.\n\n[1] Primal Wasserstein Imitation Learning, Dadashi et al. ICLR 2021", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "5ilkK57o2FH", "reviewer": "Reviewer_RnpT", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies imitation learning in partially observable environments. The proposed method is to learn a world model using a dataset of any quality, unroll the agent's policy in the latent space of the world model to simulate the agent's trajectory, and train the policy with a reward function to make the agent's trajectory close to the demonstration trajectory in the latent space.\n\nSpecifically, the world model follows the recurrent neural network to predict future latent states, and the world model is trained with ELBO objective, consisting of the KL divergence between the prior and posterior distributions of latent states, and reconstruction error for the raw observation. The agent is trained via the on-policy actor-critic algorithm, where the policy is executed in the latent space of the world model. The dense reward signal penalizes the distance on the latent state densities between the agent's trajectory and the expert demonstration trajectory.\n\nExperiments are conducted on five Atari games and the proposed is compared with behavior cloning, and generative adversarial imitation learning, showing advantages over the baselines in three out of five environments.", "summary_of_the_review": "Overall, this paper proposed a technically solid method for imitation learning, but the evaluation can be better if conducted in more environments, with more datasets, and more design choices for the reward function.", "strength_and_weaknesses": "Strengths:\nThe method is well-motivated to study the covariate shift issue in imitation learning, which is an important and interesting problem.\nThe proposed method is technically solid. The learning of the world model and on-policy actor-critic is well executed to make the relatively complicated pipeline work in Atari games. \nThe paper writing is mostly clear and easy to follow.\n\nWeaknesses:\nThe performance improvement in only three out of five Atari games is not very impressive.\nThe statement \"address the problem of covariate shift\" in the abstract is not well-supported by the experiments. Regarding the experiments, it is unclear whether the bottleneck for baselines is covariate shift, and how much covariate shift can be addressed by the proposed method. \nThe evaluation and analysis can be improved. For example, how does the quality of the dataset affect the quality of the world model, and how does the quality of the world model affect the final performance of imitation learning? I doubt \"the quality of the policy used to generate world model training episodes did not affect final performance\". See below for more detailed questions and comments. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject