Overcoming Knowledge Barriers: Online Imitation Learning from Visual Observation with Pretrained World Models

Thank you for your thoughtful feedback. We appreciate your recognition of our improvement and raise the score. We would like to address your remaining questions: **Re real world use cases**: - We apologize if our example of EEF to joint torque action space conversion appeared ambitious. The core point is the action space may change from the demonstration to the local setup for imitation. Still taking open-x as an example, most datasets are collected under EEF position action space. Thus, it will create a mismatch if we want to imitate a policy in joint velocity space, which is favourable for RL [2]. We would like to further point out that even if the high-level action space is the same as the EEF position, the low-level controller can also be different. Thus, even if the policy perfectly imitates the action sequence in the dataset, we can still result in different outcomes. Our approach can theoretically handle both of these cases. - Regarding the second point, in ILfO the goal is to imitate the expert's observation, i.e. the results of the actions. For this purpose, it is not always necessary to accurately infer the true actions. As shown in Figure 5 left, when AIME-NoB reach the expert performance, the action mse is still on the scale of 0.1 which is not so accurate. On the other hand, our method is purely data-driven. If we have a new form of robot, our method can be deployed by just collecting some data to pretrain a world model, and use the observation-only demonstrations. The other hard-designed methods are mainly developed for robot arms. Thus, when a new form of robot comes into play, the algorithm may need to be redesigned to fit the new robot. In this viewpoint, our method is more scalable. **Re EKB and DKB**: We are glad to hear that you find the definition helpful! We would like to first clarify: - Regarding $\hat p$, we use the notation to indicate it is an empirical distribution [3]. And since it is formed based on the ground truth action, we also refer it as a ground truth distribution. We are sorry if it causes confusions. Do you have some suggestions about a better notation to use? - $\pi_{\mathrm{demo}}$ is the policy that collects the demonstration dataset. It is defined in Section 2 (Line 131). Regarding "quantitatively measure to close the loop", we did the ablation in Figure 3. In terms of the theory framework, each line in the figure means: - **Expert:** $\mathcal{R}(\pi_{\mathrm{demo}})$. - **MBBC (oracle):** $\mathcal{R}(\pi_{\omega^*}^{\mathrm{AIME}})$ with Eq. 3 as the $J_{\mathrm{policy}}$. - **AIME:** $\mathcal{R}(\pi_{\psi^*}^{\mathrm{AIME}})$ with Eq. 3 as the $J_{\mathrm{policy}}$ and Eq. 2 as the $J_{\mathrm{model}}$. - **AIME-NoEKB:** $\mathcal{R}(\pi_{\psi^*}^{\mathrm{AIME-NoEKB}})$ with Eq. 3 as the $J_{\mathrm{policy}}$ and **Eq. 4** as the $J_{\mathrm{model}}$. (only change the model learning to overcome EKB) - **AIME-NoB:** $\mathcal{R}(\pi_{\psi^*}^{\mathrm{AIME-NoB}})$ with **Eq. 7** as the $J_{\mathrm{policy}}$ and Eq. 4 as the $J_{\mathrm{model}}$. (further change the policy learning objective to overcome DKB) The EKB and DKB for AIME can be measured by: - $\text{EKB} = \mathcal{R}(\pi_{\omega^*}^{\mathrm{AIME}}) - \mathcal{R}(\pi_{\psi^*}^{\mathrm{AIME}})$ - $\text{DKB} = \mathcal{R}(\pi_{\mathrm{demo}}) - \mathcal{R}(\pi_{\omega^*}^{\mathrm{AIME}})$ Then given $\mathcal{R}(\pi_{\psi^*}^{\mathrm{AIME-NoEKB}}) \approx \mathcal{R}(\pi_{\omega^*}^{\mathrm{AIME}})$, we argue that EKB is overcome by **Eq. 4**. Then, given $\mathcal{R}(\pi_{\psi^*}^{\mathrm{AIME-NoB}}) = \mathcal{R}(\pi_{\mathrm{demo}})$, we argue the DKB is further overcome by **Eq. 7**. We could potentially also have the oracle $\mathcal{R}(\pi_{\omega^*}^{\mathrm{AIME-NoB}})$ with Eq. 7 as the $J_{\mathrm{policy}}$ to compute the EKB and DKB for AIME-NoB. But given AIME-NoB already reaches the expert performance, we can infer $\mathrm{EKB}=\mathrm{DKB}=0$ for AIME-NoB on the walker-run task. We agree with you that it should be helpful to include the definition in the main text. For now, we keep it in the appendix to ease our discussion, and we will merge that to the main text in the final version. [2] Aljalbout *et al.*, On the role of the action space in robot manipulation learning and sim-to-real transfer, RA-L 2024 [3] Empirical distribution function - Wikipedia, https://en.wikipedia.org/wiki/Empirical_distribution_function

