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

## Abstract

Abstract Incorporating the successful paradigm of pretraining and finetuning from Computer Vision and Natural Language Processing into decision-making has become increasingly popular in recent years.
In this paper, we study Imitation Learning from Observation with pretrained models and find existing approaches such as BCO and AIME face knowledge barriers, specifically the Embodiment Knowledge Barrier (EKB) and the Demonstration Knowledge Barrier (DKB), greatly limiting their performance.
The EKB arises when pretrained models lack knowledge about unseen observations, leading to errors in action inference.
The DKB results from policies trained on limited demonstrations, hindering adaptability to diverse scenarios.
We thoroughly analyse the underlying mechanism of these barriers and propose AIME-v2 upon AIME as a solution.
AIME-v2 uses online interactions with data-driven regulariser to alleviate the EKB and mitigates the DKB by introducing a surrogate reward function to enhance policy training.
Experimental results on tasks from the DeepMind Control Suite and Meta-World benchmarks demonstrate the effectiveness of these modifications in improving both sample-efficiency and converged performance.
The study contributes valuable insights into resolving knowledge barriers for enhanced decision-making in pretraining-based approaches.
Code will be available at https://github.com/argmax-ai/aime-v2 .

## 1 Introduction

In recent years, we are going through a paradigm shift from learning from scratch to pretraining and finetuning in Computer Vision (CV) [, , ] and Natural Language Processing (NLP) [, , , , ] fields due to the increasing availability of foundation models [] and ever growing datasets.
However, how to adapt this new paradigm into decision making is still unclear, i.e. what type of models we need to pretrain and how these models can be adapted to solve downstream tasks.
Recent works [, , , , ] showcase that one can pretrain a shared latent space world model, enabling successful and efficient transfer to new tasks with either reinforcement learning [, , ] or Imitation Learning from Observation (ILfO) [, ].
Arguably, ILfO [, , , , , ], especially from videos [, , , ], is a more promising approach in this new paradigm since it does not require a handcrafted reward function which is hard to define for many real-world tasks.

However for ILfO with pretrained models, we identify two barriers, which limit the effectiveness of employing pretrained models, the Embodiment Knowledge Barrier (EKB) and the Demonstration Knowledge Barrier (DKB).
The EKB describes the shortcomings of a pretrained model when confronted with novel observations and actions beyond its training experience.
The DKB describes the known problem of generalisation from a limited number of expert demonstrations in imitation learning [].
We find current approaches like BCO [] and AIME [] typically suffer from these two knowledge barriers.
First, these algorithms depend on the pretrained model to successfully infer missing actions from observation sequences.
Thus, when the model has not seen a specific observation before, it may not contain sufficient knowledge about the embodiment to infer the correct action.
Second, policy optimisation is only guided by a limited demonstrations, which may not provide sufficient coverage of the space and lead to a brittle policy working well in some scenarios but not others.

Figure: Figure 1: Main idea of this paper. On the left, we extend the plots from Figure 4 in Zhang et al. [] by adding the MBBC variant and draw the BCO(0) and AIME separately. The purple region between the oracle version and the expert is the DKB while the orange region between the algorithm and its oracle version represents the EKB. On the right, we present the solutions proposed in this paper to solve the two barriers. The blue parts represent the original version of the algorithms that suffer from the knowledge barriers. Orange parts demonstrate the solution for EKB, where the agent is allowed to interact with the environment and use $D_{\mathrm{online}}$ together with $D_{\mathrm{body}}$ to update the world model. Purple parts show the solution for DKB, where a VIPER model is pretrained and used to label the online dataset and then used as a RL signal for policy learning.
Refer to caption: extracted/5566725/figures/AIME-v2.jpg

To better showcase the two barriers, in Figure 1 we extended Figure 4 from Zhang et al. [] by adding a model-based BC (MBBC) variant where the policy is learned by running behaviour cloning on the learned latent space of the world model.
BCO(0) [] and AIME both pretrain a model from a large embodiment dataset and use that to infer the actions for the observation-only demonstrations.
The oracle versions remove the need to infer the missing actions, thus removing the EKB.
As we can see from the figure, the two algorithms are always upper-bounded by the corresponding oracle version, and the difference between them represents the EKB.
On the other hand, even if given the true actions of the expert, imitation performance may still be impacted by a limited number of demonstrations providing insufficient coverage of the state space.
Thus, the difference between the oracle version and the expert performance represents the DKB.

In this paper, we study how to resolve these barriers to improve the performance of ILfO approaches, in particular of AIME proposed by Zhang et al. [].
For the EKB, we extend the setting from offline to online by allowing the agent to further interact with the environment to gather more data to train the world model.
While for the DKB, we introduce a surrogate reward function to allow the policy to essentially train on more data.
We demonstrate that the proposed modifications significantly improve the performance on $9$ tasks in DeepMind Control Suite (DMC) [] and $6$ tasks in Meta-World [].

## 2 Preliminary

We mostly follow the problem setup as described in Zhang et al. [].
We consider a POMDP problem defined by the tuple $\{S,A,T,R,O,\Omega\}$, where $S$ is the state space, $A$ is the action space, $T:S\times A\rightarrow S$ is the dynamic function, $R:S\rightarrow\mathbb{R}$ is the reward function, $O$ is the observation space, which is image in this paper, and $\Omega:S\rightarrow O$ is the emission function.
The goal is to find a policy $\pi:S\rightarrow A$ which maximises the accumulated reward, i.e. $\sum_{t}r_{t}$.

