Reward-free World Models for Online Imitation Learning

We thank the reviewer for the constructive review.

Weaknesses

W1: 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.

A1: We highlight the theoretical and empirical insights as follows:

1. Additional Theoretical Guarantee In Section 6 in the anonymous repository: https://anonymous.4open.science/r/rebuttal-reward-free, we have included an additional theoretical analysis of our approach. We use a sub-optimality bound to show that our training minimizes two key terms: the distance between expert and learned state-action distributions and the divergence between true and learned dynamics. The critic and policy objectives address the first term, while the consistency loss minimizes the second. Together, these ensure the value function $V^\pi$ approximates the expert value $V^{\pi_E}$ as the model learns.

2. Better Performance on Complex Tasks and Sampling Efficiency Empirically, our method demonstrates stable, expert-level performance across a range of complex tasks. By learning a latent transition model, our approach achieves greater stability and better sampling efficiency (especially in high-dimensional settings such as Dog environment) than model-free methods like IQ-Learn.

3. Broad Application of the Reward-free Setup for World Model Training Our approach offers a potential solution for task-oriented world model learning from human demonstrations in cases where rewards are difficult to obtain. This can support world model training in real-world robotics applications.

W2: 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].

A2: We would like to clarify that we did not claim that reward-based world models lead to training instability. Rather, we assert that our approach demonstrates more stable performance compared to previous methodologies. Regarding the additional empirical assessment, we provide the comparison with the model-free version of hybrid IRL (HyPE). We present the experimental results for the extra baseline in the additional materials available in Section 1 in the anonymous repository: https://anonymous.4open.science/r/rebuttal-reward-free.

Regarding the model-based version (HyPER), we find it to be highly costly in terms of model rollouts and prone to instability during training. We were unable to achieve stable, expert-level results with HyPER in our benchmarks.

W3: 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?

A3: We experimented on Walker Run task with diffusion policy, using the same 100 expert demonstrations. We find that we cannot obtain satisfactory performance with this amount of data. The training MSE loss is converged but the policy cannot reach optimal. We present our additional experimental results for diffusion policy in Section 5 in the anonymous repository: https://anonymous.4open.science/r/rebuttal-reward-free.

W4: 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.

A4: We apologize for the confusion. Our model, trained on demonstrations without reward data, is capable of estimating rewards for previously unseen state-action pairs. We will revise the writing of this part.

We appreciate the reviewer’s constructive feedback and have revised the manuscript accordingly. Below, we address the reviewer’s concerns:

1. Deeper technical insights.

A1: We have added the theoretical analysis of our approach in Section 4.2 and Appendix H.3 of the revised manuscript.

2. Additonal baseline for comparsion.

A2: We have added the additional baseline, HyPE, on 6 locomotion tasks and 3 dexterous hand manipulation tasks. The new results are presented in Figure 2, Figure 3, and Table 1 of the revised manuscript.

3. Computational overhead concern.

A3: We have added an analysis of the additional computational time consumption in Appendix F of the revised manuscript.

4. Impact of noisy environment dynamics.

A4: We have added an experiment on noisy environment dynamics, presented in Figure 15 of Appendix E.4 in the revised manuscript.

5. More detailed hyperparameter analysis.

A5: We have included an additional ablation study of the hyperparameter $\alpha$ in Figure 14 of Appendix E.3.

We will continue to revise the manuscript and welcome any further feedback or suggestions the reviewer may have.

We thank the reviewer for the constructive review.

Q1: 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;

A1: One possible explanation for the high variance in the estimated expert rewards is as follows:

There are multiple equivalent reward formulations that lead to optimal trajectories, and the maximum entropy objective selects the one with the highest entropy. Our actor-critic architecture is optimized with the maximum entropy inverse RL objective, leading to an even distribution of rewards for expert demonstrations. As a result, rewards that are closer to the expert tend to exhibit higher variance. A similar phenomenon is also observed in [1].

Q2: 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;

A2: For the sampling complexity analysis, our method converges to the optimum significantly faster than model-free methods (IQL+SAC) in high-dimensional settings. This trend is demonstrated in Figure 10 of our manuscript. A similar phenomenon is observed in the visual experiment with the Walker Walk task, as shown in Figure 5. We should also clarify that better sampling complexity does not refer to a smaller number of expert demonstrations. Instead, sampling complexity refers to the number of environment interactions required to reach optimal performance.