We would like to thank for your time to review our paper and giving constructive feedback. Please find our replies to your concerns below. > I appreciate that the authors conduct experiments on both tasks from DMControl and Meta-World, and consider tasks with varying difficulty in the case of Meta-World. However, I would have liked to see a few examples of tasks that are more challenging, i.e., where the proposed method (along with baselines) struggle a bit more. For DMControl this could be e.g. any of the Humanoid or Dog tasks, and for Meta-World this could be e.g. Stick Pull / Push or Pick Place (Shelf). Thanks for your suggestions of the harder tasks. As from previous benchmark results, baselines BCO($\alpha$), OT and PatchAIL are not well-performed, so we are not expecting them to be good on these even harder tasks. Thus, we currently mainly compare between AIME-NoB with AIME. If time permitted during the rebuttal phase, we will add these results later. All the results are available in Appendix L of the revised paper. For MetaWorld, we include four additional tasks, namely pick-place, shelf-place, stick-pull and stick-push. The results are shown in Figure 15. From the results, AIME-NoB reliably outperforms AIME. Even in the very hard tasks stick-pull and stick-push, AIME-NoB manages around 60% success rates. However, AIME-NoB doesn't performs so well on pick-place and shelf-place. We conjecture it is due to the visual difficulties of the tasks. It is known that the world models based on reconstruction loss struggle at modeling small objects, which is the small cube we need to pick-up in these two tasks. Improving the world models' ability of modeling small objects or increase the resolutions of the observations will likely improve the performance. For DMC, we conduct additional experiments on humanoid-stand, and show the results in Figure 16. Results show AIME-NoB's capability to make progress on and potentially solve complex tasks like Humanoid. To the best of our knowledge, this is the first time that an ILfO algorithm show progress on vision-based humanoid task. > I find that some of the claims / conclusions are a bit exaggerated relative to what the experimental results show. For example, the authors claim that " the model pretrained on MW-mt50 offers much better results" (L424) while the results in Figure 4(c) show a fairly small difference between the two curves (absolute ~15% increase in success rate on avg it seems). I would prefer claims to accurately reflect the evidence. We thank you for this careful observation. We agree that our language should more precisely reflect the quantitative results. We have revised the paper accordingly, with changes highlighted in blue to better align our claims with the experimental evidence. > Figure 10 indicates that the proposed method succeeds in all of the considered tasks except Cartpole Swingup. I would have expected this task to be the easiest of them all; can the authors please comment on why the method fails on this particular task? Thanks for the question. We were also very surprised by the poor result on CartPole during experiments. Thus, we have conducted a separate analysis on the CartPole at Appendix J. In a nutshell, we think it is a problem of the environment design -- the camera position cannot cover the entire moving space of the CartPole which makes the most important behaviour, the swingup, not visible from the image observation. This makes the AIME loss less effective and even harmful. Thus, we tune the hyper-parameters to lower the effect of the AIME loss and we have shown in this way AIME-NoB can solve CartPole swingup. Please refer to Appendix J for more details.

