Reward-free World Models for Online Imitation Learning

## Abstract

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.

## 1 Introduction

Imitation learning (IL) has garnered considerable attention due to its broad applications across various domains, such as robotic manipulation (Zhu et al., ; Chi et al., ) and autonomous driving (Hu et al., ; Zhou et al., ). 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 et al., ) and Implicit BC (Florence et al., ), highlight the advantages of leveraging large datasets in conjunction with relatively straightforward behavioral cloning (BC) methodologies. However, despite its wide applicability, IL methods that do not incorporate online interaction often suffer from poor generalization outside the expert data distribution, especially when encountering out-of-distribution states. Such limitations make these methods vulnerable to failure, as even minor perturbations in state can lead to significant performance degradation. This is often reflected in issues such as bias accumulation and suboptimal results (Reddy et al., ). These challenges stem from BC’s inability to fully capture the underlying dynamics of the environment and its inherent lack of exploration capabilities (Garg et al., ).

To address these shortcomings, methods like GAIL (Ho & Ermon, ), SQIL (Reddy et al., ), IQ-Learn (Garg et al., ), and CFIL (Freund et al., ) have introduced value or reward estimation to facilitate a deeper understanding of the environment, while leveraging online interactions to enhance exploration. Nevertheless, these approaches continue to face substantial challenges, particularly when applied to tasks with high-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 learning (IRL) techniques, introduces instability during training (Garg et al., ). Recent advancements in world models have demonstrated exceptional performance across a wide range of control tasks, underscoring their potential in complex decision-making and planning scenarios (Hafner et al., ; ; ; ; Hansen et al., ; ). Specifically, world models offer advantages over model-free agents in terms of sampling complexity and future planning capabilities, resulting in superior performance on complex tasks (Hansen et al., ; ; Hafner et al., ). Notably, decoder-free world models, which operate exclusively in latent spaces without reconstruction, have proven to be highly effective and efficient in modeling complex environment dynamics (Hansen et al., ; ).

Motivated by these insights, we explore the application of world models in the context of online imitation learning without rewards, enabling IL agents to develop a deeper understanding of environmental dynamics and improve their performance in tasks characterized by high-dimensional observations and complex dynamics. In this work, we present a novel approach to online imitation learning that leverages the strengths of decoder-free world models, specifically designed for complex tasks involving high-dimensional observations, intricate dynamics, and vision-based inputs. In contrast to conventional latent world models, which rely on reward and Q-function estimation, our approach completely 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-policy space. By utilizing an inverse soft-Q learning objective for the critic network (Garg et al., ), our method derives rewards directly from Q-values and the policy, effectively rendering the world model reward-free. This novel formulation mitigates key challenges in IL, including out-of-distribution errors and bias accumulation.

Through online training with finite-horizon planning based on learned latent dynamics, our method demonstrates strong performance in complex environments. We evaluate our approach across a diverse set of locomotion and manipulation tasks, utilizing benchmarks from DMControl (Tunyasuvunakool et al., ), MyoSuite (Caggiano et al., ), and ManiSkill2 (Gu et al., ), and demonstrate superior empirical performance compared to existing online imitation learning methods.

Our contributions are as follows:

- •
We introduce a novel, robust methodology that leverages world models for online imitation learning, effectively addressing the challenges posed by complex robotics tasks.
- •
We propose an innovative gradient-free planning process, operating without explicit reward modeling, within the context of model predictive control.
- •
We showcase the model’s effectiveness in inverse reinforcement learning tasks by demonstrating a positive correlation between decoded and ground-truth rewards.

## 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, ) formulated the reward learning as a min-max problem similar to GAN (Goodfellow et al., ). Inverse Soft Q-Learning (Garg et al., ) reformulated the learning objective of GAIL and integrated their findings into soft actor-critic (Haarnoja et al., ) and soft Q-learning agents for imitation learning. CFIL (Freund et al., ) introduced a coupled flow approach for reward generation and policy learning using expert demonstrations. ValueDICE (Kostrikov et al., ) proposed an off-policy imitation learning approach by transforming the distribution ratio estimation objective. Das et al. () proposed a model-based inverse RL approach by predicting key points for imitation learning tasks. SQIL (Reddy et al., ) proposed an online imitation learning algorithm with soft Q functions. Diffusion Policy (Chi et al., ) is a recent offline IL method using a diffusion model for behavioral cloning. Implicit BC (Florence et al., ) discovers that treating supervised policy learning with an implicit model generally improves the empirical performance for robot learning tasks.

##### Model-based Reinforcement Learning

Contemporary model-based RL methods often learn a dynamics model for future state prediction via data-driven approaches. PlaNet (Hafner et al., ) was introduced as a model-based learning approach for partially observed MDPs by proposing a recurrent state-space model (RSSM) and an evidence lower-bound (ELBO) training objective. Dreamer algorithm (Hafner et al., ; ; ) is a model-based reinforcement learning approach that uses a learned world model to efficiently simulate future trajectories in a latent space, allowing an agent to learn and plan effectively. TD-MPC series (Hansen et al., )(Hansen et al., ) learns a scalable world model for model predictive control using temporal difference learning objective.

Our approach employs a model-based methodology to address challenges in online imitation learning. By integrating a data-driven approach for latent dynamics learning with planning for control, the agent is able to effectively capture and leverage the underlying environment dynamics. Empirical evaluations demonstrate that our model achieves superior performance on complex online imitation learning tasks compared to existing methods.

## 3 Preliminary

We model the decision-making process in the environment as a Markov Decision Process (MDP), which can be defined as a tuple $\langle\mathcal{S},\mathcal{A},p_{0},\mathcal{P},r,\gamma\rangle$. $\mathcal{S}$ and $\mathcal{A}$ represent state and action space. $p_{0}$ is the initial state distribution and $\mathcal{P}:\mathcal{S}\times\mathcal{A}\rightarrow\Delta_{\mathcal{S}}$ is the transition probability. $r(s,a)\in\mathcal{R}$ is the reward function and $\mathcal{R}$ is the reward space. $\gamma\in(0,1)$ is the discount factor. We denote the expert state-action distribution as $\rho_{E}$ and the behavioral distribution as $\rho_{\pi}$. Similarly, we denote the expert policy as $\pi_{E}$ and the behavioral policy as $\pi$. $\Pi$ is the set of all stochastic stationary policies that sample an action $a\in\mathcal{A}$ given a state $s\in\mathcal{S}$. $\mathcal{Z}$ is the space for the latent representation of the original state observations, and $\mathcal{Q}$ is the space for all possible Q functions. $\mathcal{H}(\cdot)$ represents the entropy of a distribution.

##### Maximum Entropy Inverse Reinforcement Learning

Inverse Reinforcement Learning (IRL) focuses on recovering a specific reward function $r(s,a)$ in the reward space $\mathcal{R}$ given a certain amount of expert samples using expert policy $\pi_{E}$. Maximum entropy IRL seeks to solve this problem by optimizing ${\text{max}}_{r\in\mathcal{R}}{\text{min}}_{\pi\in\Pi}\mathbb{E}_{\rho_{E}}[r(
\mathbf{s},\mathbf{a})]-(\mathbb{E}_{\rho_{\pi}}[r(\mathbf{s},\mathbf{a})]+
\mathcal{H}(\pi))$. GAIL (Ho & Ermon, ) generalized the objective into a form including an explicit reward mapping with a convex regularizer $\psi(r)$:

$$ $\underset{r\in\mathcal{R}}{\text{max}}\underset{\pi\in\Pi}{\text{min}}\quad \mathbb{E}_{\rho_{E}}[r(\mathbf{s},\mathbf{a})]-\mathbb{E}_{\rho_{\pi}}[r( \mathbf{s},\mathbf{a})]-\mathcal{H}(\pi)-\psi(r)$ (1) $$