Q3: 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.

A3: We thank the reviewer for pointing out the need to include additional baselines for reward recovery. We will incorporate them accordingly.

[1] Freund, G. J., Sarafian, E., & Kraus, S. (2023, July). A coupled flow approach to imitation learning. In International Conference on Machine Learning (pp. 10357-10372). PMLR.

We thank the reviewer for the constructive review.

Q1: Theoretical Guarantee.

A1: Thank you for your constructive comments. To address your theoretical concern, we first borrow a bound from prior work [1]. According to Theorem 2 in [1], the sub-optimality is bounded:

Given an unknown MDP $\mathcal{M}$ and our learned MDP $\hat{\mathcal{M}}$ with transition probability $d$ and $\hat{d}$ in the latent space $\mathcal{Z}$, and let $R\_\max$ be the maximum of the unknown reward in MDP, the value is bounded by: $|V^{\pi_E}\_\mathcal{M}-V^\pi\_\mathcal{M}|\leq\frac{2R\_{\max}}{1-\gamma}D\_{TV}(\rho^\pi\_{\hat{\mathcal{M}}},\rho^{\pi_E}\_{{\mathcal{M}}})+\frac{\gamma R\_{\max}}{(1-\gamma)^2}\mathbb{E}\_{\rho^\pi\_{\hat{\mathcal{M}}}}\Big[D\_{TV}(d(z'|z,a),\hat{d}(z'|z,a))\Big]$

Our critic and policy objectives can be interpreted as a min-max optimization of Eq. (4) in our manuscript, as outlined in lines 156-158 and 185-186. This approach can be viewed as minimizing a statistical distance with an entropy term, corresponding to Eq. (2) in our manuscript. Thus, our critic and policy objectives effectively minimize the first term, $D\_{TV}(\rho^\pi\_{\hat{\mathcal{M}}}, \rho^{\pi\_E}\_{{\mathcal{M}}})$, in the bound above.

Our consistency loss minimization is minimizing the second term $\mathbb{E}\_{\rho^\pi\_{\hat{\mathcal{M}}}}\Big[D\_{TV}(d(z'|z,a),\hat{d}(z'|z,a))\Big]$ in the bound.

In the specific case of Gaussian dynamics, assuming that the transition probabilities $d$ and $\hat{d}$ are approximately Gaussian and that our estimated standard deviation $\hat{\sigma}$ is close to the actual standard deviation $\sigma$, the consistency loss minimizes the upper bound of $\mathbb{E}\_{\rho^\pi\_{\hat{\mathcal{M}}}}\left[D\_{TV}(d(z'|z,a), \hat{d}(z'|z,a))\right]$.

Using Pinsker inequality, we have: $\mathbb{E}\_{\rho^\pi\_{\hat{\mathcal{M}}}}\Big[D\_{TV}(d(z'|z,a),\hat{d}(z'|z,a))\Big]\leq\mathbb{E}\_{\rho^\pi\_{\hat{\mathcal{M}}}}\sqrt{\frac{1}{2}D\_{KL}(d(z'|z,a),\hat{d}(z'|z,a))}$

With the assumptions, we can represent the KL divergence by mean and standard deviation: $D_{KL}(d(z'|z,a),\hat{d}(z'|z,a))=\log\frac{\hat\sigma}{\sigma}+\frac{\sigma^2+(\mu-\hat\mu)^2}{2\hat\sigma^2}-\frac{1}{2}\approx\frac{(\mu-\hat\mu)^2}{2\sigma^2}$