Thanks for spending your time to offer feedback to our paper. We address your concerns below: > The logic is confusing. The authors state that in order to alleviate EKB, they additionally use online interactions. It is obvious that online trajectories will bring more benefits since the behavior policy to collect offline datasets differs from the current policy. So the improvement compared to AIME might be largely caused by the setting difference rather than algorithm improvement. We appreciate this concern but respectfully disagree for several reasons: 1. **Novel Setting as Contribution**: Extending algorithms to new settings with demonstrated improvements is a recognized contribution in reinforcement learning research. For instance, DrQ [1] extended SAC [2] to handle image inputs, while Cal-QL [3] adapted CQL [4] to online settings. Both works were considered significant contributions to the field. 2. **Non-trivial Algorithm Design**: While online trajectories may intuitively seem beneficial, proper algorithm design is crucial for realizing these benefits. Our experiments demonstrate this: the BCO(α) baseline, which also utilizes online data for fine-tuning, shows no improvement or even degraded performance (Figure 2). This highlights that simply having access to online data is insufficient without proper algorithmic design. 3. **Comprehensive Benchmarking**: While we compare against AIME to demonstrate resolution of knowledge barriers, our evaluation extends beyond this single comparison. AIME-NoB significantly outperforms state-of-the-art online ILfO algorithms like PatchAIL and OT, validating our approach's effectiveness. 4. **Fundamental Solution to EKB**: We contend that online interaction is the most scalable solution to EKB. This mirrors human learning - many complex skills (e.g., athletic movements) cannot be mastered through observation alone but require practical experience. While performance improvements in purely offline settings are possible through specialized architectures or regularization, such knowledge engineering approaches are inherently less scalable than allowing online interaction. > The setting is very complicated. To my understanding, AIME's setting includes an offline dataset ($s_{off}$,$a_{off}$) for world model training, an expert dataset ($s_{expert}$) for imitation learning. And AIME-NoB additionally assumes access to the environment to collect an online dataset ($s_{on}$,$a_{on}$). The complexity of the setting makes it hard to make fair experimental comparisons as baselines might only assume access to a part of the datasets. We acknowledge this complexity but would like to contextualize it: 1. The apparent complexity largely stems from the current state of the field, where pre-trained world models are not yet readily available. Once such models become standard resources, the training pipeline will be significantly streamlined, as the current state in NLP and CV. 2. Our setting is fundamentally comparable to standard online RL in terms of required components - the expert dataset effectively replaces the reward function in traditional RL. The additional offline dataset aligns with the growing paradigm of pre-training and fine-tuning in machine learning. 3. While some comparisons (e.g., with PatchAIL and OT) may not utilize all available data, we argue that the ability to leverage diverse data sources is increasingly important in modern machine learning. Our method's ability to effectively utilize multiple data sources should be viewed as a feature rather than a limitation. [1] Yarats *et al.*, Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels, ICLR 2021 Spotlight [2] Haarnoja *et al.*, Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor, ICML 2018 [3] Nakamoto *et al.*, Cal-QL: Calibrated Offline RL Pre-Training for Efficient Online Fine-Tuning, NeurIPS 2023 [4] Kumar *et al.*, Conservative Q-Learning for Offline Reinforcement Learning, NeurIPS 2020

We would like to thank for your time to review our paper and giving constructive feedback. Please find our replies to your concerns below. > How do you differentiate from previous work that uses online interactions and surrogate rewards? The key innovation of AIME-NoB compared to previous online ILfO approaches like PatchAIL and OT is its ability to leverage pre-trained world models. While previous methods require training models from scratch, AIME-NoB can effectively utilize models pretrained on large datasets. As demonstrated in Figure 6(c), this capability is crucial for AIME-NoB's performance, showing much better sample-efficiency over training from scratch. This aligns with the broader paradigm shift in AI toward pre-training and fine-tuning approaches, which have proven highly effective across multiple domains. > How does it compare with RL with some hand-designed reward or a learned surrogate reward? First we would like to point out that one of the most important motivations to study ILfO is that we can remove the needs of reward. The reward engineering is a notorious problem in the RL field, which prevents RL from scaling to hundreds or even thousands of different tasks. So in our setup, we deliberately remove the reward to make the algorithm more applicable in a broader scope of problems. To address your question directly, we compare AIME-NoB with Dreamer with the hand-designed reward from the environments. The results are shown in Appendix J of the revised paper. For DMC, as we can see from the results in Figure 17, AIME-NoB achieve better sample efficiency on 8 out of 9 tasks excepted the problematic environment cartpole-swingup as we shown in Appendix J. On hopper-hop and quadruped-run, AIME-NoB is surpassed by Dreamer in the end. This is because as an imitation learning algorithm, AIME-NoB's performance is limited by the quality of the expert. For the MetaWorld, since DreamerV3 is not officially benchmarked on it, we have to produce the results ourselves. we report the results at Figure 18. As the manipulation tasks typically impose challenges for exploration, we see that Dreamer with or without the help of the pretrained model struggles to accomplish the task within the 500k environment steps. In the same time, with the help of the demonstrations, it is much easier for AIME-NoB to explore the related regions in the observation which results in a better sample efficient. This marks an advantage of using demonstrations over using a scalar reward to define the task. > Does the method generalize to real-world environments, such as learning from videos? In this paper, we run all our experiments on image observations, i.e. videos as demonstrations. But indeed, we only run our experiments on the simulators. We believe AIME-NoB is well-suited for real-world applications, given its ability to leverage pre-trained models is particularly valuable for real-world scenarios where data collection is costly. Do you have some suggested real-world environments that we can test on? Would these align with your expectations for real-world validation?