We presume the existence of three datasets of the same embodiment available to our agent.
The *embodiment dataset* $D_{\mathrm{body}}$ contains trajectories $\{o_{0},a_{0},o_{1},a_{1}\dots\}$ that represent past experiences of interacting with the environment.
This dataset provides information about the embodiment for the algorithm to learn a world model.
In addition, we also allow the agent to interact with the environment to collect new data in a *replay buffer* $D_{\mathrm{online}}$.
Note that, although the simulator will give us the reward information, the agent is not allowed to use them, and we only use the reward for evaluation purpose.
The *demonstration dataset* $D_{\mathrm{demo}}$ contains a few expert trajectories $\{o_{0},o_{1},o_{2}\dots\}$ of the embodiment solving a certain task defined by $R_{\mathrm{demo}}$.
The crucial difference between this dataset and the other two datasets is that the actions are not provided anymore since they are not observable from a third-person perspective.
The goal of our agent is to learn a policy $\pi$ from $D_{\mathrm{demo}}$ which can solve the task defined by $R_{\mathrm{demo}}$ as well as the expert who generated $D_{\mathrm{demo}}$.

### 2.1 World Models

A World Model [] is a generative model which models a probability distribution over sequences of observations, i.e. $p(o_{1:T})$.
The model can be either unconditioned or conditioned on other factors, such as previous observations or actions.
When the actions taken are known, they can be considered as the condition, i.e.  $p(o_{1:T}|a_{0:T-1})$, and the model is called embodied [].

In this paper, we consider variational latent world models where the observation is governed by a Markovian hidden state.
In the literature, this type of model is also referred to as a State-Space Model (SSM) [, , , , ].
Such a variational latent world model involves four components, namely

$$ encoder $\displaystyle z_{t}=f_{\phi}(o_{t}),$ posterior $\displaystyle s_{t}\sim q_{\phi}(s_{t}|s_{t-1},a_{t-1},z_{t}),$ decoder $\displaystyle o_{t}\sim p_{\theta}(o_{t}|s_{t}),$ prior $\displaystyle s_{t}\sim p_{\theta}(s_{t}|s_{t-1},a_{t-1}).$ $$

$f_{\phi}(o_{t})$ is the encoder to extract the features from the observation; $q_{\phi}(s_{t}|s_{t-1},a_{t-1},z_{t})$ and $p_{\theta}(s_{t}|s_{t-1},a_{t-1})$ are the posterior and the prior of the latent state variable; while $p_{\theta}(o_{t}|s_{t})$ is the decoder that decodes the observation distribution from the state.
$\phi$ and $\theta$ represent the parameters of the inference model and the generative model respectively.

Typically, the model is trained by maximising the Evidence Lower Bound (ELBO) which is a lower bound of the log-likelihood, or evidence, of the observation sequence, i.e. $\log p_{\theta}(o_{1:T}|a_{0:T-1})$.
Given a sequence of observations, actions, and states, the objective function can be computed as

$$ ELBO $\displaystyle=\sum_{t=1}^{T}J_{t}^{\mathrm{rec}}-J_{t}^{\mathrm{KL}}=\sum_{t=1 }^{T}\log p_{\theta}(o_{t}|s_{t})-D_{\text{KL}}[q_{\phi}||p_{\theta}].$ (1) $$

The objective function is composed of two terms: the first term $J^{\mathrm{rec}}$ is the likelihood of the observation under the inferred state, which is usually called the reconstruction loss; while the second term $J^{\mathrm{KL}}$ is the KL divergence between the posterior and the prior distributions of the latent state.
To compute the objective function, we use the re-parameterisation trick [, ] to autoregressively sample the inferred states from the observation and action sequence.

In summary, a world model is trained by solving the optimisation problem as

$$ $\displaystyle\phi^{*},\theta^{*}=\mathop{\arg\!\max}_{\phi,\theta}\mathbb{E}_{ \{o,a\}\sim D_{\mathrm{body}},s\sim q_{\phi}}[\text{ELBO}].$ (2) $$

### 2.2 AIME

AIME is a recently proposed algorithm that uses a pretrained world model to solve ILfO in an offline setting.
Specifically, it uses the pretrained world model as an implicit inference model by solving for the best action sequence that makes the demonstration most likely under the trained world model.
The imitation can be done jointly with the action inference using amortised inference and the re-parameterisation trick by solving the following optimisation problem

$$ $\displaystyle\psi^{*}=\mathop{\arg\!\max}_{\psi}\mathbb{E}_{o\sim D_{\mathrm{ demo}},s\sim q_{\phi^{*},\theta^{*}},a\sim\pi_{\psi}}[\text{ELBO}],$ (3) $$

where $\psi$ is the parameter for policy $\pi_{\psi}(a_{t}|s_{t})$.
The resulting objective is very similar to Equation 2, with a subtle difference of the sampling path.
That is in the new objective, only the observations are sampled from the dataset and both states and actions are sampled iteratively from the learned model and the policy, respectively.

## 3 Methodology

In this section we will make thorough analysis separately on the EKB and DKB especially upon the AIME algorithm as the case study.
Based on the analysis we introduce the solution for each knowledge barriers and combine them as AIME-v2.
We depict the general framework of the solutions on Figure 1 and show the pseudocode of AIME-v2 in Algorithm 1 in Appendix A.

### 3.1 Resolving the EKB

The most natural way to solve the EKB is to allow the agent to further interact with the environment.
By allowing the agent to interact with the environment, the new experiences can minimise the error in the pretrained model in proximity of the policy $\pi_{\psi}$ and gain more embodiment knowledge relevant for the task at hand.
Torabi et al. [] proposed a modified version of BCO(0) called BCO($\alpha$) which introduced such an interaction phase.
However, from their and our empirical results, it did not resolve the EKB since there remains a gap between BCO($\alpha$) and the BC oracle when the environment is complex. In fact, as we will show in the following, the basic idea of adding online interactions is not straightforward to successfully implement in practice.