For a non-restrictive set of reward functions $\mathcal{R}=\mathbb{R}^{\mathcal{S}\times\mathcal{A}}$, the objective can be reformulated into a minimization of the statistical distance between distributions $\rho_{E}$ and $\rho_{\pi}$ (Ho & Ermon, ):

$$ $\underset{\pi}{\min}\quad d_{\psi}(\rho_{\pi},\rho_{E})-\mathcal{H}(\pi)$ (2) $$

##### Inverse Soft-Q Learning

Prior work (Garg et al., ) introduced a bijection mapping $\mathcal{T}^{\pi}:\mathbb{R}^{\mathcal{S}\times\mathcal{A}}\rightarrow\mathbb{
R}^{\mathcal{S}\times\mathcal{A}}$ between Q space $\mathcal{Q}$ and reward space $\mathcal{R}$, i.e., the inverse Bellman operator:

$$ $(\mathcal{T}^{\pi}Q)(\mathbf{s},\mathbf{a})=Q(\mathbf{s},\mathbf{a})-\gamma \mathbb{E}_{\mathbf{s}^{\prime}\sim\mathcal{P}(\cdot|\mathbf{s},\mathbf{a})}V^ {\pi}(\mathbf{s}^{\prime})$ (3) $$

where $V^{\pi}(\mathbf{s})=\mathbb{E}_{\mathbf{a}\sim\pi(\cdot|\mathbf{s})}[Q(\mathbf
{s},\mathbf{a})-\log\pi(\mathbf{a}|\mathbf{s})]$. The reward decoding is defined as $r=\mathcal{T}^{\pi}Q$. By applying the operator $\mathcal{T}^{\pi}$ over Eq.1, prior work reformulated the GAIL training objective in Q-policy space (Garg et al., ):

$$ $\begin{split}\mathcal{J}(\pi,Q)&=\mathbb{E}_{(\mathbf{s},\mathbf{a})\sim\rho_{ E}}\Big{[}Q(\mathbf{s},\mathbf{a})-\gamma\mathbb{E}_{\mathbf{s}^{\prime}\sim \mathcal{P}(\cdot|\mathbf{s},\mathbf{a})}V^{\pi}(\mathbf{s}^{\prime})\Big{]}\\ &-\mathbb{E}_{(\mathbf{s},\mathbf{a})\sim\rho_{\pi}}\Big{[}V^{\pi}(\mathbf{s}) -\gamma\mathbb{E}_{\mathbf{s}^{\prime}\sim\mathcal{P}(\cdot|\mathbf{s},\mathbf {a})}V^{\pi}(\mathbf{s}^{\prime})\Big{]}-\psi(\mathcal{T}^{\pi}Q)\end{split}$ (4) $$

which is the inverse soft-Q objective for critic learning. In the objective function, $p_{0}$ is the initial distribution of states. In this way, we can perform imitation learning by leveraging actor-critic architecture. The critic and policy can be learned by finding the saddle point in a joint optimization problem $Q^{*}=\text{argmax}_{Q\in\mathcal{Q}}\min_{\pi\in\Pi}\mathcal{J}(\pi,Q)$ and $\pi^{*}=\text{argmin}_{\pi\in\Pi}\max_{Q\in\mathcal{Q}}\mathcal{J}(\pi,Q)$. Garg et al. () proved the uniqueness of the saddle point. For a fixed Q, the optimization for policy has a closed-form solution, which is the softmax policy:

$$ $\pi_{Q}(\mathbf{a}|\mathbf{s})=\frac{\exp Q(\mathbf{s},\mathbf{a})}{\sum_{ \mathbf{a}}\exp Q(\mathbf{s},\mathbf{a})}$ (5) $$

In the actor-critic setting, we can optimize the policy $\pi$ using maximum-entropy RL objective, which approximates $\pi_{Q}$, and learn critic using:

$$ $\begin{split}\underset{Q\in\mathcal{Q}}{\max}\mathcal{J}(\pi,Q)&=\underset{Q \in\mathcal{Q}}{\max}\quad\Bigg{[}\mathbb{E}_{(\mathbf{s},\mathbf{a})\sim\rho_ {E}}\Big{[}\phi(Q(\mathbf{s},\mathbf{a})-\gamma\mathbb{E}_{\mathbf{s}^{\prime} \sim\mathcal{P}(\cdot|\mathbf{s},\mathbf{a})}V^{\pi}(\mathbf{s}^{\prime}))\Big {]}\\ &-\mathbb{E}_{(\mathbf{s},\mathbf{a})\sim\rho_{\pi}}\Big{[}V^{\pi}(\mathbf{s}) -\gamma\mathbb{E}_{\mathbf{s}^{\prime}\sim\mathcal{P}(\cdot|\mathbf{s},\mathbf {a})}V^{\pi}(\mathbf{s}^{\prime})\Big{]}\Bigg{]}\end{split}$ (6) $$

where $\phi$ is a concave function. Specifically, if we leverage $\chi^{2}$ regularization, we will have $\phi(x)=x-\frac{1}{4\alpha}x^{2}$, and $\alpha$ is a scalar coefficient to ensure stable training and prevent overfitting.

## 4 Methodology

In conventional latent world models, an explicit reward function is typically used to map latent representations and actions to observed reward data. However, in imitation learning, the model is limited to learning from a finite set of expert state-action demonstrations, making it difficult for traditional world models to directly perform imitation learning. To address this limitation, we propose a reward-free world model that learns solely from expert demonstrations and environment interactions, without requiring explicit reward data. Additionally, our model can decode dense rewards for the environment, allowing it to learn a reward function through state-only interactions and a limited set of expert demonstrations. This allows our model to effectively address inverse reinforcement learning (IRL) tasks.

Figure: Figure 1: 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 in Section 4.1. The policy prior from the world model guides the MPPI planning process along with rewards decoded from Q estimations. The detailed planning process is revealed in Algorithm 1.
Refer to caption: x1.png

### 4.1 Learning Process of a Reward-free World Model

World models used in reinforcement learning settings often contain a reward model $R(\mathbf{z},\mathbf{a})$ that requires supervised learning using explicit reward signals from the online environment interactions or the offline data. However, if our learning objective is able to form a bijection between Q space $\mathcal{Q}$ and reward space $\mathcal{R}$, it would be natural to decode the reward from the Q value instead of learning another separate mapping for the reward, which also enables the world model to perform imitation learning with expert demonstrations without explicit reward signal. An overview of our proposed method is shown in Figure 1. The detailed training algorithm is shown in Algorithm 2.

##### Model Components

We introduce our approach for imitation learning as Inverse Soft-Q Learning for Model Predictive Control, or IQ-MPC as an abbreviation. Our architecture consists of four components:

$$ Encoder: $\displaystyle\quad\mathbf{z}=h(\mathbf{s})$ (7) Latent dynamics: $\displaystyle\quad\mathbf{z}^{\prime}=d(\mathbf{z},\mathbf{a})$ (8) Value function: $\displaystyle\quad\hat{q}=Q(\mathbf{z},\mathbf{a})$ (9) Policy prior: $\displaystyle\quad\mathbf{\hat{a}}=\pi(\mathbf{z})$ (10) $$

where $\mathbf{s}$ and $\mathbf{a}$ are states and actions, $\mathbf{z}$ is latent representations. The policy prior $\pi$ guides the model predictive planning process, along with rewards decoded from the value function $Q$. We maintain two separate replay buffers $\mathcal{B}_{E}$ and $\mathcal{B}_{\pi}$ 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 $\mathcal{B}=\mathcal{B}_{E}\cup\mathcal{B}_{\pi}$. We sample trajectories with short horizons of length $H$ from the replay buffers.

##### Model Learning

We learn the encoder $h$, latent dynamics $d(\mathbf{z},\mathbf{a})$ and Q function $Q(\mathbf{z},\mathbf{a})$ jointly by minimizing the objective for prediction consistency and critic learning:

