SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xingyuan Zhang $^{1,2*}$ , Philip Becker-Ehmck $^{1}$ , Patrick van der Smagt $^{1,3}$ , Maximilian Karl $^{1}$ $^{1}$ Machine Learning Research Lab, Volkswagen Group,  $^{2}$ Technical University of Munich,  $^{3}$ Eötvös Loránd University Budapest  
{xingyuan.zhang, philip.becker-ehmck, maximilian.karl}@volkswagen.de
Unlike most reinforcement learning agents which require an unrealistic amount of environment interactions to learn a new behaviour, humans excel at learning quickly by merely observing and imitating others. This ability highly depends on the fact that humans have a model of their own embodiment that allows them to infer the most likely actions that led to the observed behaviour. In this paper, we propose Action Inference by Maximising Evidence (AIME) to replicate this behaviour using world models. AIME consists of two distinct phases. In the first phase, the agent learns a world model from its past experience to understand its own body by maximising the evidence lower bound (ELBO). While in the second phase, the agent is given some observation-only demonstrations of an expert performing a novel task and tries to imitate the expert's behaviour. AIME achieves this by defining a policy as an inference model and maximising the evidence of the demonstration under the policy and world model. Our method is "zero-shot" in the sense that it does not require further training for the world model or online interactions with the environment after given the demonstration. We empirically validate the zero-shot imitation performance of our method on the Walker and Cheetah embodiment of the DeepMind Control Suite and find it outperforms the state-of-the-art baselines. Code is available at: https://github.com/argmax-ai/aime.
In recent years, deep reinforcement learning (DRL) has enabled intelligent decision-making agents to thrive in multiple fields [1, 2, 3, 4, 5, 6]. However, one of the biggest issues of DRL is sample inefficiency. The dominant framework in DRL is learning from scratch [7]. Thus, most algorithms require an incredible amount of interactions with the environment [1, 2, 3].
In contrast, cortical animals such as humans are able to quickly learn new tasks through just a few trial-and-error attempts, and can further accelerate their learning process by observing others. An important difference between biological learning and the DRL framework is that the former uses past experience for new tasks. When we try a novel task, we use previously learnt components and generalise to solve the new problem efficiently. This process is augmented by imitation learning [8], which allows us to replicate similar behaviours without direct observation of the underlying muscle movements. If the DRL agents could similarly harness observational data, such as the abundant online video data, the sample efficiency may be dramatically improved [9]. The goal of the problem is related to the traditional well-established Learning from Demonstration (LfD) field from the robotics community [10, 11], but instead of relying on knowledge from the engineers and researchers, e.g. mathematical model of robot's dynamic or primitives, we aim to let the robots learn by itself.
Figure 1: Overview of AIME algorithm. In phase 1, both observations and actions are provided by the embodiment dataset and the agent learns a variational world model to model the evidence of observations conditioned on the actions. Then the learnt model weights are frozen and transferred to phase 2. In phase 2, only the observations are provided by the demonstration dataset, so the agent needs to infer both states and actions. The action inference is achieved by the policy model which samples actions given a state. The grey lines indicate the world model parameters are frozen in phase 2. Both phases are optimised toward the same objective, i.e. the ELBO.
However, directly learning a model from observation-only sequences [12, 13] is insufficient for both biological and technical systems. Without knowing the actions that lead to the observations, the observation sequences are highly stochastic and multi-modal [14]. Trying to infer these unknown actions without prior knowledge of the world is difficult due to the problem of attributing which parts of the observations are influenced by the actions and which parts are governed by normal system evolution or noise.
Therefore, in this work, we hypothesise that in order to make best use of observation-only sequences, an agent has to first understand the notion of an action. This can be achieved by learning a model from an agent's past experiences where both the actions and their consequences, i.e. observations, are available. Given such a learnt model which includes a causal model of actions and their effects, it becomes feasible for an agent to infer an action sequence leading to given observation-only data.
In this work, we propose a novel algorithm, Action Inference by Maximising Evidence (AIME), to try to replicate the imitation ability of humans. The agent first learns a world model from its past experience by maximising the evidence of these experiences. After receiving some observation-only demonstrations of a novel task, the agent tries to mimic the demonstrator by finding an action sequence that makes the demonstration most likely under the learnt model. This procedure is shown in Figure 1.
Our contribution can be summarised as follows:
- We propose AIME, a novel method for imitation from observation. AIME first learns a world model by maximising the evidence of its past experience, then considers the policy as an action inference model and imitates by maximising the evidence of demonstration.
- We conduct experiments with a variety of datasets and tasks to demonstrate the superior performance of AIME compared with other state-of-the-art methods. The results showcase the zero-shot transferability of a learnt world model.
Consider an MDP problem defined by the tuple  $\{S, A, T, R\}$ , where  $S$  is the state space,  $A$  is the action space,  $T: S \times A \to S$  is the dynamic function and  $R: S \to \mathbb{R}$  is the reward function. A POMDP adds partial observability upon an MDP with two components: the observation space  $O$  and the emission function  $\Omega: S \to O$ . The six components of a POMDP can be categorised into three groups:  $S$ ,  $A$  and  $T$  define the embodiment of our agent,  $O$  and  $\Omega$  define the sensors of our agent and
$R$  itself defines the task. The goal is to find a policy  $\pi : S \to A$  which maximises the accumulated reward, i.e.  $\sum_{t} r_{t}$ .
In this paper, we want to study imitation learning within a fixed embodiment across different tasks. We presume the existence of two datasets for the same embodiment:
- Embodiment dataset  $D_{\mathrm{body}}$  contains trajectories  $\{o_0, a_0, o_1, a_1 \ldots\}$  that represent past experiences of interacting with the environment. This dataset provides information about the embodiment for the algorithm to learn a model. For example, in this paper, the dataset is a replay buffer filled while solving some tasks with the same embodiment. But in general, it may be any collection of past experiences of the embodiment.
- Demonstration dataset  $D_{\mathrm{demo}}$  contains a few expert trajectories  $\{o_0, o_1, o_2, \ldots\}$  of the embodiment solving a certain task defined by  $R_{\mathrm{demo}}$ . The crucial difference between this dataset and the embodiment dataset 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 use information in  $D_{\mathrm{body}}$  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}}$ . For simplicity, we assume that the two datasets share the same observation space  $O$  and the emission model  $\Omega$ .
In this section, we describe our proposed method, AIME, in detail. AIME consists of two phases. In the first phase, the knowledge of the embodiment is learnt through a form of world model; while in the second phase, this knowledge is used to imitate the expert.
3.1 Phase 1: Model Learning
In the first phase, we need to learn a model to understand our embodiment. We achieve this by learning a world model. As an analogy to a language model, we define a world model as a probability distribution over sequences of observations. The model can be either unconditioned or conditioned on other factors such as previous observations or actions. For phase 1, the model needs to be the conditional distribution, i.e.  $p(o_{1:T}|a_{0:T - 1})$ , to model the effect of the actions. When given an observation sequence, the likelihood of this sequence under the model is referred to as evidence.
In this paper, we consider variational 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) [15, 16, 17, 18, 19, 20]. Such a variational world model involves four components, namely
$$
\operatorname {e n c o d e r} z _ {t} = f _ {\phi} (o _ {t}),
$$
$$
\text {p o s t e r i o r} s _ {t} \sim q _ {\phi} \left(s _ {t} \mid s _ {t - 1}, a _ {t - 1}, z _ {t}\right),
$$
$$
\text {p r i o r} s _ {t} \sim p _ {\theta} \left(s _ {t} \mid s _ {t - 1}, a _ {t - 1}\right),
$$
$$
\operatorname {d e c o d e r} o _ {t} \sim p _ {\theta} \left(o _ {t} \mid s _ {t}\right).
$$
$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, a variational world model is trained by maximising the 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
$$
J \left(o _ {1: T}, s _ {0: T}, a _ {0: T - 1}\right) = J _ {\text {r e c}} \left(o _ {1: T}, s _ {0: T}, a _ {0: T - 1}\right) + J _ {\mathrm {K L}} \left(o _ {1: T}, s _ {0: T}, a _ {0: T - 1}\right), \tag {1}
$$
$$
\text {w h e r e} J _ {\mathrm {r e c}} \left(o _ {1: T}, s _ {0: T}, a _ {0: T - 1}\right) = \sum_ {t = 1} ^ {T} \log p _ {\theta} \left(o _ {t} \mid s _ {t}\right), \tag {2}
$$
$$
J _ {\mathrm {K L}} \left(o _ {1: T}, s _ {0: T}, a _ {0: T - 1}\right) = \sum_ {t = 1} ^ {T} - D _ {\mathrm {K L}} \left[ q _ {\phi} \left(s _ {t} \mid s _ {t - 1}, a _ {t - 1}, f _ {\phi} (o _ {t})\right) \right| | p _ {\theta} \left(s _ {t} \mid s _ {t - 1}, a _ {t - 1}\right) ]. \tag {3}
$$
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 [21, 22] to autoregressively sample the inferred states from the observation and action sequence.
Combining all these, we formally define the optimisation problem for this phase as
$$
\phi^ {*}, \theta^ {*} = \underset {\phi , \theta} {\operatorname {a r g m a x}} \mathbb {E} _ {\left\{o _ {1: T}, a _ {0: T - 1} \right\} \sim D _ {\text {b o d y}}, s _ {0: T} \sim q _ {\phi}} [ J \left(o _ {1: T}, s _ {0: T}, a _ {0: T - 1}\right) ]. \tag {4}
$$
3.2 Phase 2: Imitation Learning
In the second phase, we want to utilise the knowledge of the world model from the first phase to imitate the expert behaviour from the demonstration dataset  $D_{\mathrm{demo}}$  in which only sequences of observations but no actions are available. We derive our algorithm from two different perspectives.
The Bayesian derivation Since the actions are unknown in the demonstration, instead of modelling the conditional evidence in phase 1, we need to model the unconditional evidence, i.e.  $\log p_{\theta}(o_{1:T})$ . Thus, we also need to model the actions as latent variables together with the states. In this way, the reconstruction term  $J_{\mathrm{rec}}$  will stay the same as eq. (2), while the KL term will be defined on the joint distribution of states and actions, i.e.
$$
J _ {\mathrm {K L}} \left(o _ {1: T}, s _ {0: T}, a _ {0: T - 1}\right) = \sum_ {t = 1} ^ {T} - D _ {\mathrm {K L}} \left[ q _ {\phi , \psi} \left(s _ {t}, a _ {t - 1} \mid s _ {t - 1}, f _ {\phi} \left(o _ {t}\right)\right) \right| | p _ {\theta , \psi} \left(s _ {t}, a _ {t - 1} \mid s _ {t - 1}\right) ]. \tag {5}
$$
If we choose the action inference model in the form of a policy, i.e.  $\pi_{\psi}(a_t|s_t)$ , and share it in both posterior and prior, then the new posterior and prior can be factorised as
$$
q _ {\phi , \psi} \left(s _ {t}, a _ {t - 1} \mid s _ {t - 1}, f _ {\phi} (o _ {t})\right) = \pi_ {\psi} \left(a _ {t - 1} \mid s _ {t - 1}\right) q _ {\phi} \left(s _ {t} \mid s _ {t - 1}, a _ {t - 1}, f _ {\phi} (o _ {t})\right) \tag {6}
$$
$$
\text {a n d} p _ {\theta , \psi} \left(s _ {t}, a _ {t - 1} \mid s _ {t - 1}\right) = \pi_ {\psi} \left(a _ {t - 1} \mid s _ {t - 1}\right) p _ {\theta} \left(s _ {t} \mid s _ {t - 1}, a _ {t - 1}\right) \tag {7}
$$
respectively. When we plug them into the eq. (5), the policy term cancels and we will get a similar optimisation problem with phase 1 as
$$
\psi^ {*} = \underset {\psi} {\operatorname {a r g m a x}} \mathbb {E} _ {o _ {1: T} \sim D _ {\mathrm {d e m o}}, \left\{s _ {0: T}, a _ {0: T - 1} \right\} \sim q _ {\phi^ {*}, \psi}} \left[ J \left(o _ {1: T}, s _ {0: T}, a _ {0: T - 1}\right) \right]. \tag {8}
$$
The main difference between eq. (4) and eq. (8) is where the action sequence is coming from. In phase 1, the action sequence is coming from the embodiment dataset, while in phase 2, it is sampled from the policy instead since it is not available in the demonstration dataset.
The control derivation From another perspective, we can view phase 2 as a control problem. One crucial observation is that, as shown in eq. (1), given a trained world model, we can evaluate the lower bound of the evidence of any observation sequence given an associated action sequence as the condition. In a deterministic environment where the inverse dynamics model is injective, the true action sequence that leads to the observation sequence is the most likely under the true model. In general, the true action sequence may not necessarily be the most likely under the model. This is, however, a potential benefit of our approach. We are mainly interested in mimicking the expert's demonstration and may be better able to do so with a different action sequence.
Thus, for each observation sequence that we get from the demonstration dataset, finding the missing action sequence can be considered as a trajectory-tracking problem and can be tackled by planning. To be specific, we can find the missing action sequence by solving the optimisation problem
$$
a _ {0: T - 1} ^ {*} = \underset {a _ {0: T - 1}} {\operatorname {a r g m a x}} \mathbb {E} _ {o _ {1: T} \sim D _ {\mathrm {d e m o}}, s _ {0: T} \sim q _ {\phi^ {*}}} \left[ J \left(o _ {1: T}, s _ {0: T}, a _ {0: T - 1}\right) \right]. \tag {9}
$$
If we solve the above optimisation problem for every sequence in the demonstration dataset, the problem will be converted to a normal imitation learning problem and can be solved with standard techniques such as behavioural cloning. We can also view this as forming an implicit inverse dynamics model (IDM) by inverting a forward model w.r.t. the actions.
To make it more efficient, we use amortised inference. We directly define a policy  $\pi_{\psi}(a_t|s_t)$  under the latent state of the world model. By composing the learnt world model and the policy, we can form
Algorithm 1: AIME
Data: Embodiment dataset  $D_{\mathrm{body}}$ , Demonstration dataset  $D_{\mathrm{demo}}$ , Learning rate  $\alpha$
Phase 1: Model Learning
Initialise world model parameters  $\phi$  and  $\theta$
while model has not converged do
Initialise policy parameters  $\psi$
$\{o_{1:T},a_{0:T - 1}\} \sim D_{body}$ $s_0\gets 0$  for  $t = 1:T$  do  $s_t\sim q_\phi (s_t|s_{t - 1},a_{t - 1},f_\phi (o_t))$  Compute objective function  $J$  from eq. (1) Update model parameters  $\phi \leftarrow \phi +\alpha \nabla_{\phi}J,\theta \leftarrow \theta +\alpha \nabla_{\theta}J$  # Phase 2: Imitation Learning
while policy has not converged do
$o_{1:T}\sim D_{demo}$ $s_0\gets 0$    
for  $t = 1:T$  do  $\begin{array}{r}a_{t - 1}\sim \pi_{\psi}(a_{t - 1}|s_{t - 1})\\ s_t\sim q_\phi (s_t|s_{t - 1},a_{t - 1},f_\phi (o_t)) \end{array}$  Compute objective function  $J$  from eq. (1) Update policy parameters  $\psi \leftarrow \psi +\alpha \nabla_{\psi}J$
a new generative model of the state sequence by the chain of  $s_t \to a_t \to s_{t+1} \to a_{t+1} \dots \to s_T$ . Then we will get the same optimisation problem as eq. (8).
To sum up, in AIME, we use the same objective function – the ELBO – in both phases with the only difference being the source of the action sequence. We provide the pseudo-code for the algorithm in Algorithm 1 with the colour highlighting the different origins of the actions between the two phases.
To test our method, we need multiple environments sharing an embodiment while posing different tasks. Therefore, we consider Walker and Cheetah embodiment from the DeepMind Control Suite (DMC Suite) [23]. Officially, the Walker embodiment has three tasks: stand, walk and run. While the Cheetah embodiment only has one task, run, we add three new tasks, namely run backwards, flip and flip backwards, inspired by previous work [24]. Following the common practice in the benchmark [19], we repeat every action two times when interacting with the environment. For both embodiments, the true state includes both the position and the velocity of each joint and the centre of mass of the body. In order to study the influence of different observation modalities, we consider three settings for each environment: MDP uses the true state as the observation; Visual uses images as the observation; LPOMDP uses only the position part of the state as the observation, so that information-wise it is identical to the Visual setting but the information is densely represented in a low-dimensional form.
To generate the embodiment and demonstration datasets, we train a Dreamer [19] agent in the Visual setting for each of the tasks for 1M environment steps. We take the replay buffer of these trained agents as the embodiment datasets  $D_{\mathrm{body}}$ , which contain 1000 trajectories, and consider the converged policy as the expert to collect another 1000 trajectories as the demonstration dataset  $D_{\mathrm{demo}}$ . We only use 100 trajectories for the main experiments, and the remaining trajectories are used for an ablation study. The performance of the policy is measured by accumulated reward. The exact performance of the demonstration dataset can be found in Appendix D. Besides the above embodiment datasets, we also study two datasets generated by purely exploratory behaviour. First, we use a random policy that samples uniformly from the action space to collect 1000 trajectories, and we call this the random dataset. Second, we train a Plan2Explore [24] agent for 1000 trajectories and label its replay buffer as the  $p2e$  dataset. Moreover, for the Walker embodiment, we also merge all the above datasets except the run dataset to form a mix dataset. This resembles a practical setting where one possesses a lot of experience with one embodiment and uses all of it to train a single foundational world model.
Figure 2: Performances on Walker. Each column indicates one task and its associated demonstration dataset, while each row indicates the embodiment datasets used to train the model. The title of each figure is named according to  $D_{\mathrm{body}} \rightarrow D_{\mathrm{demo}}$ . Numbers are computed by averaging among 100 trials and then normalised to the percentage of the expert's performance. Error bars are showing one standard deviation. The last row and column are averaged over the corresponding task or dataset. The error bar is large for them due to aggregating performance distributed in a large range.
4.1 Benchmark results
We mainly compare our method with BCO(0) [25]. BCO(0) first trains an IDM from the embodiment dataset and then used the trained IDM to label the demonstration dataset and then uses Behavioural Cloning (BC) to recover the policy. We do not compare with other methods since they either require further environment interactions [26, 27] or use a goal-conditional setting [28] which does not suit the locomotion tasks. More details about related works can be found in Section 5. The implementation details can be found in Appendix B.
The main results of our comparison are shown in Figure 2 and Figure 3. Overall, we can see that AIME largely outperforms BCO(0) in all the environment settings on Walker and on POMDP settings on Cheetah. AIME typically achieves the lowest performance on the Visual setting, but even that is comparable with BCO(0)-MDP which can access the true states. We attribute the good performance of AIME to two reasons. First, the world model has a better data utilisation rate than the IDM because the world model is trained to reconstruct whole observation sequences, while the IDM only takes
Figure 3: Performances on Cheetah. Each column indicates one task and its associated demonstration dataset, while each row indicates the embodiment datasets used to train the model. The title of each figure is named according to  $D_{\mathrm{body}} \rightarrow D_{\mathrm{demo}}$ . runb and flipb are short hands for run backwards and flip backwards. Numbers are computed by averaging among 100 trials and then normalised to the percentage of the expert's performance. Error bars are showing one standard deviation. The last row and column are averaged over the corresponding task or dataset. The error bar is large for them due to aggregating performance distributed in a large range.
short clips of the sequence and only predicts the actions. Thus, the world model has less chance to overfit, learns better representations and provides better generalisation. Second, by maximising the evidence, our method strives to find an action sequence that leads to the same outcome, not to recover the true actions. For many systems, the dynamics are not fully invertible. For example, if a human applies force to the wall, since the wall does not move, one cannot tell how much force is applied by visual observation. The same situation applies to the Walker and Cheetah when certain joints are locked due to the singular pose. This same phenomenon is also discussed in [28].
We also find that, comparing with the Walker experiments, the performance on Cheetah is lower and the improvement offered by AIME is smaller. We think it is because the setup for Cheetah is much harder than Walker. Although the tasks sound similar from the names, e.g. flip and flip backward, due to the asymmetrical structure of the embodiment, the behaviour for solving the tasks can be quite different. The difference limits the amount of knowledge that can be transferred from the embodiment dataset to the demonstrations. Moreover, some tasks are built to be hard for imitation. For example, in the demonstration of the flip tasks, the cheetah is "flying" in the air and the actions taken there is not relevant for solving the tasks. That leaves only a few actions in the sequence that are actually essential for solving the task. We think this is more challenging for AIME since it needs to infer a sequence of actions, while BCO(0) is operating on point estimation. That is, when the first few actions cannot output reasonable actions to start the flip, then the later actions will create a very noisy