As shown in recent works in Offline RL, continuing training an actor-critic from the offline phase in the online phase requires certain measures to combat the shift of objective [, , ].
A similar story also applies when extending AIME from purely offline to online.
The most dominant problem we found is overfitting to the newly collected dataset.

As the training progresses iteratively between data collection, model training and policy training, in the early phase of training, there are only a few new trajectories available for training the model.
Because the world model is highly expressive, it may overly favor similar trajectories in the new data, leading to a high ELBO.
Normally, this may not be a big problem since, eventually, more and more data will be collected to combat this overfitting.
But since AIME also depends on the ELBO to train the policy, it quickly causes the policy training to diverge.

In order to address the overfitting issue, we need a regulariser for model learning.
Instead of designing ad-hoc methods to regularise the model in the parameter space, we adopt a data-driven approach.
From model’s perspective, the overfitting is caused by a sudden shift of the training data from a large and diverse pretraining dataset to a small and narrow replay buffer.
So one way to make the shift not as sudden is to append the pretraining dataset to the replay buffer, so that the distribution of the training data will change smoothly.
However, this causes data efficiency problems since the replay buffer is relatively small compared to the pretraining dataset so that uniformly sampling from them together limits utilisation rate of the new data.
Instead, we consider sampling separately from both datasets.
Concretely, we modify the objective in Equation 2 to

$$ $\displaystyle\phi^{*},\theta^{*}=\mathop{\arg\!\max}_{\phi,\theta}\>$ $\displaystyle\alpha\mathbb{E}_{\{o,a\}\sim D_{\mathrm{body}},s\sim q_{\phi}}[ \text{ELBO}]+(1-\alpha)\mathbb{E}_{\{o,a\}\sim D_{\mathrm{online}},s\sim q_{ \phi}}[\text{ELBO}].$ (4) $$

The amount of data we sample from the pretraining dataset is controlled by a hyper-parameter $\alpha$, which represents how much regularisation we put upon the model.
In this paper, we mainly consider setting $\alpha=0.5$, so that we sample the data evenly from both datasets.
We justify our choice in the ablation section with Figure 4(b).

This finding contradicts Rajeswar et al. [] and Hansen et al. [], where the pretrained world models do not need such a data-driven regulariser.
We conjecture that unlike AIME they mainly use their world models purely as generative models to predict states and rewards given action sequences, which is only indirectly influenced by overfitting the ELBO.

### 3.2 Resolving the DKB

Based on the discussion from the previous sections, the straightforward way of solving the DKB is also to increase the number of demonstrations available to the agent.
However, expert demonstrations are generally considered to be difficult and expensive to collect.
Thus, increasing the size of the demonstration dataset is not always feasible in real-world applications.

In order to propose a more practical solution, we need to look deeper into what is the real cause of the DKB.
The policy learning part of the AIME algorithm is essentially behaviour cloning, and it is only conducted on the demonstration dataset.
So for the states covered in the demonstration dataset, the policy is given clear guidance about what to do, while for other states, the behaviour is undefined.
AIME is solely relying on the generalisation abilities of the learned latent state and the trained policy network to extrapolate the correct behaviour.
In particular for small demonstration datasets, this can be unreliable or even impossible.
Therefore, if we can enlarge the space of the covered states, we should reduce the DKB [].

Generative Adversarial Imitation Learning (GAIL) style algorithms [, , , ] are perfect examples of this solution.
Instead of directly learning from the demonstration, they adversarially train a discriminator to assess how how closely each state resembles those in the demonstration dataset.
When learning the policy, they treat the discriminator’s output as a reward and encourage visiting states that are more likely belonging to the demonstrations.
This modification provides guidance for newly visited states in the replay buffer, allowing the space of supported states to grow over the course of training.

However, this benefit is not clearly separable in GAIL since it is always entangled with the adversarial training of the discriminator.
A recent work MAHALO [] shows further evidence of the importance of the size of the covered space.
The authors studied a similar ILfO setup with embodiment and demonstration datasets.
They compared two variants: for one they train an inverse dynamics model (IDM) from the embodiment dataset and use it to label the demonstration dataset, while for another they train a reward model from the demonstration dataset by labelling all time steps with a reward of $1$, and then use it to label the embodiment dataset.
Finally, they run the same offline RL algorithm on both labelled datasets.
The results show the second variant attains a much better performance even though the labelling from the reward model is not as meaningful as the actions from the IDM.

Based on these insights, we propose to introduce a surrogate reward providing guidance signal for the agent on the replay buffer dataset.
Due to the instability of adversarial training [, , ] and our focus on the pretraining paradigm, we opt to adopt the VIPER algorithm [].
Instead of training a discriminator, VIPER trains a video prediction model on the demonstration datasets and treats the likelihood of the video prediction model as the reward for policy learning, i.e.

$$ $\displaystyle r^{\mathrm{VIPER}}_{t}=\log p_{\nu}(o_{t}|o_{<t}).$ (5) $$

Using this reward, we train the policy with a dreamer-style actor-critic algorithm [] based on imagination in the latent space of the world model [].
In order to do this, we first need to modify the reconstruction term in Equation 1 by adding an extra term for decoding the VIPER reward, i.e. $p_{\theta}(r^{\mathrm{VIPER}}_{t}|s_{t})$.
Then, we further train a value estimator $V_{\xi}(s_{t})$ using TD($\lambda$)-return estimates, i.e.