Given a predicted latent state $\hat{z}'$ from the learned dynamics $\hat{d}$ and an actual latent state $z' = h(s')$ encoded from a state observation with unknown dynamics $d$, minimizing the L2 loss approximately minimizes the distance between the means of the Gaussian distributions. This, in turn, approximately minimizes the right-hand side of the Pinsker inequality. Consequently, our consistency loss minimizes the statistical distance between the dynamics.

In conclusion, our training objective ensures that as the dynamics model learns, it simultaneously minimizes the upper bound of the deviation between the value function $V^\pi$ and the expert value $V^{\pi_E}$. Given that the value function is computed by $V^\pi(s) = \mathbb{E}_{a \sim \pi(\cdot|s)} [Q(s,a) - \log \pi(a|s)]$, this also guarantees that the Q function can follow when the dynamics model is being optimized.

However, in the reviewer's question, the reviewer claimed that rewards learned from IRL methods can be "transferred", which we interpret as meaning transferable across different dynamics. However, IRL methods such as GAIL are dynamics-aware [2], as they match state-action distributions rather than policy distributions. Therefore, the learned rewards should, to some extent, contain information about the dynamics. If the explanation above doesn’t fully address your question, could you please provide further clarification?

Q2: Comparison to baselines:

A2: We provide the comparison with the model-free version of hybrid IRL (HyPE). We present the experimental results for the extra baseline in the additional materials available in Section 1 in the anonymous repository: https://anonymous.4open.science/r/rebuttal-reward-free.

Regarding the model-based version (HyPER), we find it to be highly costly in terms of model rollouts and prone to instability during training. We were unable to achieve stable, expert-level results with HyPER in our benchmarks. We have chosen to leave this as part of our future work.

We draw upon the concept of "reward ambiguity" in inverse reinforcement learning (IRL) as discussed in [3] to reinforce our first point in response to C1 from our previous reply. As noted in [3], a reward function $r(s_t,a_t,s_{t+1})$ retains its optimality under the following transformation: $\tilde r(s_t,a_t,s_{t+1})=r(s_t,a_t,s_{t+1})+\gamma\Phi(s_{t+1})-\Phi(s_t)$ where $\Phi(s)$ represents an arbitrary potential function. Given a transition function $\mathcal{P}(s,a):\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{S}$, the transformed state-action reward can be expressed as: $\tilde r(s_t,a_t)=r(s_t,a_t)+\gamma\Phi(\mathcal{P}(s_t,a_t))-\Phi(s_t)$ If an $(s_t,a_t,s_{t+1})$ tuple arises from dynamics that deviate from $\mathcal{P}(s,a)$, it can affect policy optimality. Consequently, the state-action rewards generated by AIL and many previous IRL methods are inherently dependent on the dynamics.

We sincerely appreciate the reviewer’s thoughtful feedback and are more than happy to provide further clarification if needed. Please feel free to reach out with any additional questions or comments.

[3] Luo, F. M., Cao, X., Qin, R. J., & Yu, Y. (2022). Transferable Reward Learning by Dynamics-Agnostic Discriminator Ensemble. arXiv preprint arXiv:2206.00238.

We thank the reviewer for the constructive review.

The reviewer primarily raised concerns in three areas: writing, mathematical formulations, and other aspects such as baselines and related work. However, we respectfully disagree with the reviewer’s claims regarding the alleged inaccuracies in the mathematical formulations, including $V^\pi(s)$ being mistaken for $Q(s, a)$, errors in the reward term, and the arrangement of expectations and sums. We address these concerns in detail in responses A2, A3, A4, and A5.

Writing Concerns

Q1: Problematic writing.

Q1.1: “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.

A1.1: We acknowledge that this statement is not entirely accurate and will revise the manuscript accordingly.

Q1.2: “ 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.

A1.2: We thank the reviewer for bringing this issue to our attention. The statement "In the objective function, $p_0$ is the initial distribution of states" should be removed. We will revise the manuscript accordingly.

Q1.3: “, and α is a scalar coefficient to ensure stable training and prevent overfitting.” thats a bit unspecific. More explanation is needed or add a citation.

A1.3: We thank the reviewer for highlighting the potentially unclear statement. Regarding the hyperparameter $\alpha$, we have conducted additional ablation studies to evaluate its impact. The results have been included in Section 3 in the anonymous repository: https://anonymous.4open.science/r/rebuttal-reward-free.

Q1.4: 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.

A1.4: We apologize for the unclear statement. Previous world models typically utilize an explicit reward model to fit reward data. In reinforcement learning, this often involves regressing the reward signal retrieved from the environment, whereas in online imitation learning, it is commonly achieved through adversarial training to distinguish between expert and behavioral demonstrations. In contrast, our model eliminates the need for an explicit reward model. Instead, we decode the reward directly from an adversarially trained critic. This represents a key distinction between our approach and prior work. We will revise this section to clarify accordingly.

Mathematical Concerns

Q2: 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.

A2: We apologize for any confusion, but we believe there may be a misunderstanding regarding this objective. The value function $V^\pi(s)$ mentioned in the question should not be $Q(s,a)$ in the objective, as the additional entropy term $\mathcal{H}(\pi)$ in Eq. 1 plays a crucial role in the formulation.

Regarding the derivation of Eq. 4, the objective is derived from Eq. 1 using the inverse Bellman operator defined in Eq. 3 and $V^\pi$, which is introduced in the paragraph following Eq. 3.

The objective in Eq. 1 is:

$\mathbb{E}\_{\rho\_E}[r(s,a)] - \left[\mathbb{E}\_{\rho\_\pi}[r(s,a)] + \mathcal{H}(\pi)\right] - \psi(r)$

which is equivalent to:

$\mathbb{E}\_{\rho_E}[r(s,a)] - \left[\mathbb{E}\_{\rho\_\pi}[r(s,a)] - \mathbb{E}_{\rho\_\pi} \log(\pi(a|s))\right] - \psi(r)$

The first term in Eq. 1, $\mathbb{E}\_{\rho\_E}[r(s,a)]$, can be transformed into the first term in Eq. 4, $\mathbb{E}\_{\rho\_E}[Q(s,a) - \gamma \mathbb{E}\_{s' \sim \mathcal{P}(\cdot|s,a)} V^\pi(s')]$, via the inverse Bellman operator.

In the second term, $\left[\mathbb{E}\_{\rho_\pi}[r(s,a)] - \mathbb{E}\_{\rho_\pi} \log(\pi(a|s))\right]$, since we have a $\log \pi$ term in addition to the expectation $\mathbb{E}\_{\rho\_\pi} r(s,a)$, by employing the inverse Bellman operator and the definition of the value function $V^\pi(s) = \mathbb{E}\_{a \sim \pi(\cdot|s)} [Q(s,a) - \log \pi(a|s)]$, we can derive:

$\mathbb{E}\_{\rho\_\pi}[r(s,a)] - \mathbb{E}\_{\rho\_\pi} \log(\pi(a|s))= \mathbb{E}\_{\rho\_\pi}[Q(s,a) - \gamma \mathbb{E}\_{s' \sim \mathcal{P}(\cdot|s,a)} V^\pi(s')]-\mathbb{E}_{\rho\_\pi} \log(\pi(a|s))= \mathbb{E}\_{(s,a) \sim \rho\_\pi} [V^\pi(s) - \gamma \mathbb{E}\_{s' \sim \mathcal{P}(\cdot|s,a)} V^\pi(s')].$

Finally, we convert the regularizer $\psi(r)$ into $\psi(\mathcal{T}^\pi Q)$ using the inverse Bellman operator.

Q3: “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).

