Action Inference by Maximising Evidence: Zero-Shot Imitation from Observation with World Models
Authors: Xingyuan Zhang
Philip Becker-Ehmck
Patrick van der Smagt
Maximilian Karl
Venue: NeurIPS 2023 poster (2023)
Keywords: Imitation Learning
World Models
Latent Variable Model
Transfer Learning
Variational Inference
TL;DR: We propose Action Inference by Maximising Evidence (AIME), a novel algorithm for imitation from observation by transferring learnt world models.

Abstract: 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 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.

Action Inference by Maximising Evidence: Zero-Shot Imitation from Observation with World Models
Xingyuan Zhang!. Philip Becker-Ehmek' Patrick van der Smagt' Maximilian Karl'
MMachine Learning Research Lab, Volkswagen Group, Techhical University of Munich, 'Eötvös Loránd University Budapest
(xingyuan. zzang, tchang,phiiip... philip. beckeprbeeerrr-. boocker-emmcke maximi ian karl karl
Abstract
Unlike most reinforcement learning agents which require unrealistic amount 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 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-onle demonstrations of an expert performing a novel task and tries to imitate the expert's behaviour. AIME achieves this by defining policy as 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 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-ard baselines. Code is available at: https://gt //ggtt//bb thub com/ ar gmax- ai/ aime
Introduction
In recent years, deep reinforcement learning (DRL) has enabled intelligent decision -maiin agents to thrive in multiple fields [1] 2 3 5 6). However, one of the biggest issues of DRL is sample inefficiency. The dominant framework in DRL is learning from scratch 77]. Thus, most algorithms require an incredible amount of interactions with the environment
In contrast, cortical animals such as humans are able to quickly learn new tasks through just a few trial-and-erron attempts, and can further accelerate their learning process by observing others. An important difference between biological learning and the DRI. framework is that the former uses past experience for new tasks. When we try novel task, use previously learnt components and generalise to solve the new problem efficiently. This process is augmented by imitation learning (8), which allows 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 [ The goal of the problem is related to the traditional well-established Learning from Demonstration (LfD) field from the robotics community (TO) 11,, but instead of relying knowledge from the engineers and researchers, e.g. mathematical model of robot's 's dynamic or primitives, we aim to let the robots learn by itself.
"Corresponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Phase 1: Model Learning
Phase 2: Imitation Learning
Transfer Parameters gb.
pc+1
2
Pc-1
OL12
Embodiment dataset
Demonstration dataset fo 02. ...
R2
Gt
[0 R, O2.
8t
Se+1
Pr+2
$
Sc-1
5++2
Pa
[OUT. Su.T. Ro4- (0rr + +JI (Oum, Sacr, 9or-11)
Figure I: Overview of AIME algorithm. In phase I, both observations and actions are provided by the embodiment dataset and the agent learns 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 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 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 infer these unknown actions without prior knowledge of the world is difficult due to the problem of attributing which parts of the observations 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 agent's past experiences where both the actions and their consequences, i.e. observations, available. Given such a learnt model which includes a causal model actions and their effects, becomes feasible for agent to infer action sequence leading 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 world model from its past experience by maximising the evidence of these experiences. After receiving some observation-ons demonstrations of novel task, the agent tries to mimic the demonstrator by finding action sequence that makes the demonstration most likely under the learnt model. This procedure is shown Figure]
Our contribution can be summarised as follows:
We propose AIME, novel method for imitation from observation. AIME first learns 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 variety of datasets and tasks demonstrate the superior performance of AIME compared with other state-of-the-arr methods. The results showcase the zero-shot transferability of learnt world model.
2 Problem formulation
Consider an MDP problem defined by the tuple R)), where S is the state space, A is the action space, T: S x - S is the dynamic function and R: - R is the reward function. A POMDP adds partial observability upon MDP with two components: the observation space O and the emission function Ω + O. The six components of POMDP can be categorised into three groups: S, A and T define the embodiment of our agent, O and Ω define the sensors of our agent and

---

R itself defines the task. The goal is to find a policy : S - A which maximises the accumulated reward, i.c. E, Tt.
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 Doady contains trajectories {oo- 40, 01, .1... that represent past experiences of interacting with the environment. This dataset provides information about the embodiment for the algorithm to learn model. For example, in this paper, the dataset is replay buffer filled while solving some tasks with the same embodiment. But in general, may be any collection of past experiences of the embodiment.
Demonstration dataset Demo contains a few expert trajectories {oo. 01, 02, 02... of the embodiment solving certain task defined by Rdemo- The crucial difference between this dataset and the embodiment dataset is that the actions are not provided anymore since they are not observable from third-person perspective.
The goal of our agent is to use information in Doddy to learn a policy æ from Daemo which can solve the task defined by Rdemo well as the expert who generated Daemo- For simplicity, we assume that the two datasets share the same observation space O and the emission model Ω.
3 Methodology
In this section, we describe our proposed method, AIME, in detail. AIME consists of phases. In the first phase, the knowledge of the embodiment is learnt through 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 world model. As analogy language model, we define world model probability distribution over sequences of observations. The model can be either unconditioned or conditioned on other factors such previous observations actions. For phase the model needs to be the conditional distribution, ie.. p(01.T ag.r-1): 1). to model the effect of the actions. When given an observation sequence, the likelihood of this sequence under the model referred as evidence.
In this paper, we consider variational world models where the observation governed by Markovian hidden state. In the literature, this type of model is also referred to state-space model (SSM) 20). Such a variational world model involves four components, namely
15 16
17
18
encoder = fp((t)
posterior St ~ (0t 11-1,
prior 8, ~ 170(81100-1,
decoder Or ~ po st))
fa(ot) is the encoder to extract the features from the observation; (0(5)18111,0 01-1, =)) and (o(se(8t-1,4t-11 are the posterior and the prior of the latent state variable; while the decoder that decodes the observation distribution from the state. ó and 0 represent the parameters of the inference model and the generative model respectively.
is
Typically, a variational world model is trained by maximising the ELBO which is a lower bound of the log- likelihooddd or evidence, of the observation sequence, i.e. log po (o(11r (01:r/00:7 (0::7 Given sequence of observations, actions, and states, the objective function can be computed as
J(LTT, 80.T, 10) 1)= = (ooe 80.7, (0:7-1) + JKL(O1:T 80.T, a0:r-1),
(1)
T
where Jree 80:T, 11) =
Σ
log Po (orls)))))
(2)
[=1
J (0.r = Σ
Dxi.
(3)
3

---

The objective function is composed of two terms: the first term Jzee is the likelihood of the observation under the inferred state, which is usually called the reconstruction loss; while the second term JKI. is the KI. divergence between the posterior and the prior distributions of the latent state. To compute the objective function, we use the re- ee-prammterisation trick [2 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
0 argmax E (OL (oL Dona [J (01:21 80:27 00-r-1)]-
(4)
0.8
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 Daemo 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 unknown in the demonstration, instead of modelling the conditional evidence in phase 1, we need to model the unconditional evidence, i.e. log (o(0).)) Thus, also need model the actions as latent variables together with the states. In this way, the reconstruction term Jree will stay the same as eq. while the KL term will be defined on the joint distribution of states and actions, i.e.
(OI,, S0:T; 00:T-11
DRL. Raa.lgo.w(es (st, 11-1 So (1 (-19, (5)
If choose the action inference model in the form of policy, i.e. n. (arls)), and share in both posterior and prior, then the new posterior and prior can be factorised as
( a 1e-1/st-1, =
fo(o.))
(6)
and pa.e (8t, a1-1/s1-1) = To - (ooe-)))o(eell (=
respectively. When we plug them into the eq. (5), the policy term cancels and we will get a similar optimisation problem with phase
argmax E Emrr Daimmo ( (50eT aa..-a------- [J(o So:T, a0:r-1)]-
(8)
The main difference between eq. and eq. (8) 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 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, shown in eq. ( given trained world model, lower bound of the evidence of any observation sequence given an associated action sequence as the condition. In 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.
evaluate the
Thus, for each observation sequence that we get from the demonstration dataset, finding the missing action sequence can be considered as a trajectory trackinn problem and can be tackled by planning. To be specific, we can find the missing action sequence by solving the optimisation problem
do. argmax Eo, y Do [J(o.T So:,, Gu:r-1))-
40.T
we solve the above optimisation problem for every sequence in the demonstration dataset, the problem will be converted 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 forward model w.r.t. the actions.
To make more efficient, we use amortised inference. We directly define a policy ny (st) under the latent state of the world model. By composing the learnt world model and the policy, can form

---

Algorithm 1: AIME
Data: Embodiment dataset Dhody: Demonstration dataset Demos Learning rate # Phase 1: Model Learning
Initialise world model parameters and
while model has not converged do
{01:T, 00:T-1)
body
80 4 0
for
8:~0 (0 0,,,,, fo(o))
Compute objective function J from eq.
Update model parameters & +- Φ + a $5 0 + 0 aV
Phase 2: Imitation Learning
Initialise policy parameters v/r while policy has not converged do
01:T Daemo
so +- 0
for = T:
a T-1
~ (0(3) St-1, 0t-11 (o(o) 8t
Compute objective function J from eq. Update policy parameters vir f- v + V
generative model of the state sequence by the chain of st - at - 8141 a++1 Then we will get the same optimisation problem as eq. (8)
+ ST-
To sum up, in AIME, we use the same objective function the ELBO in both phases with the only difference being the souree 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.
4 Experiments
To test our method, we need multiple environments sharing an embodiment while posing different tasks. Therefore, consider Walker and Cheetah embodiment from the DeepMind Control Suite (DMC Suite) 233 Officially, the Walker embodiment has three tasks: stand, walk and run. While the Cheetah embodiment only has one task, run, 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, 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 the observation, so that information wise is identical to the Visual setting but the information is densely represented in low- -iw-dimensional form.
consider three settings for
To generate the embodiment and demonstration datasets, we train Dreamer [19] agent in the Visual setting for each of the tasks for IM environment steps. We take the replay buffer of these trained agents the embodiment datasets Dhody, which contain 1000 trajectories, and consider the converged policy as the expert to collect another 1000 trajectories as the demonstration dataset Daemo. 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 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 mix dataset. This resembles a practical setting where one possesses lot of experience with one embodiment and uses all of train a single foundational world model.
5

---

stand
random
random
random
(CC1O) Visus
method
BCOKO-LPONDP
(CDIO)-MDP
AIME-LPOMDP AIME-MDP
p70 + stand
p7e +> walk
p2o +> nun
p7e ~3
Misus BCOK0)-PPOMP BCOIOI-MDP AIME Visus
method
AIME LPONDP
NIME-MPP
stand -> stand
stand ->> walk
stand -> run
stand -> all
BCO/O)
method
BCOKO)LPOMMP BCD/0)-MDP
WME-Visus AIME LPOMOP
walk Stand
walk walk
wolk
walk all
BCD(0)-Visus LPONDP
method
8CDI0I-MDP AIME-Viuul LPONDP AIME-NDP
stand
methed
BCO(O) LPONOP
BCOIOI-MDP
AIME-Viual
AIME-LPOMOP
AIME MDP
mix >> stand
mix >> walk
mix -> Tun
mix >>
RCO00)-Muua/
method
BCCKO)-PPOMDP
BCD(O)-NDP
AIME-Visul
AIME-LPONDP
AIME-NDP
all >> stand
all -> walk
all
811 all
BCOY0)-Viuail BCONO)-
methed
SCO(O) NDP AIME-Visue
AIME-LPONDP AIME-NDP
25
50
75
100
0
25
50
100
25
50
75
100
25
50 performance
75
100
performance
performence
performance
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 cach figure is named according to Dbody + Daemo- Numbers are computed by averaging among 100 trials and then normalised to the percentage of the expert's '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 large range.
4.1 Benchmark results
We mainly compare our method with BCO(O) 25 BCO(O) 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 goal-conditional setting (28) which does not suit the
locomotion tasks. More details about related works can be found in Section 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 (CO(0)-MDP which can access the true states. We attribute the good performance of AIME to two reasons. First, the world model has better data utilisation rate than the IDM because the world model is trained to reconstruct whole observation sequences, while the IDM only takes

---

rendom => run
rendom w> runb
rendom- ->> Fip
rendom-> Riib Σ
random ->> al
AC BCOJDI LPOMDP
metrod
ROOIIIMEEP AIM Macal
APNE-LITTMEP MDP
pie .> run
p2e . runb
pde m> Fip
pite flibb
płe
MCsa BCCOOI LPOMOP
methos
BCO1O MDP
Macal
AINE
Nuh runb
ran-
run flgb
BCO101 Moual BCOHD-LLOWED
PCO1O MDP
flp
runb
eGOTOL Minua BCO1D ROWER
Mf
AIRE LHOMPP MET
1ip =>
Sp-> rubb
App ~> flip
fii flpb
flip
BCOJDI LPOMDP DC00 MEP
Masal AIRE- LPOMP AIKE MDP
1lpb-> run
1ppb -2 runb
tipb -2 hip
kipb - Ripb
ripo -2 al
DC BOOM
method
0C0 MDP
AIME Mical
AMME-LPOMOP
bll -1 Fun
all - runb
all +> Alp
- fllb
all > al
BC0101 Misual
meethoo
(CO(0) MDP
ARME-AN Mieua
25
50
75
100
0
25 performance
50
75
100
25
50 performance
75
100
0
25
50 performance
75
100
25
50
75
100
performance
performance
Figure 3: Performances Cheetah. Each column indicates task and its associated demonstration dataset, while each row indicates the embodiment datasets used to train the model. The title of each figure named according to Doedy + Deemo- 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 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 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 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 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 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 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(O) 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 very noisy
7

---

gradient since none of them can explain the "flying". In general, poorly modelled regions of the world may lead 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 50Hz which is much higher than the 20Hz 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, (CO(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 large margin, which indicates the strong generalisability of 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 diverse set of experiences, which could be more valuable in real-world scenarios.
Influence of observation modality Compared with BCO()), 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 to the fact that low-dimension signals have denser information and offer a smoother landscape in the evidence space than the pixels 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.c. the p2e dataset on Walker and the random dataset Cheetah. We conjecture this the fact the world model is not trained well with the default hyper-parre 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 more realistic setting where we want to use experience from multiple tasks to transfer to a new task.
Sample elficiency and scalability To test these properties, we vary the number of demonstrations within 21, 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(O) struggles with low-data scenarios and typically needs at least 10 to 20 demonstrations to surpass the performance of 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.c. ELBO, consists of two terms, the reconstruction term Jree and the KL. term JKL. 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 []] removing either term will negatively impact the results. When we compare the two variants, only using the KL term 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 KI. 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
8

---

120
120
Nomalised Retums
100
Norallsed Returss
BC0101
Normalizd Returts
BCO10)
random
16
20-
25
10
Tor
107
107
10
10"
10
10
10
10'
Number of Demonstrations
Number Demonstrations
Number Demonstrations
Figure 4: Ablation of the number of demonstrations on mix + run transfer on the Walker embodiment. The performance shown as the normalised returns over seeds and 100 trials for each seed. The shaded region represents one standard deviation.
100
Normalised Returns
BCQ(O)
60
AIME
AIME w/o Rec
AIME w/o KL
IIDM
AIME w/ IDM
MDP
LPOMDP
Visual
Figure Ablation studies on - 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.
is to Compared with the BCO(O) baseline, AIME consists of two distinct modifications: 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 components. First, remove the SSM, we train 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 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(O), the IIDM does provide assistance in the nigh-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 better performance comparing with BCO(0), but it still performs worse than AIME, especially on the POMDP settings.
Components
Related works
Imitation learning from observations Most previous works on imitation learning from only ob- servation can be roughly categorised into two groups, one based on IDMs 25 30 28 and one based on generative adversarial imitation learning (GAIL) B1 26 27]. The core component of the first group is to learn an IDM that maps state transition pair to the action that caused the transition. 25 use 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 trained to be conditioned on future frame as goal instead of only the next frame. During deployment, the task is communicated on the fly by the user in the form of -ramess 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 long-term goal of observation is and also not practical set the next observation as the goal and demonstrate that a high frequency by the user. Different from these methods, our approach uses 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 's and the agent's s observation trajectories. Then the discriminator serves reward function, and the agent's 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 B3]
don't '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. [3 is hybrid method that uses both of the components and focus on setting that the demonstrations are coming from different embodiment.
Reusing learnt components decision-making Although transferring pre-trained models has become a dominant approach in natural language processing (NLP) R4 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 [88 9 77. 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
41
42
43 444. However, the transferability of these world models is not well-studied.
learns policy by using 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 the model by observing the expert while (24) communicates the task by ground truth reward function which is less accessible in real-world setting.
24
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 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 gap compared with the LPOMDP setting where the wwwddmmnsiooall 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 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 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 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.
Acknowledgments and Disclosure of Funding
We want to acknowledge Elie Aljalbout for the insightful discussion during the initial stage of the project and Botond Cseke for mathematical support.
10

---

References
[1] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Pe- tersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Guman-level control through deep reinforcement learning. Nature, 5 518(7540):529-555 533, February 2015.
[2] David Silver, Aja Huang. Chris J. Maddison, Arthur Guez, Laurent Sifre, George van den Driess- che, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam,, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, llya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, :29(7587):484-489 January 2016.
[3] Oriol Vinyals, Igor Babuschkin, Wojciech M. Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David H. Choi, Richard Powell, Timo Ewalds, Petko Georgiev, Junhyuk Oh, Dan Horgan, Manuel Kroiss, Ivo Danihelka, Aja Huang, Laurent Sifre, Trevor Cai, John P. Agapiou, Max Jaderberg, Alexander S. Vezhnevets, Rémi Leblond, Tobias Pohlen, Valentin Dalibard, David Budden, Yury Sulsky, James Molloy, Tom L. Paine, Caglar Gulcehre. Ziyu Wang, Tobias Pfaff, Yuhuai Wu, Roman Ring. Dani Yogatama, Dario Wünsch, Katrina McKin- ney. Oliver Smith, Tom Schaul, Timothy Lillicrap, Koray Kavukcuoglu, Demis Hassabis, Chris Apps, and David Silver. Grandmaster level in StarCraft using multi-agent reinforcement learning. Nature, 575(7782):350-354
[4] Suyoung Choi, Gwanghyeon Ji, Jeongsoo Park, Hycongjun Kim, Jubyeok Mun, Jeong Hyun Lee, and Jemin Hwangbo. Learning quadrupedal locomotion deformable terrain. Science Robotics, 8(74))eade2256, 2023.
[5] OpenAI, Ilge Akkaya, Marcin Andrychowicz, Maciek Chociej, Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, Jonas Schneider, Nikolas Tezak, Jerry Tworek, Peter Welinder, Lilian Weng, Qiming Yuan, Wojeiech Zaremba, and Lei Zhang. Solving rubik's cube with a robot hand, 2019.
[6] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal. Katarina Slama, Alex Gray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[7] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Marc G Belle- mare. Reincarnating reinforcement learning: Reusing prior computation accelerate progress. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[8] Marco lacoboni. Imitation, empathy, and mirror neurons. Annual review of psychology, 60:653-70, 10 2008
[9] Bowen Baker, Ilge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (VPT): Learning to act by watching unlabeled online videos. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[10] Yezhou Yang. Yi Li, Cornelia Fermuller, and Yiannis Aloimonos. Robot learning manipulation action plans by "watching"; unconstrained videos from the world wide web. Proceedings of the AAAI Conference on Artificial Intelligence. 29(1), Mar. 2015.
[11] Y. Kuniyoshi, M. Inaba, and H. Inoue. Learning by watching: extracting reusable task knowledge from visual observation of human performance. IEFE Transactions on Robotics and Automation, 10(6):799-8)) 822, Dec 1994.

---

[12] Younggyo Seo. Kimin Lee, Stephen L James, and Pieter Abbeel. Reinforcement learning with action-free pre-training from videos. Kamalika Chaudhuri, Stefanie Jegelka. Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 19561-19579, PMLR, 17-23 Jul 2022.
[13] Alejandro Escontrela, Ademi Adeniji, Wilson Yan, Ajay Jain, Xue Bin Peng, Ken Goldberg, Youngwoon Lee, Danijar Hafner, and Pieter Abbeel. Video Prediction Models Rewards for Reinforcement Learning, 2023.
[14] Mohammad Babaeizadeh, Chelsea Finn, Dumitru Erhan, Roy H. Campbell, and Sergey Levine. Stochastic variational video prediction. In International Conference on Learning Representa- tions, 2018.
[15] Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances in neural information processing systems, 28, 2015.
[16] Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. In International Conference on Learning Representations, 2017.
[17] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pages 2555-2565, 2019.
[18) Philip Recker-Ehmck, Jan Peters, and Patrick Van Der Smagt. Switching linear dynamics for variational Bayes filtering. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceed- ings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 553-562, PMLR, 09-15 Jun 2019.
[19] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representa- tions, 2020.
[20] Alexej Klushyn, Richard Kurle, Maximilian Soelch, Botond Cseke, and Patrick van der Smagt. Latent matters: Learning deep state-space models. Advances in Neural Information Processing Systems, 1e:11022400444 2021.
[21] Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
[22] Danilo Jimenez Rezende. Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Eric P. Xing and Tony Jebara, editors, Proceedings of the 31st International Conference on Machine Learning, volume 32 of Proceedings of Machine Learning Research, pages PMLR.
Bejing, China, 22-22 24 Jun 2014.
[23] Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap. Nicolas Heess, and Yuval Tassa. dmm_ontrol, Software and tasks for continuous control. Software Impucts, 6: 100022 2020.
[24] Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In /CML, 2020.
[25] Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In International Joint Conferences on Artificial Intelligence, pages 4950-4957, 2018.
[26] Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observa- tion, 2018
12

---

[27] Chenhao Li, Marin Vlastelica, Sebastian Blaes, Jonas Frey, Felix Grimminger, and Georg Martius. Learning agile skills via adversarial imitation of rough partial demonstrations. In 6th Annual Conference on Robot Learning, 2022.
[28] Deepak Pathak, Parsa Mahmoudich, Guanghao Luo, Pulkit Agrawal, Dian Chen, Yide Shentu, Evan Shelhamer, Jitendra Malik, Alexei A. Efros, and Trevor Darrell, Zero-shot visual imitation. In ICLR, 2018.
[29] Nicklas A Hansen, Hao Su, and Xiaolong Wang. Temporal difference learning for model predictive control. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 8387-8406. PMLR, 17-23 Jul 2022.
[30] Ashvin Nair, Dian Chen, Pulkit Agrawal, Phillip Isola, Pieter Abbeel, Jitendra Malik, and Sergey Levine. Combining self-supervised learning and imitation for vision-based rope manipulation. In 2017 IEEE International Conference on Robotics and Automation (ICRA), page 2146-2151
IEEE Press. 2017.
[31] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
[32] Ashley Edwards, Himanshu Sahni, Yannick Schroecker, and Charles Isbell. Imitating latent policies from observation. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceed- ings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 1763. PMLR, 09-15 Jun 2019.
[33] Karl Schmeckpeper, Oleh Rybkin, Kostas Danilidis, Sergey Levine, and Chelsea Finn. Re- inforcement learning with videos: Combining offline observations with interaction. In Jens Kober, Fabio Ramos, and Claire Tomlin, editors, Proceedings of the 2020 Conference on Robot Learning, volume 155 of Proceedings of Machine Learning Research, pages 339-354. PMLR, 16-18 Nov 2021.
[34] Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pages 4171-4186. Association for Computational Linguistics, 2019.
[35] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
[36] Rishi Bommasani, Drew A. Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S. Bernstein, Jeannette Bohg. Antoine Bosselut, Emma Brunskill, Erik Brynjolfsson, Shyamal Buch, Dallas Card, Rodrigo Castellon, Niladri Chatterji, Annie Chen, Kathleen Creel, Jared Quincy Davis, Dora Demszky, Chris Donahue, Moussa Doumbouya, Esin Durmus, Stefano Ermon, John Etchemendy, Kawin Ethayarajh, Li Fei-Fei, Chelsea Finn, Trevor Gale, Lauren Gillespie, Karan Goel, Noah Goodman, Shelby Grossman, Neel Guha, Tatsunori Hashimoto, Peter Henderson, John Hewitt, Daniel E. Ho, Jenny Hong, Kyle Hsu, Jing Huang, Thomas Icard, Saahil Jain, Dan Jurafsky, Pratyusha Kalluri, Siddharth Karamcheti, Geoff Keeling, Fereshte Khani, Omar Khattab, Pang Wei Koh, Mark Krass, Ranjay Krishna, Rohith Kuditipudi. Ananya Kumar, Faisal Ladhak, Mina Lee, Tony Lee, Jure Leskovec, Isabelle Levent, Xiang Lisa Li, Xuechen Li, Tengyu Ma, Ali Malik, Christopher D. Manning, Suvir Mirchandani, Eric Mitchell, Zanele Munyikwa, Suraj Nair, Avanika Narayan, Deepak Narayanan, Ben Newman, Allen Nie, Juan Carlos Niebles, Hamed Nilforoshan, Julian Nyarko, Giray Ogut, Laurel Orr, Isabel Papadimitriou, Joon Sung Park, Chris Piech, Eva Portelance, Christopher Potts, Aditi Raghunathan, Rob Reich, Hongyu Ren, Frieda Rong, Yusuf Roohani, Camilo Ruiz, Jack Ryan, Christopher Ré, Dorsa Sadigh, Shiori Sagawa, Keshav Santhanam, Andy Shih, Krishnan Srinivasan, Alex Tamkin, Rohan Taori, Armin W. Thomas, Florian Tramer, Rose E. Wang, William Wang, Bohan Wu, Jiajun Wu, Yuhuai Wu, Sang Michael Xie, Michihiro Yasunaga
13

---

Jiaxuan You, Matei Zaharia, Michael Zhang, Tianyi Zhang, Xikun Zhang, Yuhui Zhang, Lucia Zheng, Kaitlyn Zhou, and Percy Liang. On the opportunities and risks of foundation models. 2021.
[37] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 16000-16009, 2022.
[38] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1126-1135. PMLR, 06-11 Aug 2017.
[39] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In
S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Desa-Bianchh, and R. Garnett, editors, Advances in Neurul Information Processing Systems, volume 31. Curran Associates, Inc.
[40] Philip Recker-Ehmck, Maximilian Karl, Jan Peters, and Patrick van der Smagt. Learning to fly via deep model-based reinforcement learning. arXiv preprint 17Xiv:2003.08870 2020.
[41] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021.
[42] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models, 2023.
[43] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blažej Osiński, Roy Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In International Conference on Learning Representations, 2020.
[44] Julian Schrittwieser, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839)6604-66 609, 2020.
[45] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Viti,o, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Gamett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[46] Junyoung Chung, Çaglar Gülgehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. 2014.
[47] Maximilian Seitzer, Arash Tavakoli, Dimitrije Antic, and Georg Martius. On the pitfalls of heteroscedastic uncertainty estimation with probabilistic neural networks. In International Conference on Learning Representations, April 2022.
[48] Djork-Amé Clevert, Thomas Unterthiner, and Sepp Hochreiter, Fast and accurate deep network learning by exponential linear units (elus), 2015.
[49] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference machine learning, pages 1861-1870. PMLR, 2018
[50] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey F. Hinton. Layer normalization, 2016.
[51] Cong Lu, Philip Ball, Tim G. J. Rudner, Jack Parker-Holder, Michael A. Osborne, and Yee Whye Teh. Challenges and Opportunities in Offline Reinforcement Learning from Visual Observations. Transactions on Machine Learning Research, 2023.
[52] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. Advances Neural Information Processing Systems, 34, 2021.
14

---

A Computational resources
On a GTX 1080Ti graphics card, AIME typically requires 10 hours of training for phase and 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-ar performances by splitting the latent state to be a combination of deterministic and stochastic components. The RSSM implemen- tation is largely following Dreamer-v1 [9] with continuous stochastic and deterministic variables. Although newer versions of Dreamer initially choose not to them for the sake of simplicity. We use slightly larger state space for our experiment with 512 deterministic and 128 stochastic dimensions and find it generally cases 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 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
42 offer some new tricks to improve performance, we
GRU cell [46). For
the default hyper-parameters. we do not use any free nats [1),, KL scaling [19] and KL balancing
41 tricks in the literature 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 3-nll remedy this problem, and since 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 implemented 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 learning rate of lc 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 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 ppper--rrrmeters 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 oowddieesioonal observations are also with a fixed variance of 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 3e-4, and we use free nats of 1.0 and KL balancing of 0.8 to mitigate the collapse and unstable problem of the KI. term. For the LPOMDP setting, we also set the KL. scaling parameter / = 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 is due to learnt policies being stuck in some local minima that they are unable to escape.
15

---

For the BCO(O) baseline, the IDM and policy are built by using the same network architecture with the world model make fair comparison. The observations are first processed by the encoder network, and then get stacked to deal with the temporal information. An MIP is used to decode the stacked representation to the output distribution. We stack consecutive in this work. We did grid search about the width, depth of the MLP and also the number of stacking frames and didn't t find any increase of the performance. Following the original paper, we split the datasets by 7: 3, and choose the finial 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 the Bayesian derivation, factorised both the posterior and prior of the joint distribution of state and action with shared policy network, as in eq. (6) and eq. (7). Alternatively, we can re-factorise the posterior with IDM, that is
9 (st, at- 1t-11 fo((:)) g0 ootl://.......ollooooo.......a.. folo)))
(10)
One thing that needs to noticed here is that the IDM not in the familiar form do (4t-1 8t-1 8)). This is because the latent state in the world model is action dependent the familiar form is non- casual in the world model. But we should highlight here that this non-casua structure a result of the model we used in phase 1 since want to reuse the knowledge learnt there. For example, one can also factorise the joint posterior as
g
go(a (s)
(11)
However, in this case, the model is different for phase 1. In this section, stick using the factorisation in eq. (TO). Since a new IDM component is introduced, the objective of both phase and phase need to be modified. For phase since actions available in the dataset, the IDM can be treated as a decoder and trained by maximising likelihood. That is, we add (IM((((T) = Σ log (ar St, fo fo((+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, ie. ( = Σi= - DKL (qo (aels:. fo(0(+1)) T (ag )s
One caveat about this formulation is that, in phase I, the IDM forms 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 shown in Table and Table 2
Table 1: Average expert return of each demonstration dataset of Walker.
D
Ddemo
Average return
stand
957.87 (max: 1000)
walk
943.79 (max: 1000)
604.10 (max: 1000)
Table 2: Average expert return of each demonstration dataset of Cheetah,
D Demoo
Average return
run
888.65 (max: 1000)
run backwards
218.50 (max: 500)
flip
485.79 (max: 500)
flip backwards
379.91 (max: 500)
16

---

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(O)
AIME
walker-random walker-medium replayy
2.11 0.91
12.36 + 4.69
6.56
10.18 + 4.33
walker-mix
5.32
8.49 3.60
cheeda
0.01 + 0.01
9.48 + 4.72
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-D4RI. datasets [51] to showcase that AIME 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 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 the main text, we also mix the two embodiment datasets for each embodiment to form a mix dataset.
The results these datasets shown in Table 3 We and AIME is generally low, but AIME still outperform BCO(O) which proves AIME can also handle datasets generated by model-free methods. The low performance is due to more constrained setup of the task, i.e. less amount of embodiment data and less diversity. Except the cheetah-medium wwedimeedmmmmie-ee 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 that random datasets do not help much in learning model, and intuitively the medium replay dataset is better but still does not contain enough information to solve the task.
that the performance of both BCO())
We also conduct experiments 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 mix dataset. Moreover, we also merge this mix dataset with the mix dataset in the second experiment to form a total_ _ix dataset. We treat these five datasets as the embodiment dataset and the expert dataset the demonstration dataset. For the Cheetah embodiment, the benchmark provides medium and expert datasets with distraction level of easy, medium, and hard. We subsample the medium datasets get 200 trajectories from each level, then merge that with the mix dataset in the second experiment to form total mix dataset. Then the algorithms are using this total mix dataset as the embodiment dataset and the expert dataset as the demonstration dataset.
from the result from Table although we still outperform the BCO(O) baseline, AIME impacted significantly by the distractions. This behaviour 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.
As
F Additional plots
In this section, we will present some additional plots to complement the main text and provide further insights.
17

---

Table 4: Result on V-D4RL distracting datasets. Embodiment datasets are marked on the left, and the demonstration datasets are chosen 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(O)
AIME
walker-easy
2.10 + 0.88
4.73 + 2.54
walker-medium
2.15 0.87
3.94 0.99
walker-hard
2.95
4.16 1.98
walker-mix
2.12 0.86
3
walker-total_mix
2.15 0.71
12.66 + 4.51
cheetah-total_me _mi
16.61 + 7.28
32.40 + 14.52
we also provide detailed profile plots in Fig Figure 7 as recommended in (2). We can see that AIME is normally more stable w.r.t the performance by having a smaller decay region. It is clearly shown on such tasks as walk + walk and run + run on Walker where BCO(0)-MDP has some trails with very low performance, while all variants of AIME maintain decent performance.
Additional to Figure and Figure 3
130
(2e ui stand
p2e - walk
-o run
130
150
Bibot POMEE LOMES
BOX0 NEP
AINE-UPOND
Walk siand
walk walk
walk->> run
130
clab 22035
mik stand
mik - Wilk
mik->
130
25
40
50
30
200
20
43
50
82 100
22
40
62
60
200
vo expert per performance 16)
Relabive ecpert parformnnce (p)
aovvvee expert parfor mance 150
Figure 6: Performance distributions of each method on Walker tasks.
We also present some representative training curves of AIME's phase from our experiments in Figure 8 The first three figures show the transfer from the mix dataset to the run task in the three settings which are the typical success cases of AIME. During the course of training, ELBO is maximised towards convergence and the MSE between the generated actions and the true actions decreases. We can also see that for the MDP and LPOMDP settings, the converged ELBO is lower than the ELBO when evaluated with the true action sequence, indicating there is still space for improvement. However, for the Visual setting, the converged ELBO exceeds the one with true actions, which should be attributed to the over-fitting of the world model from phase 1. The last three figures
18

---

random - mi
randoms =2 nubb
anndom = 4p
random-2 -2 Aph
130
# 400spp
p20 runb
fip
Aic flab
30
oun
nu anb
Fun--o Ap
run Rip
120
Method
BOOOD LPOND
AIF (MNNO
runb 105
1u10 Hpp
Funt->>
AIN NDP
flp
Ap
1ip-> lip
Tip -- fllb
120
9n
rubb
#2055
50
Se-adive
Rebertre
Redatiuv
Reall
Figure 7: Performance distributions of each method on Cheetah tasks.
show the transfer from the random dataset to the three tasks in the Visual settings which we consider as failure cases. For the stand and walk tasks, none of the metrics are converging. For the run task, we can observe a severe over- wer--fitinn starting from the beginning of the training, and the MSE keeps increasing. We conjecture these are due to the less well-trained world models.
19

---

AIME-MDP mix training
term
action normalised
AIME
14:00
-2000
3.02
-2030
AIME
10:0
AUME
-2209
AIME
3.5
0.32
lction
action
frue action
3.45
26
460
266
460
200
<6
AIME-LPOMDP mix training
ELBO
term
KL term
action
nommalied returns
2.50
ARM
AIME
2200
3232
AIME
AIME
AIME
0.60
actioon
action
8.55
230
240
200
266
4E0
56
-3998
220
300
300
200
406
306
283
400
50
AIME-Visual mik training curve
ELBO
FDC term
KL term
action mise
norm alised returs
NIH
AMM
-1488
AIME
-4.330
A.Mt
ARM
0.38
-3.62
trie artion
true action
True, action
400
263
486
200
AIME-Visual random >> stand training curve
ELBO
rec term
KL term
action mse
normellsed retu
AIRE
AINE
-5.000
18.
1.00
-5.13
28-
AMME
AMM
AIME
1.0
true aciion
-2882
true adiion
Erue attioo
20
20
60.3
AIME Visual randem > walk training curve
FLBO
FBC Term
KL term
action mse
normalised retums
-2.000
ANE
AME
AIME
AINE
True action
Erue adiion
-700
26-
-200
AIME
12
-000
Erue action
1.0
206
460
403
200
402
400
200
460
AIME-Visual random > run training curve
ELBO
58715
rec torm
KL torm
action mise
norm commalise retums
-3,070
-3.022
AIME
AMME anttaaation
frue juction
Thie, actiio
1203
1.40
-3.008
5.8.80
1.35
AIME
650
202
420
Figure 8: Samples of training curve in phase 2 of AIME, The first three showcase the typical successful training curves, while the remaining three demonstrate the failure cases. The true action is referring to evaluating the trajectories with the true action sequence.
20


Reviews:
[{"review_id": "y1e4QEcjol", "reviewer": "Reviewer_Cs3Z", "summary": "The paper proposes action inference by maximising evidence as a way for an MBRL to replicate most likely actions using appropriate world models. The algorithm has two phases: 1) Learn the world model based on a replay buffer, and 2) imitate the expert's behaviour by inferring the policy that maximizes the evidence of the demonstration under the policy and world model. Experimental results on the Walker and Cheetah embodiments of the DeepMind Control Suite demonstrate that this zero-shot imitation performance outperforms the current state-of-the-art approaches.", "questions": "-\tDoes the world model training in Phase 1 have to converge because imitation learning can happen? Is this primarily for changes in the task? E.g., going from walking to hopping but with the same agent in the same environment? \n-\tWhat if the space of actions change between phase 1 and phase 2? Will AIME still work? It doesn’t seem like it. \n-\tCan action trajectories be learnt directly instead of one step action policies? \n-\tHow sensitive is the AIME method to the quality and diversity of the observation-only demonstrations provided? What happens if the demonstrations are not representative of the task at hand or are of poor quality?\n-\tWhat is the computational complexity of the AIME method? How well does it scale to larger and more complex tasks?\n-\tdoes the AIME method handle situations where the world model learned from past experiences is not accurate or complete? What are the implications if the world model is imperfect?\n-\tHow well does the AIME method perform in terms of transfer learning? Can the world model learned in one context be effectively applied to another for the policy?\n-\tHow robust is the AIME method to changes in the environment or task? Can it adapt to new situations without requiring additional training?", "limitations": "The authors point out the following limitations that limit the scope of the current results: \n\n- The AIME method performs well with visual input, but there is a significant performance gap when compared to the LPOMDP setting where low-dimensional signals are observed. This is attributed to the loss surface of the pixel reconstruction loss not being smooth enough to allow the gradient method to find an equally good solution.\n- The study only considers the simplest setting where both the embodiment and sensor layout are fixed across tasks. This is a limitation as humans observe others in a third-person perspective and can imitate animals whose bodies are not similar to humans. Relaxing these assumptions could allow for transfer across different embodiments and even directly from human videos.\n- For some tasks, even humans cannot achieve zero-shot imitation by only watching others. This could be due to the task's complexity or completely unfamiliar skills. Even with proper instruction, humans still need to practice in the environment and learn something new to solve some tasks. This suggests the need for an online learning phase as an extension to the AIME framework.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "-\tThe paper addresses a major issue in deep reinforcement learning (DRL), namely sample inefficiency. By suggesting a method that can harness observational data, the authors propose a way to improve the sample efficiency of DRL agents.\n- The paper introduces a new method, Action Inference by Maximising Evidence (AIME), for imitation learning. This method is designed to mimic the human ability to learn quickly from observation, which is an interesting contribution.\n- AIME's two-phase learning process, involving the creation of a world model and then using it for imitation learning, is a unique approach. This process allows the agent to understand its own body and the likely actions that led to observed behaviors, which is a crucial aspect of learning.\n-  The method is capable of \"zero-shot\" learning, meaning it does not require further training for the world model or online interactions with the environment after being given the demonstration. This is a significant advantage in terms of efficiency and practicality.\n- The authors provide empirical validation of their method on the Walker and Cheetah embodiments of the DeepMind Control Suite. They demonstrate that their method outperforms state-of-the-art baselines, which strengthens the credibility of their approach.", "weaknesses": "-\tThe paper validates the AIME method using the Walker and Cheetah embodiments of the DeepMind Control Suite, which are simulated environments. It's unclear how well the method would perform in real-world scenarios, where conditions can be more complex and unpredictable.\n-\tThe AIME method assumes that the agent can learn a perfect world model from its past experiences. This may not always be possible due to the complexity and unpredictability of many environments. The performance of the method could be affected if the world model is not accurate.\n-\tThe effectiveness of the AIME method may heavily depend on the quality of the observation-only demonstrations provided. If these demonstrations are not representative of the task at hand, or if they are of poor quality, the performance of the method could be significantly affected.\n-\tThe paper does not discuss the computational complexity of the AIME method or its scalability to larger and more complex tasks. If the method is computationally intensive, it may not be practical for use in real-time applications or on larger scales.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Does the world model training in Phase 1 have to converge because imitation learning can happen? Is this primarily for changes in the task? E.g., going from walking to hopping but with the same agent in the same environment?\n\nNo, it is not necessary to train the model until converge to enable imitation. It emerges gradually during training. From the results shown in Fig. 1 in the general response, the imitation ability is established in very early phase of the training process long before convergence. \n\n> What if the space of actions change between phase 1 and phase 2? Will AIME still work? It doesn’t seem like it.\n\nIn general, we do not care about what action space is used to collect the demonstrations for phase 2, since we only need the observations. AIME will derive a policy in the original action space that the world model is trained on during phase 1. This is actually a benefit of our method, since it allows you to have different action space for the agent and the demonstrator. For example, if you want to demonstrate a task on a robot arm, the low-level action space is not intuitive for a human operator. In our setup we can easily use another more intuitive interface like smart phones [1] to do the demonstration and let the agent to imitate on the low-level action space. But if you are referring to completely changing the action space also in deployment, which is not likely to happen in real life, it will require an extra function to map the action from the old action space to the new action space.\n\n> Can action trajectories be learnt directly instead of one step action policies?\n\nYes, they can. As we show in the control derivation from Sec. 3.2, we can infer the action trajectories directly by a planning algorithm. But it needs to be done individually for each trajectory. That is why we introduce amortised inference to improve the efficiency. And also one can define the amortised inference model in many different forms, for example $\\pi(a_{t:t+T}|s_t)$. We use the one step action policy form to keep it simple and comparable to the baselines. \n\n> How sensitive is the AIME method to the quality and diversity of the observation-only demonstrations provided? What happens if the demonstrations are not representative of the task at hand or are of poor quality?\n\nWe have to clarify that, in the setting of imitation, the task is defined by the demonstrations. We aim to replicate the behaviour of the demonstration rather than solving the original task that the demonstrator is trying to solve. This can be viewed as an alternative way of defining the task, i.e. by showing how the task is done, rather than defining a reward function. \n\nAbout the diversity of the demonstrations, we would like to point you to Fig. 4 of the paper, where we limited the number of demonstrations, which also limited the diversity. We can see this indeed has an effect on the performance, but AIME is less sensitive to this than the BCO(0) baseline.\n\n> What is the computational complexity of the AIME method? How well does it scale to larger and more complex tasks?\n\nFor the training complexity, in Appendix A, we show phase 1 requires 10 - 20 hours while phase 2 requires 5 - 10 hours on an old 1080ti GPU. For solving multiple tasks, you only need to run phase 1 once and run phase 2 multiple times for different tasks. Moreover, we would like to mention that 5 -10 hours training time is for training 500 epochs, but from Fig. 8 of Appendix E, normally AIME does not need that long for converging. This makes AIME require even less compute.\n\nFor the inference complexity, we are training a RSSM world model from the Dreamer papers. So we can do real-time inference whenever their method can. In a recent paper [2], they successfully run the model to control four different robots in real-time. Thus, our method is also applicable for the real-time scenarios. \n\nFor the scalability, we also run experiments on the mix dataset, where the dataset is 3 times larger than other datasets in this paper. The improved results suggest AIME can handle a larger scale of data. \n\n> Does the AIME method handle situations where the world model learned from past experiences is not accurate or complete? What are the implications if the world model is imperfect?\n\nWe didn't assume the world model learnt on the first phase is perfect. And actually, all the pretrained world models considered in the experiments are not perfect, since the training dataset did not cover all the dynamic of the embodiment. \n\nWhen the world model is imperfect, it could degenerate the imitation performance, like in the convergence experiment, or in the worst case leads to failure like divergence. We show a few of these failure cases in the Appendix E.\n\n> How well does the AIME method perform in terms of transfer learning? Can the world model learned in one context be effectively applied to another for the policy?\n\nThis transfer ability is exactly what we show with AIME. In the majority of the experiments, we do cross task transfer. For example, we learn a world model on the stand task, then imitate the policy of the run task. The results showcase the strong transfer ability of the world model. \n\n> How robust is the AIME method to changes in the environment or task? Can it adapt to new situations without requiring additional training?\n\nIn this paper, we mainly assume the same environment in both phases. In general, the robustness to environment change can be injected through diverse training dataset or domain-knowledge-based data augmentation. In term of tasks, that is exactly what we show in the paper. The pretrained world model exhibit the zero-shot ability for imitation, without further environment interactions to finetune the model.\n\n[1] Mandlekar *et al.,* Scaling Robot Supervision to Hundreds of Hours with RoboTurk: Robotic Manipulation Dataset through Human Reasoning and Dexterity, IROS 2019\n\n[2] Wu *et al.*, DayDreamer: World Models for Physical Robot Learning, CoRL 2022"}, {"review_id": "1LL3slRghP", "reviewer": "Reviewer_eCVB", "summary": "This paper presents an algorithm named AIME to learn the world model and apply it to downstream tasks. In the first stage, AIME learns a world model from a dataset with actions to maximize the likelihood via EBLO. While in the second stage, given observation-only demonstrations, AIME optimizes the action sequence to imitate the expert’s behavior. The empirical result shows that AIME outperforms previous methods in DMC tasks.", "questions": "Why the performance seems to have a very large variance in Figure 3?\n\n", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Optimizing actions from observation-only trajectory via ELBO is somewhat novel compared to behavior cloning methods.\n\n2. AIME performs better than previous methods even with local and image-based observations.", "weaknesses": "1. The major concern is the problem setting of AIME. In my view, more general setting the agent can only get state trajectory in the world-model learning stage while can obtain action-labeled data in the second stage. Then the agent has much more data in training (e.g., human data without actions) and only need a small amount of action-labeled data for fast adaptation. The authors should clarify the significance of the problem settings studied in this paper.\n\n2. Since the world model applies action-labeled data in the first stage, the world model will be related to the policies that generate the dataset. I wonder if the model can handle datasets with a mixture of policies or low-quality policies.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The major concern is the problem setting of AIME. In my view, more general setting the agent can only get state trajectory in the world-model learning stage while can obtain action-labeled data in the second stage. Then the agent has much more data in training (e.g., human data without actions) and only need a small amount of action-labeled data for fast adaptation. The authors should clarify the significance of the problem settings studied in this paper.\n\nYou raise a good point. We agree that having a large amount of observation-only dataset, like YouTube videos, and then adapt to some action-label datasets is a setting that has attracted a lot of attention in recent years. However, the reversed setting, as we used in this paper, also has a lot of potential. By emphasising embodiment, we require the agent to have a dataset about the embodiment in the first place. \n\nIt quite fits for robotics learning [1, 2], where people experiment with the same robot for years and naturally have access to a lot of embodiment data. These years, there are more and larger embodiment datasets getting collected [3] and open sourced [4, 5]. Moreover, it also holds true for well-studied simulator benchmarks where embodiment datasets can be easily access through open sourced benchmarks [6, 7, 8]. Besides, using embodiment data before others is also true for lifelong learning of humans. Thinking about infants who first randomly explore what their body can do before they learn some complex motor skills like walking. Last but not least, learning a model for certain embodiments is much easier than learning a general model for all the embodiments, which we don't even have in nature. \n\nThus, although the setting that first has embodiment datasets is less popular at the moment, it still has a great potential. We will update the paper to make the motivation of this setting more clear. \n\n> Since the world model applies action-labeled data in the first stage, the world model will be related to the policies that generate the dataset. I wonder if the model can handle datasets with a mixture of policies or low-quality policies.\n\nYes, the dataset for training the world model matters a lot, and it is also one of the conclusion from the paper. In the experiment section, we train the models on multiple datasets. The random dataset can represent the low-quality policies setting, and all the other datasets are collected by a mixture of policies since they stem from a replay buffer where each trajectory is collected by a different policy, and also the mix dataset contains multiple behaviours. We observe that in general, both AIME and BCO(0) don't work well on the low-quality random dataset since the experience from there hardly explains the observations seen in the demonstrations, i.e. lack of information, and the performance generally improved on p2e dataset where the quality of exploration improves. And AIME works well on the mixture of policies as we showed in other experiments.  \n\n> Why the performance seems to have a very large variance in Figure 3?\n\nThis is because the methods do well on some tasks but fail on others, when you average them all, you get a very huge variance. We agree it is not an ideal visualization, but we think it is still valuable to report the aggregated result. There are reward profile figures in the appendix E where you can find more details of each test trajectory.\n\n[1] Thrun and Mitchell, Lifelong robot learning, Robotics and Autonomous Systems, 1995\n\n[2] Singh, Transfer of learning by composing solutions of elemental sequential tasks, Machine Learning, 1992\n\n[3] Brohan *et al.*, RT-1: Robotics Transformer for Real-World Control at Scale, arXiv 2212.06817\n\n[4] Dasari *et al.*, RoboNet: Large-Scale Multi-Robot Learning, CoRL 2019\n\n[5] Ebert *et al.*, Bridge Data: Boosting Generalization of Robotic Skills with Cross-Domain Datasets, arXiv 2109.13396\n\n[6] Fu *et al.*, D4RL: Datasets for Deep Data-Driven Reinforcement Learning, arXiv 2004.07219\n\n[7] Qin *et al.*, NeoRL: A Near Real-World Benchmark for Offline Reinforcement Learning, NeurIPS 2022 Datasets and Benchmarks Track\n\n[8] Baker *et al.*, Video PreTraining (VPT): Learning to Act by Watching Unlabeled Online Videos, arXiv 2206.11795"}, {"review_id": "xCTfKqj0XX", "reviewer": "Reviewer_nqaR", "summary": "This paper presents an imitation learning approach by first training a world model to predict next observations conditioned on (given) actions, and in a second phase training a policy that amortizes action inference by maximizing the likelihood of observations under a dataset of expert demonstrations. The authors compare their method to BCO(0) on the DMC walker and cheetah environments.", "questions": "How is the performance measured? Is this accumulated reward (normalized to the expert performance)?\n\nThere are some questioning data points in the evaluations, for instance:\n- AIME on Cheetah run->flip is able to learn from Visual, but not at all from LPOMDP/MDP, same for flip->run\n- On flipb->run on Cheetah is impossible to learn, but running backwards is\n- BCO(0)-MDP seems to outperform AIME-MDP on xxx->flipb\nDo you have any insights on these results, whether this is due to the task at hand, the collected dataset, ...?\nMoreover, on Cheetah the resulting agent reaches an overall performance of under 50%. Any idea why this performance gap is there? \n", "limitations": "The authors adequately address the limitations in the conclusion.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is well written, the idea is clearly explained and the benchmark seems well executed.", "weaknesses": "The experimental results are hard to parse and have some anomalies (see my questions). Also the strongest claims \"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.\" are also mainly based on the Walker experiment, but are less outspoken on the Cheetah dataset. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The experimental results are hard to parse and have some anomalies (see my questions). Also the strongest claims \"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.\" are also mainly based on the Walker experiment, but are less outspoken on the Cheetah dataset.\n\nWe agree that we did not put enough analysis of the cheetah results. But the conclusion in the paper generally apply to both embodiments, just the improvement on cheetah embodiment is smaller. We will clarify this and add more analysis about the cheetah experiments in the revised version.\n\n> How is the performance measured? Is this accumulated reward (normalized to the expert performance)?\n\nThanks for bringing this up. Yes, you are right, the performances in all figures are reported based on the normalised accumulated reward. We will add the definition of performance to Sec. 4 in the revised version. \n\n> AIME on Cheetah run->flip is able to learn from Visual, but not at all from LPOMDP/MDP, same for flip->run\n\nWe are also confused about this. But we have to point out the RSSM (the dreamer model) is mostly optimised toward visual setting, and we find it harder to tune for other settings. And for the particular settings you mentioned, there is some small similarity between the run and flip tasks, since they all require some subtle movement of the front leg, we conjecture that the change of visual observation can be more pronounced for these subtle movements than the proprioception.\n\n> On flipb->run on Cheetah is impossible to learn, but running backwards is\n\nThis is due to the low performance of the run backward expert. In Appendix D, we show the average return of the expert on each task. Run's expert gets a return of $888.65$ while run backward's expert gets only $218.50$. This makes imitating run backward the easiest task among all. \n\nMoreover, we would like to point out that although run and run backward, flip and flip backwards sounds quite similar, due to the asymmetrical structure of the embodiment, the behaviour for solving the tasks can be quite different. \n\n> BCO(0)-MDP seems to outperform AIME-MDP on xxx->flipb Do you have any insights on these results, whether this is due to the task at hand, the collected dataset, ...?\n\nThe flips are the hardest tasks in the experiments, since the majority of the time in the expert demonstrations, 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 put additional challenge for AIME since it needs to infer a sequence of action, while BCO(0) is operated on point estimation. For example, when the first a few actions cannot output reasonable actions to start the flip, then the later actions will create very noisy gradient since none of them can explain the \"flying\".\n\n> Moreover, on Cheetah the resulting agent reaches an overall performance of under 50%. Any idea why this performance gap is there?\n\nThis is a good question. We think this gap could be the result of different control frequency between the two embodiment. Although it is not well-stated in the literature, with an action repeat of 2 as following other papers, the Walker is running on 20Hz while Cheetah is running on 50Hz. Running in a higher frequency makes each action have smaller influence on the environment, and observations as a consequence. This smaller change makes both the IDM and world model harder to train. There are also works [1, 2] using an action repeat of 4 for the Cheetah environment to get better results, which brings the control frequency to 25Hz, thus closer to the Walker one. \n\n[1] Hafner *et al.*, Learning Latent Dynamics for Planning from Pixels, ICML 2019\n\n[2] Hansen *et al.*, Temporal Difference Learning for Model Predictive Control, ICML 2022"}, {"review_id": "MnXQXW0emt", "reviewer": "Reviewer_MApq", "summary": "This paper considers the setting where we have access to a dataset containing states and actions for pre-training and must then learn to solve a task given a new observation-only dataset from a downstream task. The key assumption here is that we have access to this action labelled dataset, and once we have it, we are better off training a world model than simply learning to label the new observation only trajectories with action labels (a la BCO and subsequently VPT). The proposed method is one of the first (as far as I know) to use forward dynamics models to implicitly model the action distribution and provides a nice alternative to IDM approaches for what is becoming an increasingly relevant problem setting. Given the relevance of the topic area, and the fact this exact thing has not been done before, I am voting for the paper to be accepted. My score would be increased if some additional experiments could be conducted, since these particular ones are relatively similar and also low impact in terms of their ambition.", "questions": "How was the IDM baseline tuned? For example, in the VPT paper they mention different architecture choices made a big difference for performance. ", "limitations": "Limitations are discussed thoroughly in the main body.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The strengths of this work are clear, it is a very relevant problem setting and this method is distinct vs. previous methods like BCO which rely on an IDM. In reality, this paper is essentially \"Implicit action learning with forward dynamics models\", and that has not been done before as far as I am aware. The method itself is fairly clearly presented, and the experiments are relatively clear with sufficient ablation studies. Finally, it is great to see limitations adequately discussed in the main body, which is surprisingly rare.", "weaknesses": "Note that I have voted to accept, the following comments are not red flags but would likely improve the paper, and maybe make it possible to increase to a higher score.\n1. The experiments are fairly mundane, and while scientific best practices appear to have been followed, there is a huge gap between what the authors claim to be working towards (\"a single foundational world model\") and what is actually shown (two DMC environments). It would be fantastic to see an example of this method at larger scale, even if the results are not state-of-the-art and there is only a single seed, for compute reasons. For example, this could be done using the dataset from VPT. The Minecraft images could be resized to make them smaller, and then it would be possible to use the DreamerV3 codebase (which was tested on MineCraft and runs on a single GPU) then see if it is possible to learn from the unlabelled MineCraft videos. If this works, it would drastically increase the impact of the paper, beyond being something mildly interesting for people who care about this specific topic, to something that catches people's eye across the field.\n2. It seems slightly fishy to use a world model generated dataset to compare a world model based method and an IDM approach. It is possible there is some bias in the Dreamer or P2E data that makes it easier for a more similar approach to do well. Given how brittle many of these methods are, this could make a difference. Would it be possible to instead consider some open-source benchmarks such as VD4RL (Lu et al), which would make it \"fair\" across different approaches? For example you could use the random or mixed datasets for the embodied and then expert for the demonstration.\n3. All of these experiments are within the same single environment, with 3 different reward functions. There is no variability in terms of the dynamics or observation space. To take a tiny step towards a \"foundational world model\" could you consider some variation, such as using the distracting control suite or varied dynamics in the simulator? This is also available in VD4RL so could be used there too. My guess is the world model approach would actually do better and it would make the results more interesting.\n4. This is a very active area of research so the citations seem light, for example: \n- Edwards et al. \"Imitating latent policies from observation\". ICML 2019\n- Seo et al. \"Reinforcement learning with action-free pre-training from videos\". ICML 2022\n- Schmeckpeper et al. \"Reinforcement learning with videos: Combining offline observations with interaction\". CoRL 2020", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The experiments are fairly mundane. It would be fantastic to see an example of this method at larger scale. For example, this could be done using the dataset from VPT. The Minecraft images could be resized to make them smaller, and then it would be possible to use the DreamerV3 codebase (which was tested on MineCraft and runs on a single GPU) then see if it is possible to learn from the unlabelled MineCraft videos. If this works, it would drastically increase the impact of the paper, beyond being something mildly interesting for people who care about this specific topic, to something that catches people's eye across the field.\n\nWe agree doing the MineCraft experiment is a great way to improve the impact of this paper, and we have also thought about that. Actually, we had contacted the authors of DreamerV3 about the possibility to open source the pretrained world model immediately after the paper came out. But the authors said they won't do it due to the extra effort to make their model loadable with the open sourced code base. And if we train the world model ourselves, as stated in the appendix of DreamerV3, it will take about 17 days to train their model. And the VPT labelled dataset is almost as twice large as the final reply buffer of the DreamerV3 agent, so it will probably need a month to train a model and not speak about hyperparameter tuning. As a small lab, we don't have enough resource for an experiment at this scale. But we think this is a great suggestion, and we would like to find ways to try it out in the future. \n\n> It seems slightly fishy to use a world model generated dataset to compare a world model based method and an IDM approach. It is possible there is some bias in the Dreamer or P2E data that makes it easier for a more similar approach to do well. Given how brittle many of these methods are, this could make a difference. Would it be possible to instead consider some open-source benchmarks such as VD4RL (Lu et al), which would make it \"fair\" across different approaches? For example you could use the random or mixed datasets for the embodied and then expert for the demonstration.\n\nThat is a good suggestion. We conduct experiments with the V-D4RL main datasets. Please kindly check the experiment setup and results from the general response. \n\nWe 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 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 VD4RL have only 200 trajectories, which is much less than the 1000 trajectories in our paper. Moreover, it is already shown from Fig. 2 and Fig. 3 that random datasets does not help much about learning a model, and intuitively the medium_replay dataset is better but still does not contain enough information to solve the task.\n\nWe would like to point out that due to the limited time in the rebuttal period, these are only preliminary results. \n\n> All of these experiments are within the same single environment, with 3 different reward functions. There is no variability in terms of the dynamics or observation space. To take a tiny step towards a \"foundational world model\" could you consider some variation, such as using the distracting control suite or varied dynamics in the simulator? This is also available in VD4RL so could be used there too. My guess is the world model approach would actually do better and it would make the results more interesting.\n\nThanks for you suggestion. We conduct additional experiments with the distracting datasets. Please kindly check the experiment setup and results from the general response. \n\nAs we can see from the result, although still outperform the BCO(0) baseline, AIME is largely influenced by the distractions. This behaviour is desired 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 dynamic part of the world model and allowing the encoders and the decoders also get finetuned in the second phase. But due to the limited time in the rebuttal, we are unable to test this idea. \n\n> This is a very active area of research so the citations seem light.\n\nThanks for your suggestion. We will cite these papers in the updated version.\n\n> How was the IDM baseline tuned? For example, in the VPT paper they mention different architecture choices made a big difference for performance.\n\nThat is a good point. We are sorry that we forgot to add the implementation details of the BCO baseline to the appendix. In order to make a fair comparison in this paper, the IDM and policy are built by using the same network architecture with the world model. Especially for the visual setting, the IDM uses the same CNN as the world model, and the temporal information is handled by stacking the representation from each frame. Then the action is predicted as a TanhGaussian distribution with an MLP. 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. We will add these details to the Appendix B for the revised version.\n\nWe did not conduct experiments on other architecture design, like using a heavy transformer model to handle longer context length. But from Appendix D.1 of the VPT paper, the most important design choice is using the 3D convolution before the CNN to process each frame. We think that is quite similar to stacking the image directly when you have a short context length and don't have a transformer afterward who needs tokens from each frame. And, according to [1], stacking representations is actually a better design choice than stacking frames. \n\n[1] Shang *et al.*, Reinforcement Learning with Latent Flow, NeurIPS 2021\n"}]

Author rebuttal:
We thank all reviewers for their time and insightful feedbacks. 

We conduct three new experiments suggested by the reviewers, the results are provided in the pdf:

**We evaluate multiple checkpoints during the course of the world model pretraining to address the convergence concerns from reviewer Cs3Z. To be specify, we retrain the model on walker-mix dataset for 2000 epochs (prolong from 1000 epochs in the paper) and save a checkpoint every 100 epochs. Then, all the 20 saved models are evaluated by imitating the run policy.**

From the results shown in Fig. 1, the imitation ability is established in very early phase of the training process long before converge. 

**We conduct experiments on the VD4RL [1] main datasets, as suggested by reviewer MApq, to prove that AIME can also work with dataset generated by model-free methods. We use their random and medium_replay datasets from Walker and Cheetah as embodiment datasets. Besides, we also merge these two datasets to form a mix dataset. The models are trained on the three datasets for each embodiment. Then we treat the expert datasets in the benchmark as demonstration dataset.** 

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 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 VD4RL have only 200 trajectories, which is much less than the 1000 trajectories in our paper. Moreover, it is already shown from Fig. 2 and Fig. 3 that random datasets does not help much about learning a model, and intuitively the medium_replay dataset is better but still does not contain enough information to solve the task.

**We conduct experiments on the VD4RL [1] distracting datasets, as suggested by reviewer MApq, 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 level to 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 with 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, although still outperform the BCO(0) baseline, AIME is largely influence by the distractions. This behaviour is desired 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 dynamic part of the world model and allowing the encoders and the decoders also get finetuned in the second phase. But due to the limited time in the rebuttal, we are unable to test this idea. 

We would like to mentioned that due to the limit time during rebuttal, we only run AIME and BCO(0) baselines with the default parameters we used in the paper. Further hyper-parameter tuning could achieve a better result.

For individual questions raised by each reviewer, please find the responses below.

[1] Lu *et al.*, Challenges and Opportunities in Offline Reinforcement Learning from Visual Observations, TMLR 2023

Meta-review: Reviewers unanimously vote to accept the paper. The paper presents a methd for zero-shot imitation wherein the agent first learns a world model (i.e., forward dynamics model) and then given demonstrations of a target task uses its world model to infer actions. In the second stage, no finetuning or action information is required -- the demonstrations simply consist of a sequence of observations. Results are presented on a few locomotion tasks. The diversity of the tasks are a concern amongst the reviewer and therefore their ratings are lukewarm. 

While I am going with the reviewers and recommending acceptance, I encourage the authors to include the following discussion in the introduction: 

- There is a rich history in robotics where the "robot" first learns some basic skills and then uses them to interpret video/observation only demonstrations (see [1]-[4]) -- including them / discussing them is important to set the context for the work. Similarly, consider including other papers that the reviewers have pointed out, e.g., MApq pointed out: 
Edwards et al. "Imitating latent policies from observation". ICML 2019
Seo et al. "Reinforcement learning with action-free pre-training from videos". ICML 2022
Schmeckpeper et al. "Reinforcement learning with videos: Combining offline observations with interaction". CoRL 2020

- Inverse Dynamics Models (IDM) have been used exactly in this setting, and the authors cite and compare against it. However, more advanced IDM based approaches such as in [6] can easily be applied in locomotion tasks -- however authors argue they can't be in Lines 171-72. Even in locomotion tasks, given a pair of frames $o_t, o_{t+1}$, the frame $o_{t+1}$ can be considered to be the goal and can the trained goal-conditioned models can be applied. Please either provide a comparison with [6] or discuss in detail why it is not applicable. Further, the discussion of [5],[6] and other IDM based methods should be included in the introduction. 

The strength of this in my opinion is not that the agent's embodiment is learned first -- IDM approaches also do that. It is really that the authors are learning forward models -- and therefore its key to compare against strong implementations of IDM such as in [6]. 

[1] Yezhou Yang, Yi Li, Cornelia Fermuller, and Yiannis Aloimonos. Robot learning manipulation action plans by ”watching” unconstrained videos from the world wide web. In AAAI, 2015.
[2] Cynthia Breazeal and Brian Scassellati. Robots that imitate humans. Trends in cognitive sciences, 2002.
[3] Yasuo Kuniyoshi, Masayuki Inaba, and Hirockika Inoue. Teaching by showing: Generating robot programs by visual observation of human performance. In International Symposium on Industrial Robots, 1989.
[4] Yasuo Kuniyoshi, Masayuki Inaba, and Hirochika Inoue. Learning by watching: Extracting reusable task knowledge from visual observation of human performance. IEEE Transactions on Robotics and Automation, 1994
[5] Ashvin Nair, Dian Chen, Pulkit Agrawal, Phillip Isola, Pieter Abbeel, Jitendra Malik, and Sergey Levine. Combining self-supervised learning and imitation for vision-based rope manipulation. ICRA, 2017.
[6] Pathak, Deepak, Parsa Mahmoudieh, Guanghao Luo, Pulkit Agrawal, Dian Chen, Yide Shentu, Evan Shelhamer, Jitendra Malik, Alexei A. Efros, and Trevor Darrell. "Zero-shot visual imitation." ICLR 2018
Decision: Accept (poster)