$$ $\displaystyle V^{\lambda}_{\xi}(s_{t})$ $\displaystyle=(1-\lambda)\sum_{n=1}^{\infty}{\lambda^{n-1}V^{(n)}_{\xi}(s_{t})}$ (6) $\displaystyle\text{with}\;V^{(n)}_{\xi}(s_{t})$ $\displaystyle=\sum_{t^{\prime}=t+1}^{t+n}{\gamma^{t^{\prime}-t-1}r^{\mathrm{ VIPER}}_{t^{\prime}}}+\gamma^{n}V_{\xi}(s_{t+n}).$ $$

Using this estimate, we optimise our value function by minimising the MSE

$$ $\displaystyle\xi^{*}=\mathop{\arg\!\min}_{\xi}(V_{\xi}(s_{t})-V^{\lambda}_{\xi ^{\prime}}(s_{t}))^{2}.$ (7) $$

As is common practice, we use a target value network with parameters $\xi^{\prime}$ to stabilise training, whose parameters are updated using Polyak averaging with a learning rate $\tau$ in every iteration.

Using this value estimate, we extend the policy objective of Equation 3 to

$$ $\displaystyle\psi^{*}=\mathop{\arg\!\max}_{\psi}\mathbb{E}_{o\sim D_{\mathrm{ demo}},s\sim q_{\phi,\theta},a\sim\pi_{\psi}}[\text{ELBO}]+\beta\mathbb{E}_{\{ o,a\}\sim D_{\mathrm{online}},s\sim q_{\phi},a^{\prime}\sim\pi_{\psi},s^{ \prime}\sim p_{\theta}}[V_{\xi^{\prime}}^{\lambda}(s^{\prime})],$ (8) $$

where $\beta$ is a hyper-parameter for balancing the two terms.
We set $\beta=0.1$ by default based on the difference of default learning rate in AIME and Dreamer.

## 4 Experiments

We aim to answer the following questions: a) How does the proposed AIME-v2 compare with state-of-the-art methods on common benchmarks? b) How well does the proposed modification resolve the EKB and the DKB? c) How do different choices of hyper-parameters influence the results? In order to answer these questions, we design our experiments on DMC and Meta-World benchmarks.

### 4.1 Datasets

For the DMC benchmark, we choose 9 tasks across 6 embodiments following Liu et al. [] and use their published dataset as the demonstration datasets.
Each dataset contains only 10 trajectories to reflect the scarcity of expert demonstrations.
For the embodiment dataset, in order not to leak the task information from the pretraining phase, we follow Rajeswar et al. [] and run a Plan2Explore [] agent for each embodiment with 2M environments steps and use its replay buffer as the embodiment dataset.
Different to them taking the model directly from the Plan2Explore agent as the pretrained model, we follow Zhang et al. [] to retrain the model for $200\mathrm{k}$ gradient steps to get a better model.

For Meta-World benchmark, we use the data and model from Hansen et al. [].
The embodiment dataset was created from the replay buffer datasets.
The open-sourced replay buffer datasets contain $40\mathrm{k}$ trajectories for each of the $50$ tasks with only state information.
In order to fit to our image observation setup, we render the images by resetting the environment to the initial state of each trajectory and then executing the action sequence.
The details of resetting to the initial state can be found in Appendix D.

Following the idea of not leaking too much about the task information, inspired by the common practice in offline RL benchmarks [], we use the first $200$ trajectories from each replay buffer and form a dataset with $10\mathrm{k}$ trajectories in total.
We call this dataset MW-mt50.
To further study the out-of-distribution transfer ability of the pretrained model, we follow the difficulty classification of the tasks from [] and only use the $39$ easy and medium tasks to generate the datasets and the $11$ tasks hard and very hard tasks as hold-out tasks.
We uniformly sample $250$ trajectories from te first $10\mathrm{k}$ trajectories from each of the $39$ tasks and form a dataset with $9750$ trajectories in total.
We refer to this dataset as MW-mt39.
Hence, MW-mt39 contains some expert behaviour solving the tasks, while MW-mt50 consists of mostly exploratory behaviour.

As the demonstration datasets, we use the single-task policies open-sourced by TD-MPC2 and collect $50$ trajectories for each tasks.
We ensure that every trajectory in the demonstration dataset is successful.
Since there are $500$ steps in a DMC trajectory and only $100$ steps in a Meta-World trajectory, the resulting datasets are roughly the same size.

Figure: Figure 2: Benchmark results and 9 DMC tasks. Return are calculated by running the policy 10 times with the environment and taking the average return. The results are averaged across 3 seeds with the shade region representing 95% CI.
Refer to caption: x1.png

### 4.2 Implementation

For the world model, we use the RSSM architecture [] with the hyper-parameters in Hafner et al. [] for DMC tasks.
In addition, we use the KL Balancing trick from Hafner et al. [] to make the training more stable.
For Meta-World, since the visual scene is more complex, we use the M size model from Hafner et al. [], but still with the continuous latent variable to be aligned with other models used in this paper.
The policy network is implemented with a two-layer MLP, with $128$ neurons for each hidden layer.
All the models are trained with Adam optimiser [].
More details about the hyper-parameters can be found in Appendix C.

For the VIPER model, in the original paper, the authors first pretrain a VQ-GAN [] and then train a GPT-style auto-regressive model in the quantised space for prediction.
For simplicity of the implementation, in this paper, we consider training an unconditioned latent world model as in Seo et al. [] to model the VIPER reward.
We use the same RSSM architecture of the model learning for DMC, only removing the condition of the actions, and we train the VIPER model for each task separately.
Especially during training, we find training such a powerful model from scratch on a small dataset can easily result in over-fitting.
Thus, we empirically choose to train the model only for $500$ gradient steps for DMC models and $1000$ gradient steps for Meta-World models.
We show evidence of overfitting in Appendix E.
Due to the large scale of the ELBO, we also apply symlog [] when computing the VIPER reward.
Another difference with the original VIPER paper is that we do not use intrinsic motivation as the exploration bonus as the authors suggested, since the AIME loss for policy learning already provides task-related guidance for exploration.
We only apply an entropy regulariser to the policy as is common practice.
We further show the synergy between AIME and VIPER in Appendix F.