$$ $\mathcal{L}=\sum_{t=0}^{H}\lambda^{t}\Bigg{(}\mathbb{E}_{(\mathbf{s}_{t}, \mathbf{a}_{t},\mathbf{s}^{\prime}_{t})\sim\mathcal{B}}\|\mathbf{z}_{t+1}- \text{sg}(h(\mathbf{s}^{\prime}_{t}))\|_{2}^{2}\Bigg{)}+\mathcal{L}_{iq}$ (11) $$

where sg is the stop gradient operator and $\mathcal{L}_{iq}$ is the inverse soft-Q critic objective, which is a modification for horizon $H$ and latent representation $\mathbf{z}$ from Eq.7 based on Eq.6:

$$ $\begin{split}\mathcal{L}_{iq}(Q,\pi)=\sum_{t=0}^{H}\lambda^{t}\Bigg{[}-& \mathbb{E}_{(\mathbf{s}_{t},\mathbf{a}_{t},\mathbf{s}^{\prime}_{t})\sim \mathcal{B}_{E}}\Big{[}Q(\mathbf{z}_{t},\mathbf{a}_{t})-\gamma\bar{V}^{\pi}(h( \mathbf{s}^{\prime}_{t}))\Big{]}+\mathbb{E}_{\mathbf{s}_{0}\sim\mathcal{B}_{E} }\Big{[}(1-\gamma)V^{\pi}(\mathbf{z}_{0})\Big{]}\\ &+\mathbb{E}_{(\mathbf{s}_{t},\mathbf{a}_{t},\mathbf{s}^{\prime}_{t})\sim \mathcal{B}}\frac{1}{4\alpha}\Big{[}Q(\mathbf{z}_{t},\mathbf{a}_{t})-\gamma \bar{V}^{\pi}(h(\mathbf{s}^{\prime}_{t}))\Big{]}^{2}\Bigg{]}\end{split}$ (12) $$

Compared to Eq.6, the key difference is the second term of the objective, which computes the original value difference $\mathbb{E}_{(\mathbf{s}_{t},\mathbf{a}_{t},\mathbf{s}^{\prime}_{t})\sim
\mathcal{B}_{\pi}}[V^{\pi}(\mathbf{z}_{t})-\gamma V^{\pi}(\mathbf{z}^{\prime}_
{t})]$ using only the representation of the initial state $\mathbf{s}_{0}$. This reformulation, derived in Theorem 1 (Appendix F.1), yields more stable Q estimation, as confirmed by the ablation in Appendix E.3. We also apply $\chi^{2}$ regularization, as noted in Garg et al. (). We leverage $\lambda\in(0,1]$ as a constant discounting weight over the horizon, guaranteeing the influence to be smaller for states and actions farther ahead. Note that $\lambda$ here differs from the environment discount factor $\gamma$. All value functions in the objective are computed from Q and policy network via $V^{\pi}(\mathbf{z})=\mathbb{E}_{\mathbf{a}\sim\pi(\cdot|\mathbf{z})}[Q(\mathbf
{z},\mathbf{a})-\beta\log\pi(\mathbf{a}|\mathbf{z})]$, where $\beta$ is the entropy coefficient. We will further discuss the selection of $\beta$ in the policy learning part. Especially, $\bar{V}^{\pi}(h(\mathbf{s}^{\prime}))$ is the value function computed by the target Q network $\bar{Q}$. $\mathbf{z}$ is retrieved by rolling out dynamics model from the latent representation of the first state:

$$ $\mathbf{z}_{t+1}=d(\mathbf{z}_{t},\mathbf{a}_{t}),\quad\mathbf{z}_{0}=h( \mathbf{s}_{0})$ $$

We update the Q, encoder, and dynamics network by minimizing Eq.11 while keeping policy prior $\pi$ fixed.

##### 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:

$$ $\mathcal{L}_{\pi}=\sum_{t=0}^{H}\lambda^{t}\Bigg{[}\mathbb{E}_{(\mathbf{s}_{t} ,\mathbf{a}_{t})\sim\mathcal{B}}\Big{[}-Q(\mathbf{z}_{t},\pi(\mathbf{z}_{t}))+ \beta\log(\pi(\cdot|\mathbf{z}_{t}))\Big{]}\Bigg{]}$ (13) $$

$\beta$ is an entropy coefficient which is a fixed scalar. Hansen et al. () experimented on adaptive entropy coefficient and observed no performance improvement on model predictive control compared to a fixed scalar. Therefore, we also choose not to leverage a learnable $\beta$ for simplicity. We prove in Theorem 2 that this policy update can achieve $\pi^{*}=\text{argmax}_{\pi\in\Pi}\min_{Q\in\mathcal{Q}}\mathcal{L}_{iq}(Q,\pi)$ to find the saddle point.

##### Balancing Critic and Policy Training

We observe unstable training processes in some tasks due to the imbalance between the critic and the policy. When the discriminative power of the critic is too strong, the policy prior $\pi$ may fail to learn properly. In those cases the Q value difference between expert batch and behavioral batch $\mathbb{E}_{(\mathbf{s},\mathbf{a})_{(0:H)}\sim\mathcal{B}_{E}}Q(\mathbf{z}_{t
},\mathbf{a}_{t})-\mathbb{E}_{(\mathbf{s},\mathbf{a})_{(0:H)}\sim\mathcal{B}_{
\pi}}Q(\mathbf{z}_{t},\mathbf{a}_{t})$ will not converge. To mitigate this issue, we choose to use the Wasserstein-1 metric for gradient penalty (Gulrajani et al., ; Garg et al., ) in addition to the original inverse soft-Q objective, enforcing Lipschitz condition for the gradient:

$$ $\mathcal{L}_{pen}=\sum_{t=0}^{H}\lambda^{t}\Bigg{[}\mathbb{E}_{(\mathbf{\hat{s }}_{t},\mathbf{\hat{a}}_{t})\sim\mathcal{B}}\Big{(}\|\nabla Q(\mathbf{\hat{z}} _{t},\mathbf{\hat{a}}_{t})\|_{2}-1\Big{)}^{2}\Bigg{]}$ (14) $$

In Eq.14, $\mathbf{\hat{s}}$ and $\mathbf{\hat{a}}$ are sampled from the straight line between samples from expert buffer $(\mathbf{s},\mathbf{a})\sim\mathcal{B}_{E}$ and behavioral buffer $(\mathbf{s},\mathbf{a})\sim\mathcal{B}_{\pi}$ by linear interpolation. $\nabla$ is the gradient with respect to the interpolated input $\mathbf{\hat{z}}_{t}$ and $\mathbf{\hat{a}}_{t}$. By incorporating this additional objective, we can enforce unit gradient norm over the straight lines between state-action distribution $\rho_{\pi}$ and $\rho_{E}$. We show the ablation study regarding this regularization term in Appendix E.3.

Figure: Algorithm 1 IQ-MPC (*inference*)

### 4.2 Planning with Policy Prior

Similar to TD-MPC (Hansen et al., ) and TD-MPC2 (Hansen et al., ), we utilize the Model Predictive Control (MPC) framework for local trajectory optimization over the latent representations and acquire control action by leveraging Model Predictive Path Integral (MPPI)(Williams et al., ) with sampled action sequences $(\mathbf{a}_{t},\mathbf{a}_{t+1},...,\mathbf{a}_{t+H})$ of length $H$. Instead of planning with explicit reward models like TD-MPC and TD-MPC2, we estimate the parameters $(\mu^{*},\sigma^{*})$ using derivative-free optimization with reward information decoded from the critic’s estimation:

$$ $\begin{split}\mu^{*},\sigma^{*}=\underset{(\mu,\sigma)}{\text{argmax}} \underset{(\mathbf{a}_{t},...,\mathbf{a}_{t+H})\sim\mathcal{N}(\mu,\sigma^{2}) }{\mathbb{E}}\Bigg{[}\gamma^{H}V^{\pi}(\mathbf{z}_{t+H})+\sum^{H-1}_{h=t} \gamma^{h}(V^{\pi}(\mathbf{z}_{h})-\gamma V^{\pi}(\mathbf{z}_{h+1}))\Bigg{]} \end{split}$ (15) $$

