SlowGuess/NeurIPS_2023_OCR huggingface dataset:
gradient since none of them can explain the "flying". In general, poorly modelled regions of the world may lead to noisy gradients for the time steps before it. On the other hand, we can also find most variants achieve a good performance on the run backward demonstration dataset, which is mainly due to low expert performance (see Appendix D) for the task that makes imitation easy. Last but not least, since we follow the common practise for the benchmark [19], the Cheetah embodiment is operated on  $50\mathrm{Hz}$  which is much higher than the  $20\mathrm{Hz}$  used in Walker. Higher frequency of operation makes the effect of each individual action, i.e. change in the observation, more subtle and harder to distinguish, which poses an additional challenge for the algorithms.
Influence of different datasets As expected, for almost all the variants of methods, transferring within the same task is better than transferring between different tasks. In these settings, BCO(0)-MDP is comparable with AIME. However, AIME shines in cross-task transfer. Especially when transferring between run and walk tasks and transferring from stand to run on Walker, AIME outperforms the baselines by a large margin, which indicates the strong generalisability of a forward model over an inverse model. We also find that AIME makes substantially better use of exploratory data. On Walker, AIME largely outperforms baselines when using the p2e dataset as the embodiment dataset and outperforms most variants when using the random dataset as the embodiment dataset. Moreover, when transferring from the mix dataset, except for the MDP version, AIME outperforms other variants that train the world model on just any single individual task dataset of the mixed dataset. This showcases the scalability of a world model to be trained on a diverse set of experiences, which could be more valuable in real-world scenarios.
Influence of observation modality Compared with BCO(0), AIME is quite robust to the choice of observation modality. We can see a clear ladder pattern with BCO(0) when changing the setting from hard to easy, while for AIME the result is similar for each modality. However, we can still notice a small difference when comparing LPOMDP and Visual settings. Although these observations provide the same information, we find AIME in the LPOMDP setting performs better than in the Visual setting in most test cases. We attribute it to the fact that low-dimension signals have denser information and offer a smoother landscape in the evidence space than the pixels so that it can provide a more useful gradient to guide the action inference. Surprisingly, although having access to more information, AIME-MDP performs worse than AIME-LPOMDP on average. The biggest gaps happen when transferring from exploratory datasets, i.e. the p2e dataset on Walker and the random dataset on Cheetah. We conjecture this to the fact the world model is not trained well with the default hyper-parameters, but we defer further investigation to future work.
4.2 Ablation studies
In this section, we conduct some ablation studies to investigate how AIME's performance is influenced by different components and design choices. We will mainly focus on using the mix embodiment dataset and transfer to run task, which represents a more realistic setting where we want to use experience from multiple tasks to transfer to a new task.
Sample efficiency and scalability To test these properties, we vary the number of demonstrations within  $\{1,2,5,10,20,50,100,200,500,1000\}$ . We also include BC with the true action as an oracle baseline. The results are shown in Figure 4. BCO(0) struggles with low-data scenarios and typically needs at least 10 to 20 demonstrations to surpass the performance of a random policy. In contrast, AIME demonstrates continual improvement with as few as 2 trajectories. And surprisingly, thanks to the generalisation ability of the world model, AIME even outperforms oracle BC when the demonstrations are limited. These demonstrate the superior sample efficiency of the method. Moreover, the performance of AIME keeps increasing as more trajectories are provided beyond 100, which showcases the scalability of the method.
Objective function The objective function, i.e. ELBO, consists of two terms, the reconstruction term  $J_{\mathrm{rec}}$  and the KL term  $J_{\mathrm{KL}}$ . To investigate the role that each term plays in AIME, we retrain two variants of AIME by removing either of the terms. As we can see from Figure 5, removing either term will negatively impact the results. When we compare the two variants, only using the KL term is better in settings with low-dimensional signals, while using only the reconstruction term yields a slightly better result for the high-dimensional image signal. But on all settings, the performance of using only the KL term is very close to the one that use both terms. This suggests that the latent state in the world model has already mostly captured the essential part of the environment. Although it is
Figure 4: Ablation of the number of demonstrations on mix  $\rightarrow$  run transfer on the Walker embodiment. The performance is shown as the normalised returns over 3 seeds and 100 trials for each seed. The shaded region represents one standard deviation.
Figure 5: Ablation studies on mix  $\rightarrow$  run transfer on the Walker embodiment. Numbers are computed by averaging among 3 seeds and 100 trials for each seed, and then normalised to the percentage of the expert's performance. Error bars are showing one standard deviation.
still worse than using both terms, it sheds some light on the potential of incorporating decoder-free models [29] into the AIME framework.
Components Compared with the BCO(0) baseline, AIME consists of two distinct modifications: one is to use an SSM to integrate sequences and train the policy upon its latent representation; the other is to form an implicit IDM via gradients rather than training an IDM explicitly. We design two baselines to investigate the two components. First, to remove the SSM, we train a forward dynamics model directly on the observations of the embodiment dataset and use that as an implicit IDM for imitation on the demonstration dataset. We term this variant IIDM. Second, we train a separate IDM upon the trained latent state of the world model and use that to guide the policy learning in phase 2. The detailed derivation of the IDM formulation can be found in Appendix C. Figure 5 clearly demonstrates the significance of the latent representation for performance. Without the latent representation, the results are severely compromised across all settings. However, when compared to BCO(0), the IIDM does provide assistance in the high-dimensional Visual setting, where training an IDM directly from the observation space can be extremely challenging. While having IDM on the latent representation leads to a better performance comparing with BCO(0), but it still performs worse than AIME, especially on the POMDP settings.
5 Related works
Imitation learning from observations Most previous works on imitation learning from only observation can be roughly categorised into two groups, one based on IDMs [25, 9, 30, 28] and one based on generative adversarial imitation learning (GAIL) [31, 26, 27]. The core component of the first group is to learn an IDM that maps a state transition pair to the action that caused the transition. [25, 9] use the IDM to label the expert's observation sequences, then solve the imitation learning problem with standard BC. [30, 28] extend the IDM to a goal-conditioned setting in which the IDM is trained to be conditioned on a future frame as the goal instead of only the next frame. During deployment, the task is communicated on the fly by the user in the form of key-frames as goals. The
setup mainly suits for the robot manipulation tasks in their paper since the user can easily specify the goals by doing the manipulation himself, but not suits for the locomotion tasks, in which it is not clear what a long-term goal of observation is and also not practical set the next observation as the goal and demonstrate that in a high frequency by the user. Different from these methods, our approach uses a forward model to capture the knowledge of the embodiment. In the second group of approaches, the core component is a discriminator that distinguishes the demonstrator's and the agent's observation trajectories. Then the discriminator serves as a reward function, and the agent's policy is trained by RL [31]. As a drawback, in order to train this discriminator the agent has to constantly interact with the environment to produce negative samples. Different from these methods, our method does not require further interactions with the environment, enabling zero-shot imitation from the demonstration dataset. Besides the majority, there are also works [32, 33] don't strictly fit to the two groups. [32] also use forward model like us by learning a latent action policy and a forward dynamic based on the latent action. However, it still needs online environment interactions to calibrate the latent actions to the real actions. [33] is hybrid method that uses both of the components and focus on a setting that the demonstrations are coming from a different embodiment.
Reusing learnt components in decision-making Although transferring pre-trained models has become a dominant approach in natural language processing (NLP) [34, 35, 36] and has been getting more popular in computer vision (CV) [37, 36], reusing learnt components is less studied in the field of decision-making [7]. Most existing works focus on transferring policies [38, 9, 7]. On the other hand, the world model, a type of powerful perception model, that is purely trained by self-supervised learning lies behind the recent progress of model-based reinforcement learning [39, 17, 19, 40, 41, 42, 43, 44]. However, the transferability of these world models is not well-studied. [24] learns a policy by using a pre-trained world model from exploration data and demonstrates superior zero-shot and few-shot abilities. We improve upon this direction by studying a different setting, i.e. imitation learning. In particular, we communicate the task to the model by observing the expert while [24] communicates the task by a ground truth reward function which is less accessible in a real-world setting.
6 Discussion & conclusion
In this paper, we present AIME, a model-based method for imitation from observations. The core of the method exploits the power of a pre-trained world model and inverses it w.r.t. action inputs by taking the gradients. On the Walker and Cheetah embodiments from the DMC Suite, we demonstrate superior performance compared to baselines, even when some baselines can access the true state. The results showcase the zero-shot ability of the learnt world model.
Although AIME performs well, there are still limitations. First, humans mostly observe others with vision. Although AIME works quite well in the Visual setting, there is still a gap compared with the LPOMDP setting where the low-dimensional signals are observed. We attribute this to the fact that the loss surface of the pixel reconstruction loss may not be smooth enough to allow the gradient method to find an equally good solution. Second, in this paper, we only study the simplest setting where both the embodiment and sensor layout are fixed across tasks. On the other hand, humans observe others in a third-person perspective and can also imitate animals whose body is not even similar to humans'. Relaxing these assumptions will open up possibilities to transfer across different embodiments and even directly from human videos. Third, for some tasks, even humans cannot achieve zero-shot imitation by only watching others. This may be due to the task's complexity or completely unfamiliar skills. So, even with proper instruction, humans still need to practise in the environment and learn something new to solve some tasks. This motivates an online learning phase 3 as an extension to our framework. We defer these topics to future work.
We hope this paper demonstrates the great potential of transferring a learnt world model, incentivises more people to work in this direction and encourages researchers to also share their learnt world model to contribute to the community.
We want to acknowledge Elie Aljalbout for the insightful discussion during the initial stage of the project and Botond Cseke for mathematical support.
A Computational resources
On a GTX 1080Ti graphics card, AIME typically requires 10 hours of training for phase 1 and 5 hours of training for phase 2 with MDP and LPOMDP setups. The required time nearly doubled when running with Visual settings, due to heavier visual backbones and rendering. We conduct experiments on a shared local cluster which uses A100 and RTX8000 GPUs. The newer GPUs can slightly improve the training speeds but not much since the main computational bottleneck is the recurrent structure. In terms of one GTX 1080Ti, it will require roughly 50 GPU days to produce the benchmark results.
B Implementation and training details
We implement all listed methods in PyTorch [45].
For the world model, we use RSSM [17, 19], which offers state-of-the-art performances by splitting the latent state to be a combination of deterministic and stochastic components. The RSSM implementation is largely following Dreamer-v1 [19] with continuous stochastic and deterministic variables. Although newer versions of Dreamer [41, 42] offer some new tricks to improve performance, we initially choose not to use them for the sake of simplicity. We use a slightly larger state space for our experiment with 512 deterministic and 128 stochastic dimensions and find it generally eases the policy training process to collect the datasets. For the Visual setting, the encoder and decoder are implemented with CNNs. The decoder output a Gaussian distribution with the mean output by CNN and a fixed variance of 1. For the low-dimensional settings, the encoder is implemented as an identity function while the decoders are Gaussian distributions with both the mean and variance parameterised by MLPs. The deterministic part of the state is implemented as a GRU cell [46]. For the default hyper-parameters, we do not use any free nats [19], KL scaling [19] and KL balancing [41] tricks in the literature to relax the constraint of the KL term. When decoding low-dimensional signals, we sometimes observed the decoder yielding a degenerate solution as found in [47]. We use their  $\beta$ -nll to remedy this problem, and since it re-weights the reconstruction term, we re-weight the KL term accordingly to maintain the balance.
Without further mention, all the CNN encoders and decoders above are implemented as in [39, 17], while all the MLPs are with 2 hidden layers and 128 units of each layer with ELU [48] activation function. All the components are trained with Adam optimiser with a learning rate of  $1\mathrm{e} - 3$ . For the stochastic policy, the output distribution is modelled by a TanhGaussian distribution [49] with both the mean and variance parameterised by neural networks.
For AIME, we consider 100 gradient steps as an epoch. For phase 1, the model is trained for 1000 epochs, while for phase 2 we train the policy for 500 epochs. Both the final model and policy are from the last epoch without any early stopping criteria.
When training the world model on the Cheetah dataset, we find the default hyper-parameters cannot stably train a good world model. Thus, we adapt the implementation to exactly the same network structure as the origin repository. Specifically, the decoders of low-dimensional observations are also with a fixed variance of 1 and all the MLPs are widened to 512 neurons in each hidden layer and equipped with Layer Normalisation [50]. For the hyper-parameters, the learning rate is decreased to  $3\mathrm{e} - 4$  , and we use free nats of 1.0 and KL balancing of 0.8 to mitigate the collapse and unstable problem of the KL term. For the LPOMDP setting, we also set the KL scaling parameter  $\beta = 0.0002$  to relax the constraint. One thing that needs to mention is, while the tricks about the KL term are helpful for model training, they hurt the results in phase 2. It could be because, in phase 2, the model is frozen, so that no-more stability issues will be encountered. So in this case it is better to optimise the policy with the true ELBO.
To be strict with our setup of the two phases, we retrain the world model after data collection for all the experiments. However, one can also directly use the world model from the trained dreamer agent. We empirically find these models yield similar results with the world model retrained afterwards on the same reply buffer. One caveat is that, although it is tempting to also reuse the trained policy as initialisation in phase 2, we found it is actually harmful to the performance. We conjecture that it is due to learnt policies being stuck in some local minima that they are unable to escape.
For the BCO(0) baseline, the IDM and policy are built by using the same network architecture with the world model to make a fair comparison. The observations are first processed by the encoder network, and then get stacked to deal with the temporal information. An MLP is used to decode the stacked representation to the output distribution. We stack 5 consecutive in this work. We did a grid search about the width, depth of the MLP and also the number of stacking frames and didn't find any increase of the performance. Following the original paper, we split the datasets by  $7:3$ , and choose the final model based on the validation loss.
C AIME with IDM
In this section, we introduce an alternative variant of AIME which also uses IDM. Recall that in the Bayesian derivation, we factorised both the posterior and prior of the joint distribution of state and action with a shared policy network, as in eq. (6) and eq. (7). Alternatively, we can re-factorise the posterior with IDM, that is
$$
q _ {\phi , \theta} \left(s _ {t}, a _ {t - 1} \mid s _ {t - 1}, f _ {\phi} \left(o _ {t}\right)\right) = q _ {\phi} \left(a _ {t - 1} \mid s _ {t - 1}, f _ {\phi} \left(o _ {t}\right)\right) q _ {\phi} \left(s _ {t} \mid s _ {t - 1}, a _ {t - 1}, f _ {\phi} \left(o _ {t}\right)\right). \tag {10}
$$
One thing that needs to be noticed here is that the IDM is not in the familiar form of  $q_{\phi}(a_{t - 1}|s_{t - 1},s_t)$ . This is because the latent state in the world model is action dependent so the familiar form is non-casual in the world model. But we should highlight here that this non-casual structure is a result of the model we used in phase 1 since we want to reuse the knowledge learnt there. For example, one can also factorise the joint posterior as
$$
q _ {\phi} \left(s _ {t}, a _ {t - 1} \mid s _ {t - 1}, f _ {\phi} \left(o _ {t}\right)\right) = q _ {\phi} \left(a _ {t - 1} \mid s _ {t - 1}, s _ {t}\right) q _ {\phi} \left(s _ {t} \mid s _ {t - 1}, f _ {\phi} \left(o _ {t}\right)\right). \tag {11}
$$
However, in this case, the model is different for phase 1. In this section, we stick to using the factorisation in eq. (10). Since a new IDM component is introduced, the objective of both phase 1 and phase 2 need to be modified. For phase 1, since actions are available in the dataset, the IDM can be treated as a decoder and trained by maximising likelihood. That is, we add  $J_{\mathrm{IDM}}^{\mathrm{p1}}(\tau^{(T)}) = \sum_{t=0}^{T-1} \log q_{\phi}(a_t | s_t, f_{\phi}(o_{t+1}))$  to the objective function. For phase 2, since actions are not available, the IDM serves as the posterior and guides the prior policy through a KL divergence, i.e.  $J_{\mathrm{IDM}}^{\mathrm{p2}}(\tau^{(T)}) = \sum_{t=0}^{T-1} -D_{KL}[q_{\phi}(a_t | s_t, f_{\phi}(o_{t+1}))||\pi_{\psi}(a_t | s_t)]$ .
One caveat about this formulation is that, in phase 1, the IDM forms a loop on the graphical model. In order to stabilise the training process, we detach the gradient from the IDM to the rest of the network.
D Dataset details
Here we provide extra information about the datasets. The expert return which we normalised against is shown in Table 1 and Table 2.
Table 1: Average expert return of each demonstration dataset of Walker.
Ddemo	Average return
stand	957.87 (max: 1000)
walk	943.79 (max: 1000)
run	604.10 (max: 1000)
Table 2: Average expert return of each demonstration dataset of Cheetah.
Ddemo	Average return
run	888.65 (max: 1000)
run backwards	218.50 (max: 500)
flip	485.79 (max: 500)
flip backwards	379.91 (max: 500)
Table 3: Result on V-D4RL main datasets. Embodiment datasets are marked on the left, and the demonstration datasets are chosen to be the expert dataset for each task in the origin environment. Values are averaged over 100 trajectories and reported as accumulated reward divided by 10, as suggested in the V-D4RL paper.
BCO(0)	AIME
walker-random	2.11 ± 0.91	12.36 ± 4.69
walker-medium_replay	6.54 ± 6.56	10.18 ± 4.33
walker-mix	5.32 ± 5.23	8.49 ± 3.60
cheetah-random	0.01 ± 0.01	9.48 ± 4.72
cheetah-medium_replay	15.47 ± 7.38	31.47 ± 16.14
cheetah-mix	16.08 ± 6.61	40.27 ± 11.52
E Experiments on V-D4RL datasets
We provide here some additional results of AIME on V-D4RL datasets [51] to showcase that AIME can also work with datasets collected by non-model-based methods. V-D4RL provides multiple different datasets for the Walker and Cheetah embodiments from DMC Suite, and it is original designed for offline RL with visual inputs. The datasets are collected by running a few model-free RL methods and either keep the replay buffer or rollout from a policy checkpoints. Since our setting requires a bit more exploration in the embodiment datasets to understand the embodiment, we choose to use their random and medium_replay datasets as the embodiment datasets. The expert datasets are used as the demonstration datasets. Same as what we did for the Walker embodiment in the main text, we also mix the two embodiment datasets for each embodiment to form a mix dataset.
The results on these datasets are shown in Table 3. We can see that the performance of both BCO(0) and AIME is generally low, but AIME still outperform BCO(0) which proves AIME can also handle datasets generated by model-free methods. The low performance is due to a more constrained setup of the task, i.e. less amount of embodiment data and less diversity. Except the cheetah-medium_replay having 400 trajectories, the other three datasets provided by V-D4RL have only 200 trajectories, which is much less than the 1000 trajectories in the main experiments. Moreover, it is already shown from Figure 2 and Figure 3 that random datasets do not help much in learning a model, and intuitively the medium_replay dataset is better but still does not contain enough information to solve the task.
We also conduct experiments on the V-D4RL distracting datasets, to test the performance of AIME on distracting datasets. For the Walker embodiment, the benchmark provides random datasets with a distraction level of easy, medium, and hard. We also merge these three levels to form a mix dataset. Moreover, we also merge this mix dataset with the mix dataset in the second experiment to form a total_mix dataset. We treat these five datasets as the embodiment dataset and the expert dataset as the demonstration dataset. For the Cheetah embodiment, the benchmark provides medium and expert datasets with a distraction level of easy, medium, and hard. We subsample the medium datasets to get 200 trajectories from each level, then merge that with the mix dataset in the second experiment to form a total_mix dataset. Then the algorithms are using this total_mix dataset as the embodiment dataset and the expert dataset as the demonstration dataset.
As we can see from the result from Table 4, although we still outperform the BCO(0) baseline, AIME is impacted significantly by the distractions. This behaviour is expected since the world model is trained with reconstruction loss. It is not easy to handle observations with distractions. A potential solution to this problem is to freeze only the dynamics part of the world model and allowing encoders and decoders to fine-tune their parameters in the second phase. We leave these improvements for our future works.
F Additional plots