A3: We apologize for the confusion. The transformation can be described in the following details:

1. In Eq. 12, we use expectations over batches sampled from the buffers $\mathcal{B}\_E$ and $\mathcal{B}\_\pi$ to approximate the expectations over $\rho\_E$ and $\rho\_\pi$ in Eq. 6.

2. In Eq. 12, we compute this objective using latent representations $z = h(s)$ instead of the actual states $s$ in Eq. 6, as stated in lines 236-237.

3. Regarding the derivation of the term $\mathbb{E}_{s_0 \sim \mathcal{B}_E}[(1-\gamma)V^\pi(z_0)]$ in Eq. 12 from its original form in Eq. 6, we refer to Appendix F.1 of our manuscript. This is also mentioned in lines 247-248 of the main text.

4. As mentioned in lines 193-194, we employ $\chi^2$ regularization in Eq. 6. Therefore, we use $\phi(x) = x - \frac{1}{4\alpha}x^2$, where the additional $-\frac{1}{4\alpha}x^2$ corresponds to the third term $\mathbb{E}_{(s,a,s') \sim \mathcal{B}} \frac{1}{4\alpha}[Q(z_t, a_t) - \gamma \bar{V}^\pi(h(s'_t))]^2$ in Eq. 12. Empirically, we find it better to compute this term using both expert and behavioral batches, so we sample from $\mathcal{B} = \mathcal{B}\_E \cup \mathcal{B}\_\pi$ rather than only from $\mathcal{B}\_E$ in Eq. 12.

5. As noted in lines 253-254, we compute the target value function $\bar{V}^\pi$ using the target Q network $\bar{Q}(z,a)$ in Eq. 12, which is maintained via soft updates.

6. In Eq. 12, we sum the inverse soft-Q loss over a horizon $H$ using the factor $\lambda$, whereas in Eq. 6, it only represents the loss for a single timestep, as mentioned in lines 236-237 and 249-251.