where $\mu,\sigma\in\mathbb{R}^{H\times m},m=\dim\mathcal{A}$. $(\mathbf{z}_{t},...,\mathbf{z}_{t+H})$ are computed by unrolling with $(\mathbf{a}_{t},..,\mathbf{a}_{t+H})$ using dynamics model $d_{\theta}$. Eq.15 is solved by iteratively computing soft expected return $\phi_{\Gamma}$ of sampled actions from $\mathcal{N}(\mu,\sigma^{2})$ and update $\mu,\sigma$ based on weighted average with $\phi_{\Gamma}$. We describe the detailed planning procedure in Algorithm 1. Eq.15 is an estimation of the soft-Q learning objective (Haarnoja et al., ) for RL with horizon $H$. After iteration, we execute the first action sampled from the normal distribution $a_{t}\sim\mathcal{N}(\mu^{*}_{t},(\sigma^{*}_{t})^{2}\mathrm{I})$ in the environment to collect a new trajectory for behavioral buffer $\mathcal{B}_{\pi}$.

## 5 Experiments

We conduct experiments for locomotion and manipulation tasks to demonstrate the effectiveness of our approach. We choose to leverage the online version of IQ-Learn+SAC (referred to as IQL+SAC in the experiment plots) (Garg et al., ) and CFIL+SAC (Freund et al., ) as our baselines for comparison studies. The empirical results regarding state-based and visual experiments 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 Section 5.2. We conduct ablation studies for our model, which are discussed in Section 5.3. The details of the environments and tasks can be found in Appendix D.

Figure: Figure 2: Locomotion Results Our method demonstrates much stabler performance near expert level compared to baseline methods. In the plots, blue lines refers to the online version of IQ-L+SAC (Garg et al., ), orange lines refers to the CFIL+SAC method (Freund et al., ) 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.
Refer to caption: extracted/5935868/walker_run_comparison.png

### 5.1 Main Results

#### 5.1.1 State-based Experiments

##### Locomotion Tasks

We benchmark our algorithm on DMControl (Tunyasuvunakool et al., ), evaluating tasks in both low- and high-dimensional environments. Our method outperforms baselines in performance and training stability. We use 100 expert trajectories for low-dimensional tasks (Hopper, Walker, Quadruped, Cheetah), 500 for Humanoid, and 1000 for Dog (both high-dimensional). Each trajectory contains 500 steps, sampled using trained TD-MPC2 world models (Hansen et al., ). Performance is averaged over 3 seeds per task. The results are demonstrated in Figure 2. We also conducted high-dimensional experiments on various tasks in the Dog environment, with results provided in Appendix E.1.

Figure: Figure 3: Manipulation Results in MyoSuite Our IQ-MPC shows stable and outperforming results in 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 is unable to train after 20K time steps. Thus, we interpolate the rest of the time steps with a straight line in the plot.
Refer to caption: extracted/5935868/key_turn_comparison.png

**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.**
| Method | IQL+SAC | CFIL+SAC | IQ-MPC(Ours) |
| --- | --- | --- | --- |
| Key Turn | 0.72$\pm$0.04 | 0.65$\pm$0.08 | 0.87$\pm$0.03 |
| Object Hold | 0.00 $\pm$ 0.00 | 0.01$\pm$0.01 | 0.96$\pm$0.03 |
| Pen Twirl | 0.00 $\pm$ 0.00 | 0.00$\pm$0.00 | 0.73$\pm$0.05 |

##### Manipulation Tasks

We consider manipulation tasks with a dexterous hand from MyoSuite (Caggiano et al., ) to show the capability and robustness of our IQ-MPC model in high-dimensional and complex dynamics scenarios. 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 IQ-Learn+SAC 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, we refer to Figure 3. Table 1 shows the success rate results. We take the mean for 3 seeds regarding the performance for each task. We have conducted additional experiments on ManiSkill2 (Gu et al., ), for which we refer to Appendix E.2.

Figure: Figure 4: Reward Recovery. The IQ-MPC model successfully recovers rewards in the inverse RL setting. The recovered rewards exhibit a positive correlation with ground truth rewards. This experiment is conducted on the Cheetah Run task using state-based observations from DMControl.
Refer to caption: extracted/5935868/reward_recovery.png

#### 5.1.2 Visual Experiments

We further investigate the capability of handling visual tasks for our IQ-MPC model. We conduct the experiments on locomotion tasks in DMControl with visual observations. We demonstrate that our IQ-MPC model can cope with visual modality inputs by only replacing the encoder with a 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 a modification of IQL+SAC as 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 a series of visual experiments in the DMControl environment. We demonstrate the results in Figure 5.

### 5.2 Reward Recovery

We evaluate the ability to recover rewards using a trained IQ-MPC model, which demonstrates our model’s capability of handling inverse RL tasks. We observe a positive correlation between ground-truth rewards and our recovered rewards. We conduct this experiment on the DMControl Cheetah Run task and decode rewards via $r(\mathbf{z},\mathbf{a})=Q_{\theta}(\mathbf{z},\mathbf{a})-\gamma\mathbb{E}_{
\mathbf{z}^{\prime}\sim d_{\theta}(\cdot|\mathbf{z},\mathbf{a})}V^{\pi}(
\mathbf{z}^{\prime})$. We evaluate over 5 trajectories sampled from a trained IQ-MPC. The results are revealed in Figure 4.