### 4.3 Benchmark Results

Figure: Figure 3: Benchmark results and $6$ Meta-World tasks. Trajectories are only counted as success when it success at the last time steps and the success rates are calculated with $10$ policy rollouts. The results are averaged across $3$ seeds with the shade region representing $95$% CI.
Refer to caption: x2.png

The benchmark results of DMC are shown in Figure 2.
AIME-v2 significantly outperforms the PatchAIL baseline in $7$ out of $9$ tasks in terms of sample efficiency.
Benefiting from the pretrained world model, AIME-v2 typically can reach expert performance within $200$k environment steps.
Compared with BCO($\alpha$), updating the model is regularised and is benefitting more from the online interaction to resolve the EKB.
Compared with the original offline AIME, AIME-v2 reliably improves performance, especially in hard tasks such as walker-run and hopper where offline AIME did not manage to make any progress.

However, there are still two tasks for which AIME-v2 does not show much progress, namely cartpole-swingup and quadruped-run.
For cartpole-swingup, we observe that the policy learns to move the cart out of the scene so that the static image yields a high likelihood from the video prediction model.
A similar phenomenon was also discussed in the original VIPER paper [].
For quadruped-run, we conjecture that it is due to visual difficulties of a reconstruction-based model.
When the quadruped is initialised on the ground, due to the symmetric structure of the robot, it is impossible to figure out which action corresponds to which leg, and it easily leads the action inference process to diverge.
We additionally show AIME-v2 can work on these tasks with the help of the true reward in Appendix F.

The benchmark results of Meta-World are shown in Figure 3.
We choose 4 tasks, namely disassemble, assembly, hand-insert and push from hard and very hard difficulties, and 2 tasks, namely sweep and hammer from the medium difficulty.
While PatchAIL does not work on these tasks at all, AIME and AIME-v2 can make progress on them.
AIME-v2 with either pretrained model outperforms AIME in all the tasks.
On hard and very hard tasks, AIME with the mt50 model is better than AIME with the mt39 model.
This is because they contain an unseen novel object by the mt39 model creating a large EKB.
But in the online setting of AIME-v2, the two models are mostly on par.
Moreover, using the mt50 models is better than using the mt39 models on average, which may imply covering diverse bahaviour is more valuable than knowing the expert directly.

### 4.4 Ablation Results

Figure: (a) Demonstration efficiency.
Refer to caption: x3.png

We conduct our ablation studies on walker-run task from DMC.

How well does the proposed methods resolve knowledge barriers?
In order to show how well AIME-v2 resolves the two knowledge barriers, we conduct a similar experiment as Zhang et al. [] by providing the agent with different numbers of demonstrations.
The result is shown in Figure 4(a).
As we discussed before, MBBC as an oracle method that circumvents the EKB is a strict upper bound for AIME.
And AIME-v2 which addresses both the EKB and DKB achieves much better results and is an upper bound for MBBC.
From AIME-v2 can achieve near-expert performance with as few as $5$ demonstrations for this challenging task.

Influence of the data regulariser ratio $\alpha$.
We set the regulariser ratio $\alpha$ from $[0.0,0.25,0.5,0.75]$ and plot the results in Figure 4(b).
As we can see from the result, as long as we enable the regulariser, i.e. set $\alpha>0$, we get reliable improvements over the course of training.
But if we disable the regulariser by setting $\alpha=0$, the learning only starts once there is enough new data in the replay buffer.
As we discussed in Section 3.1, without the regulariser, in the early stage of the training, the model can easily overfit to the replay buffer, and it explains the early flattening phase of the training.
As the training progresses, more and more data is available from the replay buffer, and it can establish the regulariser on its own, which explains the dramatic growing phase of the curves.

We also plot the MSE between the inferred actions and the true actions during the training process.
From that we can see that a higher regulariser ratio offers more stable inference of the actions in the early phase of training.

Influence of the value gradient loss weight $\beta$.
We set the weight $\beta$ from $[0.0,0.01,0.05,0.1,0.5,1.0]$ and plot the results in Figure 4(c).
As the result shows, having a small $\beta$ slows learning progress toward convergence.
On the other hand, setting $\beta$ to a much larger value will improve the sample-efficiency without causing instability.
For the sample efficiency, since we only have $10$ demonstrations, DKB dominates over EKB as shown in Figure 4(a).
Thus, having a larger $\beta$ will make the learning much faster.
In terms of the stability, as we discussed in 3.2, AIME loss and the value gradient loss operate on different regions of the environment states.
This could make their influence on the policy independent of each other.

## 5 Related Work

Imitation Learning from Observation.
ILfO [, , , , , , ] has become more popular in recent years due to their potential to utilise internet-scale videos for behaviour learning.
Most of the previous works [, , ] study the problem only with the true state as observation.
Recent works [, , , ] have started to shift toward image observations as a more general setting.
Our work is a continuation of this journey.

Pretrained Models for Decision-Making.
Inspired by the tremendous progress made in recent years in CV and NLP fields with the power of pretrained models, the decision-making community is also trying to follow the trend.
Most recent works focus on the use of Large Language Model (LLM) for decision-making.
A prompted model is used for producing trajectories and plans [, , , ], code [, , , , ] or for modifying the reward [, ].
There are also other people studying the benefit of pretrained visual models for visuomotor tasks [, , , ] while others try to train large policy networks directly with transformers [] and huge datasets [, , ].
However, there is only little attention being put on pretrained world models [, , ], which are natively developed by the model-based decision-making community and perfectly fit into the pretraining and finetuning paradigm.
Our work explores this overlooked domain and showcases its potential.

