Reward-free World Models for Online Imitation Learning

Authors: Shangzhe Li, Zhiao Huang, Hao Su
Primary area: reinforcement learning
Keywords: World Models, Imitation Learning, Inverse Reinforcement Learning

Abstract:
Imitation learning (IL) enables agents to acquire skills directly from expert demonstrations, providing a compelling alternative to reinforcement learning. However, prior online IL approaches struggle with complex tasks characterized by high-dimensional inputs and complex dynamics. In this work, we propose a novel approach to online imitation learning that leverages reward-free world models. Our method learns environmental dynamics entirely in latent spaces without reconstruction, enabling efficient and accurate modeling. We adopt the inverse soft-Q learning objective, reformulating the optimization process in the Q-policy space to mitigate the instability associated with traditional optimization in 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.

Decision: reject
Reviewer scores: 5, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729708077889

thread_type:
official_review

<thread_object>
creator: reviewer_yqjv

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1729931683293

thread_type:
official_review

<thread_object>
creator: reviewer_erye

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730667280276

thread_type:
official_review

<thread_object>
creator: reviewer_d7rz

<content>
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.

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.

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

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730905915733

thread_type:
official_review

<thread_object>
creator: reviewer_4dpd

<content>
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.

strengths:
- the idea is interesting. 
- I believe that (when done right), this is a scaleable approach.

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.

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
</content>
</thread_object>

</thread>
<thread>
creation date:
1732578333846

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Reviewer-author discussion

comment:
Dear Reviewers,

We sincerely thank you for taking the time to review our paper and for providing thoughtful and constructive feedback. We have put considerable effort into addressing your concerns in our rebuttal, and we hope we have clarified all the issues raised.

As the discussion phase deadline approaches, please don’t hesitate to reach out if further clarification or additional experiments are needed.

Thank you once again for your time and engagement!

Best regards,

Authors of Submission 8939
</content>
</thread_object>

</thread>
<thread>
creation date:
1732944146649

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
Dear Reviewers,

We sincerely thank you for your valuable feedback and constructive suggestions, which have significantly enhanced the quality of our work.

We highlight the key strengths recognized by the reviewers:

1. **Novelty (Reviewer D7Rz and eRye):** The work introduces a novel reward-free formulation of world models, integrating an inverse soft-Q objective to address diverse imitation learning tasks.  

2. **Solid Experimental Results (Reviewer D7Rz, eRye and yQJv):** The approach is thoroughly evaluated across a variety of benchmarks, including DMControl, MyoSuite, and ManiSkill2, demonstrating stable, expert-level performance. Comprehensive ablation studies provide insightful analysis of the proposed method's components.

3. **Clarity and Structure (Reviewer D7Rz and yQJv):** The paper is well-written, with a clear and structured presentation, making the approach and its underlying formulation easy to understand.  

The reviewers' concerns primarily focus on the following three aspects:

1. Mathematical concerns and additional theoretical guarantees. (Reviewer 4dPD and D7Rz) 
2. Additional experimental requirements, including comparisons with new baselines, ablation studies, robustness in noisy environments, reward correlation analysis, and computational overhead evaluation. (All the reviewers) 
3. Unclear writing and incomplete related work coverage. (Reviewer 4dPD, D7Rz and yQJv)  

To address these concerns, we have conducted further analysis and revised the manuscript as summarized below:

1. **Mathematical Concerns and Theoretical Guarantees**:  
   We have clarified the mathematical concerns raised by Reviewer 4dPD and provided additional theoretical analysis of our method. The theoretical analysis is detailed in Section 4.2 and Appendix H.3 of the revised manuscript.  

2. **Additional Experiments**:  
   We have added baseline comparisons for six locomotion tasks and three manipulation tasks, as shown in Figure 2, Figure 3, and Table 1. Furthermore, we conducted additional experiments addressing ablation studies, robustness in noisy environments, reward correlation, and computational overhead. These results are presented in Appendix E, Appendix F, and Appendix G of the revised manuscript.  

3. **Writing and Related Works**:  
   We have revised the manuscript to improve clarity and incorporated related works suggested by the reviewers. These updates are reflected in Sections 2, 3, and 4 of the revised manuscript.  

We are grateful for your thorough review process and have worked diligently to address the concerns raised. We hope these revisions and the accompanying detailed responses adequately address your feedback, further strengthening the manuscript. Thank you again for your insightful suggestions and thoughtful engagement with our work.  

Best regards,  
Authors of Submission 8939
</content>
</thread_object>

</thread>
<thread>
creation date:
1733181970206

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Reviewer-author discussion

comment:
Dear Reviewers,

We sincerely thank you for taking the time to review our paper and for providing thoughtful and constructive feedback. We have put considerable effort into addressing your concerns in our rebuttal, and we hope we have clarified all the issues raised.

As the discussion phase deadline (December 2nd) approaches, please don’t hesitate to reach out if further clarification is needed.

Thank you once again for your time and engagement!

Best regards,

Authors of Submission 8939
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.