Figure: Figure 5: 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 IQL+SAC with an additional convolutional encoder as IQL+SAC (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.
Refer to caption: extracted/5935868/walker_walk_rgb_comparison.png

### 5.3 Ablation Studies

In this section, we will show the ablation studies of our model, including the ablation over expert trajectory numbers. Regarding the ablations for objective formulation and gradient penalty selection, we refer to Appendix E.3.

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 is observed with 5 expert trajectories in the Hopper Hop task. We reveal the empirical results for this ablation in Figure 6.

Figure: Figure 6: Ablation on Expert Trajectory Numbers We show the performance of our IQ-MPC model with different numbers of expert trajectories. We can still have stable expert-level performance with only 10 expert demonstrations for the Hopper Hop task in locomotion and 5 expert demonstrations for the Object Hold task with a dexterous hand.
Refer to caption: extracted/5935868/ablation_hopper_traj_num.png

## 6 Conclusions and Broader Impact

We propose an online imitation learning approach that utilizes reward-free world models to address tasks in complex environments. By incorporating latent planning and dynamics 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 high-dimensional locomotion 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.

## Appendix A Hyperparameters and Architectural Details

This section will show the detailed hyperparameters and architectures used in our IQ-MPC model.

### A.1 World Model Architecture

All of the components are built using MLPs with Layernorm (Ba, ) and Mish activation functions (Misra, ). We leverage Dropout for Q networks. The amount of total learnable parameters for IQ-MPC is 4.3M. We depict the architecture in a Pytorch-like notation:

The exact parameters above represent the situation when the state dimension is 91, and the action dimension is 39.

Additionally, we also show the convolutional encoder used in our visual experiments:

### A.2 Hyperparameter Details

The detailed hyperparameters used in IQ-MPC are as 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 1 as the weight.
- •
We leverage $\lambda=0.5$ in a horizon.
- •
We apply the same heuristic discount calculation as TD-MPC2 (Hansen et al., ), using 5 as the denominator, with a maximum discount of 0.995 and a minimum of 0.95.
- •
We iterate 6 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 are generated by the policy prior $\pi$, 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-4$.
- •
The entropy coefficient $\beta=1e-4$.
- •
We found no significant improvement by adding the Wasserstein-1 gradient penalty (Eq.14) in locomotion tasks. Therefore, we only apply gradient penalty to manipulation tasks.
- •
We use $\alpha=0.5$ for $\chi^{2}$ divergence $\phi(x)=x-\frac{1}{4\alpha}x^{2}$.
- •
We use soft update coefficient $\tau=0.01$.

## Appendix B Training Algorithm

For completeness, we show the pseudo-code for IQ-MPC training in Algorithm 2.

Figure: Algorithm 2 IQ-MPC (*training*)

## Appendix C Task Visualizations

We visualize each task using the random initialization state of an episode. Regarding the locomotion tasks in DMControl, we show them in Figure 7. Figure 8 shows the visualizations of manipulation tasks with dexterous hands in MyoSuite.

Figure: Figure 7: Locomotion Visualizations The visualizations for DMControl environments, including Hopper, Cheetah, Walker, Quadruped, Humanoid, and Dog.
Refer to caption: extracted/5935868/1_walker.png

Figure: Figure 8: Manipulation Visualizations with Dexterous Hands The visualizations for MyoSuite tasks, including Key Turn, Object Hold, and Pen Twirl.
Refer to caption: extracted/5935868/1_key_turn.png

Figure: Figure 9: Manipulation Visualizations with Robot Arms The visualizations for ManiSkill2 tasks, including Pick Cube and Lift Cube.
Refer to caption: extracted/5935868/1_pick.png

## Appendix D Environment and Task Details

### D.1 Locomotion Environments

We experiment on 6 locomotion environments in DMControl. The details of the corresponding environments are shown in Table 2. Regarding the visual inverse RL tasks, we take RGB image observations with the shape of $64\times 64\times 9$ for inputs. Each observation consists of 3 RGB frames.

**Table 2: Environment Details for State-based Experiments in DMControl. We show the environment details for experiments on DMControl with state-based observations. High-dimensional tasks have higher hard levels compared to normal tasks for imitation learning.**
| Environment | Observation Dimension | Action Dimension | High-dimensional? |
| --- | --- | --- | --- |
| Hopper | 15 | 4 | No |
| Cheetah | 17 | 6 | No |
| Quadruped | 78 | 12 | No |
| Walker | 24 | 6 | No |
| Humanoid | 67 | 24 | Yes |
| Dog | 223 | 38 | Yes |

### D.2 Manipulation Environment

We experiment on 5 manipulation tasks in ManiSkill2 and MyoSuite. Among these tasks, 2 of them are in ManiSkill2, for which we describe the task details in Table 3, and 3 of them are in MyoSuite, for which we describe the task details in Table 4.

**Table 3: Task Details for Experiments in ManiSkill2. We show the environment details for experiments on ManiSkill2. The ManiSkill2 benchmark is built for large-scale robot learning and features extensive randomization and diverse task variations.**
| Task | Observation Dimension | Action Dimension |
| --- | --- | --- |
| Lift Cube | 42 | 4 |
| Pick Cube | 51 | 4 |

**Table 4: Task Details for Experiments in MyoSuite. We show the environment details for experiments on MyoSuite. The MyoSuite benchmark is designed for physiologically accurate, high-dimensional musculoskeletal motor control, featuring highly complex object manipulation using a dexterous hand.**
| Task | Observation Dimension | Action Dimension |
| --- | --- | --- |
| Object Hold | 91 | 39 |
| Pen Twirl | 83 | 39 |
| Key Turn | 93 | 39 |

## Appendix E Additional Experiments

### E.1 Additional High-dimensional Locomotion Experiments

To show the robustness of our model in high-dimensional tasks, we conduct locomotion experiments on the Dog environment with different tasks such as standing, trotting, and walking, in addition to the running task in Section 5.1.1. The dog environment is a relatively complex environment due to 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 10.

Figure: Figure 10: Additional High-dimensional Locomotion Experiments Our IQ-MPC shows stable and expert-level performance on different tasks in the Dog environment, which demonstrates our model’s capability in handling high-dimensional tasks. In the plots, the color settings are the same as those in Figure 2.
Refer to caption: extracted/5935868/dog_stand_comparison.png

### E.2 Additional Manipulation Experiments

We also evaluate our method on simpler manipulation tasks in ManiSkill2 (Gu et al., ). We show stable and comparable results in the pick cube task and lift cube task. IQL+SAC (Garg et al., ) also performs relatively well in these simple settings. Figure 11 shows the episode rewards results in ManiSkill2 tasks, and Table 5 demonstrates the success rate of each method.

Figure: Figure 11: Manipulation Results in ManiSkill2 Our IQ-MPC shows stable and comparable results in ManiSkill2 manipulation experiments. In the plots, the color settings are the same as those in Figure 2.
Refer to caption: extracted/5935868/pick_cube_comparison.png

**Table 5: Manipulation Success Rate Results in ManiSkill2 We evaluate the success rate of IQ-MPC on pick and lift tasks in the ManiSkill2 environment. We show outperforming empirical results compared to IQL+SAC, and CFIL+SAC. We show the results by averaging over 100 trajectories and evaluating over 3 random seeds.**
| Method | IQL+SAC | CFIL+SAC | IQ-MPC(Ours) |
| --- | --- | --- | --- |
| Pick Cube | 0.61$\pm$0.13 | 0.00$\pm$0.00 | 0.79$\pm$0.05 |
| Lift Cube | 0.85 $\pm$ 0.04 | 0.01$\pm$0.01 | 0.89$\pm$0.02 |

### E.3 Additional Ablation Studies

In this section, we demonstrate the results of ablating over objective formulation and gradient penalty.

##### Objective Formulation

We observe performance improvement using the reformulated objective Eq.12 as we mentioned in the Model Learning part of Section 4.1. In details, we changed the value temporal difference term $\mathbb{E}_{(\mathbf{s}_{t},\mathbf{a}_{t},\mathbf{s}^{\prime}_{t})\sim
\mathcal{B}_{\pi}}[V^{\pi}(\mathbf{z}_{t})-\gamma V^{\pi}(\mathbf{z}^{\prime}_
{t})]$ into a form only containing value from initial distribution $\mathbb{E}_{\mathbf{s}_{0}\sim\mathcal{B}_{E}}[(1-\gamma)V^{\pi}(\mathbf{z}_{0
})]$. This technique is also mentioned in the original IQ-Learn paper (Garg et al., ). We have given the theoretical proof for mathematical equivalence in Theorem 1. In this section, we provide the empirical analysis regarding the effectiveness of this technique in the context of our IQ-MPC model. We observe stabler Q estimation leveraging this technique. Moreover, in this case, the difference in 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 as 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.11) during training. The results in Humanoid Walk task are shown in Figure 12.

Figure: Figure 12: 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.
Refer to caption: extracted/5935868/ablation_human_Q.png

##### 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 Q 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 13 and Table 6.

Figure: Figure 13: Ablation on Gradient Penalty. We show the improvement by adopting the Wasserstein-1 gradient penalty by demonstrating the effect over the convergence of Q-difference, which is the difference between Q estimation on 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.
Refer to caption: extracted/5935868/ablation_pick_qdiff_grad_pen.png

**Table 6: Ablation on Gradient Penalty with Success Rate We evaluate the success rate of IQ-MPC with and without gradient penalty on ManiSkill2 Pick Cube task. We show the results by averaging over 100 trajectories and evaluating over 3 random seeds.**
| Gradient Penalty? | Yes | No |
| --- | --- | --- |
| Success Rate | 0.79$\pm$0.05 | 0.51$\pm$0.11 |

## Appendix F Additional Theoretical Supports

We first give a proper definition of distributions involving latent state representations:

###### Definition 1 .

Define a latent state distribution $\tilde{p}^{\pi}_{t}=h_{*}p^{\pi}_{t}$ as a pushforward distribution of original state distribution $p^{\pi}_{t}$ for policy $\pi$ given an encoder mapping $h:\mathcal{S}\rightarrow\mathcal{Z}$.

###### Definition 2 .

