Reward-free World Models for Online Imitation Learning
Authors: Shangzhe Li
Zhiao Huang
Hao Su
Venue: ICML 2025 poster (2025)
Keywords: world models
imitation learning
Primary area: reinforcement_learning

Abstract: Imitation learning (IL) enables agents to acquire skills directly from expert demonstrations, providing a compelling alternative to reinforcement learning. However, prior online IL approaches struggle with complex tasks characterized by high-dimensional inputs and complex dynamics. In this work, we propose a novel approach to online imitation learning that leverages reward-free world models. Our method learns environmental dynamics entirely in latent spaces without reconstruction, enabling efficient and accurate modeling. We adopt the inverse soft-Q learning objective, reformulating the optimization process in the Q-policy space to mitigate the instability associated with traditional optimization in the reward-policy space. By employing a learned latent dynamics model and planning for control, our approach consistently achieves stable, expert-level performance in tasks with high-dimensional observation or action spaces and intricate dynamics. We evaluate our method on a diverse set of benchmarks, including DMControl, MyoSuite, and ManiSkill2, demonstrating superior empirical performance compared to existing approaches.

Reward-free World Models for Online Imitation Learning
Shangzhe Li"
Zhiao Huang?
Hao Su 2
Abstract
behavioral cloning (BC) methodologies. However, despite its wide applicability. IL methods that do not incorporate on- line interaction often suffer from poor generalization outside the expert data distribution, especially when encountering out-of nut-of-distribuutin states. Such limitations make these meth- ods vulnerable to failure, as even minor perturbations in state can lead to sigmificant performance degradation. This is often reflected in issues such as bias accumulation and suboptimal results (Reddy al., 2019). These challenges stem from BC's inability to fully capture the underlying dynamics of the environment and its inherent lack of explo- ration capabilities (Garg et al., 2021).
Imitation learning (IL) enables agents to acquire skills directly from expert demonstrations, pro- viding a compelling alternative to reinforcement learning. However, prior online IL approaches struggle with complex tasks characterized by high- dimensional inputs and complex dynamics. In this work, we propose a novel approach on- line imitation learning that leverages reward-free world models. Our method learns environmental dynamics entirely in latent spaces without recon- struction, enabling efficient and accurate model- ing. We adopt the inverse soft-Q learning objec- tive, reformulating the optimization process in the Q-policy space mitigate the instability associ- ated with traditional optimization in the reward- policy space. By employing learned latent dy- namics model and planning for control, our ap- proach consistently achieves stable, expert-level performance in tasks with high- nigh-dimensional ob- servation or action spaces and intricate dynam- ics. We evaluate our method on diverse set of benchmarks, including DMControl, MyoSuite, and ManiSkil12, demonstrating superior empirical performance compared to existing approaches.
To address these shortcomings, methods like GAIL (Ho Ermon, 2016), SQIL (Reddy et al., 2019). IQ-Learn (Garg et al., 2021), and CFIL (Freund et al., 2023) have introduced value reward estimation to facilitate deeper understand- ing of the environment, while leveraging online interactions to enhance exploration. Nevertheless, these approaches con- tinue to face substantial challenges, particularly when ap- plied to tasks with nigh-dimensional observation and action spaces or complex dynamics. Additionally, framing online IL as a min-max optimization problem within the reward- policy space, often inspired by inverse reinforcement learn- ing (IRL. techniques, introduces instability during training (Garg et al., 2021). Recent advancements in world models have demonstrated exceptional performance across a wide range of control tasks, underscoring their potential in com- plex decision-makee 201 9a;b; 2020; 2023; Hansen et al., 2022; 2023). Specifi- cally, world models offer advantages over model-free agents in terms of sampling complexity and future planning capa- bilities, resulting in superior performance on complex tasks (Hansen et al., 2022: 2023; Hafner et al., 2019a). Notably, decoder-free free world models, which operate exclusively in la- tent spaces without reconstruction, have proven to be highly effective and efficient in modeling complex environment dynamics (Hansen et al., 2022; 2023).
1. Introduction
and planning scenarios (Hafner et al.,
Imitation learning (II.) has garnered considerable attention due to its broad applications across various domains, such as robotic manipulation (Zhu et al., 2023; Chi et al., 2023) and autonomous driving (Hu et al., 2022; Zhou et al., 2021). Unlike reinforcement learning, where agents learn through reward signals, IL involves learning directly from expert demonstrations. Recent advances in offline IL, including Diffusion Policy (Chi al., 2023) and Implicit BC (Flo- rence et al., 2022), highlight the advantages of leveraging large datasets in conjunction with relatively straightforward
Motivated by these insights, we explore the application of world models in the context of online imitation learning without rewards, enabling II. agents to develop a deeper un- derstanding of environmental dynamics and improve their performance in tasks characterized by high-dimensional ob- servations and complex dynamics. In this work, we present novel approach to online imitation learning that leverages
'oout China University of Technology HHillbot Inc. " UUniversity of California, San Diego Correspondence to: Shangzhe Li <shangzhe @d...
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Reward-free World Models for Online Imitation Learning
the strengths of decoder-free world models, specifically designed for complex tasks involving high-dimensional ob- servations, intricate dynamics, and vision-based inputs. In contrast to conventional latent world models, which rely on reward and Q-function estimation, our approach com- pletely eliminates the need for explicit reward modeling. We propose a framework for reward-free world models that redefines the optimization process within the Q-policy space, addressing the instability associated with min-max optimization in the reward-police space. By utilizing an in- verse soft-Q learning objective for the critic network (Garg et al., 2021), our method derives rewards directly from Q- values and the policy, effectively rendering the world model reward-free. Moreover, by performing imitation learning online, our model addresses key challenges in IL, such as out-of-distribution errors and bias accumulation.
Soft Q-Learning (Garg al., 2021) reformulated the learn- ing objective of GAIL and integrated their findings into soft actor-critic (Haarnoja et al., 2018) and soft Q-learning agents for imitation learning. CFIL (Freund et al., 2023) in- troduced a coupled flow approach for reward generation and policy learning using expert demonstrations. ValueDICE (Kostrikov et al., 2019) proposed an off off-policy imitation learning approach by transforming the distribution ratio esti- mation objective. (Das et al., 2021) proposed a model-based inverse RL approach by predicting key points for imita- tion learning tasks. SQII. (Reddy et al., 2019) proposed an online imitation learning algorithm with soft Q func- tions. Diffusion Policy (Chi et al., 2023) is a recent offline IL method using a diffusion model for behavioral cloning. Implicit BC (Florence et al., 2022) discovers that treating supervised policy learning with an implicit model generally improves the empirical performance for robot learning tasks, Hybrid inverse reinforcement learning (Ren et al., 2024) proposed a new methodology leveraging a mixture of online and expert demonstrations for agent training, achieving ro- bust performance in environments with stochasticity. Prior works (Englert et al., 2013; Hu et al., 2022; Igl et al., 2022) explored the potentials of model-based imitation learning real-world robotics control and autonomous driving. Ef- ficientlmitate (Yin et al., 2022) combined EfficientZero (Ye et al., 2021) with adversarial imitation learning, achieving excellent results in DMControl (Tassa et al., 2018) () imitation learning tasks. V-MAIL (Rafailov et al., 2021) introduced a model-based approach for imitation learning using vari- ational models. CMIL (Kolev et al., 2024) proposed an imitation learning approach with conservative world models for image-based manipulation tasks. Ditto (DeMoss et al., 2023) developed an offline imitation learning approach with Dreamer V2 (Hafner et al., 2020) and adversarial imitation learning. DMIL (Zhang et al., 2023) utilized a discriminator to simultaneously evaluate both the accuracy of the dynam- ies and the suboptimality of model rollout data relative to real expert demonstrations in the context of offline imitation learning.
Through online training with finite-horizon planning based on learned latent dynamics, our method demonstrates strong performance in complex environments. We evaluate our ap- proach across a diverse set of locomotion and manipulation tasks, utilizing benchmarks from DMControl (Tunyasuvu- nakool et al., 2020), MyoSuite (Caggiano et al., 2022), and ManiSkil12 (Gu al., 2023), and demonstrate superior em- pirical performance compared to existing online imitation learning methods.
Our contributions are as follows:
We introduce a novel, robust methodology that lever- ages world models for online imitation learning, ef- fectively addressing the challenges posed by complex robotics tasks.
We propose an innovative gradient-free planning pro- cess, operating without explicit reward modeling, within the context of model predictive control.
We showcase the model's effectiveness in inverse rein- forcement learning tasks by demonstrating a positive correlation between decoded and ground-truth rewards.
Model-based Reinforcement Learning Contemporary model-based RL methods often learn dynamics model for future state prediction via data- approaches. PlaNet (Hafner et al., 2019b) was introduced as a model-basse learning approach for partially observed MDPs by propos- ing a recurrent state-space model (RSSM) and an evidence lower-bound (ELBO) training objective. Dreamer algo- rithm (Hafner et al., 2019a; 2020; 2023) a model-based reinforcement learning approach that uses a learned world model to efficiently simulate future trajectories in latent space, allowing an agent to learn and plan effectively. TD- MPC series (Hansen et al., 2022se et al., 2023) learns a scalable world model for model predictive control using temporal difference learning objective.
2. Related Works
Our work builds upon literature in Imitation Learning (IL) and Model-based Reinforcement Learning.
Imitation Learning Recent works regarding IL leveraged deep neural architectures to achieve better performance. Generative Adversarial Imitation Learning (GAIL) (Ho & Ermon, 2016) formulated the reward learning as min-max problem similar to GAN (Goodfellow et al., 2014). Model- based Adversarial Imitation Learning (MAIL) (Baram et al., 2016) extended the GAIL approach to incorporate a for- ward model trained by data-driven methodology. Inverse
2

---

Reward-free World Models for Online Imitation Learning
Our approach employs a model-based methodology ad- dress challenges in online imitation learning. By integrating data-driven approach for latent dynamics learning with planning for control, the agent is able to effectively cap- ture and leverage the underlying environment dynamics. Empirical evaluations demonstrate that our model achieves superior performance on complex online imitation learning tasks compared existing methods.
operator T over Eq. prior work reformulated the GAIL training objective in Q-policy space (Garg et al., 2021):
[TT, Q) = E (s.n)~pp
((s.aa) - -P( x.n) V (s')
ppp
V*(s)) -7lpg' - V (s')
QT)Q))
(4)
3. Preliminary
which is the inverse soft-Q objective for critic learning. In this way, we can perform imitation learning by lever- aging actor-critic architecture. The critic and policy can be learned by finding the saddle point in a joint opti- mization problem Q' or min I (IT, Q) and " argmin maxqeo [ Q). (Garg al., 2021) proved the uniqueness of the saddle point. For fixed Q. the optimization for policy has closed-form solution, which is the softmax policy:
We model the decision-making process in the environment as a Markov Decision Process (MDP), which can be defined as a tuple (S,A, Po: P, r, r)), S and A represent state and action space. Po the initial state distribution and P S x A + As is the transition probability. r(s, a) E R is the reward function and R is the reward space. Y E (0, 1) is the discount factor. We denote the expert state- action distribution as PE and the behavioral distribution as Pr. Similarly, we denote the expert policy as गæ and the behavioral policy as т. II is the set of all stochastic stationary policies that sample an action a E A given a state S. Z is the space for the latent representation of the original state observations, and Q is the space for all possible Q functions. H((-) represents the entropy of distribution.
exp Q(s, a) Σ. exp Q a)
To
In the actor-critic setting. we optimize the policy x using maximum-entropp RL objective, which approximates xq. and learn critic using:
max (IT, Q) QEQ
Maximum Entropy Inverse Reinforcement Learning Inverse Reinforcement Learning (IRL) focuses on recov- ering a specific reward function r(s,a) in the reward space R given certain amount of expert samples us- ing expert policy Tg. Maximum entropy IRL (Ziebart et al., 2008) seeks to solve this problem by optimiz- ing mii [r(s, - (Ep, [r(s, a)] + H(T)). GAIL (Ho & Ermon, 2016) generalized the objective into a form including an explicit reward mapping with a convex regularizer (r)::
max QEo Biss)) (o(((s,a) [1s0)))p(o((((ssa)) - P( (e
(s (T())= (Be
(6)
where & is a concave function. Specifically, if we leverage
x² regularization, we will have o(x) = x ..= 1² The scalar coefficient a controls the strength of x² regulariza- tion in the inverse soft-Q objective. Intuitively, the additonal regularization term penalizes the magnitude of the estimated reward. Prior work (Al-Hafez et al., 2023) interpreted the objective with this regularizer as minimizing the squared Bellman error, establishing connection between inverse soft-Q learning and SQIL (Reddy et al., 2019). A detailed empirical analysis on hyperparameter a is shown in Ap- pendix E.3.
maxmin a)) E,,, -()) (1) rERTEll
For a non-restrictive set of reward functions R = S the objective can be reformulated into a minimization of the statistical distance between distributions pr and Par (Ho & Ermon, 2016):
min (p: PF:) H(T))
Inverse Soft-Q Learning Prior work (Garg et al., 2021) introduced bijection mapping T 5 + SxA between Q space Q and reward space R, i... the inverse Bellman operator:
4. Methodology
In reinforcement learning, world models typically regress explicit reward signals provided by the environment. In imitation learning, prior approaches (Kolev et al., 2024; DeMoss et al., 2023) aim to train a reward model through adversarial objectives alongside a separate critic network trained on temporal difference objectives. In contrast, we
(T)))ss, a) = Q(s, a) - -E (s') (3)
where (s) = Ea (Q(s,a) (s(ss, a)- - logg T (a|s)]. The re- ward decoding is defined as = TTQ. By applying the
3

---

Reward-free World Models for Online Imitation Learning
Interaction
Environment
Latent Planning
Learned Model Components
Observation
Reward-free World Model Training
Behavioral
Bulfer
Reward-fres World Model
Expert
Buffer
Figure I. IQ-MPC We demonstrate the training workflow for IQ-MPC. The reward-free world model leverages both expert and behavioral data for training, using objectives Section 4.1. The policy prior from the world model guides MPPI planning process along with rewards decoded from Q estimations. The detailed planning process is revealed in Algorithm 1.
eliminate the need for separate reward model retrieving rewards directly from the learned critic. To this end, we propose reward-free world model that learns exclusively from reward-free expert demonstrations and environment interactions, without training a dedicated reward model. Furthermore, since our model can decode dense rewards from the critic, can solve inverse reinforcement learning tasks using reward-free interactions and a limited set of expert demonstrations that include only states and actions.
Value function:
q = (z,,a)
Policy prior: â T())
(10)
where and a are states and actions, z latent representa- tions. The policy prior и guides the model predictive plan- ning process. along with rewards decoded from the value function Q. We maintain two separate replay buffers BE and B2 for expert and behavioral data storage respectively. Behavioral data are collected during the learning process. For simplicity, we denote the sampling process from the joint buffer as B Bg B,, We sample trajectories with short horizons of length H from the replay buffers.
4.1. Learning Process of a Reward-free World Model
World models used in reinforcement learning settings often contain reward model R22, a) that requires supervised learning using explicit reward signals from the online en- vironment interactions or the offline data. However, if our learning objective able to form a bijection between Q space Q and reward space R, it would be natural to decode reward from the Q value instead of learning another sep- arate mapping for the reward, which also enables the world model perform imitation learning with expert demonstra- tions without explicit reward signal. An overview of our proposed method is shown in Figure 1. The detailed train- ing algorithm is shown in Algorithm 2. We also provide a theoretical analysis our training objective, as detailed in Section 4.2 Appendix H.3.
Model Learning We learn the encoder h, latent dynamics d(z,, a) and Q function Q(z, a) jointly by minimizing the objective for prediction consistency and critic learning:
L
Σ Xa ( sg(h(s())1 + Lig 1=0
(11)
where sg is the stop gradient operator and Lig is the inverse soft-Q critic objective, which modification for horizon H and latent representation from Eq.7 based on Eq.6:
(ig(2,n)
H
Σ
x
We introduce our approach for im- itation learning as Inverse Soft-Q Q Learning for Model Predictive Control, or IQ-MPC as an abbreviation. Our architecture consists of four components:
Model Components
1=
+ Eso~Be
(1 VT
Q(zz, as) -
Encoder: z h(ss)
(7)
+ B(s 4a
Latent dynamics: z' d(z, a)
(8)
(12)
4

---

Reward-free World Models for Online Imitation Learning
Compared to Eq.6, the key difference is the second term of the objective, which computes the original value difference a,,,,, (ot.,, (Z) - -V -z)( using only the represen- tation of the initial state Su. This reformulation, derived in
Algorithm 1
IQ-MPC (inference)
Require: learned network parameters
,0,0% o0: initial parameters for N
Lemma H.3 (Appendix H. yields more stable Q estima- tion, as confirmed by the ablation in Appendix E.3. We also apply x" regularization, as noted (Garg et al., 2021). We leverage (0. 1] as a constant discounting weight over the horizon, guaranteeing the influence to be smaller for states and actions farther ahead. Note that À here differs from the environment discount factor 7. All value functions in the objective are computed from Q and policy network via (*(z) = -- (Q (2, a) 3 log я (alz)]. where 3 is the entropy coefficient. We will further discuss the selection of / in the policy learning part. Especially, V (h(s')) is
N, Na: number of sampic/policy trajectories
St, H: current state, rollout horizon
1: Encode state 2, f (se)
2: for each iteration 1..J do
3: Sample N trajectories of length from ((1-11(00)1)))
Sample N, trajectories of length H using TT. dg // Estimate trajectory returns or using do, Qo, no. starting from 2, and initialize Φr = 0:
5:
for all N + N2 trajectories (ar, a..1..........
do
6:
for step t 0..H do
the value function computed by the target Q network Q. z is retrieved by rolling out dynamics model from the latent representation of the first state:
7:
Z++1 f dg (Z), ai)
4 Latent transition
4,+1 ~ (0(122+41)
Compute V (Zc+1) via Q0(2e+1,
B To (ā6+1 1 2(+1)
Z[+1 d(zz, at), Zo h(so)
10:
r(z,, a1) Q (2.2 ar) (1a)-770
Reward decoding
We update the Q, encoder, and dynamics network by mini- mizing Eq. 11 while keeping policy prior fixed.
11:
Φr = Φr + + (r(( ar) - Blog Te (ac (arlz)]))
12:
end for
13:
Φr = Ørr + H (zz)
Terminal value
14:
end for
Policy Prior Learning We choose to learn the policy prior network with the maximum entropy reinforcement learning. We minimize the following maximum entropy RL objective using data sampled from both the expert buffer and the behavioral buffer:
15: // Update parameters μe, a for next iteration:
16: 1 1²,00 3 + MPPI update with ør.
17:
end for
18: return a ~ N(u², " (0")21)
dient:
H
E FiBLaaieBBB [ -((2,, T (2.)) (0gg(((--TTT///)///)
H
x2 B B(B. (IV(()((((VV1))))) (IVQ(2,, &0)le - -1)² (14)
Epen
Σ 1=0
(13)
33 is an entropy coefficient which is fixed scalar. (Hansen et al., 2023) experimented on adaptive entropy coefficient and observed no performance improvement on model pre- dictive control compared to fixed scalar. Therefore, we also choose not to leverage learnable 33 for simplicity. We prove in Theorem H.4 that this policy update can achieve n = argmax=en mingeo (i,(Q,T) T) to find the saddle point.
In 14, S and are sampled from the straight line between samples from expert buffer (s. BE and behavioral buffer (s, a) ~ B, by linear interpolation. V is the gradient with respect to the interpolated input 2, and â, By incorpo- rating this additional objective, we can enforce unit gradient norm over the straight lines between state action distribu- tion Pa and pg. We show the ablation study regarding this regularization term in Appendix E.3.
Balancing Critic and Policy Training We observe un- stable training processes in some tasks due to the imbal- ance between the critic and the policy. When the dis- criminative power of the critic too strong, the policy prior π may fail to learn properly. In those cases the Q value difference between expert batch and behavioral batch ((..a) ((z,, a.) E.eea) (...a))....BBB... (22,, ac) will not converge. To mitigate this issue, we choose to use the Wasserstein-1 metric for gradient penalty (Gulrajani et al., 2017; Garg et al.. 2021) in addition to the original inverse soft-Q objective, enforeing Lipschitz condition for the gra-
4.2. Theoretical Analysis on the Learning Objective
In this section, we demonstrate theoretically that the learn- ing objectives of IQ-MPC effectively minimize the value dif- ference between the current policy and the expert, ensuring that Q-value estimation can follow as the latent dynamics model learns. We begin by utilizing the following lemma established in (Kolev et al., 2024):
Lemma 4.1 (Bounded Suboptimality). Given an unknown latent MDP M and our learned latent MDP M with tran- sition probabilities d and d in the latent state space Z and
5

---

Reward-free World Models for Online Imitation Learning
action space A, and letting Rmax denole the maximum reward of the unknown MDP, the difference between the expected return of the current policy, năo and that of the expert policy, T.. is bounded by:
objective (Haarnoja et al., 2017) for RL with horizon H. After iteration, we execute the first action sampled from the normal distribution de ~ ~V (at)² in the environment to collect new trajectory for behavioral buffer B,.
2Rmax
ITM -n'l== -
1-> DrV MMMMPPPMM PM)
5. Experiments
11
We conduct experiments for locomotion and manipulation tasks demonstrate the effectiveness of our approach. We choose to leverage the online version of IQ-Lear++SA (referred to as IQL+SAC in the experiment plots) (Garg et al., 2021), CFIL+SAC (Freund et al., 2023), and HyPE (Ren et al., 2024) as our baselines for comparison stud- ies. The results presented below for our IQ-MPC model are obtained through planning. We provide an analysis of the computational overhead of our model in Appendix F, The empirical results regarding state-based and visual ex- periments are shown in Section 5.1. We experiment on the reward recovery capability of our IQ-MPC model, for which we reveal the results in Appendix G. We conduct ablation studies for our model, which are discussed in Section 5.2. The details of the environments and tasks can be found in Appendix D. We also analyze the training time and the ro- bustness of our model under noisy environment dynamics. The corresponding results are presented in Appendix F and E.4.
7 R.uax
(z, (1 7)2 Bp dt) d(z' z a
+
72
Our critic and policy objectives can be interpreted as a min- max optimization of Eq.4. Moreover, this approach can be viewed as minimizing a statistical distance with an entropy term, corresponding to Eq.2. Thus, on one hand, our critic and policy objectives effectively minimize T1 in the bound provided in Lemma 1. On the other hand, optimizing our consistency loss in Eq. 11 is approximately minimizing the second term T2 in the bound. Our training objective ensures that as the dynamics model learns, it simultaneously minimizes the upper bound of the deviation between the expected return for current policy π and the expert expected return. A more detailed analysis on T2 is given in Appendix H.3.
4.3. Planning with Policy Prior
5.1. Main Results
Similar to TD-MPC (Hansen et al., 2022) and TD-MPC2 (Hansen et al., 2023), we utilize the Model Predictive Con- trol (MPC) framework for local trajectory optimization over the latent representations and acquire control action by lever- aging Model Predictive Path Integral (MPPI,WWlliams et al., 2015) with sampled action sequences (a,, at+1. a(+H) of length H. Instead of planning with explicit reward mod- els like TD-MPC and TD-MPC2, we estimate the parame- ters (p",a") using derivative free optimization with reward information decoded from the critic's 's estimation:
5.1.1. STATE-BASED EXPERIMENTS
Locomotion Tasks We benchmark our algorithm DM- Control (Tunyasuvunakool et al., 2020), evaluating tasks in both low- and high-dimensional environments. Our method outperforms baselines in performance and training stabil- ity. We use 100 expert trajectories for low-dimensional tasks (Hopper, Walker, Quadruped, Cheetah). 500 for Hu- manoid, and 1000 for Dog (both high-dimensional), Each trajectory contains 500 steps, sampled using trained TD- MPC2 world models (Hansen et al., 2023). Performance is averaged over 3 seeds per task. The results are demon- strated in Figure 2. Our method is comparable to HyPE in the Quadruped Run and Chectah Run tasks, while outper- forming all other baselines in the remaining tasks. We also conducted high-dimensional experiments on various tasks in the Dog environment, with results provided in Appendix E.
,r, 0² = argmax (u.O)
E
[ (T(+T)
(a. (a..... B (1)
H 1
+
> ~ (r(zeth, - logo (ag
h=
(15)
where μ, a E "xm computed by unrolling with (ar,.., ac+H) using dynamics model do. The reward r(z, a) is computed by Q(z, a) V'
dim A. (Z Z(+H) ) are .m ,
Manipulation Tasks We consider manipulation tasks with dexterous hand from MyoSuite (Caggiano al., 2022) to show the capability and robustness of our IQ-MPC model in high-dimensional and complex dynamics scenar- ios. We leverage 100 expert trajectories with 100 steps sampled from trained TD-MPC2 for each task. We evaluate the episode reward and success rate of our model along with
15 is solved by iteratively com- puting soft expected return Ør of sampled actions from N( o²) and update μ., a based on weighted average with ør. We describe the detailed planning procedure in Al- gorithm 1. Eq. 15 is an estimation of the soft-Q learning
6

---

Reward-free World Models for Online Imitation Learning
Malk
Condraped
My
Cheeron
Figure Locomotion Results Our method demonstrates much stabler performance near expert level compared baseline methods. the plots, blue lines refers the online version of IQL+SAC (Garg et al., 2021), orange lines refers to the HyPE method (Ren al., 2024), purple lines refers to the CFIL+SAC (Freund et al., 2023) baseline and red lines refers to our IQ-MPC model. The dotted green lines are the mean episode reward for the expert trajectories used during training.
12 9er
11080
Figure 3. Manipulation Results in MyoSuite Our IQ-MPC shows stable and outperforming results MyoSuite manipulation experiments with dexterous hands. In the plots, the color settings are the same as those in Figure 2. In the Pen Twirl task, the CFIL+SAC agent unable to train after 20K time steps. Thus, we interpolate the rest of time steps with a straight line plot.
SQ-Lear++SAC, HyPE, and CFIL+SAC. We show superior empirical performance in three different tasks, including object holding, pen twirling. and key turning. Regarding the results for episode reward, refer to Figure Table I shows the success rate results. We take the mean for 3 seeds regarding the performance for each task. We have conducted additional experiments on ManiSkil12 (Gu et al., 2023), for which we refer to Appendix E.2.
shallow convolution network and keeping the rest of the model unchanged. We sample the expert data using trained TD-MPC2 models with visual inputs. We take 100 expert trajectories for each task. The expert trajectories contain actions and RGB frame observations. We leverage mod- ification of IQL+SAC our baseline. We add the same convolutional encoder as our IQ-MPC for processing visual inputs and keep the rest of the architecture the same. We perform superior to the baseline model in series of visual experiments in the DMControl environment. We demon- strate the results in Figure 4.
5.1.2. VISUAL EXPERIMENTS
We further investigate the capability of handling visual tasks for IQ-MPC model. We conduct the experiments locomotion tasks in DMControl with visual observations. We demonstrate that our IQ-MPC model can cope with vi- sual modality inputs by only replacing the encoder with a
5.2. A Abltion Studies
In this section. we will show the ablation studies of our model, including the ablation expert trajectory num-

---

Reward-free World Models for Online Imitation Learning
Method
IQLASAC
CFIL+SAC
HyPE
IQ-MPC((Ors)
Key Turn
0.72+t0.04
0.65=#0.08
0.55 0.09
0.87=
Object Hold
0.00 + 0.00
0.01+00.01
0.13 + 0.10
0.96
Pen Twirl
0.00 + 0.00
0.00++0.00
0.00 + 0.00
0.73+-0.08
Table 1. Manipulation Success Rate Results in MyoSuite We evaluate the success rate of IQ-MPC on the Key Turn, Object Hold, and Pen Twirl tasks in MyoSuite. Our IQ-MPC demonstrates strength in handling complex manipulation tasks with dexterous hands and musculoskeletal motor control. We show the results by averaging over 100 trajectories and evaluating over 3 random seeds.
M W
Colece hooee hoor
juy
Figure 4. Results for Visual Experiments Our IQ-MPC (red lines) shows stable and expert-level results in visual observation tasks. In the plots, we denote the ASA with an additional convolutional encoder as IQLASAC (Visual) (blue lines). Our model outperforms IQL+SAC (Visual) in the Cheetah Run and Walker Run, and it has comparable performance in the Walker Walk task. The expert trajectories used for training are sampled from TD-MPC2 trained on visual observations.
bers. Regarding the ablations for objective formulation, gradient penalty selection, and hyperparameter a, we refer to Appendix E.3.
700
Upioode Reward on Hepper Hoop
Petccarko
500
We ablate over the expert trajectories used for IQ-MPC training. We demonstrate our results with 100, 50, 10, and 5 expert trajectories in the Hopper Hop task and Object Hold task. We show that our world model can still reach expert-level performance with only a small amount of expert demonstrations but with slower convergence. The instability observed with 5 expert trajectories in the Hopper Hop task. We reveal the empirical results for this ablation in Figure 5.
100
reportt vigettatie
Erpes
6. Conclusions and Broader Impact
We propose an online imitation learning approach that uti- lizes reward-free world models to address tasks in complex environments. By incorporating latent planning and dynam- ies learning, our model can have a deeper understanding of intricate environment dynamics. We demonstrate stable, expert-level performance on challenging tasks, including dexterous hand manipulation and sigh-diimnsional locomo- tion control. In terms of broader impact, our model holds potential for real-world applications in manipulation and locomotion, particularly for tasks that involve visual inputs and complex environment dynamics.
-2500
-5000
-1500
-10000
-33900
Tireectep 10
Figure 5. Ablation Expert Trajectory Numbers. Performance of IQ-MPC with varying numbers of expert trajectories. Stable expert-level performance is achieved with only 10 expert demon- strations for Hopper Hop (top) and 5 for Object Hold (bottom).
8

---

Reward-free World Models for Online Imitation Learning
Impact Statement
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. Advances in neural informa- tion processing systems, 27, 2014.
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which feel must be specifically highlighted here.
Gu, J., Xiang, F., Li, X., Ling, Z.,, Liu, X., MU, T., Tang, Y., Tao, S., Wei, X., Yao, Y., et al. Maniskil12: A unified benchmark for generalizable manipulation skills. arXiv preprint arXiv:2302.04654 2023.
References
Al-Hafez, F., Tateo, D., Arenz, O., Zhao, G., and Peters, J. Ls-iq: Implicit reward regularization for inverse re- inforcement learning, arXiv preprint arXiv:2303.0053 2023.
Gulrajani, I,, Ahmed, F., Arjovsky, M., Dumoulin, V., and Courville, A. C. Improved training of wasserstein gans. Advances neural information processing systems, 30, 2017.
Ba, J. L. Layer normalization. 1607v:660..06450, 2016.
arXiv preprint
Haarnoja, T., Tang, H., Abbeel, P., and Levine, S. Rein- forcement learning with deep energy- policies. In International conference on machine learning, pp. 1352- 1361. PMLR, 2017.
Baram, N., Anschel, O., and Mannor, S. Model- based adversarial imitation learning. arXiv preprint arXiv:1612.02179, 2016.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, Soft actor-critic: Off-policy maximum entropy deep reinforce- ment learning with a stochastic actor. In International conference machine learning, pp. 1861-1870. PMLR, 2018.
Caggiano, V., Wang, H., Durandau, G., Sartori, M., and Kumar, V. Myosuite-a a contact-rich simulation suite for musculoskeletal motor control. arXiv preprint arXiv:2205. 2022.
Chi, C., Feng, S., Du, Y., Xu, Z.,, Cousineau, E., Burchfiel, B., and Song, S. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
Hafner, D., Lillicrap, T., Ba, and Norouzi, M. Dream control: Learning behaviors by latent imagination. arXiv preprint arXiv: 1912.01033 2019a.
Hafner, D., Lillicrap, T., Fischer, I,, Villegas, R., Ha, D.,
Das, N., Bechtle, S., Davchev, T., Jayaraman, D., Rai, A., and Meier, F. Model-based inverse reinforcement learning from visual demonstrations. In Conference on Robot Learning, pp. 1930-1942. PMLR, 2021.
Lee, H., and Davidson, J. Learning latent dynamics for planning from pixels. In International conference on machine learning, pp. 2555-2565. PMLR, 2019b.
Hafner, D., Lillicrap, T., Norouzi, M., and Ba, J. Mas- tering atari with discrete world models. arXiv preprint arXiv:2010.02/99: 2020.
DeMoss, B., Duckworth, P., Hawes, N., and Posner, I. Ditto: Offline imitation learning with world models. arXiv preprint arXiv:2302.03006 2023.
Hafner, D., Pasukonis, J., Ba, J., and Lillicrap, T. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04/04 2023.
Englert, P., Paraschos, A., Peters, J., and Deisenroth, M. Model-based imitation learning by probabilistic trajectory matching. In 2013 IEEE International Conference on Robotics and Automation, pp. 1922-1927, 2013. doi: 10.1 11110//CA..211.6660888.
Hansen, N., Wang, X., and Su, H. Temporal difference learning for model predictive control. arXiv preprint arXiv:2203.0499 2022.
Florence, P., Lynch, C., Zeng, A., Ramirez, O. A., Wahid, A., Downs, L., Wong, A., Lee, J., Mordatch, I., and Tompson, J. Implicit behavioral cloning, In Conference on Robot Learning, pp. 158-168, PMLR, 2022.
Hansen, N., Su, H., and Wang, X. Td-mpe2: Scalable, ro- bust world models for continuous control. arXiv preprint arXiv:2310. 16828, 2023.
Freund, G. J., Sarafian, E., and Kraus, S. A coupled flow ap- proach to imitation learning. In International Conference on Machine Learning, pp. 10372. PMLR, 2023.
Ho, J. and Ermon, S. Generative adversarial imitation learn- ing. Advances in neural information processing systems, 29, 2016.
Garg, D., Chakraborty, S., Cundy, C., Song, J., and Ermon, S. Iq-learn: Inverse soft-q learning for imitation. Ad- vances in Neural Information Processing Systems, 34: 4028 4028-4039, 2021.
Hu, A., Corrado, G., Griffiths, N., Murez, Z., Gurau, C., Yeo, H., Kendall, A., Cipolla, R., and Shotton, J. Model-based imitation learning for urban driving. Advances in Neural Information Processing Systems, 35:20703-20716 2022.

---

Reward-free World Models for Online Imitation Learning
Igl, M., Kim, D., Kuefler, A., Mougin, P., Shah, P., Shiarlis, K., Anguelov, D., Palatucci, M., White, B., and White- son, S. Symphony: Learning realistic and diverse agents for autonomous driving simulation. In 2022 Interna- tional Conference on Robotics and Automation (ICRA), pp. 2445-2451, IEEE, 2022.
Zhou, J., Wang, R., Liu, X., Jiang, Y., Jiang, S., Tao, J., Miao, J., and Song, S. Exploring imitation learning for autonomous driving with feedback synthesizer and differ- entiable rasterization. In 2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 1450-1457. IEEE, 2021.
Kolev, V., Rafailov, R., Hatch, K., Wu, J., and Finn, C. Effi- cient imitation learning with conservative world models. arXiv preprint arXivv2405.131993 2024.
Zhu, Y., Joshi, A., Stone, P., and Zhu, Y. Viola: Imita- tion learning for vision- manipulation with object proposal priors. In Conference on Robot Learning, pp. 1199-1210. PMLR, 2023.
Kostrikov, I., Nachum, O., and Tompson, Imitation learn- ing via off-policy distribution matching. arXiv preprint arXivv 191 192.00033. 2019.
Ziebart, B. D., Maas, A. L., Bagnell, J. A., and Dey, A. K. Maximum entropy inverse reinforcement learn- ing. In AAAI Conference on Artificial Intelligence, 2008. URL https: org org/CorruaID:333211
Misra, D. Mish: A self regularized non mon-monntonic activa- tion function. arXiv preprint arXiv: 08681, 2019.
Rafailov, R., Yu, T., Rajeswaran, A., and Finn, C. Visual adversarial imitation learning using variational models. Advances in Neural Information Processing Systems, 34: 3016-3028, 2021.
Reddy, S., Dragan, A. D., and Levine, Sqil: Imitation learning via reinforcement learning with sparse rewards. arXiv preprint arXiv: /905. 110, 2019.
Ren, J., Swany, G., Wu, Z. S., Bagnell, J. A., and Choud- hury, Hybrid inverse reinforcement learning. arXiv preprint arXiv:2402.088: 088448, 2024.
Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., Casas, D. d. L., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq. A., et al. Deepmind control suite. arXiv preprint arXiv:1801.00690. 2018
Tunyasuvunakool, S., Muldal, A., Doron, Y., Liu, S., Bohez, S., Merel, J., Erez, T., Lillicrap, T., Heess, N., and Tassa, Y. dm amm_ontrol: Software and tasks for continuous control. Software Impacts, 6:100022, 2020.
Williams, G., Aldrich, A., and Theodorou, E. Model pre- dictive path integral control using covariance variable importance sampling. arXiv preprint arXiv:1509.01199, 2015.
Ye, W., Liu, S., Kurutach, T., Abbeel, P., and Gao, Mas- tering atari games with limited data. Advances in neural information processing systems, 34:25476-25488, 2021.
Yin, Z.HH., Ye, W., Chen, Q. and Gao, Y. Planning for sample efficient imitation learning. Advances in Neural Information Processing Systems, 35:2577-2589, 2022.
Zhang, W., Xu, H., Niu, H., Cheng, P., Li, M., Zhang, H., Zhou, G., and Zhan, X. Discriminator- srriidd model- based offline imitation learning. In Conference on Robot Learning, pp. 1266 1276. PMLR, 2023.
10

---

Reward-free World Models for Online Imitation Learning
A. Hyperparameters and Architectural Details
This section will show the detailed hyperparameters and architectures used our IQ-MPC model.
A.1. World Model Architecture
All of the components are built using MLPs with Layernorm (Ba, 2016) and Mish activation functions (Misra, 2019). We leverage Dropout for Q networks. The amount total learnable parameters for IQ-MPC 4.3M. We depict the architecture
Pytorch-like notation:
Architen ieetttuuur: (
(_encoder: MeduleDict:
(state): Sequential
(0): NozmedLinear (in featufrrrrreesttteeee dim, out features=256, bias=True, act-Xish)
(1) NermedLinear (in features-256, out features=51,,, bias=True, act=SimNorm.
Ldynamics): Sequential
(0): NormedLincar (in_ffaturre=s22Hcciioddddi fut_foatures=512, bias-True, act-Mish)
(1): NormedLinear fin_feeaur,,=512, out_features=512, bias=True, act=Mish)
(2): NormedLinear (i__faauree55511 out_features-512, bias-True, act=SimNorm)
(pi): Sequent
bias=Irue, act=Mish) bias=True, act=Mish) (2): Linear (in features=512, out futt.features2-aatiinn din, biss=True)
(0): NormedLinear (in_teutue,s=b1,, out_ att_feaauues=112
(1) NermedLinear (in features=512, 12attess22, out features==-
(Qs): Vectorized ModuleList
(0-4) : x Sequential
(0) NozmedLinear Set
fet_feaauues=512, bias=True, dropout
=0.01, act=Mish)
(1): : (inalliineatt__fatuueeeeeeee out_ fettfeauure-=512 bias=True, act=Mish)
(2): Linear (in fcaturrsr2- cut_features=1, blas-True)
(_target_ Vectorized Modulel.ist
(0-4) : x Sequential
(0) Normediinear in httarrrriinftttrrsss11aaattiiio
512 blas-True, dropout
=0.01, act=Mish)
(1) : NormedLinear features=512, out_ features=512, bias=True, act=Mish)
(2) : Linear (in, features=512, out out_featuresstt bias=True)
Learnable parameters: 4,274,259
The exact parameters above represent the situation when the state dimension is 91, and the action dimension is 39.

---

Reward-free World Models for Online Imitation Learning
Additionally, we also show the convolutional encoder used in our visual experiments:
(_encoder) MeduleDict
(rgo): Sequentiai
(0) ShistAug()
(1): : PixelPreprocess
(2) : (9nv2d 32, kernel_size- 7), stride- (2, 2))
(3) (elld[nnlace-rrree)
(4) Conv2d 32, sernel_size-(5, 5), stride- (2, 2))
(5): ReLU(inplace-True)
(6) : Conv2d(32, 32, termel_size= (3, 3), stride= (2, 2))
(7): (elU(inplace-Tre))
(8) C32v2d 32, kernel_sizo=13, 3), stride=(1, 1))
(9): : Flatten(star__dimel, did_di==-1)
(10) MnNorrmdimm22
A.2. Hyperparameter Details
The detailed hyperparameters used in IQ-MPC are follows:
The batch size during training is 256.
We balance each part of the loss by assigning weights. For inverse soft Q loss, we assign 0.1. For consistency loss, we assign 20. For the policy and gradient penalty, we assign as the weight.
We leverage À 0.5 in horizon.
We apply the same heuristic discount calculation as TD-MPC2 (Hansen et al., 2023), using as the denominator, with maximum discount of 0.995 and minimum of 0.95.
We iterate times during MPPI planning.
We utilize 512 samples as the batch size for planning.
We select 64 samples via top-k selection during MPPI iteration.
During planning, 24 of the trajectories generated by the policy prior x, while normal distributions generate the rest.
Planning horizon H 3.
The temperature coefficient is 0.5.
We set the learning rate of the model to 3e
The entropy coefficient 3 = le - 4.
We found no sigmificant improvement by adding the Wasserstein-1 gradient penalty (Eq.14) in locomotion tasks. Therefore, we only apply gradient penalty to manipulation tasks.
We use a = 0.5 for x2 divergence x(x) = I 4ey 2
We use soft update coefficient 0.01.
12

---

Reward-free World Models for Online Imitation Learning
B. Training Algorithm
For completeness, we show the pseudo-code for IQ-MPC training in Algorithm
Algorithm 2 IQ-MPC (training)
Require: 0,0 randomly initialized network parameters
71, T, B,, Bg: learning rate, soft update coefficient, horizon discount coefficient, behavioral buffer, expert buffer for training steps do
// Collect episode with IQ-MPC from So
for step = 0...P do
Compute a, with (T(-Iha(s) using Algorithm (s), rt) env.step(a.)
Planning with IQ-MPC
Bx B U (1,,ar, s() s!
Add to behavioral buffer
St+1
end for
// Update reward-free world model using collected data in B and Bg:
for num updates per step do
(st; ar, rt. s().H ~ x U Bg
Combine behavioral and expert batch
Zo h00)))
4 Encode first observation
// Unroll for horizon H
for O...H do
1 do do(2t. at)
d. = QZz at)
end for
Compute critic and consistency loss L(zo.H, 90:H, h(sb. (.. A) Compute policy prior loss L (Zo:H, X)
Equation 11
Equation 13
if use gradient penalty then Compute gradient penalty Lpen (20:11, a0:H; A)
Equation 14
else
ppn 0 end if
0 - 1 Vo Lx -ppen)
Update online network
0 + (1 - -)0 +70
Update target network
end for
end for
13

---

Reward-free World Models for Online Imitation Learning
C. Task Visualizations
visualize each task using the random initialization state of an episode. Regarding the locomotion tasks in DMControl, we show them in Figure 6. Figure shows the visualizations of manipulation tasks with dexterous hands in MyoSuite.
Figure Locomotion Visualizations The visualizations DMControl environments, including Hopper. Cheetah, Walker, Quadruped, Humanoid, and Dog.
Figure 7. Manipulation Visualizations with Dexterous Hands The visualizations for MyoSuite tasks, including Key Turn, Object Hold, and Pen Twirl.
Figure 8. Manipulation Visualizations with Robot Arms The visualizations for ManiSkil12 tasks, including Pick Cube and Lift Cube.
14

---

Reward-free World Models for Online Imitation Learning
D. Environment and Task Details
D.I. Locomotion Environments
We experiment on locomotion environments in DMControl. The details of the corresponding environments are shown in Table Regarding the visual inverse RL tasks, we take RGB image observations with the shape of 64 64 for inputs. Each observation consists of RGB frames.
Environment
Observation Dimension
Action Dimension
High-dimensional?
Hopper
15
4
No
Cheetah
17
No
Quadruped
78
12
No
Walker
24
No
Humanoid
67
24
Yes
Dog
223
38
Yes
Table Enviromment Details for State-based Experiments DMControl. show the environment details experiments on DMControl with state-based observations. High-dimensional tasks have higher hard levels compared to normal tasks for imitation learning.
D.2. Manipulation Environment
We experiment on manipulation tasks in ManiSkil12 and MyoSuite. Among these tasks, 2 of them are in ManiSkil12, for which we describe the task details Table 3, and them are MyoSuite, for which we describe the task details Table
Task
Observation Dimension
Action Dimension
Lift Cube Pick Cube
42
4
51
4
Table Task Details Experiments ManiSkil12. show environment details experiments on ManiSkill2, ManiSkill2
benchmark is built for large-scale robot learning and features extensive randomization and diverse task variations.
Task
Observation Dimension
Action Dimension
Object Hold
91
39
Pen Twirl
83
39
Key Turn
93
Table Task Details for Experiments in MyoSuite. We show the environment details for experiments MyoSuite. The MyoSuite benchmark is designed for physiologically accurate, high-dimensional musculoskeletal motor control, featuring highly complex object manipulation using dexterous hand.
15

---

Reward-free World Models for Online Imitation Learning
E. Additional Experiments
E.1. Additional High-dimennional Locomotion Experiments
To show the robustness our model in nigh-dimensiioal tasks, we conduct locomotion experiments on the Dog environment with different tasks such standing, trotting, and walking, in addition to the running task in Section 5.1.1. The dog environment relatively complex environment due its high-dimensional observation and action spaces. We leverage 500 expert trajectories sampled from trained TD-MPC2 for each experiment. We show the results in Figure
Tro Vinel
Too
Well
1-1
Thandar
Figure 9. Additional High-dimensionall Locomotion Experiments Our IQ MPC shows stable and expert-level performance on different tasks in the Dog environment, which demonstrates model's capability in handling high-dimensional tasks. In the plots, the blue lines and orange lines represent the results from IQL+SAC (Garg et al., 2021) and CFIL+SAC (Freund et al., 2023), respectively, while the red lines represent the results from our IQ-MPC.
E.2. Additional Manipulation Experiments
We also evaluate our method on simpler manipulation tasks ManiSkil12 (Gu al., 2023). We show stable and comparable results in the pick cube task and lift cube task. IQL+SAC (Garg et al., 2021) also performs relatively well in these simple settings. Figure 10 shows the episode rewards results in ManiSkill2 tasks, and Table demonstrates the success rate of each method.
CFIL-SAC
CLICAAC
300
Epioode Reaa
0.3
0.2
64
2.4
2.0
100
Timrsttoo (50
Timestco
Figure 10. Manipulation Results in ManiSkill2 Our IQ-MPC shows stable and comparable results in ManiSkil12 manipulation experi- ments. In the plots, the color settings are the same as those in Figure
Method
IQL+SAC
CFILASAC
IQ-MPC(Ours)
Pick Cube
0.61=0.13
0.00+1.000
0.79:10.05
Lift Cube
0.85 + 0.04
0.01
0.89+00.02
Table 5. Manipulation Success Rate Results in ManiSkil12 We evaluate the success rate of IQ-MPC on pick and lift tasks in the ManiSkil12 environment. We show outperforming empirical results compared IQL+SAC, and CFIL+SAC. We show the results by averaging over 100 trajectories and evaluating over 3 random seeds.
16

---

Reward-free World Models for Online Imitation Learning
E.3. Additional Ablation Studies
In this section, we demonstrate the results of ablating over objective formulation, gradient penalty, and hyperparameter a.
Objective Formulation We observe performance improvement using the reformulated objective Eq. we mentioned in the Model Learning part of Section 4.1. In details, we changed the value temporal difference term =,, [V* (z:) >V*(2,)) into a form only containing value from initial distribution BB [(1 (zo)] This technique is also mentioned in the original IQ-Learn paper (Garg et al., 2021). We have given the theoretical proof for mathematical equivalence in Lemma H.3. In this section, we provide the empirical analysis regarding the effectiveness of this technique in the context our IQ-MPC model. We observe stabler Q estimation leveraging this technique. Moreover, in this case, the difference Q estimation between the expert batch and the behavioral batch can converge more easily, especially for high-dimensional cases like the Humanoid Walk and Dog Run task. The better convergence of Q estimation difference shows that the Q function faces difficulty in distinguishing between expert and behavioral demonstrations, which implies that the policy prior behaves similarly expert demonstrations. The stable Q estimation results in a better learning behavior for the latent dynamics model, which is observed by measuring the prediction consistency loss (The first term in Eq. 1 1) during training. The results in Humanoid Walk task are shown in Figure 11.
Q Estimation Humanoid Walk
0 Estimation Differance Humanoid Walk
Olantec with Maol
116
104
400
Timeatop IHI
Timestop
Prodiction Consistoncy
Humanoid
Epicado Reward Humanoid Walk
cluntie win viautoonaaiion
Olanctive inlulioorration Obacthv OPhernnc
Oblachi with Tarrgori OPferncce
13bue
0.000
-00
0.93
3.00
2.8
Timesico
Figure 11. Ablation on Objective Formulation We show that the Q estimation and training dynamics are stabler by utilizing objective with initial distribution compared to leveraging objective with temporal difference. Moreover, we obtain stable expert-level performance leveraging the objective with the initial distribution. We depict the stability by showing plots regarding Q estimation and prediction consistency. The red lines are the stabler results using the objective with initial distribution while orange lines are the results with temporal difference objective. The ablation experiments are conducted on the Humanoid Walk task.
17

---

Reward-free World Models for Online Imitation Learning
Gradient Penalty We ablate over the Wasserstein-1 metric gradient penalty in Eq.14 with our experiments. This training technique balances the discriminative power of the network to ensure stable policy learning. We show improvement in training stability on the Pick Cube task in the ManiSkill2 environment. By leveraging gradient penalty, we observe stable convergence regarding the difference in Q estimation between expert and behavioral batch. This behavior results in stabler policy learning, especially in tasks with low dimensional state or action space, where expert and behavioral demonstrations can be easily distinguished. The ablation results are shown in Figure 12 and Table 6.
Estimation Difference on Pick Cube
Episode Reward Pick Cube
Tainng with Greclent farety
Thining Groclett fenety
500
Ville
Tirestop (N)
Timestep (10
Figure 12, Ablation on Gradient Penalty. We show the improvement by adopting the Wasscrstein-1 gradient penalty by demonstrating the effect over the convergence of Q-difference, which is the difference between estimation expert and behavioral demonstrations. The converging Q-difference implies stable policy learning and reasonable discriminative power of the Q network. We also demonstrate the effectiveness of the gradient penalty by episode reward during training. The red lines represent results with gradient penalty while orange lines represent results without it.
Gradient Penalty?
Yes
No
Success Rate
0.79-10.05
0.51
Table Ablation on Gradient Penalty with Success Rate We evaluate the success rate of IQ-MPC with and without gradient penalty on ManiSkil12 Pick Cube task. We show the results by averaging over 100 trajectories and evaluating over 3 random seeds.
Hyperparameter Selection We perform ablation study the selection of the hyperparameter a in Eq. The hyperparameter controls the strength of the x2 regularization applied to the inverse soft-Q objective. Intuitively, the last term in Eq.12 12 serves as penalty on the magnitude of the estimated reward. Therefore, smaller values of a result in a larger penalty on the estimated reward magnitude, which helps enforce training stability and prevents Q estimation from exploding. In contrast, larger values a encourage more aggressive estimation of the reward and Q value, increasing the chances of training instability. We experiment with the effect of this hyperparameter in the Humanoid Walk task and conclude that 0.5 is the optimal choice. We present our results in Figure 13.
E.4. Experiments on Noisy Environment Dynamics
In this section, we evaluate the robustness of our IQ-MPC model under noisy and stochastic environment dynamics. HyPE (Ren al., 2024) has demonstrated relatively robust performance when subjected minor noise perturbations in environment transitions. Although our model is primarily designed for fully deterministic settings, observe that it exhibits degree of robustness when handling stochastic environment dynamics.
For our experiments, we adopt the same environment settings as HyPE (Ren et al., 2024), introducing trembling noise probability, Peremble. into the environment transitions. Specifically, we assess the impact of Ptremble on our IQ-MPC model in the Walker Run task. The results indicate that our model maintains certain robustness even in the presence of transition noise. We present the results the Figure 14.
18

---

Reward-free World Models for Online Imitation Learning
-00
1000
322
100
364
0..0
200
150
100
0.0
0.5
3.50
0.75
1,00
122 1.50
1.75
2.00
0.00
0.25
0.50
6.72
1,39
325
1.50
2.35
2.00
Figure 13. Ablation Hyperparameter Selection. In ablation study, found that larger values of may lead higher estimations, resulting in suboptimal and unstable training behavior. Conversely, smaller values a lead to lower Q estimations, which correspond to stronger regularization and may also cause suboptimal performance.
Walker Run
Cheetah Run
1728
LOC4
0.03
0.00
000
Eo cod8 Rewarr
600
800
400
200
100 Timestep
323
150
375
Timestep
Figure Experiments on Noisy Environment Dynamics We evaluate our model's performance the Walker and Cheetah Run task under different values of Peremble, where larger Peremble indicates greater stochasticity in the environment dynamics. Specifically, we experiment with Peremble 0.01 and peremble 0.02, observing only slight performance degradation. This suggests that our model exhibits a degree of robustness to noisy environment dynamics.
E.5. Additional Experiments with Few Expert Demonstrations on Visual Tasks
In this section, we demonstrate that our approach can also learn visual tasks using only 10 expert demonstrations. Results with 10 demonstrations are shown Figure 15. While the convergence slower, successful learning is still achievable under this low-data regime.
Watorr Vaau
Chediah 3an Vbud
Vitid
12:00 Novrr
30
Figure 15. Additional Experiments with Few Expert Demonstrations on Visual Tasks We show that our IQ-MPC method can successfully learn from only 10 expert demonstrations DMControl visual tasks. We present results on the Walker Run, Chectah Run, and Walker Walk tasks, and include performance curves with 100 export demonstrations as a reference.
19

---

Reward-free World Models for Online Imitation Learning
F. Training Time Evaluation
We evaluate the computational overhead of our approach in comparison to model-free baselines (IQL+SAC (Garg et al., 2021), CFIL+SAC (Freund et al., 2023). HyPE (Ren et al., 2024)) and the model- based baseline (HyPER (Ren al., 2024)). Additionally, we assess the training time of IQ-MPC with model predictive control enabled and when interacting solely with the policy prior. The experiments are conducted the Humanoid Walk task, with training time reported in seconds. All baselines are trained using a single RTX 2080 Ti GPU. The results are shown in Figure 16.
Training Time on Humanoid Walk
I5-RPCC wwth MPCI
NPC (with Policy Prigi
350000
IQL-SAC
CFIL+GAC
300000
250000
Time
150000
100000
Timesteg
Figure 16. Computational Overhead We evaluate the computational cost during training of our model the Humanoid task. Leveraging a policy prior for direct interaction, instead of relying on MPC, accelerates the training process but may introduce greater instability. Our model requires less computational time compared to HyPER, although its training remains slower than model-free baselines.
G. Reward Correlation Analysis
We evaluate the ability to recover rewards using a trained IQ-MPC model, which demonstrates our model's capability of han- dling inverse RL tasks. We observe a positive correlation between ground-truth rewards and our recovered rewards. We con- duct this experiment on the DMControl Cheetah Run task and decode rewards via r(z, a) = Qo a) " (=2'2do()) V/ (z'). We evaluate over 5 trajectories sampled from a trained IQ-MPC. The results are revealed in Figure 17.
We further analyze the correlation between the decoded rewards and ground-truth rewards in the Hopper Hop, Cheetah Run, Quadruped Run, and Walker Run tasks. Specifically, we compute the Pearson correlation between the estimated and ground-truth rewards in these settings, using IQL+SAC (Garg et al., 2021) as the comparison baseline. The results are presented in Table 7.
In Figure 17, we observe that the variance of the estimated rewards higher when the ground- truth reward high. One possible explanation for this high variance in the estimated expert rewards is as follows:
There multiple equivalent reward formulations that result in optimal trajectories, and the maximum entropy objective selects the one with the highest entropy. Our actor-critic architecture, optimized with the maximum entropy inverse RL objective, leads to more evenly distributed reward structure for expert demonstrations, Consequently, rewards closer to the expert tend to exhibit higher variance, a phenomenon also observed in (Freund et al., 2023).
H. Addiionna Theoretical Supports
We first give a proper definition of distributions involving latent state representations: Definition H.1. Defin latent state distribution P." h.p? pushforward distribution of original state distribution p. for policy æ given an encoder mapping h S Z.
Definition H.2. Define a latent state-action distribution with policy и as p" on Z x A from an original state-action distribution pa on A with an encoder mapping h S Z.
20

---

Reward-free World Models for Online Imitation Learning
Reward Recovery Correlation
1.0
Normalized Recovered Reward
0.8
0.4
0.2
Ground Reward
Figure Reward Recovery. The IQ-MPC model successfully recovers rewards in the inverse RL setting, showing a positive correlation with ground-truth rewards. This experiment is conducted on the Cheetah Run task with state-based observations from DMControl.
Method
IQL+SAC
IQ-MPC (Ours)
Hopper Hop
0.49
0.88
Cheetah Run
0.79
0.87
Walker Run
0.65
0.91
Quadruped Run
0.88
0.93
Table 7. Pearson Correlations of Reward Recovery We evaluate the Pearson correlation between the decoded rewards from IQL+SAC and IQ-MPC in the Hopper Hop, Cheetah Run, Quadruped Run, and Walker Run tasks. Our results demonstrate that IQ-MPC achieves higher correlation with ground-truie couud--ruuhh rewards when trained on these tasks.
H.1. Objective Equivalence
In this section, we will provide proof for the reformulation of the second term in Eq. 12 for completeness. We borrow the proof from (Garg et al., 2021) and slightly modify fit our setting with latent representations instead of actual states. The proof is demonstrated in Lemma H.3. In 1q 12, use the mean over encoded latent representation batch sampled from the expert buffer Bg to approximate the mean over initial distribution po latent representation.
Lemma H.3 (Objective Equivalence). Given latent transition model d(z' |z, a), latent state distribution PT for time step I and a latent state-action distribution p", we have:
E (2,a)~ (V) y V (z')l = (1 - -)) VV (zo)]
Proof. We decompose the left-hand side into a summation:
E (x,aa~<pp V (z) - ((( V= (z'))
= (1 - 7)
7 E' (V") -IE V* (z')]
00
- 7)
Σ 7 Baain [v* 7) Σ Ex 1/ [2*
(1 7) Ezo (z)))
21

---

Reward-free World Models for Online Imitation Learning
H.2. Policy Update Guarantee
We prove that policy update objective Eq. can search for the saddle point in optimization, which increases (TT Q) with Q fixed, following (Garg et al., 2021). For simplicity, prove it with horizon H 1, and it's generalizable objective with discounted finite horizon.
Theorem H.4 (Policy Update). Updating the policy prior via maximum entropy objective increases (Tg Q) with Q fixed. We assume entropy coefficient 1.
Proof. For a fixed Q:
V (z) (12)) -a/2(=(
- DEL (ox((T((1)) m((12))
exp Q(z,),,)
+ log Σ exp((((z,a))
Ea exp(Q(z,a))
Policy update with maximum entropy objective is optimizing:
(2,)) π argmin, DEL (x(-12) exxp(( exp(p ((, a))
Assume that we have an updated policy n' via gradient descent with learning rate E:
exp Q(z,)) :) Ea exp(Q zz, a))
= π - VV (1(2)))(((()))))
We can obtain (*(z) < V= (z). In regions where o(x) is monotonically non-decreasing and Q fixed, we can have (ig( Q) Lig ( Q)-
H.3. Analysis on the Consistency Loss
We provide more detailed analysis regarding the relationship between minimizing the consistency loss in Eq. and minimizing T2 in the bound provided by Lemma 4. S.eeciffiilll consistency loss directly minimizes the upper bound of T2 under following assumptions:
Assumption H.5. The latent dynamics d 2 x A + Ag is approximately a Gaussian distribution on latent space 2 with N V(ua,
Assumption H.6. The standard deviation of our learned latent dynamics da is close to the actual standard deviation aa- Considering T2 in Lemma 4.1 and neglecting the constant coefficient, according to Pinsker Inequality, we have:
1 P
B
E
With Assumption H.5 and H.6, we can represent the KI. divergence by mean and standard deviation of actual and learned latent dynamics:
(xz a)) = llog δa + 03 + (aa--au Paa²
1
(μa paa)
202
2
20
Given predicted latent state 2' from the learned dynamics d and an actual latent state x' = h(ss) encoded from the next state observation with unknown dynamics d, minimizing the L. loss approximately minimizes the distance between the means of the learned and actual latent dynamics distributions. This, in turn, minimizes the right-hant hand side of the Pinsker inequality under our assumptions. Consequently, our consistency loss minimizes the statistical distance between the dynamics.
22

---

Reward-free World Models for Online Imitation Learning
I. Limitations and Future Work
One limitation our method is its sensitivity to the size of the observation and action spaces. Empirically, we find that in scenarios with low-dimensional observations or actions, the discriminator can become overly powerful, potentially leading training instability, discussed in Section 4.1. While we address this issue by incorporating a Wasserstein-1 gradient penalty, it may not be sufficient in all settings. Future work could explore more robust stabilization techniques tailored to tasks with small observation or action space dimensions. Additionally, applying our method to real-world robotic tasks would be a valuable direction to assess its practical effectiveness and generalization capabilities.
23


Reviews:
[{"review_id": "3bO6YoDZj9", "reviewer": "Reviewer_R4Uu", "summary": "The paper proposes a reward-free world model approach, IQ-MPC, for online imitation learning, addressing challenges in high-dimensional and complex tasks. The method integrates decoder-free latent dynamics models with inverse soft-Q learning, eliminating explicit reward modeling. By reformulating the optimization in Q-policy space and leveraging model predictive control (MPC), the approach achieves stable, expert-level performance on benchmarks including DMControl, MyoSuite, and ManiSkill2.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "We sincerely thank you for your review and your recognition of its strengths, including our novel approach, comprehensive evaluation, and clear presentation. Below, we provide detailed responses to your comments.\n\n### 1: Limited Scope on Visual Tasks and Noisy Dynamics Test\n\nThe evaluated scope on visual tasks is limited (only 3 tasks). The noisy dynamics test (Figure 14) is restricted to one task and minor noise with no comparison to baselines, which may be insufficient to validate robustness claims.\n\n**Answer:**\nWe have conducted additional experiments with visual observations on the Quadruped Walk task, as well as noisy environment dynamics experiments on the Cheetah Run and the dexterous hand manipulation task Key Turn. The table below compares the performance degradation in terms of success rate under noisy environment dynamics between our approach and IQL+SAC:\n\n| $p_{tremble}$ | IQL+SAC        | HyPE         | IQ-MPC         |\n|----------------|----------------|--------------|----------------|\n| 0              | 0.72 ± 0.04   | 0.55 ± 0.09 | 0.87 ± 0.03   |\n| 0.005          | 0.54 ± 0.13   | 0.43 ± 0.14 | 0.79 ± 0.07   |\n| 0.01           | 0.28 ± 0.16   | 0.33 ± 0.13 | 0.61 ± 0.12   |\n\nThe results demonstrate that our approach is more robust to noisy environment dynamics compared to IQL+SAC and HyPE. Additional experimental results for Quadruped Walk with visual observations and Cheetah Run with noisy dynamics are available in the anonymous repository: https://anonymous.4open.science/r/reward-free-C1D5/.\n\nRegarding concerns about the noise level, our noisy setting aligns with HyPE [1], which applied $p_{tremble}$ values of 0.01 and 0.025 for locomotion tasks, comparable to our noise levels in Figure 14. We hope the experimental results provided above sufficiently address your concerns.\n\n[1] Ren, J., Swamy, G., Wu, Z. S., Bagnell, J. A., & Choudhury, S. (2024). Hybrid Inverse Reinforcement Learning. *arXiv preprint arXiv:2402.08848.*\n\n---\n\n### 2: MPC Computation Overhead\n\nMPC’s computation overhead (Appendix F) is acknowledged but not fully quantified against baselines.\n\n**Answer:**\nWe have provided a comparative analysis of computational overhead between our approach (both with and without MPC implementation) and other baseline methods (IQL+SAC, CFIL+SAC, HyPE, HyPER) in Figure 15 of our manuscript. Should any further clarification or additional comparative data be required to support your review, we would be happy to provide more detailed explanations or conduct supplementary analyses. Please feel free to share any specific aspects you would like us to elaborate on, and we will gladly address them accordingly.\n\n---\n\n### 3: Number of Expert Trajectories\n\nWhy do the authors use such a large number of expert trajectories (e.g., 100 expert trajectories for visual tasks in DMC, totaling 100 × 500 expert transition steps)? Given the total number of online training steps is 1M and the common choice in this setting is just 10 expert trajectories [1][2], this seems unusually high. While Figure 5 shows the method’s robustness to fewer expert trajectories, this ablation is only conducted on two tasks. What is the reasoning behind the initial choice?\n\n**Answer:**\nOur approach can also effectively learn using only 10 expert trajectories across all three visual tasks, consistent with the settings in [1][2]. The results using 10 expert trajectories are available in our anonymous repository: https://anonymous.4open.science/r/reward-free-C1D5/. We initially chose to use 100 demonstrations to maintain consistency in the number of expert trajectories across tasks of varying difficulty, as using only 10 expert demonstrations for all tasks would be challenging.\n\n---\n\n### 4: Real-World Experiments\n\nDo the authors plan to extend the proposed method to real-robot experiments? I am eager to see the results in real-world settings, as the paper only presents experiments in simulated environments.\n\n**Answer:**\nYes, we plan to extend our method to real-world experiments in future work. Given our method's promising performance in simulation, we believe it has strong potential to tackle real-world robotic tasks. We appreciate your interest and look forward to sharing future results.", "claims_and_evidence": "1. IQ-MPC outperforms baselines in tasks with high-dimensional observations and complex dynamics and mitigates instability during training. The compelling experimental results across locomotion and manipulation tasks over other baselines support this claim.\n\n2. The reward-free world model learns effectively from expert demonstrations without explicit rewards, and can decode reliable rewards from the learned Q-values. The theoretical bijection between Q value and reward spaces is established. Moreover, Figure 16 and Table 7 show positive correlations between decoded and ground-truth rewards.", "methods_and_evaluation_criteria": "Methods: The integration of latent dynamics with inverse soft-Q learning is novel, and the gradient-free MPC planning aligns with recent trends (e.g., TD-MPC).\n\nEvaluation: The benchmarks are diverse, covering low- and high-dimensional, state-based and visual-based, locomotion and manipulation tasks, which aligns with the problem of complex IL. The use of multiple baselines provides a fair comparison. However, The evaluation of noisy dynamics is limited to one task (Walker Run) with minor perturbations and lack comparison with other baselines, insufficient for a robust assessment.", "theoretical_claims": "The paper includes theoretical analysis in Section 4.2 and Appendix H, focusing on the learning objective’s suboptimality bound (Lemma 4.1) and consistency loss (Appendix H.3). Lemma 4.1 bounds the suboptimality of the policy based on distribution discrepancy between agent and expert and dynamics modeling error. The proof in Appendix H is mathematically sound, aligning with standard RL theory. Appendix H.3 links consistency loss to minimizing KL divergence between true and learned dynamics under Gaussian assumptions, but this analysis seems relatively trivial.", "experimental_designs_or_analyses": "The experimental design is generally sound. The diverse task set and ablation on expert trajectory numbers  enhance the evaluation’s robustness. \n\nHowever, the evaluated scope on visual tasks is limited (only 3 tasks). The noisy dynamics test (Figure 14) is limited to one task and minor noise with no comparison to baselines, may be insufficient to validate robustness claims.", "relation_to_broader_scientific_literature": "The key contributions of the paper are closely related to imitation learning, model-based rl and inverse rl, particularly in the areas of decoder-free world models and inverse soft Q-learning.", "essential_references_not_discussed": "I don't think there are any essential references missing, but I'm not completely sure. I would still recommend considering the opinions of other reviewers.", "other_strengths_and_weaknesses": "Strengths:\n1. Novel integration of inverse soft-Q learning with reward-free, decoder-free world models.\n2. Comprehensive evaluation and analysis.\n3. The paper is well-structured and explains its methodology clearly.\n\nWeakness:\n1. MPC’s computation overhead (Appendix F) is acknowledged but not fully quantified against baselines.\n2. Robustness claim on noisy setting needs more experimental backing.", "other_comments_or_suggestions": "The limitation statement of the proposed method is not found in the paper. It is recommended that the authors explicitly discuss the limitations of the proposed approach.", "questions_for_authors": "1. Why do the authors use such a large number of expert trajectories in the experiment (e.g., 100 expert trajectories for visual tasks in DMC, totaling 100 × 500 expert transition steps)? Given that the total number of online training steps is 1M and the common choice in this setting is just 10 expert trajectories [1][2], this seems unusually high. Although Figure 5 demonstrates the method’s robustness to the number of expert trajectories, showing that performance remains solid with just 10 expert trajectories, this ablation study is only conducted on two tasks. I am curious about the reasoning behind the initial choice.\n2. Do the authors plan to extend the proposed method to real-robot experiments? I am eager to see the results in real-world settings, as the paper only presents experiments in simulated environments.\n\n[1] Rafailov R, Yu T, Rajeswaran A, et al. Visual adversarial imitation learning using variational models[J]. Advances in Neural Information Processing Systems, 2021, 34: 3016-3028.\n\n[2] Wan S, Wang Y, Shao M, et al. Semail: eliminating distractors in visual imitation via separated models[C]//International Conference on Machine Learning. PMLR, 2023: 35426-35443.", "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": "", "extra_scores": {}}, {"review_id": "J2wKvzKqwS", "reviewer": "Reviewer_VeS8", "summary": "The paper proposes a world-model-based approach for online imitation learning. The method aims to address limitations in current imitation learning techniques by leveraging world models to improve training stability and performance. The authors demonstrate that their approach achieves better training efficiency and stability across a wide range of environments with diverse observation spaces. The proposed method is empirically compared against three state-of-the-art baselines, showing superior results.\n\n### Update after rebuttal:\nThe authors have provided clear and helpful responses to my concerns. The rationale for using planning over RL is now better justified, and the discussion of limitations and potential generalization was appreciated. Code availability further improves confidence in reproducibility.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "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": "We sincerely thank you for your review and your recognition of its strengths, including its clarity, the simplicity and effectiveness of our approach, and the solid theoretical grounding with comprehensive empirical validation. Below, we provide detailed responses to your comments.\n\n### 1: Rationale for Using Planning Instead of Reinforcement Learning\n\nThe rationale behind using planning instead of reinforcement learning (RL) for sampling interactions needs further elaboration. Why was planning chosen over RL for sampling environment interactions? How does this choice impact performance in different settings?\n\n**Answer:**\nPlanning is particularly effective for complex tasks with high-dimensional observation and action spaces, such as the Dog Run task. On simpler tasks with low-dimensional observation and action spaces, like the Walker Run task, the performance difference between planning and directly using the policy prior for action sampling is not as significant. Moreover, the planning process is initialized with the RL policy prior, which means it generally enhances performance.\n\n---\n\n### 2: Discussion of Failure Cases and Limitations\n\nA deeper discussion of failure cases and limitations would strengthen the contribution.\n\n**Answer:**\nThank you for the suggestion. We have discussed the potential instability caused by an imbalance between the policy and critic in the paragraph titled Balancing Critic and Policy Training on Page 5 of our manuscript. While we mitigate this issue using gradient penalty techniques, which help achieve strong performance across various tasks, the potential for instability may persist in other settings. We will provide a more detailed explanation and additional examples in the revised manuscript.\n\n---\n\n### 3: Generalization to Unseen Tasks or Domains\n\nCould the proposed method generalize to unseen tasks or domains beyond the evaluated environments?\n\n**Answer:**\nTD-MPC2 [1] has demonstrated the capability to generalize to unseen tasks after multi-task training. However, we did not conduct multi-task training for our IQ-MPC method. As a result, its generalization to unseen tasks would be limited. Expanding our approach to enable better generalization through multi-task training is an interesting direction for future research.\n\n---\n\n### 4: Code Availability and Reproducibility\n\nThe lack of a code appendix makes it difficult to assess the reproducibility of the results. Do the authors plan to release the code to improve result reproducibility?\n\n**Answer:**\nTo ensure reproducibility, we have provided our code in the following anonymous repository for review: https://anonymous.4open.science/r/reward-free-C1D5/. We also plan to make our code publicly available in the future.\n\n---\n\n### 5: Additional Related Work  \n\nTo the best of my knowledge, the related work has been discussed to a sufficient extent. However, the authors may want to consider Altmann et al. (2024), *\"Discriminative Reward Co-Training\"* as additional related work.  \n\n**Answer:**  \nThank you for the suggestion. We will incorporate this additional related work in the revised manuscript.\n\n\n[1] Hansen, N., Su, H., & Wang, X. (2023). TD-MPC2: Scalable, Robust World Models for Continuous Control. *arXiv preprint arXiv:2310.16828.*", "claims_and_evidence": "The paper claims that using a world model enhances the performance of online imitation learning. This claim is supported by extensive empirical evaluations demonstrating improved stability and efficiency compared to three baselines. However, further clarification on why planning was chosen over reinforcement learning (RL) for environment interaction sampling would strengthen the justification for the approach.", "methods_and_evaluation_criteria": "The methodology is sound, employing a well-defined world model framework. The evaluation criteria include performance metrics across multiple environments, making the study comprehensive. However, additional discussion on the choice of baselines and potential limitations in generalization to unseen tasks would be beneficial.", "theoretical_claims": "The theoretical aspects appear well-grounded, and the method aligns with established principles in imitation learning and model-based reinforcement learning.", "experimental_designs_or_analyses": "The experimental setup is robust, involving diverse environments to test generalizability. The comparisons with state-of-the-art baselines further support the effectiveness of the proposed method.", "relation_to_broader_scientific_literature": "The paper builds on prior work in imitation learning and model-based RL, contributing a novel integration of world models for stability improvements.", "essential_references_not_discussed": "To the best of my knowledge, related work has been discussed in sufficient extent. However, the authors may want to consider Altmann et al. 2024, \"Discriminative Reward Co-Training\" as additinal reated work.", "other_strengths_and_weaknesses": "Strengths:\n- The paper is well-written and easy to follow.\n- The approach is simple yet effective in addressing limitations in imitation learning.\n- Theoretical grounding is solid, and empirical validation is comprehensive.\n\nWeaknesses:\n- The rationale behind using planning instead of RL for sampling interactions needs further elaboration.\n- A deeper discussion of failure cases and limitations would strengthen the contribution.\n- The lack of a code appendix makes it difficult to assess the reproducibility of the results.", "other_comments_or_suggestions": "Providing a code appendix and open-sourcing the implementations would significantly improve the reproducibility of the results.\n\nA discussion on the computational cost of the approach compared to baselines would be useful.", "questions_for_authors": "Why was planning chosen over reinforcement learning for sampling environment interactions? How does this choice impact performance in different settings?\n\nCould the proposed method generalize to unseen tasks or domains beyond the evaluated environments?\n\nDo the authors plan to release the code to improve result reproducibility?", "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": "", "extra_scores": {}}, {"review_id": "hUvR4J6H0s", "reviewer": "Reviewer_79r3", "summary": "This paper gives an approach for online imitation learning with reward free world models that learn dynamics efficiently in latent space. The method is evaluated in DMControl, MyoSuite, and ManiSkill2 against several baselines.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "We sincerely thank you for your review and your recognition of its strengths, including our novel approach, theoretical support, strong performance, and clear presentation. Below, we provide detailed responses to your comments.\n\n### 1: Harder Evaluation Tasks in ManiSkill2\n\nThe experiments could be improved by evaluating on harder tasks in ManiSkill2. How does the method compare to baselines in more complex manipulation tasks?\n\n**Answer:**\nWe have conducted experiments on ManiSkill2 using the Lift Cube and Pick Cube tasks in Appendix E.2. Additionally, we extended our evaluation to more challenging tasks in ManiSkill2, and the results are shown below:\n\n| Method      | IQL+SAC        | CFIL+SAC      | IQ-MPC (Ours)  |\n|--------------|----------------|---------------|----------------|\n| Pick YCB    | 0.10 ± 0.05   | 0.00 ± 0.00  | 0.31 ± 0.06   |\n| Stack Cube  | 0.23 ± 0.08   | 0.01 ± 0.01  | 0.57 ± 0.05   |\n\nThese results demonstrate that our approach achieves superior performance in more complex tasks compared to the baselines.\n\n---\n\n### 2: Generalization to Out-of-Distribution (OOD) States\n\nIs there empirical evidence supporting generalization to OOD states?\n\n**Answer:**\nTo provide empirical evidence for OOD state generalization, we compare our performance with a Behavior Cloning (BC) baseline trained on the same expert dataset. Since BC is fully offline and tends to struggle with OOD states, it serves as a suitable comparison. The evaluation was conducted on MyoSuite dexterous hand manipulation tasks, with results presented below in terms of success rate:\n\n| Method       | BC             | IQ-MPC (Ours)  |\n|---------------|----------------|----------------|\n| Key Turn     | 0.31 ± 0.15   | 0.87 ± 0.03   |\n| Object Hold  | 0.26 ± 0.11   | 0.96 ± 0.03   |\n| Pen Twirl    | 0.04 ± 0.02   | 0.73 ± 0.05   |\n\nThese results indicate that our approach significantly outperforms BC, serving as indirect empirical evidence of our method's superior OOD generalization capabilities.\n\n---\n\n### 3: Sample Efficiency Trade-Offs for Reward-Free Methods\n\nWhat are the sample efficiency trade-offs for reward-free methods?\n\n**Answer:**\nThe primary advantage of our approach in terms of sample efficiency arises from its model-based methodology. However, we are not entirely clear on the specific concerns regarding the sample efficiency trade-offs mentioned in your question. Would you kindly clarify further?\n\n---\n\n### 4: Evidence Supporting Reward Recovery\n\nThe method's claim of recovering rewards for inverse reinforcement learning is not well-supported in the main text, with only one environment shown in the reward correlation plot.\n\n**Answer:**\nIn Appendix G of our manuscript, we further support our reward recovery claim by computing the Pearson correlations for reward recovery across four different tasks, as presented in Table 7. This complements the reward correlation plot in Figure 16, providing additional evidence for the validity of our approach.", "claims_and_evidence": "Several key claims are made in this paper.\n\n1. The author claims a novel approach that leverages reward free models for online imitation learning.\n\nThis claim is well supported and the method learns dynamics in latent space efficiently. \n\n2. The method consistently achieves stable, expert-level performance.\n\nWhile the method tends to perform better than baselines on the tasks shown, it has considerable variance in some of the tasks and does not reach expert performance, for example in Figure 3 pen twirl and Figure 10 pick cube and lift cube. The evaluation tasks chosen are on the easy side and on the manipulation tasks the method does not perform consistently better than baselines. \n\n3. The method recovers the rewards for inverse reinforcement learning. \n\nThis is a main claim in the contribution but is not supported in the main text with some evidence in the appendix. The plot for reward recovery correlation is only shown for one environment.", "methods_and_evaluation_criteria": "The benchmark environments make sense but are on the easy side. The experiments could be made better with harder evaluation tasks for example in ManiSkill2.", "theoretical_claims": "The theoretical claims are sound.", "experimental_designs_or_analyses": "The paper compares with relevant baselines across multiple environments. There are a lot of ablation studies to isolate the contribution of different components and test time analysis.", "relation_to_broader_scientific_literature": "The paper is relevant to online imitation learning, which is an active area of research.", "essential_references_not_discussed": "The essential references are discussed.", "other_strengths_and_weaknesses": "The paper addresses a key limitation in prior work, namely the need for explicit reward modeling. It uses novel approaches for optimization and planning and contains theoretical support to back up the claims. The method demonstrate high performance on the baseline tasks. The paper is clear and methodology is motivated. \n\nSome ablations in the paper are performed only on a single task which can be made better. The method can be evaluated on harder tasks which is a main weakness.", "other_comments_or_suggestions": "No other comments or suggestions", "questions_for_authors": "1. How does the method compare to baselines in more complex manipulation tasks?\n\n2. Is there empirical evidence for generalizing to out of distribution states?\n\n3. What are the sample efficiency trade offs for reward free method?", "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": "", "extra_scores": {}}, {"review_id": "rGIFsI93zE", "reviewer": "Reviewer_5SNt", "summary": "This paper presents an approach called IQ-MPC for online imitation learning using reward-free world models. It can alleviate some issues in IL, such as handling complex, high-dimensional inputs and intricate environmental dynamics without explicit reward modeling. Key ideas include leveraging decoder-free world models to learn environmental dynamics purely in latent spaces, using an inverse soft-Q learning objective to stabilize learning, and integrating MPC with a gradient-free planning strategy. The main findings demonstrate superior performance across benchmarks like DMControl, MyoSuite, and ManiSkill2, significantly outperforming baselines in terms of stability, sample efficiency, and robustness to high-dimensional and visual input scenarios.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "We sincerely thank you for your review and your recognition of its strengths, including its clear positioning within the literature, effective extensions of prior works, and practical advantages in complex tasks. Below, we provide detailed responses to your comments.\n\n### 1: Lack of Comparison with Model-Based IL Algorithms\n\nThe authors mention many model-based imitation learning (IL) algorithms in the related works but do not compare them as baselines. Only XXX + SAC are considered, which seems unusual and requires further explanation.\n\n**Answer:**\nWe have conducted experiments with HyPER [1], the model-based version of HyPE. However, we were unable to achieve reasonable scores using this algorithm on most of our tasks. Therefore, we opted to use the model-free version, HyPE, as our baseline. Additionally, we compared our approach with EfficientImitate [2], a model-based IL algorithm that combines EfficientZero with an adversarial imitation learning approach. The results of the comparison on MyoSuite dexterous hand manipulation tasks are shown below:\n\n| Method           | EfficientImitate   | IQ-MPC (Ours)   |\n|-------------------|--------------------|-----------------|\n| Key Turn          | 0.81 ± 0.05       | 0.87 ± 0.03    |\n| Object Hold       | 0.56 ± 0.08       | 0.96 ± 0.03    |\n| Pen Twirl         | 0.31 ± 0.07       | 0.73 ± 0.05    |\n\nThese results demonstrate the effectiveness of our approach compared to EfficientImitate in complex dexterous hand manipulation tasks.\n\n---\n\n### 2: Incremental Contribution and Limited Generalizability\n\nThe work appears somewhat incremental as it heavily relies on IQ-Learn and TD-MPC, and it may not generalize easily to other model-based algorithms.\n\n**Answer:**\nWhile our method builds upon IQ-Learn and TD-MPC, our contribution lies in developing a novel approach for online imitation learning using reward-free world models. Furthermore, the methodology we propose for learning the Q function and eliminating the need for a reward model can potentially be applied to other latent world models. However, further investigation is necessary to determine whether our approach generalizes effectively to other model-based algorithms. We consider this a promising direction for future research.\n\n[1] Ren, J., Swamy, G., Wu, Z. S., Bagnell, J. A., & Choudhury, S. (2024). Hybrid Inverse Reinforcement Learning. *arXiv preprint arXiv:2402.08848.*\n\n[2] Yin, Z. H., Ye, W., Chen, Q., & Gao, Y. (2022). Planning for Sample Efficient Imitation Learning. *Advances in Neural Information Processing Systems, 35*, 2577-2589.", "claims_and_evidence": "The authors claim that IQ-MPC achieves stable, expert-level performance in tasks with high-dimensional observation or action spaces and complex dynamics. This claim is supported by empirical evaluations on diverse benchmarks, where IQ-MPC outperforms existing methods. The paper presents quantitative results demonstrating improved performance metrics, thereby providing clear evidence to support its claims.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria align well with the problem at hand. Leveraging reward-free world models allows the framework to focus on modeling environmental dynamics without relying on explicit reward signals, which is suitable for imitation learning scenarios.\n\nI think the authors can also take the inference time into consideration, and also try some universal world models for diverse tasks to prove some scalability.", "theoretical_claims": "The theorems are derived from other works.", "experimental_designs_or_analyses": "The authors mention lots of model-based IL algorithms in the related works but do not compare them as baselines. Only XXX + SAC are considered, which is wired and requires more explanation.", "relation_to_broader_scientific_literature": "The paper clearly situates itself within existing literature on imitation learning (GAIL, SQIL, IQ-Learn, CFIL) and model-based reinforcement learning (Dreamer, TD-MPC). It effectively extends prior works by removing explicit reward dependencies in world models and providing a novel integration of inverse Q-learning with MPC for imitation learning. The paper notably complements the findings from IQ-Learn (Garg et al., 2021) and TD-MPC series (Hansen et al., 2022, 2023), extending these frameworks into a cohesive IL solution with clear practical advantages, especially in complex control tasks.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Weakness: I think the work is a little incremental since it heavily relies on IQ-Learn and TD-MPC and cannot easily be applied to other model-based algorithms.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The AC read the paper, reviews and discussions. 3/4 reviewers suggested weak accept, which the AC agrees with. 1/4 reviewer suggested accept. The content is ok in terms of addressing challenges in high-dimensional and complex tasks, with MPC based algorithm. The tasks are not very challenging, not sure if the Dreamer line of algorithms can also do it for free.
Decision: Accept (poster)