## 6 Discussion

In this paper, we identify two knowledge barriers, namely the EKB and the DKB, limiting the performance of state-of-the-art ILfO methods using pretrained models.
We thoroughly analyse the underlying cause of each barrier and propose practical solutions.
Specifically, we propose to use online interaction with a data-driven regulariser to solve the EKB and surrogate reward labelling to reduce the DKB.
Combining these solutions, we propose AIME-v2, a method based on the recently proposed AIME and showcase its efficiency compared to state-of-the-art ILfO methods.
We also conducted a series of ablation studies about how each knowledge barrier is addressed by the proposed solution and how their hyper-parameters influence the performance.

However, the proposed solutions still have drawbacks.
First, the data-driven regulariser is not practical when the model is pretrained on huge datasets like the foundation models popular in the fields of CV and NLP.
Reducing the amount of data needed for the regulariser could greatly improve the usability of the method.
Second, although having pretrained models is beneficial, having too many pretrained components can be detrimental for model selection.
Especially in AIME-v2, the world model and the VIPER model share a very similar interface.
Designing a shared model that can serve both interfaces could ease the use of the method.
Last but not least, due to the high demand of computing resources, we only study the pretrained world model on a very small scale.
It will be an interesting direction to study these model at larger scales.

We hope our work can shed some light on the future development of ILfO method and bring more attention to the great potential of pretrained world models.

## Appendix A Algortihm

Figure: Algorithm 1 AIME-v2

## Appendix B Compute Resources

All the experiments are run on a local cluster with a few A100 and RTX8000 instances.
All the experiments are tuned to use less than 10GB of GPU memory so that they can run in A100 MIG.
World models pretraining requires about 24 GPU hours, while VIPER models require negligible time for training.
Each DMC experiment requires about 40 GPU hours while each Meta-World experiment requires about 24 GPU hours.

## Appendix C Hyper-parameters

Here, we document the detailed hyper-parameters for all the trained models in Table 1.

**Table 1: AIME-v2 hyper-parameters use for each benchmark.**
|  | DMC | Meta-World |
| --- | --- | --- |
| World Model |  |  |
| CNN structure | Ha and Schmidhuber [] | Hafner et al. [] |
| CNN width | $32$ | $48$ |
| MLP hidden size | $512$ | $640$ |
| MLP hidden layer | $2$ | $3$ |
| MLP activations | LayerNorm + Swish |  |
| Deterministic latent size | $512$ | $1024$ |
| Stochastic latent size | $30$ |  |
| Free nats | $1.0$ |  |
| KL balancing | $0.8$ |  |
| Policy |  |  |
| Hidden size | $128$ |  |
| Hidden layer | $2$ |  |
| Activation | ELU |  |
| Distribution | Tanh-Gaussian |  |
| Value network |  |  |
| Hidden size | $128$ |  |
| Hidden layer | $2$ |  |
| Activation | ELU |  |
| Target EMA decay | $0.95$ |  |
| Training |  |  |
| Batch Size | $50$ | $16$ |
| Horizon | $50$ | $64$ |
| Total Env steps | $1\mathrm{M}$ | $500\mathrm{k}$ |
| Update ratio | $0.1$ |  |
| Gradient clip | $100$ |  |
| Policy entropy regulariser weight | $1\mathrm{e}{-4}$ |  |
| Model learning rate | $3\mathrm{e}{-4}$ |  |
| Policy learning rate | $3\mathrm{e}{-4}$ |  |
| Value network learning rate | $8\mathrm{e}{-5}$ |  |
| Discount factor $\gamma$ | $0.99$ |  |
| TD-Lambda parameter $\lambda$ | $0.95$ |  |
| Imagine horizon | $15$ |  |
| AIME-v2 specific |  |  |
| Policy pretraining iterations | $2000$ |  |
| Data-driven regulariser ratio $\alpha$ | $0.5$ |  |
| Value gradient loss weight $\beta$ | $0.1$ |  |

## Appendix D Details for Resetting Meta-World Tasks

To generate the image observation datasets from the TD-MPC2 replay buffer [], we modify the Meta-World codebase to reset the environment to the initial state of the trajectory from the first observation.
Luckily, the starting position of the robot arm is always the same for each task, so that we do not need to apply inverse kinematics to solve for the initial pose of the robot arm.
For the object and the target position, for most of the tasks, the internal reset position can be computed by making a constant shift on the object position and the target position in the observations.
There are, however, also a few edge cases which we handle differently.

In button-press-topdown and button-press-topdown-wall, the object’s true position only appears in the observation upon the second time step, presumably due to some simulator delay in the resetting process.
So for these two tasks, the initial state is reset by the second observation.

For basketball and box-close, it seems like there is some internal collision detection that will alter the object and robot position after the task is reset, so computing the exact reset value from the observation is not possible.
For these two tasks, we instead resort to a search-based method.
To be specific, we use a gradient-free optimiser from [] to search over the resetting space of the object and find the reset position that minimise the L2 distance with the true observation.

More details of the implementation can be found in the code.

## Appendix E Overfitting of the VIPER Models

Figure: Figure 5: Correlation of the VIPER reward and the real reward with models trained with different number of gradient steps. Each point represents one trajectory. We can clearly see the model gradually overfitting and losing the correlation with the real reward when training for more than $1\,000$ gradient steps.
Refer to caption: x6.png