Define a latent state-action distribution with policy $\pi$ as $\tilde{\rho}^{\pi}$ on $\mathcal{Z}\times\mathcal{A}$ from an original state-action distribution $\rho^{\pi}$ on $\mathcal{S}\times\mathcal{A}$ with an encoder mapping $h:\mathcal{S}\rightarrow\mathcal{Z}$.

###### Definition 1 .

###### Definition 2 .

### F.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. () and slightly modify it to fit our setting with latent representations instead of actual states. The proof is demonstrated in Theorem 1. In Eq.12, we use the mean over encoded latent representation batch sampled from the expert buffer $\mathcal{B}_{E}$ to approximate the mean over initial distribution $\tilde{p}_{0}$ on latent representation.

###### Theorem 1 (Objective Equivalence) .

Given a latent transition model $d(\mathbf{z}^{\prime}|\mathbf{z},\mathbf{a})$, a latent state distribution $\tilde{p}^{\pi}_{t}$ for time step $t$ and a latent state-action distribution $\tilde{\rho}^{\pi}$, we have:

$$ $\mathbb{E}_{(\mathbf{z},\mathbf{a})\sim\tilde{\rho}_{\pi}}[V^{\pi}(\mathbf{z}) -\gamma\mathbb{E}_{\mathbf{z}^{\prime}\sim d(\cdot|\mathbf{z},\mathbf{a})}V^{ \pi}(\mathbf{z}^{\prime})]=(1-\gamma)\mathbb{E}_{\mathbf{z}_{0}\sim\tilde{p}_{ 0}}[V^{\pi}(\mathbf{z}_{0})]$ $$

###### Proof.

We decompose the left-hand side into a summation:

$$ $\begin{split}&\mathbb{E}_{(\mathbf{z},\mathbf{a})\sim\tilde{\rho}_{\pi}}[V^{ \pi}(\mathbf{z})-\gamma\mathbb{E}_{\mathbf{z}^{\prime}\sim d(\cdot|\mathbf{z}, \mathbf{a})}V^{\pi}(\mathbf{z}^{\prime})]\\ &=(1-\gamma)\sum^{\infty}_{t=0}\gamma^{t}\mathbb{E}_{\mathbf{z}\sim\tilde{p}^{ \pi}_{t},\mathbf{a}\sim\pi(\mathbf{z})}[V^{\pi}(\mathbf{z})-\gamma\mathbb{E}_{ \mathbf{z}^{\prime}\sim d(\cdot|\mathbf{z},\mathbf{a})}V^{\pi}(\mathbf{z}^{ \prime})]\\ &=(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}\mathbb{E}_{\mathbf{z}\sim\tilde{p}^{ \pi}_{t}}[V^{\pi}(\mathbf{z})]-(1-\gamma)\sum^{\infty}_{t=0}\gamma^{t+1} \mathbb{E}_{\mathbf{z}\sim\tilde{p}^{\pi}_{t+1}}[V^{\pi}(\mathbf{z})]\\ &=(1-\gamma)\mathbb{E}_{\mathbf{z}_{0}\sim\tilde{p}_{0}}[V^{\pi}(\mathbf{z}_{0 })]\end{split}$ $$

∎

###### Theorem 1 (Objective Equivalence) .

###### Proof.

### F.2 Policy Update Guarantee

We prove that policy update objective Eq.13 can search for the saddle point in optimization, which increases $\mathcal{L}_{iq}(\pi,Q)$ with $Q$ fixed, following Garg et al. (). For simplicity, we prove it with horizon $H=1$, and it’s generalizable to objective with discounted finite horizon.

###### Theorem 2 (Policy Update) .

Updating the policy prior via maximum entropy objective increases $\mathcal{L}_{iq}(\pi,Q)$ with $Q$ fixed. We assume entropy coefficient $\beta=1$.

###### Proof.

For a fixed $Q$:

$$ $\begin{split}V^{\pi}(\mathbf{z})&=\mathbb{E}_{a\sim\pi(\cdot|\mathbf{z})}[Q( \mathbf{z},\mathbf{a})-\log(\pi(\mathbf{a}|\mathbf{z}))]\\ &=-D_{KL}\Big{(}\pi(\cdot|\mathbf{z})\Big{\|}\frac{\exp Q(\mathbf{z},\cdot)}{ \sum_{\mathbf{a}}\exp(Q(\mathbf{z},\mathbf{a}))}\Big{)}+\log\Big{(}\sum_{ \mathbf{a}}\exp(Q(\mathbf{z},\mathbf{a}))\Big{)}\end{split}$ $$

Policy update with maximum entropy objective is optimizing:

$$ $\pi^{*}=\text{argmin}_{\pi}D_{KL}\Big{(}\pi(\cdot|\mathbf{z})\Big{\|}\frac{ \exp Q(\mathbf{z},\cdot)}{\sum_{\mathbf{a}}\exp(Q(\mathbf{z},\mathbf{a}))}\Big {)}$ $$

Assume that we have an updated policy $\pi^{\prime}$ via gradient descent with learning rate $\xi$:

$$ $\pi^{\prime}=\pi-\xi~{}\nabla_{\pi}D_{KL}\Big{(}\pi(\cdot|\mathbf{z})\Big{\|} \frac{\exp Q(\mathbf{z},\cdot)}{\sum_{\mathbf{a}}\exp(Q(\mathbf{z},\mathbf{a}) )}\Big{)}$ $$

We can obtain $V^{\pi}(\mathbf{z})<V^{\pi^{\prime}}(\mathbf{z})$. In regions where $\phi(x)$ is monotonically non-decreasing and $Q$ is fixed, we can have $\mathcal{L}_{iq}(\pi^{\prime},Q)>\mathcal{L}_{iq}(\pi,Q)$.
∎

###### Theorem 2 (Policy Update) .

###### Proof.

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 3
presentation: 1
questions: - Equation (4): There seem to be an error in the expectation over the policy state-action distribution. What is V(s) should be Q(s, a). Or some part of the derivation is missing.
- “Compared to Eq.6, the key difference is the second term of the objective, which computes the original value difference E(st,at,s ′ t )∼Bπ [V π (zt) − γV π (z ′ t )] using only the representation of the initial state s0.” This sentence does not explain the transition from Eq. 6 to 12. And still What is V(s) should be Q(s, a).
- In Eq. 12, I think that the sum should be inside the expectation.
- In Eq. 15, the V π (zt) − γV π (z ′ t) is used for the reward. This can not be a valid reward signal, since it is just equivalent to a reward shaping term, which is policy invariant. 
- On what task was the ablation study done, high-dim or low dim? 
- I believe there is an important work missing [1]. This work provides a bridge between SQIL and IQ-Learn, showing that the underlying problems are equivalent, a provide a solution to the instability also encountered during this paper. 


Minor Things:
- Line 212: typo → Q twice


[1] Al-Hafez et. al. LS-IQ: Implicit Reward Regularization for Inverse Reinforcement Learning
soundness: 1
strengths: - the idea is interesting. 
- I believe that (when done right), this is a scaleable approach.
summary: This paper introduces an extension of the TD-MPC framework for imitation, specifically within the inverse reinforcement learning context. Rather than learning an explicit reward function, the authors propose learning implicit rewards in the Q-function space. They use the implicitly learned reward function, along with learned world models, to perform MPC in the latent space. The results demonstrate state-of-the-art performance when compared to other baselines.
weaknesses: - the paper is written poorly. Many conclusion done lack causality and need further clarification (see below).
- mathematically there seem to be many issues (see questions)
- I believe many experiments use badly tuned baselines (see questions)



### Problematic Causality in writing:

- “By utilizing an inverse 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. This novel formulation mitigates key challenges in IL, including out-of-distribution errors and bias accumulation.” → Inverse Q-learning is not the novel formulation that mitigates these problems, these are general advantange of inverse RL. 
- “ In the objective function, p0 is the initial distribution of states. In this way, we can perform imitation learning by leveraging actor-critic architecture.” → unclear why the initial state distribution should allow to leverage actor-critic architectures. Also, p0 is not mentioned anywhere in the equations. 
- “, and α is a scalar coefficient to ensure stable training and prevent overfitting.” thats a bit unspecific. More explanation is needed or add a citation. 
- Section 4, first paragraph. The storyline is not very clear. “In conventional latent world models, an explicit reward function is typically used to map latent representations and actions to observed reward data.  However, in imitation learning, the model is limited to learning from a finite set of expert state-action demonstrations, making it difficult for traditional world models to directly perform imitation learning. To address this limitation, we propose a reward-free world model that learns solely from expert demonstrations and environment interactions, without requiring explicit reward data. ” The authors first note that rewards are often learned alongside world models during training. They then argue that in imitation learning, conventional world models are typically used without rewards and only on expert data (not true), making imitation challenging. To address this, the authors propose their solution: training a world model without rewards.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: The "Weakness" section contains all questions I had. 

Other clarifications:
1. Line 155: there is no p_o in the equation above
2. Line 200-204: I failed to understand the argument at all. One is always free to learn dynamics models without reward?
soundness: 2
strengths: This paper proposes a new approach to interactive imitation learning that is model-based and reward-free. This results in sample-complexity wins (from the model) and stability (from the lack of adversarial reward training). The paper also shows that the methods scale across different environments, including visual imitation learning tasks.

Originality: The work presents a unique integration of IQ-Learn’s inverse soft-Q learning with model predictive control, enabling a reward-free formulation that sidesteps adversarial reward estimation.

Quality: Empirical results support the approach’s benefits in diverse settings, achieving stable, expert-level performance on a range of tasks from DMControl and MyoSuite to ManiSkill2.

Clarity: The paper is easy to follow, and the algorithm does a good job of clearly explaining training and inference procedures. The results section is well organized. 

Significance: The proposed method is practical and the reward-free model-based learning is likely to scale well.
summary: The paper looks at the problem of model-based inverse reinforcement learning. The authors propose a "reward-free" approach by combining reward-free model-free imitation learning (IQ-learn) with learned dynamics models. The proposed method, IQ-MPC, is shown to be more sample efficient in comparison to IQ-learn across a range of continuous control tasks. 

The contributions are:
1. A novel model-based, reward-free IRL algorithm
2. Empirical validation across a range of  both state-based and vision-based environments, including DMControl for locomotion, MyoSuite for dexterous hand manipulation, and ManiSkill2 for object handling
3. Ablations showing that the method is performant even as the number of expert trajectories is reduced.
weaknesses: **Fundamental technical concern:** The nice thing about recovering reward functions in IRL is that rewards transfer while Q doesn't. The Q function entangles reward with environment dynamics. I would imagine as the model changes across iterations, the Q function in the worst case may fail to transfer. I would love to be convinced by math that this is not the case (it may not have manifested in the limited set of experiments). Can the authors provide theoretical analysis demonstrating how their IQ-MPC handles Q-function transfer across model iterations, or to discuss potential failure modes?

**Comparison to baselines:** The paper would benefit from broader comparisons, especially with other model-based and model-free IRL approaches, to support claims of stability and efficiency. For example, how does it compare to hybrid inverse reinforcement learning (both model-free and model-based) https://arxiv.org/pdf/2402.08848v1 which have SOTA results on Mujoco and D4RL environments. Could the authors run against the baselines in https://arxiv.org/pdf/2402.08848v1?

**Computational overhead:** IQ-MPC has a more expensive inference time process than the baselines, making the comparisons unfair. Perhaps the planner can be distilled into the same policy class to have a fair comparison with model-free baselines. Additionally, including the number of interactions in the model would be helpful. One would expect IQ-MPC to have lower interactions with environment, but perhaps a larger (world + model) interactions compared to IQ-Learn. 

**Demonstrations efficiency:** The high number of expert trajectories required (100-500) contrasts with prior work, such as IQ-Learn’s ability to learn with as few as 5-10 trajectories in MuJoCo. It would be good to resolve this discrepancy, perhaps running IQ-Learn in the same mujoco setting as IQ-Learn and hybrid inverse reinforcement learning baselines. 

**Limited novelty:** IQ-MPC appears as a blend of IQ-Learn’s inverse soft-Q learning and TD-MPC’s model predictive control framework, without substantial originality beyond this combination. While the reward-free setup is interesting, the paper could be strengthened by examining distinct advantages from this integration. This could be done via an apples-to-apples comparison with reward-based model-based IRL. 

**Scope of “State-Only” data:** The mention of “state-only” data in the appraoch is misleading, as state-action pairs are needed to learn the Q-function and policy. Clarifying this would prevent confusion.

**Incomplete related work:** The paper does not cite fundamental papers like Ziebart’s max-entropy IRL framework or recent model-based IRL methods that directly relate to the current work. Including these citations would better contextualize IQ-MPC within the IRL landscape. Adding recent model-based IRL techniques to related work would help readers understand how IQ-MPC compares and advances beyond existing methods.
- Ziebart, B. D., Maas, A. L., Bagnell, J. A., & Dey, A. K. (2008, July). Maximum entropy inverse reinforcement learning. In Aaai (Vol. 8, pp. 1433-1438).
- Ren, J., Swamy, G., Wu, Z. S., Bagnell, J. A., & Choudhury, S. (2024). Hybrid inverse reinforcement learning. arXiv preprint arXiv:2402.08848.
- P. Englert, A. Paraschos, J. Peters, and M. P. Deisenroth. Model-based imitation learning by probabilistic trajectory matching. In 2013 IEEE International Conference on Robotics and Automation, pp.1922–1927. 2013. doi:10.1109/ICRA.2013.6630832.
- A. Hu, G. Corrado, N. Griffiths, Z. Murez, C. Gurau, H. Yeo, A. Kendall, R. Cipolla, and J. Shotton.
Model-based imitation learning for urban driving. arXiv preprint arXiv:2210.07729, 2022.
- R. Kidambi, J. Chang, and W. Sun. Mobile: Model-based imitation learning from observation alone, 2021.
- N. Baram, O. Anschel, and S. Mannor. Model-based adversarial imitation learning. Conference on Neural Information Processing Systems, 2016.
- M. Igl, D. Kim, A. Kuefler, P. Mougin, P. Shah, K. Shiarlis, D. Anguelov, M. Palatucci, B. White, and S. Whiteson. Symphony: Learning realistic and diverse agents for autonomous driving simulation. arXiv preprint arXiv:2205.03195, 2022.
- Z.-H. Yin, W. Ye, Q. Chen, and Y. Gao. Planning for sample efficient imitation learning. Conference on Neural Information Processing Systems, 2022.
- R. Rafailov, T. Yu, A. Rajeswaran, and C. Finn. Visual adversarial imitation learning using variational models, 2021.
- B. DeMoss, P. Duckworth, N. Hawes, and I. Posner. Ditto: Offline imitation learning with world models. 2023.
- W. Zhang, H. Xu, H. Niu, P. Cheng, M. Li, H. Zhang, G. Zhou, and X. Zhan. Discriminator-guided model-based offline imitation learning. Conference on Robot Learning, 2023.

Edit: Increased score during rebuttal period

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Please refer to weakness.

I have made every effort to review this paper; however, I would appreciate it if the area chair could consider that its content is not closely aligned with my expertise when evaluating all the reviews.
soundness: 3
strengths: Novelty in Reward-Free Approach: The idea of removing reward dependence by operating in latent space addresses a significant challenge in imitation learning, making the model highly adaptable to complex environments with intricate dynamics.