7. The $V^\pi(s)$ in Eq. 12 should not be $Q(s,a)$, and the reasons are shown in the answer A2.

Q4: In Eq. 12, I think that the sum should be inside the expectation.

A4: These two formulations are equivalent with only notation difference. We can always interchange the sum and expectation under conditions like Fubini's Theorem. In our notation, at each timestep $t$, we sample a batch from the buffer and compute the inverse soft-Q loss for that timestep. We then sum the losses from $t=0$ to $t=H$. This notation is aligned with our actual implementation.

We appreciate the reviewer’s constructive feedback and have revised the manuscript accordingly. Below, we address the reviewer’s concerns:

1: Problematic Causality in writing

1.1: Misleading claim about inverse soft-Q learning in IRL.

A1.1: We have revised the relevant section of the introduction, specifically lines 65–66, as highlighted in blue in the updated manuscript.

1.2: Problem with statement containing $p_0$.

A1.2: To address this issue, we revised lines 162–163 in the updated manuscript, with the changes highlighted in blue.

1.3: Lack of explanation for $\alpha$.

A1.3: We addressed this issue by revising lines 179–183 in the updated manuscript. Additionally, we conducted an ablation study on $\alpha$, with the results presented in Figure 14 of Appendix E.3.

1.4 Misleading statement in Section 4.

A1.4: We addressed this issue by revising the first paragraph of Section 4, with the changes highlighted in blue.

2. Important work missing.

A2: We mentioned LS-IQ's new perspective on the regularizer in lines 182–184 in the revised manuscript.

We will continue to revise the manuscript and welcome any further feedback or suggestions the reviewer may have.

We sincerely thank the reviewer for the thoughtful feedback. We are pleased to hear that many of the previous concerns have been resolved. In the response, the reviewer raised two points that require further revision or clarification:

1. Further clarification on the term $V^\pi(z_t) - \gamma V^\pi(z_{t+1})$.
2. Additional concerns regarding writing and figure labels.

We provide our detailed responses below:

A1: Regarding the term $V^\pi(z_t) - \gamma V^\pi(z_{t+1})$, the reviewer suggested it represents a reward shaping term that is policy-invariant. However, we respectfully hold a different view and suggest that this term is not policy-invariant.

As noted in our previous response, the term $V^\pi(z_t) - \gamma V^\pi(z_{t+1})$ evaluates to $r(z_t, a_t) + \alpha \mathcal{H}(\pi(\cdot | z_t))$. In Eq. (15) of our manuscript, this expression can be understood as the expectation over sequences of $a_t$ sampled from the Gaussian distribution $\mathcal{N}(\mu, \sigma^2)$, applied to the term $r(z_t, a_t) - \alpha \log(\pi(a_t | z_t))$.

Thus, the MPPI planning process effectively optimizes

over multiple steps, where the mean $\mu$ and standard deviation $\sigma$ are parameters optimized during the MPPI planning process. Additionally, in the term $V^\pi(z_t) - \gamma V^\pi(z_{t+1})$, the next state $z_{t+1}$ is unrolled using the action sequence sampled from $\mathcal{N}(\mu, \sigma^2)$ and the model, as described in lines 360–361 of the revised manuscript.

Consequently, the expression

depends on the action sequences sampled from $\mathcal{N}(\mu, \sigma^2)$. The Gaussian distribution $\mathcal{N}(\mu, \sigma^2)$ can be interpreted as a new "policy" optimized during the planning process. Therefore, the term $V^\pi(z_t) - \gamma V^\pi(z_{t+1})$ is dependent on this newly optimized "policy," making it neither policy-invariant nor a purely reward-shaping term. While it shares a similarity in the mathematical form with the reward shaping term $\Phi(z_t) - \gamma \Phi(z_{t+1})$, its underlying meaning is fundamentally different.

If this explanation does not fully address your concerns, could you please clarify your definition of "policy invariance" and "reward-shaping term"? This would allow us to better understand and respond to your feedback.

A2: We appreciate the reviewer’s observation that newly added lines should be removed and that the figure labels need to be larger. We fully agree with these suggestions and will implement the revisions. However, since the manuscript update deadline has passed, we will incorporate these improvements into the final version.

We thank the reviewer again for the valuable comments and are happy to provide further clarifications if needed. Please feel free to reach out with any additional questions or concerns.