To better illustrate the overfitting problem for VIPER models and justify our choice of training fewer iterations, we train the VIPER models for a varying number of gradient steps and evaluate the correlations between the VIPER reward and the true reward on both the expert dataset from PatchAIL, where the VIPER model is trained on, and the replay buffer dataset from Zhang et al. [].

Specifically, we train the same VIPER model with $\{100,500,1000,2000,5000,100000\}$ gradient steps and plot the result in Figure 5.
As we can clearly see, when training with less than or equal to $1\,000$ gradient steps, VIPER reward has a very nice correlation with the true reward, with the middle-range performance even like a linear correlation.
The best model could be selected from $500$ and $1000$ gradient steps.
However, as we train the model for longer, the VIPER reward for the expert trajectories is boosted even higher, and as a side effect, it also relatively boosts up the VIPER reward for low-performance trajectories.
This is because, when overfitting the expert trajectories, the model increases the marginal likelihood of all the observations in the expert trajectories to a higher value, which also includes a few frames of the robot lying on the ground at the very beginning of each trajectory after reset.
For these low-performance trajectories, the robot remains mainly stuck around the initial position and struggles on the ground.
This artifact of the overfitted VIPER reward creates a sharp local maximum in the low-performance region that the agent can hardly get away from.

## Appendix F Additional Experiments

Figure: Figure 6: Additional ablation on DMC tasks by exploring the synergy between AIME and VIPER model. Return are calculated by running the policy 10 times with the environment and taking the average return. The results are averaged across 3 seeds with the shade region representing 95% CI.
Refer to caption: x7.png

Synergy between AIME and VIPER model.
We also find there is a synergy between AIME and VIPER model.
As we shown in Appendix E, one inherent problem of VIPER reward is that it not only incentivises the expert behaviour as the optimal, but also a stationary behaviour with very low reward as a local maxima.
In order to work with the VIPER reward, the agent needs to have the ability to escape from the local maxima region.
AIME offers the IL loss to imitate the expert demonstrations and can achieve decent performance even when pretrained offline, which helps to escape the local maxima.
To better show the synergy, we provide additional ablation results with the VIPER reward in Figure 6.
In the experiments, we include two other variants: the AIME-v2 w/o AIME is to remove the AIME IL loss from the online policy learning, so that the policy is pretrained by AIME loss but finetuned with only RL loss from the VIPER reward; while the VIPER is following the implementation in the original VIPER paper with RL loss on both the VIPER reward for the task and intrinsic reward for exploration.
From the result, we can clearly see that without the help of AIME IL loss, VIPER reward cannot reliably motivate the agent to learn good behaviours.
Even when in walker tasks, the w/o AIME variant can solve the tasks to certain extends, it depends strongly on the random seeds.
In conclusion, the good performance of AIME-v2 cannot be achieved by either AIME IL loss or VIPER RL loss alone but a combination of both.

Figure: Figure 7: Additional benchmark results on 3 DMC tasks with an additional variant of AIME-v2 with the true reward. Return are calculated by running the policy 10 times with the environment and taking the average return. The results are averaged across 3 seeds with the shade region representing 95% CI.
Refer to caption: x8.png

Improving AIME-v2 with better rewards.
We show additional results on the $3$ not so well-performing DMC tasks, namely cartpole-swingup, cheetah-run and quadruped-run, in 7.
In the plot we add a new variant using the true reward from the environment to replace the VIPER reward.
As the results show, if we had a better estimation of the surrogate reward, AIME-v2 could also achieve good performance on these tasks.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - How do you differentiate from previous work that uses online interactions and surrogate rewards?
- How does it compare with RL with some hand-designed reward or a learned surrogate reward?
- Does the method generalize to real-world environments, such as learning from videos?
soundness: 4
strengths: - Clearly identify the shortcomings of existing IFO methods and propose reasonable solutions to reduce the gaps.
- Extensive experiments on the different design choices, e.g. reward functions, and comparisons with baselines demonstrate solid improvement.
- Paper writing is clear and easy to follow.
summary: This paper studies two barriers in imitation learning from observations, namely the Embodiment Knowledge Barrier (EKB) and Demonstration Knowledge Barrier (DKB). EKB refers to the gap when generalizing to new observations and DKB refers to the gap when generalizing from limited number of expert demonstrations.

It proposes to use online interaction to reduce EKB and introduces a weighted loss between new interactions and pre=training data. For DKB, it proposes to use a surrogate reward, such as a discriminator in AIL.

Experimental results suggest AIME-NoB outperforms baselines with significant margin in DMC and MetaWorld.
weaknesses: - The ideas to reduce EKB and DKB seems a simple combination of previous methods, such as using online interaction with weighted sampling, adversarial training, etc. It's unclear the distinctions from previous methods (such as AIL, e.g.) is significant enough.
- These environments aren't hard to come up with a hand-designed reward or learning a surrogate reward. Lack of comparisons with Dreamer-like methods using hand-designed or surrogate reward.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: NA
soundness: 2
strengths: - AIME-NoB alleviates EKB through online interaction. To address DKB, the most natural way is to collect more expert demonstrations, which is expensive in terms of robot manipulation. Therefore, AIME-NoB circumvents this burden by providing reward signals using surrogate models and optimizes the policy with reward-labeled online trajectories.
- Experiments in DMC and MetaWorld show that AIME-NoB can bring significant improvement to AIME.
summary: This paper extends AIME to AIME-NoBarries, which addresses EKB with online interaction and DKB with online trajectories labeled with surrogate rewards. The authors evaluate the effectiveness of AIME-NoB in DMC and MetaWorld.
weaknesses: - The logic is confusing. The authors state that in order to alleviate EDK, 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.
- The setting is very complicated. To my understanding, AIME's setting includes an offline dataset $(s_{\mathbf{off}},a_{\mathbf{off}})$ for world model training, an expert dataset $(s_{\mathrm{expert}})$ for imitation learning. And AIME-NoB additionally assumes access to the environment to collect an online dataset $(s_{\mathrm{on}},a_\mathrm{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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: I would like the authors to address my comments listed in "weaknesses" above during the rebuttal. I have one additional question:

- 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?
soundness: 3
strengths: - The paper is generally well written and easy to follow. Illustrations (Figure 1 in particular) are useful for understanding the proposed method, and Section 2 provides sufficient background for an unfamiliar reader to appreciate the contributions.
- Experiments are reasonably extensive, covering a variety of tasks from DMControl and Meta-World, as well as several (what appears to be) baselines appropriate for the problem setting. Empirical results are strong compared to baselines.
- It is a purely empirical paper, but there are sufficient ablations to understand how the algorithm behaves in different settings (e.g. number of demos), and which components are the main drivers of performance (e.g. choice of surrogate reward).
summary: This paper studies the problem of imitation learning (IL) from (visual) observation, i.e., demonstrations do not contain action information using pretrained world models. The paper identifies two key bottlenecks in current IL algorithms, namely (1) OOD observations, and (2) OOD task configurations, which the authors refer to as the Embodiment Knowledge Barrier (EKB) and the Demonstration Knowledge Barrier (DKB), respectively. The key technical contribution is an algorithmic extension of the method Action Inference by Maximizing Evidence (AIME); the core idea is to finetune the learned IL policy via limited online interaction and a RL objective, with surrogate rewards derived from the demonstration dataset. Experiments are conducted on DMControl and Meta-World from visual observations, and the proposed method compares favorably to both AIME without finetuning as well as other IL methods that can be finetuned online (BCO, OT, PatchAIL).
weaknesses: - 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).
- 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.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: See suggestions in the weaknesses
soundness: 3
strengths: - Writing is fairly straightforward to understand 
- Results show clear improvements over relevant methods
- The paper includes fairly extensive ablations/empirical analysis
summary: **Problem setting**: have access to a demonstration data set without actions, an offline data set of suboptimal experience, and the ability to sample experience online without reward feedback

**Proposed approach**: (1) balance the offline and online datasets when updating the model, (2) combine AIME and AIL objectives (though also considers other alternatives to AIL)

**Experiments**: comparisons to relevant prior methods on DMC and MetaWorld tasks with visual observations
weaknesses: Overall, the technical contribution seems like a fairly basic combination of prior works. While many methods build upon prior works, this particular combination seems closer to an implementation choice than a significant technical contribution. The method is also more complex than prior methods as it introduces one hyperparameters and combines multiple objectives. 

Other weaknesses
- It’s a bit unclear how this connects to a real-world problem. In robotics, you typically either collected data from a robot, where you can get actions or something close to them, or videos of humans/animals, which lacks actions but also has significant domain shift. Connecting the problem statement to a real world problem would be helpful to motivate the paper’s significance.
- EKB and DKB are defined quite informally and then are used extensively, sometimes in a hand-wavy way. I’m not sure if introducing them is helpful for understanding. I think it would be better to either remove them or define them more formally and measure the extent to which they contribute to poor performance.
- I appreciate that the experiments are on visual observations, though they would be stronger if they included more complex tasks such as dextrous tasks, longer horizon tasks, or tasks with greater diversity/generalization

Some relevant related works to consider citing/discussing:
- Mobile (https://arxiv.org/abs/2102.10769), though uses low-dim observations
- VMAIL (https://arxiv.org/abs/2107.08829), though assumes actions in demos

Meta-review: (a) The paper introduces AIME-NoB, an algorithm designed to improve Imitation Learning from Observation (ILfO) using pre-trained world models. The authors define two barriers in imitation learning from observation (ILfO) methods: the Embodiment Knowledge Barrier (EKB) and the Demonstration Knowledge Barrier (DKB).  The EKB describes the limitation of a pretrained model when confronted
with novel observations and actions beyond its training experience. The DKB describes the generalization from a limited number of expert demonstrations in imitation learning. The proposed approaches to address these limitations are: (1) balance the offline and online datasets when updating the model, (2) combine AIME and AIL objectives. 

(b) Strengths:
- Problem Definition: The paper identifies and analyzes the EKB and DKB limitations in current ILfO methods using pretrained models. The clarity improved during rebuttal.
- Effective Solutions: The authors propose practical solutions for each knowledge barrier that show significant performance improvements.
- Ablation Studies: The authors perform thorough ablation studies and sensitivity tests to analyze the contribution of different components, such as the surrogate reward and data regularizer, to the overall performance.
- Clear Presentation: Multiple reviewers praised the paper's clear and accessible writing style, which helped in understanding the proposed approach. The use of figures (e.g. Figure 1) was also cited as being helpful.
- Strong results on DMC and meta-world.

(c) Weaknesses:
- Limited novelty: Several reviewers questioned the novelty of the proposed approach. The core method combines online interaction, a data-driven regularizer, and a surrogate reward, all of which have been previously used in some form.
- Incremental Contribution: Some reviewers felt that combining existing techniques, although effective, was closer to an implementation choice than a significant technical contribution.
- Realistic impact: The reviewers questioned the direct applicability to real-world robotics problems since actions are typically available, or can be estimated in robotics, or real-world videos often have a large domain shift.

(d) My decision is to reject the paper. 

I think the paper does a good job at motivating the design choices above, and has thorough experiments with strong results, especially with the new additions during the rebuttal, which show that the design choices significantly improve over prior methods. However, in terms of whether this paper makes significant contributions at least in (1) enough novel contributions in method or, (2) significant realistic impact. As the reviewers pointed out, it seems to be a weaker case here.

Decision: Reject