Comprehensive Evaluation: Experiments conducted on a variety of benchmarks, such as DMControl, MyoSuite, and ManiSkill2, demonstrate the model's capabilities across tasks, including both state-based and visual IL tasks.
summary: The manuscript presents a novel approach to online imitation learning (IL) that leverages reward-free world models, specifically focusing on latent-space representations rather than explicit rewards. The proposed model, IQ-MPC, utilizes inverse soft-Q learning in the policy space, sidestepping traditional optimization in the reward-policy space, which has previously led to instability in online IL applications. IQ-MPC aims to demonstrate robust, expert-level performance across high-dimensional and complex control tasks using a combination of latent dynamics and model predictive control (MPC) in the absence of explicit reward functions.
weaknesses: 1. It is recommended that the author further explain the experimental results in 5.2, such as why the model restores rewards closer to the ground-truth rewards when the medium, and why the model prediction variance increases sharply when the real rewards are high; encourage the author to conduct further experiments to illustrate the extent to which the deviation of reward prediction affects the results of MPPI;

2. The author mentioned that the main advantages of the world model are sampling complexity and its future planning ability. The manuscript focuses on its excellent planning ability but lacks comparison with the baseline in sampling complexity. It is recommended to expand the experiment in 5.3 and compare IQ-MPC with the baseline in the case of a small amount of expert data;

3. Section 3 gives a general mapping of the inverse Bellman equation to recover rewards. Experiments can be added to reflect the performance of other baseline methods in reward recovery.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. A number of prior works seem to observe the instability of IQ-Learn in the regime of low expert demonstrations. The proposed approach seems to handle this instability better, despite following a similar framework. I am curious to hear why the authors posit this to be the case? 
2. I am curious how adding noise to the environment dynamics will affect the performance of all methods and baselines? It seems that prior work has seen these Q-function based learning methods (IQ-Learn in particular) to suffer from stochastic transition dynamics [1].
3. I appreciate the detailing of the hyperparameters used in the appendix. How sensitive is the performance of IQ-MPC to changes in these hyperparameters?
4. As mentioned in the weaknesses section, it would be nice to see a comparison to other online IL baselines such as those in [1], as well as offline BC baselines such as those in [2]. In particular, do offline BC/other online IL methods already come close to expert performance in the case of ~100 expert demonstrations? How about in the case of image-based inputs? 


[1] Ren, Juntao, et al. "Hybrid inverse reinforcement learning." ICML (2024).

[2] Chi, Cheng, et al. "Diffusion policy: Visuomotor policy learning via action diffusion." The International Journal of Robotics Research (2023): 02783649241273668.
soundness: 3
strengths: * This paper is well structured, with a clearly written formulation and decent benchmarks/baselines across which empirical results are shown. The experiments are also thoughtful, covering a variety of questions such as design ablations, impact of number of expert demonstrations, and correlation with reward. The environments for such evaluations are also chosen so as to include both locomotion and manipulation tasks.
* The explanation of the proposed soft-Q learning objective within a learned model is clear and well-grounded in prior work.
* The proposed method shows compelling sample efficiency gains over selected baselines, and addresses the instability of prior interactive IL approaches. The authors also document clearly their use of various techniques to improve training stability.
* The method demonstrates good performance in both state-based and image-based environments.
* The proposed method performs reasoning entirely in latent space, and is not limited by a reconstruction objective when learning the world model, which allows it to handle high-dimensional input such as images without suffering huge computational overhead from the reconstruction objective.
summary: This paper broadly considers problems in the regime of interactive imitation learning by learning with learned world models. It differs from previous model-based imitation learning techniques (e.g. model based inverse reinforcement learning) in that the model does not explicitly learn a reward function, and instead does optimization within the space of Q-functions. In essence, this work seems to lift the insights from IQ-Learn [1] to the space of models-based learning.

[1] Garg, Divyansh, et al. "Iq-learn: Inverse soft-q learning for imitation." Advances in Neural Information Processing Systems 34 (2021): 4028-4039.
weaknesses: ### High Level Feedback:
* The paper introduces a novel framework of reward-free world model, but the main contributions seem to build heavily upon existing work such as IQ-Learn [1] and TD-MPC [2, 3]. I recommend making more explicit highlights of any new theoretical insights or mathematical contributions beyond the removal of reward modeling, or highlighting the transferrable empirical insights from the methodology section.
* The authors claim that reward-based world models lead to training instability could be supported more strongly in the empirical results. In particular. Relatedly, it has been shown that IQ-Learn [1] is rather unstable in prior works, so it would be nice to see how this reward-free world model compares to other works focused on sample-efficiency in online IL (both model-free and model-based), such as [4]. 
* The method mentions the use of MPPI for planning in latent space, which could be potentially computationally demanding due to the large number of sampled trajectories. It would be nice to see a more detailed analysis or mentioning of the computational overhead of the proposed method compared to baselines, especially during real-time execution. It could be helpful for the authors to provide some mentioning of empirical metrics for the computational overhead, perhaps in comparison to baselines.
* While the empirical results demonstrates efficiency comparisons with other online IL methods, it does not show the performance of offline IL methods. In the realizable setting (with 100 expert demonstrations), how does an offline IL method such as Diffusion Policy [5] perform? 

### Low Level Technical Feedback:
* In the first paragraph to Section 4, the authors claim in lines 202-204 that they propose learning "a reward-free world model... without requiring explicit reward data". But the next sentence from lines 204-206 proceed to state, "our model can decode dense rewards for the environment, allowing it to **learn a reward function** through...". This seems self-contradictory, so I recommend authors clean up the language such that it is consistent.

[1] Garg, Divyansh, et al. "Iq-learn: Inverse soft-q learning for imitation." Advances in Neural Information Processing Systems 34 (2021): 4028-4039.

[2] Hansen, Nicklas, Xiaolong Wang, and Hao Su. "Temporal difference learning for model predictive control." arXiv preprint arXiv:2203.04955 (2022).

[3] Hansen, Nicklas, Hao Su, and Xiaolong Wang. "Td-mpc2: Scalable, robust world models for continuous control." arXiv preprint arXiv:2310.16828 (2023).

[4] Ren, Juntao, et al. "Hybrid inverse reinforcement learning." ICML (2024).

[5] Chi, Cheng, et al. "Diffusion policy: Visuomotor policy learning via action diffusion." The International Journal of Robotics Research (2023): 02783649241273668.

Meta-review: The paper proposes IQ-MPC, a model-based, reward-free imitation learning framework that combines inverse soft-Q learning (IQ-Learn) with model predictive control (MPC). Instead of explicitly learning reward functions, IQ-MPC operates in the Q-function space and employs latent dynamics models for planning. The framework is evaluated across diverse domains (e.g., DMControl, MyoSuite, ManiSkill2), showing sample improved efficiency and scalability for both state-based and visual imitation learning tasks. Ablation studies highlight the model's robustness to reduced expert trajectories.

Reasons to accept
- IQ-MPC does not rely on explicit reward functions, addressing key challenges of instability and reward bias in imitation learning.
- The approach performs well across state-based and vision-based tasks.
- Comprehensive experiments on various benchmarks show satisfactory performance and improved sample efficiency compared to baselines.
- The paper provides a clear explanation of the proposed method, training process, and inference procedures.
- Thoughtful ablations examine the impact of design choices and varying numbers of expert trajectories.

Reasons to reject
- The method heavily builds on existing work (IQ-Learn and TD-MPC) with limited new theoretical insights. The contribution is mainly in combining these frameworks.
- Broader comparisons with other model-based and model-free IRL approaches are missing, particularly to recent state-of-the-art hybrid methods.
- Computational overhead and interaction efficiency comparisons are underexplored.
- IQ-MPC requires significantly more expert demonstrations (100-500) compared to prior works like IQ-Learn, which can learn with as few as 5-10.
- Many related works were missing. (addressed by the rebuttal)
- The reviewers were concerned with the quality of writing. (mostly fixed by the revised paper)
- During the AC-reviewer discussion phase, multiple reviewers shared the concern that the proposed method resembles a potential-based reward shaping term, and optimizing such a term should not alter the optimal policy learned.

While this paper studies a promising research direction and presents an interesting approach, I believe its weaknesses outweigh its strengths. Consequently, I recommend rejecting the paper.

Decision: Reject