Imitating Human Behaviour with Diffusion Models
Authors: Tim Pearce
Tabish Rashid
Anssi Kanervisto
Dave Bignell
Mingfei Sun
Raluca Georgescu
Sergio Valcarcel Macua
Shan Zheng Tan
Ida Momennejad
Katja Hofmann
Sam Devlin
Venue: ICLR 2023 poster (2023)
Keywords: imitation learning
behavioral cloning
behavioral cloning
diffusion models
generative models

Abstract: Diffusion models have emerged as powerful generative models in the text-to-image domain. This paper studies their application as observation-to-action models for imitating human behaviour in sequential environments. Human behaviour is stochastic and multimodal, with structured correlations between action dimensions. Meanwhile, standard modelling choices in behaviour cloning are limited in their expressiveness and may introduce bias into the cloned policy. We begin by pointing out the limitations of these choices. We then propose that diffusion models are an excellent fit for imitating human behaviour, since they learn an expressive distribution over the joint action space. We introduce several innovations to make diffusion models suitable for sequential environments; designing suitable architectures, investigating the role of guidance, and developing reliable sampling strategies. Experimentally, diffusion models closely match human demonstrations in a simulated robotic control task and a modern 3D gaming environment.

Published as a conference paper at ICLR 2023
IMITATING HUMAN BEHAVIOUR WITH DIFFUSION MODELS
Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, Sam Devlin Microsoft Research
ABSTRACT
Diffusion models have emerged as powerful generative models in the text-to- image domain. This paper studies their application els for imitating human behaviour in sequential environments. Human behaviour is stochastic and multimodal, with structured correlations between action dimen- sions. Meanwhile, standard modelling choices in behaviour cloning are limited in their expressiveness and may introduce bias into the cloned policy. We begin by pointing out the limitations of these choices. We then propose that diffusion models are an excellent fit for imitating human behaviour, since they learn an ex- the joint action space. We introduce several innovations
ieeevation-to-action mod-
pressive distribution to make diffusion models suitable for sequential environments; designing suitable architectures, investigating the role of guidance, and developing reliable sampling strategies. Experimentally. diffusion models closely match human demonstrations in a simulated robotic control task and a modern 3D gaming environment. Code: pt //ttt/gttttttt// www//oooooootot//oooooooooooooct /Imiittttttttttttaaeeaeaeaaaataa ing Ruman-Behavi Behavlour --D-----iutttt Diffuaion.
INTRODUCTION
To enable Human-A AI collaboration, agents must learn to best respond to all plausible human be- haviors (Dafoe et al., 2020; Mirsky et al., 2022). In simple environments, suffices to generate all possible human behaviours (Strouse et al., 2021) but as the complexity of the environment grows this approach will struggle to scale. If we instead assume access to human behavioural data, collab- orative agents can be improved by training with models of human behaviour (Carroll et al., 2019).
In principle, human behavior can be modelled via imitation learning approaches which an agent trained to mimic the actions of demonstrator from offline dataset of observation and action tuples. More specifically, Behaviour Cloning (BC), despite being theoretically limited (Ross et al., 2011), 1), has been empirically effective in domains such as autonomous driving (Pomerleau, 1991), robotics (Florence et al., 2022) and game playing (Ye et al., 2020; Pearce and Zhu, 2022).
Popular approaches to BC restrict the types of distributions that be modelled to make learning simpler. A common approach for continuous actions is to learn a point estimate, optimised via Mean Squared Error (MSE), which can be interpereted Another popular approach is to discretise the action space into a finite number bins and frame as a classification problem. These both suffer due the approximations they make (illustrated in Figure I), either encouraging the agent to learn an 'average' policy or predicting action dimensions independently resulting in "uncoordinated' behaviour (Ke al., 2020).
isotropic Gaussian of negligible variance.
p(ajo)
Gauseian
Discretised
K-Means
M-Means
Diffusion
Figure 1: Expressiveness of a variety of models for behaviour cloning in a single-step. arcade claw game with two simultancous, continuous actions. Existing methods fail to model the full action distribution, p(ajo), whilst diffusion models excel at covering multimodal & complex distributions.

---

Published as a conference paper at ICLR 2023
Such simplistic modelling choices can be successful when the demonstrating policy is itself of re- stricted expressiveness (e.g. when using trajectories from a single pre-trained policy represented by simple model.) However, for applications requiring cloning of human behaviour, which con- tains diverse trajectories and multimodality at decision points, simple models may not be expressive enough to capture the full range and fidelity of behaviours (Orsini et al., 2021).
For these reasons, we seek to model the full distribution of actions observed. In particular, this pa- per focuses on diffusion models, which currently lead in image, video and audio generation (Saharia et al., 2022; Harvey et al., 2022; Kong et al., 2020), and avoid issues of training instability in gen- erative adversarial networks (Srivastava et al.. 2017), or sampling issues with energy based models (Florence et al., 2022). By using diffusion models for BC we are able to: 1) more accurately model complex action distributions (as illustrated in Figure 1): 2) significantly outperform oate-of-the-art methods (Shafiullah et al., 2022) on a simulated robotic benchmark; and 3) scale to modelling hu- man gameplay in Counter-Strike Strike: Global Offensive a modern, 3D gaming environment recently proposed as platform for imitation learning research (Pearce and Zhu, 2022).
To achieve this performance, we contribute several innovations to adapt diffusion models to se- quential environments. Section 3.2 shows that good architecture design can significantly improve performance. Section 3.3 then shows that Classifier- Free Guidance (CFG), which is a core part of text-to-image models, surprisingly harms performance in cett-vattiontooacctii models. Finally, Section 3.4 introduces novel, reliable sampling schemes for diffusion models. The appendices in- clude related work, experimental details, as well as further results and explanations.
2 MODELLING CHOICES FOR BEHAVIOURAL CLONING
In this section we examine common modelling choices for BC. For illustration purposes, we created a simple environment to highlight their limitations. We simulated an arcade toy claw machine, as shown in Figures 1, 4. An agent observes a top-down image of toys (o) and chooses a point in the image, in 2D continuous action space, a E IR22. If the chosen point is inside the boundaries of a valid toy, the agent successfully obtains the toy. To build a dataset of demonstrations, we synthetically generate images containing one more toys, and uniformly at random pick a single demonstration action a that successfully grabs a toy (note this toy environment uses synthetic rather than human data). The resulting dataset is used to learn p(alo). To make training quicker, we restrict the number of unique observations o to seven, though this could be generalised.
MSE. A popular choice for BC in continuous action spaces approximates p(a o) by point- that is optimised via MSE. This makes a surprisingly strong baseline in the literature despite its simplicity. However, MSE suffers from two limitations that harm its applicability modelling the full, complex distributions of human behaviour. 1) MSE outputs a point-estimate. This precludes it from capturing any variance or multimodality present in p(alo). 2) Due its opti- misation objective, MSE learns the 'average' of the distribution. This can bias the estimate towards more frequently occurring actions, or can even lead to nut-of-distribution actions (e.g. picking the action between two modes). The first can be partially mitigated by instead assuming a Gaussian distribution, predicting variance for each action dimension and sampling from the resulting Gaus- sian. However, due to the MSE objective, the learnt mean is still the average of the observed action distribution. These limitations are visualised in Figure 1.
our goal of
Discretised. A second popular choice is to discretise each continuous action dimension into B bins, and frame it as a classification task. This has two major limitations. 1) Quantisation errors arise since the model outputs single value for each bin. 2) Since each action dimension is treated in- dependently. the marginal rather than the joint distribution is learnt. This can lead to issues during sampling whereby dependencies between dimensions haviour. This can be observed in Figure 1 where points outside of the true distribution have been sampled in the bottom-right corner. This can be remedied by modelling action dimensions autore- gressively, but these models bring their own challenges and drawbacks (Lin et al., 2021).
ignored, leading to 'uncoordinated" be-
K-Means. Another method that accounts for dependencies between action dimensions, first clusters the actions across the dataset into K bins (rather than Bia)) using K-Means. This discretises the joint-action distribution, rather than the marginal as in 'Discretised'. Each action is then associated with its nearest cluster, and learning can again be framed as a classification task. This approach
2

---

Published as a conference paper at ICLR 2023
avoids enumerating all possible action combinations, by placing bins only where datapoints exist. However two new limitations are introduced: 1) Quantisation errors can more severe than for "Discretised' due only K action options. 2) The choice of K can be critical to performance (Guss et al., 2021). ). If K' is small, important actions may be unavailable to an agent; K' is large, learning becomes difficult. Additionally, since the bins chosen with respect to the entire dataset they can fall outside of the target distribution for a given observation. This is observed in Figure 1.
E-Means+Resddusl Shafiullah et al. (2022) extended K-Means by learning dependent residual that is added to the bin's center and optimised via MSE. This increases the fidelity of the learnt distribution. However, MSE. 1) The distribution p(alo) is still modelled by finite number of point estimates (maximum of K). 2) The residual learns the "average' of the actions that fall within each bin, In addition, it still requires careful choice of the hyperparameter K.
observation-
still carries through some issues of K-Means and
Diffusion models. The limitations of these existing modelling choices all arise from approximations made in the form of p(alo), combined with the optimisation objective. But it is precisely these approximations that make training models straightforward (optimising network via MSE or cross- entropy is trivial!).. So how can we learn p(a)o) while avoiding approximations? We propose that recent progress in generative modelling with diffusion models provides an answer. Diffusion models able to output expressive conditional joint distributions, and are powerful enough to scale to problems as complex as text-to-image or video generation. Thus, we will show they provide benefits they make coarse approximations about the action
when imitating human demonstrations distribution, and avoid many of the limitations of the previously discussed choices. This can be observed in Figure where only diffusion provides an accurate reconstruction of the true action distribution.
3 OBSERVATIION-TO-ACTT DIFFUSION MODELS
Diffusion models have largely been developed for the image domain. This section first introduces the underlying principles of diffusion models, then studies several aspects that require consideration to apply them to sequential environments.
3.1 DIFFUSION MODEL OVERVIEW
Diffusion models are generative models that map Gaussian noise to some target distribution in an iterative fashion, optionally conditioned some context (Dhariwal and Nichol, 2021). Beginning from, ay ~ (O, I), a sequence ar-1. ar-2 ... ag is predicted, each a slightly denoised version of the previous, with ag 'clean' sample. Where T is the total number of denoising steps (not the environment timestep as is common in sequential decision making).
This paper uses denoising diffusion probabilisitic models (Ho et al., 2020). During training. noisy inputs can be generated as: a, Vā, a +I- o, z, for some variance schedule ok,, random noise (O, I). neural network, e(-), is trained predict the noise that added input, by minimising:
CDPPM == Fo [l](o,, a,, -00-1]] =z(²]
where the expectation over all denoising timesteps, U(1 T), and observations and actions are drawn from a demonstration dataset, D.
At sampling time, with further variance schedule parameters a, & a, inputs are iteratively denoised:
1- ar E(o,, a... + o.z.
a Va, a,
-1-ar ar
3.2 ARCHITECTURAL DESIGN
This section explores neural network architectures for observation-to-action diffusion models. Specifically, the requirements of the network are: Input: Noisy action ar-1 R glall denoising
timestep т, observation (possibly with history); Output: Predicted noise mask, â Rial
3

---

Published as a conference paper at ICLR 2023
Desiising neiwork architecture choices
Text-to-lmaee diffusion
Observation intervatonnto-cctiion to-action diffusion
Bane MLP
MLP Stm
Transkene
Obsrevacion
Larggge masiel
-... -00, (00.1
Imags, x pisle)
Elicddrr
Eavirmmeet
plalo)
Action sampiinng choices
Denriving aetwork
Dilfision Br
Dilfucionn
DillisinnKKi
**
denoising timestops, guidance
deuniing Limnstepe,
guilassee
Figure 2: Diffusion BC generates action vector conditioned observation (which may be
image). By contrast, text-to-image diffusion models generate an image conditioned on a vector.
While U-Nets have become standard components of text-to-image diffusion models, their only makes sense for large, spatial input and outputs, while we require generation of an action vector of modest dimensionality. Therefore, now describe three architectures of varying complexity. Section 4 empirically assesses these three architectures, finding performance improvements in the order: Basic MLP MLP Sieve Transformer.
Basic MLP. This architecture directly concatenates all relevant inputs together, 0, This input is fed into multi- multi-layer perceptron (MLP).
MLP Sieve. This uses three encoding networks to produce embeddings of the observation, denois- ing timestep, and action: o", t², t,, E emmbe dim These are concatenated together as input to denoising network, (of, t., af, af_1). The denoising network is a fully-connected architecture, with residual skip connections, and with the raw denoising timestep т and action a,-1 repeatedly con- catenated after each hidden layer. To include longer observation history, previous observations are passed through the same embedding network, and embeddings are concatenated together.
Transformer. This creates embeddings as for MLP Sieve. A multi-headed attention architecture (Vaswani et al., 2017) (as found in modern transformer encoder networks) is then used as the de- noising network. At least three tokens used input, o tt af, and this can be extended to incorporate longer history of observations (only the current t", a are needed since the diffusion process is Markovian).
Sampling rate. The MIP Sieve and Transformer are carefully designed so the observation encoder separate from the denoising network. At test time, this means only single forward pass is required for the observation encoder, with multiple forward passes run through the lighter denoising network. This results manageable sampling time-1 in the experiment playing video game from pixels (section 4.2), we were able to roll out our diffusion models at 8Hz on an average gaming GPU (NVIDIA GTX 1060 Mobile). Table provides detailed breakdown, Note that sampling time is more severe issue in text-to-image diffusion models, where forward passes of the heavy U-Net architectures are required for all denoising timesteps.
3.3 WHY CLASSIFIER- FREE GUIDANCE FAILS
Classifier-Free Guidance (CFG) has become core ingredient for text-to-image models, allowing one to trade-off image typicality with diversity (Ho and Salimans, 2021). In CFG, a neural network is trained both a conditional and unconditional generative model. During sampling, by introducing a "guidance weight' w, one places higher weight (w >> 0) on the prediction conditioned on some context (here, o), and negative weight the unconditional prediction,
2, (1 Wconodd (a,-1, 0, WCCuncomd 12
One might anticipate that CFG would also be beneficial in the sequential setting, with larger producing trajectories of higher likelihood, but at the cost the diversity. Surprisingly, we find that CFG can actually encourage less common trajectories, and degrade performance.

---

Published as a conference paper at ICLR 2023
s(alo)
Unique p(ajn)
Weight
Weight
Weight
Figure 3: We vary the CFG 'weight' parameter (w value in Eq. 3) during sampling in the Arcade Claw environment. CFG encourages selection of actions that were specific to an observation (max- imising p(ola)). This can lead to less common trajectories being sampled more often.
In Figure we visualise p(alo) for the claw machine game, under varying guidance strengths, w. An interpretation of CFG is that encourages sampling of actions that would maximise an implicit classifier, p(ola). Hence, CFG encourages selection of actions that were unique to a particular observation (Ho al., 2022). Whilst this is useful for rext-to-image models (generate images that are more specific to a prompt), in sequential environments this leads to an agent rejecting higher- likelihood actions in favour of less usual ones that were paired with some observation. In later experiments (section 4.1) we demonstrate empirically that this can lead to less common trajectories being favoured, while degrading overall performance. Appendix provides a didactic example of when CFG fails in this way.
3.4 RELIABLE SAMPLING SCHEMES
In text-to-image diffusion, several samples are typically generated in parallel, allowing the user to select their favourite, and ignore any failures. However, when rolling out Donervation-to-action diffusion model, such manual screening not feasible. There remains a risk that bad action could be selected during a roll-out, which may send an agent toward an out-of-distribution state. Hence, propose 'Diffusion-X" and 'Diffusion-KDE' as variants of Diffusion BC, that mirror this screening process by encouraging higher-likelihen actions during sampling. For both methods, the training procedure is unchanged (only the conditional version of the model is required). The algorithms for all sampling methods are given appendix D.
Diffusion-X. The sampling process runs normal for T denoising timesteps. The denoising timestep then fixed, r 1, and extra denoising iterations continue run for M timesteps. The intuition behind this is that samples continue to be moved toward higher-likelihood regions.
Diffusion-KLD KDE. Generate multiple action samples from the diffusion model as usual (these can be done in parallel). Fit a simple kernel-density estimator (KDE) over all samples, and score the likelihood of each. Select the action with the highest likelihood.
Obs.
p(alo)
M
M
32
Diff. KDB
Figure 4: This figure shows the predictive distributions of diffusion models using various sampling schemes. When M = 0 this is Diffussio BC' and when M > this is 'Diffusion-X: Diff. KDE refers to Diffusion-KDEi

---

Published as a conference paper at ICLR 2023
The effect of these sampling modifcations is demonstrated in Figure 4. While Diffusion BC gener- ates a small number of actions that fall outside of the true p(alo) region, Diffusion-X and Diffusion- KDE avoid these bad actions. Note that the two distinct modes in the figure are recovered by both sampling methods, suggesting that multimodality is not compromised, though the diversity within each mode is reduced. Both techniques are simple implement, and experiments in Section 4 show their benefit.
4 EXPERIMENTS
This section empirically investigates the efficacy of diffusion models for BC. We assess our method in two complex sequential environments, which have large human datasets available, with the aim of answering several questions: Q1) How do diffusion models compare to existing baseline methods for BC, in terms of matching the demonstration distribution? Section 4.1 compares to four popular modelling choices in BC, as well as recent state-of-the-ee models. Section 4.2 provides further focused comparisons. Q2) How is performance affected by the architectures designed in section 3.2. CFG, and the sampling schemes in section 3.4? We provide full ablations in Section 4.1, and targeted ablations over sampling schemes in Section 4.2. Q3) Can diffusion models scale to complex environments efficiently? Section 4.2 tests environment where the observation is high-resolution image, the action space mixed continuous & discrete, and there are strict time constraints for sampling time.
Evaluation. To evaluate how closely methods imitate human demonstrations, we compare the be- haviours of our models with those of the humans in the dataset. To do so, we will compare both a high-level, analysing observable outcomes (e.g. tasks completions or game score), as well as at a low-level comparing the stationary distributions over states or actions. Both of these are important in evaluating how humanlike our models are, and provide complimentary analyses. Appendix B provides details on our evaluation metrics, which are introduced less formally here.
Baselines. Where applicable we include Human to indicate the metrics achieved by samples drawn from the demonstration dataset itself. We then re-implement five baselines. Three correspond to popular BC modelling choices, namely MSE: a model trained via mean squared error; Discretised: each action dimension discretised into 20 uniform bins, then trained independently via cross- entropy; and K-means: set of K candidate actions are first produced by running K-means over all actions, actions are discretised into their closest bin, and the model is trained via cross-entropy. A further two baselines can be considered strong, as with K-means, but additionally learns a continuous residual on top of each bin prediction, trained via MSE, which the core innovation of Behaviour Transformers (Be'T) (Shafiullah al., 2022); and EBM: a generative energy-based model trained with a contrastive loss, proposed in (Florence al., 2022) full details about the challenges of this method given in Appendix B.4.
complex methods, namely K-means+ Residual:
One of our experiments uses the set up from Shafiullah et al. (2022), allowing us to compare to their reported results, including Behaviour Transformers (BeT): the K-mean+residual combined with large 6-layer transformer, and previous 10 observations as history; Implicit BC: the official imple- mentation of energy-based models for BC (Florence et al., 2022); SPiRL: using a VAE, originally from Pertsch et al. (2021); and PARROT: a flow-based model, originally from Singh et al. (2020).
Architectures & Hyperparameters. We trial the three architecture options described in Section 3.2, which are identical across both the five re-implemented baseline methods and our diffusion variants (except where required for output dimensionality). Basic MIP and MLP Sieve both use networks with GELU activations and 3 hidden-layers of 512 units each. Transformer uses four standard encoder blocks, each with 16 attention heads. Embedding dimension is 128 for MLP Sieve and Transformer. Appendix B gives full hyperparameter details.
4.1 LEARNING ROBOTIC CONTROL FROM HUMAN DEMONSTRATION
In this environment, an agent controls a robotic arm inside a simulated kitchen. is able to perform seven tasks of interest, such opening microwave, or turning on stove (Appendix B contains full details). The demonstration dataset contains 566 trajectories. These were collected using a virtual reality setup, with human's movements translated to robot joint actuations (Gupta et al.,

---

Published as a conference paper at ICLR 2023
2020). Each demonstration trajectory performed four predetermined tasks. There are 25 different task sequences present in the dataset, of roughly equal proportion.
This environment has become popular offline RL benchmark for learning cies (Fu et al., 2020). However, as our goal is to learn the full distribution of demonstrations, we instead follow the setup introduced by Shafiullah al. (2022), which ignores any goal conditioning and aims to train an agent that can recover the full set of demonstrating policies.
poli-
The kitchen environment's observation space is 30-dimensional continuous vector containing in- formation about the positions of the objects and robot joints. The action space a 9-dimensional continuous vector of joint actuations. All models receive the previous two observations input, allowing the agent to infer velocities. For diffusion models, we set T = 50.
The kitchen environment is challenging for several reasons. 1) Strong (sometimes non-linear) cor- relations exist between action dimensions. 2) There is multimodality in p(alo) at the point the agent selects which task complete next, and also in how it completes We show that our diffusion model learns to represent both these properties in Figure 9, which visualises relationships between all action dimensions during one rollout.
4.1.1 MAIN RESULTS
Comparing the behaviour of our agents with that of the humans demonstrations is challenging re- search problem in its own right. Table presents several metrics that provide insight into how closely models match the human demonstrations from high-level analysis of task sequences selected, to low-level statistics of the observation trajectories generated by models. We briefly summarise these, more technical descriptions can be found in Appendix B. Training and sampling times of methods are given in Table 6.
24 We first measure the proportion of rollouts for which models perform four valid tasks, which nearly all human demonstrations achieve.
Tasks Wasserstein. For each method record how many times each different task sequence was completed during rollouts. The Wasserstein distance is then computed between the resulting histogram, compared to the human distribution.
Time Wasserstein. As well as analysing which tasks were completed, we also analyse when they were completed. Figure 5 plots the distribution of the time taken to complete different numbers of tasks (normalised to exclude failures), for MSE and Diffusion-X. We then compute the Wasserstein distance between the human and agent distributions (see Appendix C for full break down.)
State Wasserstein. If we truly capture the full diversity and fidelity of human behaviour present in the dataset, this will be reflected in the state occupancy distribution. We compute the Wasserstein between agent and human distributions, with lower value indicating they are more similar.
Density & Coverage. We use the 'Density' and 'Coverage" metrics from Naeem et al. (2020) that are used to evaluate GANs. Roughly speaking 'Density' corresponds to how many states from human trajectories are close to the agent's 's states and 'Coverage' corresponds to the proportion of human states that have an agent generated state nearby.
In terms of complconn rate, our Diffusion BC approaches outperform all baselines, with the ordering: Diffusion BC-Diffusion X<<Diffusion-Ki KDE. This ordering supports our hypothesis that these sampling schemes more reliably avoid bad actions. These improvements come at the cost of increased sampling time- for MLP Sieve, sampling rate drops from 16 Hz for Diffusion BC to 14 Hz for Difffusio-XX X to 12 Hz for Diffusion- KDE (Table 6).
For all Wasserstein metrics, diffusion models again outperform other baselines, and sampling schemes are usually ordered Diffusion KDE<Diffusion BC-<Diffusion-.. This is important; while Diffusion-KDD KDE completes more tasks, it signals a tendency overfit smaller number of task sequences, generating less of the diversity from the demonstration distribution. This is quantitatively confirmed by Diffusion-Ko KDEE scoring significantly higher Density, but lower Coverage.
Architecture ablation. In terms of the different architectures designed in Section 3.2. Table I shows that metrics usually improve in order: Basic MIP.<MLPP Sieve< Sieve<TTransfommn Transformer. These improvements
7

---

Published as a conference paper at ICLR 2023
Table 1: Robotic control results. Mean + one standard error three training (100 rollouts). Methods marked with asterisk (*) are our proposed methods.
Density
Coverage
ML.P Basic Architecture "Diffusion BC, Basic MLP
0.45 0.03
1,96 0.12
12,04 2.20
0.463 0.012
0.54 =: 0.02
0.38 0.01
DDiffusion
0.03
0.481
"Diffusion-X, Basic MLP
0.58 + 0.02
1,51 = 0.14
8.61 0.14
0144 0.017
0.64 + 0.00
0.41 + 0.00
MLP Sieve Architecture
0.07
0.01
0.40
Discretised, MLP Sieve
0.18 0.02
3,43 = 0.14
11.30 = 1.29
0.651 = 0.026
0.38 = 0.02
0.31 # 0.01
K- K-Meass, MLP Sieve
0.0 0.0
5.25 = 0.0
1,469 a 0.120
0.09 1: 0.00
0.06 1: 0.00
K-Menas Recidual, MI.P Sieve
0.23++0.00
2.87 0.16
11.60 2.11
0.607 0.027
0.51 0.01
0.36 0.00
Free,
"Diffusion BC, MLP Sieve
0.68 0.02
1.31 2 0.05
6,06 1.10
0.373 0.012
0.66 1: 0.01
0.42 0.00
"Diffusien-KDD KDE, MLP Sieve
0.79 + 0.04
166 0.24
6.77 0.64
0.439 0.039
0.93 + 0.02
0.41 + 0.01
DDiffuionnX X. MLP Sieve
0.77 0.02
1,06 0,05
5.24
0.344 0.004
0.78 0.01
0.45 0.00
Transformer Architecture
MSE, Tranffomer
0.89 +: 0.02
1.47 = 0.13
5.85 0.27
0.397 = 0.034
0.81 =: 0.01
0.42 +: 0.01
Discrelise,, Transformr
0,34 0.02
2,51 0.14
6.13 0.49
0,512 0.002
0.47 1: 0.01
0.36 0.00
K-Means, Transformer
0.34 0.02
2.25 = 0.16
7,80 0.87
0.426 0.018
0.66 +: 0.02
0.38 = 0.01
"Diffusion BC, Transformer
0.77 0.01
1,35 - 0.11
4,11 0.05
0,340 - 0.003
0.74 1 0.01
0.44 1: 0.00
DDiffuionn KDE. Transfomer
0.89 0.01
1.31 = 0.03
5.28
0.418
0.97 0,02
0.43 0.01
0.13
From Shufiollah et al. (2022)
Behaviour Transformers
0.44
Implicit BC
SPIRL VAE
0.0
PARROT Normalizing Flow
0.0
Human
0.98
Human sub-sampled
0.223 1: 0.006
1.00 1: 0.01
0.56 1: 0.01
do come at the of increased training and inference time for Diffusion BC the sampling rate drops from 24 Hz for MLP Basic to 16 Hz for MLP Sieve to 4 Hz for transformer (Table 6).
4.1.2 CLASSIFIER FREE GUIDANCE ANALYSIS
Section 3.3 provided intuition for why expect CFG to fail in to-suuuueee-sseeueee models. We now test our hypotheses empirically. During training the observation embedding is randomly masked with zeros with probability 0.1. This allows us learn both conditional and unconditional generative model. We then perform rollouts under different guidance weights, w E {0, 1, 4, 8), measuring the rate of completion of 4 tasks and monitoring which task the agent completes first.
Table 2 shows that completion rate drops from 0.63 without guidance (w = 0) to 0.08 with strong guidance (w 8). Mcanwhile, CFG also creates strong bias towards selection of Bottom Burner as the first task. This is significant as the human demonstrators select Bottom Burner just 10% of the time, but this increases from 7% (no guidance, w 0) to 25% (strong guidance, w 8), showing that CFG encourages less usual trajectories.
Tatts Complere
Tix05000
Figure S: Time complete robotic control kitchen tasks. Shaded: Human demonstrations. Left: MSE Transformer, Right: Diffusion-X Transformer.
8

---

Published as a conference paper at ICLR 2023
Table 2: Effect of CFG for Diffusion BC, MLP Sieve. Mean three training (100 rollouts).
First task peformed Bottom Buruer
Tasks Microwave
Kettle
Otber/Failure
Guidance 0.0
0.63
62.3
29.7
7.3
0.7
Guidance
Guidance 22 4.0
0.45
52.7
28.3
17.0
2.0
Guidance x0 8.0
0.08
24.0
26.7
24.7
25.6
Demonstrations
4.2 MODELLING HUMAN GAMEPLAY IN VIDEO GAME
We further tested our models in the environment 'Counter-Strike: Global Offensive" (CSGO) intro- duced by Pearce and Zhu (2022) (https: : /ttt:://tttttt //githtbbb... com Tenn:/eeeaaare//e/ /Counter-rrrrreeeeee Strike. Behavioural Cl Cloning). CSGO is one of the world's most popular video games in player and spectator numbers. We use the 'aim train' environment, where player is fixed on platform in the center of map, and must defend themselves against enemies controlled by built-in AI, who rush towards the player. Success requires precise and coordinated control over a mixed continuous & discrete action space, as well as dealing with multimodality in target selection and aiming.
The demonstration dataset contains 45,000 observation/action tuples, recorded from a high-skill human player. Observations are 280 150 RGB images, and the action space is three dimensional a fixed test only the
mouse x E I, mouse I, left click E {0, 1]. The environment runs asynchronously rate, providing rigorous test of sampling speed of models. Due to this constraint, MLP Sieve architecture, which offers a good trade-off between inference speed and performance. We also exclude the slightly slower Diffusion-KDDE sampling method. T is set 20.
For baselines, we compare to MSE, discretised, and K R-Means-Residual which were the strongest baselines in the kitchen environment. We also include EBM using derivative-free optimisation. We consider two options for the observation encoder: 1) a lightweight CNN, and 2) ResNet18 with ImageNet weights. Diffusion models 20 denoising timesteps, and Discrete action dimension. For MSE, we optimise the left click action via binary cross-entropp and mouse & y via MSE. Appendix B provides further details on hyperparameters and metrics.
20 bins per
Table 3 reports results. Diffusion-X performs best across both observation encoders in terms of game score and distance to the human distribution, which was measured the Wasserstein distance on the actions predicted (grouped into sequences of different lengths to assess the temporal consistency). Training and sampling times of methods are given Table Diffusion-X runs at 18 Hz compared to 200 Hz for MSE. Training time is similar.
Table 3: Video game results, over three rollouts of 10 minutes. Methods marked with asterisk (*) are our proposed methods.
Wasserstein Distance, Human Model
Game Score
1xtimesteps
16 timesteps (1 sec) 32x 32xtimess (2 sec)
Observation encader: 6-layer CNN MLP Sieve
54.7
Discrete, MI P Sieve
7.8
12.9
35.0
55.9
DDiffusion MLP Sieve
12.9
14.3
52.8
DiifiusionXX X, MI.P Sieve
17.3
11.7
30.0
47.8
Observation encoder: ResNet18 MSE, ML.P Sieve
17.8
5.5
28.1
48.9
Discrete, MLP Sieve
14.7
53.0
-.MannsRRssdunn MLP Sieve
16.8
3.8
29.2
51.8
EBM Derivative-Free. MLP Sieve
4.3
17.2
50.0
74.1
"Diffusion BC. MLP Sieve
19.0
29.5
50.4
*Diffusion X, MLP Sieve
24.0
4.5
24.5
44.4
Baselines
Human
36,5
0.73
0.57
0.38

---

Published as a conference paper at ICLR 2023
5 DISCUSSION & CONCLUSION
This paper provided the first thorough investigation into using diffusion models for imitating human behaviour. Existing modelling choices in BC make various approximations that simplify the learn- ing problem. But these approximations come at cost: they introduce limitations and biases on the cloned policy, which we systematically identified in section 2. Recent developments in generative modelling for images have shown that diffusion models are capable of learning arbitrarily complex conditional distributions, and are stable to train. This paper has combined these two worlds, propos- ing that diffusion models are also excellent fit for learning the complex observation-to- to-action distributions found in datasets of humans behaving in sequential environments. Our key insight is that, unlike many modelling choices for BC, diffusion models make coarse approximations the target distribution, so avoid many of their shortcomings.
Diffusion models do come with several limitations. They increase inference time MSE model can sample actions at 666 Hz and 200 Hz in the kitchen and CSGO environments, while Diffusion BC samples 16 Hz and 32 Hz. They also introduce several new hyperparameters. Finally, our diffusion models address only one challenge in imitation learning that of learning complex action distributions single environment timestep. There are other open challenges they do not address such as learning correlations across environment timesteps.
This paper contributed several innovations successfully adapt diffusion models to the observation- to-action domain. Our experiments demonstrated the effectiveness of these with several key take- aways. Diffusion models offer improvements over other methods in matching the demonstrations in terms of reward and distribution. 2) Reliable sampling schemes Diffusion-X and Diffusion-KI -DE offer benefits over Diffusion BC. 3) Good architecture design important to success of Diffu- sion models, allowing trade-off of performance and sampling speed. 4) CFG should be avoided as mechanism to condition on observations for diffusion agents in sequential environments.
Experiments support hypothesis that, by avoiding coarse approximations about the action dis- tribution, diffusion models improve over existing methods at modelling human demonstrations. On complex control task, diffusion models achieved a task completion rate of 89%, exceeding cent state-of-the-aa of 44%. We also demonstrated that our models can scale to learn directly from image observations, mix of continuous and discrete actions, sampling actions 18 Hz. This demonstrates the possibility of applying our models in complex, real-world settings.
REFERENCES
Anurag Ajay. Yilun Du Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-mate making?? International Conference on Learning Representations, 2023.
Micah Carroll, Rohin Shah, Mark K Ho, Tom Griffiths, Sanjit Seshia, Pieter Abbeel, and Anca Dragan. On the utility of learning about humans for human-ai coordination. Advances in neural information processing systems, 32, 2019.
Robert Dadashi, Leonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal wasserstein imita- tion learning. In International Conference on Learning Representations, 2020.
Allan Dafoe, Edward Hughes, Yoram Bachrach, Tantum Collins, Kevin R McKee, Joel Z Leibo, Kate Larson, and Thore Graepel. Open problems in cooperative ai. arXiv preprint arXiv:2012.086300, 2020.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794 2021.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar's Z. Alaya, Aurélie Boisbunon, Stanis- las Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Foumier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy. levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78): 1-8, 2021.
10

---

Published as a conference paper at ICLR 2023
Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Conference Robot Learning. pages 158-168. PMLR, 2022.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. Datasets for data-driven reinforcement learning. arXiv preprint arXiv:2004.00 07219, 2020.
Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu. A divergence minimization perspective on imitation learning methods. Leslie Pack Kaelbling, Danica Kragic, and Komei Sugiura, editors, Proceedings of the Conference on Robot Learning, volume 100 of Proceedings of Machine Learning Research, pages 1259-1277. PMLR, 30 Oct-01 Nov 2020.
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. In Conference on Robot Learning. pages 1025- 1037. PMLR, 2020.
William Hebgen Guss, Stephanie Milani, Nicholay Topin, Brandon Houghton, Sharada Mohanty, Andrew Melnik, Augustin Harter, Benoit Buschmaas, Bjarne Jaster, Christoph Berganski, Den- nis Heitkamp, Marko Henning, Helge Ritter, Chengjie Wu, Xiaotian Hao, Yiming Lu, Hangyu Mao, Yihuan Mao, Chao Wang, Michal Opanowiez, Anssi Kanervisto, Yanick Schraner, Chris- tian Scheller, Xiren Zhou, Lu Liu, Daichi Nishio, Toi Tsuneda, Karolis Ramanauskas, and Gabija Juceviciute. Towards robust and domain agnostic reinforcement learning competitions: Minerl 2020. In Proceedings of the NeurIPS 2020 Competition and Demonstration Track, 2021.
William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. 2022. URL
11495. wwwp:s///g//oo.o.tooo.o /atp//bb///222// /ttp///ww..../////2/....oooo
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural
information processing systems, 29, 2016.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in 2020.
Neural Information Processing Systems,
Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02300: 2022.
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In International Conference Learning Representations, 2019.
Ahmed Hussein, Mohamed Medhat Gaber, Eyad Hlyan, and Chrisina Jayne. Imitation learning: A survey of learning methods. ACM Computing Surveys, 50(2):1-35, 2017.
Léonard Hussenot, Marcin Andrychowicz. Damien Vincent, Robert Dadashi, Anton Raichuk, Sabela Ramos, Nikola Momchev, Sertan Girgin, Raphael Marinier, Lukasz Stafiniak, al. Hy- perparameter selection for imitation learning. In International Conference on Machine Learning,
pages 4511-4522 PMLR, 2021.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022.
Liyiming Ke, Sanjiban Choudhury, Matt Barnes, Wen Sun, Gilwoo Lee, and Siddhartha Srinivasa. Imitation learning as f-divergence minimization. In International Workshop on the Algorithmic Foundations of Robotics, pages 313-329. Springer, 2020.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations, 2020.
11

---

Published as a conference paper at ICLR 2023
Chu- Chu-Cheng Lin, Aaron Jaech, Xin Li, Matthew R Gormley, and Jason Eisner. Limitations of au- toregressive models and their alternatives. Proceedings of the 2021 Conference of the North American Chapter the Association for Computational Linguistics: Human Language Technolo- gies, pages 5147-5173, 2021.
Reuth Mirsky, Ignacio Carlucho, Arrasy Rahman, Elliot Fosong, William Macke, Mohan Sridharan, Peter Stone, and Stefano V Albrecht. A survey of ad hoc teamwork: Definitions, methods, and open problems. arXiv preprini arXiv:2202. 2ani:2222224400 10450, 2022.
Muhammad Ferjad Naeem, Seong Joon Oh, Youngjung Uh, Yunjey Choi, and Jaejun Yoo. Reliable fidelity and diversity metrics for generative models. In International Conference on Machine Learning, pages 71 7176-7185. PMLR, 2020.
Manu Orsini, Anton Raichuk, Léonard Hussenot, Damien Vincent, Robert Dadashi, Sertan Girgin, Matthieu Geist, Olivier Bachem, Olivier Pietquin, and Marcin Andrychowicz. What matters for adversarial imitation learning? Advances Neural Information Processing Systems, 34:14656- 14668, 2021.
Tim Pearce and Jun Zhu. Counter-strike deathmatch with large-scale behavioural cloning. In 2022 IEEE Conference on Games (CoG), pages 104-111. IEEE, 2022.
Karl Pertsch, Youngwoon Lee, and Joseph Lim. Accelerating reinforcement learning with learned skill priors. In Conference on robot learning, pages 188-204. PMLR, 2021.
Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neu- ral computation, 3(1):88-97, 1991.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204.06/25, April 2022.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent. arXiv preprint arXiv:2205.0611 0615,, 2022.
Stéphane Ross, Geoffrey Gordon, and J Andrew Bagnell. reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, 2011.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kam- yar Seyed Ghasemipour, Bureu Karagol Ayan, S. Sara Mahdavi, Rapha Gontijo Lopes, Tim Sal- imans, Jonathan Ho, David J Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffu- sion models with deep language understanding. arXiv preprint arXiv:2205. 1487, May 2022.
Nur Muhammad Mahi Shafiullah, Zichen Jeff Cui, Ariuntuya Altanzaya, and Lerrel Pinto. Behav- ior transformers: Cloning & modes with stone. Advances in neural information processing systems, 2022.
Avi Singh. Huihan Liu, Gaoyue Zhou, Albert Yu, Nicholas Rhinehart, and Sergey Levine. Parrot: Data-driven behavioral priors for reinforcement learning. arXiv preprint 1.Xi::0011.10000 2020.
Akash Srivastava, Lazar Valkov, Chris Russell, Michael U Gutmann, and Charles Sutton. Veegan: Reducing mode collapse in gans using implicit variational learning. Advances in neural informa- tion processing systems, 30, 2017.
DJ Strouse, Kevin McKee, Matt Botvinick, Edward Hughes, and Richard Everett. Collaborating with humans without human data. Advances in Neural Information Processing Systems, 34: 14502-14515, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
12

---

Published as a conference paper at ICLR 2023
Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.06193, August 2022.
Deheng Ye, Guibin Chen, Peilin Zhao, Fuhao Qiu, Bo Yuan, Wen Zhang, Sheng Chen, Mingfei Sun, Xiaoqian Li, Siqin Li, al. Supervised learning achieves human-level performance in moba games: A case study of honor of kings. IEEE Transactions on Neural Networks and Learning Systems, 2020.
13

---

Published as a conference paper at ICLR 2023
This appendix is organised as follows.
Section A: Extended Related Work
Section B: Experimental details
Section C: Further Results
Section D: Sampling Algorithms
Section E: Classifier-Free Guidance Analysis
A EXTENDED RELATED WORK
Implicit BC. Florence et al. (2022) formulated BC as a conditional energy-hased modeling prob- lem. They showed that the ability of energy-based policies to represent multimodal distributions and discontinuous functions can provide competitive results with state-of-the-an offline reinforce- ment learning methods. Our paper is complimentary to this work, continuing the investigation on expressive models, studying the fit of diffusion models for BC. As well as showing that they can out- perform state-of oo--off-hhe of-the-art art baselines in terms of score, we also show they outperform in closeness to the demonstration distribution, which was not explored by Florence et al. (2022). We note that dif- fusion models avoid some of the challenges of energy-based models they remove the requirement to generate negative samples during training, and avoid some of the complexities in the sampling process that energy-based models face.
Transformer BC. Although there are multiple works using transformers for RL, this section we only consider approaches that use transformers for BC, where no interaction with the environment is allowed, and only state-action trajectories and no reward signal are available in a dataset, for which the existing literature is scarce.
Reed et al. (2022) proposed GATO, an agent that could model different input modalities, including text, images and behaviour policies. However, they focused on demonstrating multitask capabilities, rather than on on how to do BC efficiently. GATO relies on discretisation, which can lead to inac- curate actions due to quantisation errors (see Sec. 2); and on autoregressive predictions, which can lead to repetitive and out distribution trajectories (Holtzman et al., 2019).
Shafiullah et al. (2022) proposed transformer based BC method that implements the K- Residull approach discussed in See. 2. This done by using two heads, one with a cat- egorical loss for predicting the centroid of the cluster which the action belongs; and the other one with an MSE loss to predict the 'residual', that is the distance from the centroid to the action. This resulted in discrete approximation of the policy. Although having limitations, it was able to outperform other approximations in a number of environment, using both human and synthetic data.
Here, we focus generative modelling with diffusion models, evaluating the impact of differ- ent neural network architectures and different sampling schemes, and showing they are expressive enough model the policy without implicit approximations, outperforming previous approaches, including discretisation and K-Meens-Residual.
Diffusion models. One main difference with previous works on diffusion models is that in order to obtain high quality samples, they relied on guiding the sampling process, either by using an expensive classifier (Dhariwal and Nichol, 2021), by training conditional and an unconditional model together (Ho and Salimans, 2021). However, as explained in Sec. 3.2, guidance can bias sampling towards low likelihood trajectories, which is problematic for BC applications, so in order to obtain high quality and high likelihood samples, we introduce two novel quidance-free methods.
RL and Diffusion Models. Up to our knowledge, the only previous work that used diffusion models for RL is by Janner et al. (2022) who proposed Diffuser: model that predicts all states and actions of a complete trajectory; and uses classifier-like guidance for trajectory optimization and similar U-Net architectures to those from text-to-image diffusion models (Ramesh et al., 2022; Saharia et al., 2022), replacing the spatial convolutions with temporal convolutions. The main limitation of Diffuser is that sampling complete trajectories is computationally expensive, especially since the trajectory has to be sampeed at each timestep (open-loop). Hence, it was demonstrated for environments with oowdeinnnii state-action sets. Our model is able to generate trajectory by
14

---

Published as a conference paper at ICLR 2023
sampling step time (closed loop). Combined with replacing U-Nets with other architectures like MLPs and transformer, this enables learning from high- nigh-dimensional states like images.
Concurrent to our work. Wang et al. (2022) use conditional diffusion model with an MLP predict a single action at a time conditioned on the current state. However, their focus is on trajectory optimization, relying on bootstrapping to learn value function that is used for (classifier-like) guidance, and test their approach in environments with low dimensional state sets. Here, we focus on BC, for which observe guidance can bias the sampling towards low-likelihood trajectories, so we introduce Iwo novel guidance-free conditional sampling methods. In addition, we explore the impact of multiple architectures, including transformers, and demonstrate their performance even under high-dimensional state sets.
Ajay et al. (2023) provide a further concurrent work. They propose to diffuse sequences future states, using a separate rnverse-dyamml model to then infer actions. This contrasts with our own approach of directly diffusion actions, which allows us to scale up to image observations. They also explore conditioning on reward, skills and constraints, finding that CFG can help enforce the conditioning. By contrast, our investigation considered using CFG when the observation is the conditioning variable.
Imitation Learning (II.. This paper lies under the II. paradigm (Hussein et al., 2017), in particular the line of work that frames IL matching the state-action occupancy distribution in the dataset, which is related to Inverse RL (Ho and Ermon, 2016; Ghasemipour et al., 2020; Ke et al., 2020). Indeed, we diffusion models a practical BC approach that optimise a variational bound of the usual BC objective (which can be interpreted as minimising the forward KL. w.r.t. to the ex- perts distribution Ghasemipour al. (2020); Ke al. (2020)), but, discussed in Sec. 2, diffusion models are able to approximate multimodal distributions, surmounting one major limitation of pre- vious BC methods (Ke al., 2020), especially when working with real world datasets from multiple demonstrators.
Finally, similarly Dadashi et al. (2020); Hussenot et al. (2021), compute the Wasserstein distance between the occupancy distributions as an evaluation criteria.
B EXPERIMENTAL DETAILS
B. CLAW ENVIRONMENT
The claw environment consists of seven unique pictures, each with different combination of valid toys to pick. To create the dataset, pick one of the seven images random, and then pick demonstration action from the true p(alo). The final dataset consists of 20000 such demonstrations.
Figure shows full comparison of all methods in the seven different images. Each individual point represents one sample from the trained model when it given the image observation o.
We train all methods for 100 epochs with batch size of 32, using Adam optimizer and decaying learning rate starting from 0.0001. We use 50 diffusion steps for the diffusion models. We set K 10 for K-Means clustering.
B.2 KITCHEN
The seven tasks of interest are: open the microwave ("microwave"), move kettle onto a stove ('ket- tle'), turn on the bottom stove ('bottom burner'), turn on the top stove ('top burner'), turn on light-switch ("light-switeh") ) open a sliding cabinet ('slide cabinet'), open a swinging door ('hinge cabinet').
For each method we train different seeds, and roll out 100 trajectories length 280 for evaluation. K-Means Transformer only has single seed.
Time Wasserstein: For each method the histogram over the number of timesteps taken to complete N tasks is computed. For each tasks, the Wasserstein distance between the the histogram of the model's 's timesteps and the humans is computed. We then average those Wassersteins across all tasks. In the situation where method did not ever complete tasks, we leave the entry blank.
15

---

Published as a conference paper at ICLR 2023
State-based Wasserstein: We are approximating the Wasserstein between the stationary distribu- tions over the state space of the human dataset and our models: ((PAGENTSS) (s), PHUMAN(s)), To do so produce a 'dataset' for each method by concatenating the states for all 100 trajectories (similarly for the humans). We then turn these into the empirical distributions by having uniform mass over each data point. We then use the POT library's emd2 function (Flamary al., 2021) to compute the Wasserstein distance, using the Le cost function between the 30 dimensional entries.
To produce the sub-sampled human dataset, we pick 100 trajectories uniformly without replacement from the human dataset. We do this 5 times.
Density and Coverage: We the Density and Coverage metries proposed in Naeem et al. (2020), and compute them using the code provided by the authors. We use the human dataset as the set of real state samples, and our model's 's rollouts as the fake state samples. We use 10 nearest neighbours for the calculation.
Hyperparameters. Neural network architectures were identical across all methods (except for the final linear layer which allows for differing output dimensionality). Basic MLP and MLP Sieve used three hidden-layers of 512 nodes and GELU activations. Embedding dimension of observation, action, and timestep was fixed at 128. Embedding networks were single hidden-layer MLPs with 128 nodes, using leaky ReLU activations except for the timestep encoder which used Sinusoidal activations. The Transformer used four encoding blocks, each with 16 self-attention heads and an internal embedding dimension of 64. All models receive the previous two observations as history.
Other hyperparameters were kept consistent across methods as far as possible. Initial experimenta- tion showed that methods did not suffer from overfitting, and as default we trained models for 500 epochs with cosine learning rate decay. MLP models used a learning rate of 1e-3 and batchsize of 512, while transformer models used a learning rate of 5e-4 and batchsize of 1024. We set K = 64 for K-means and discretised used 20 bins per action dimension.
We made extensive efforts to bring performance of K-means+residual inline with that reported in Shafiullah et al. (2022). Our best settings used a fixed learning rate of Ie-4, with K 64, though this still gave a task completion rate of 0.34, which falls slightly below 0.44 as originally reported. The difference might be explained by the larger network used by Shafiullah et al. (2022) - layers, and an observation history of 10 steps. Note we also tested hyperparameters as reported in their paper (50 epochs. batchsize of 64), but this provided worse performance for us.
For diffusion models, we set T = 50 and standard 3 schedules linearly decaying in [lc-4, 0.02]. (Note that the variance schedule parameters 0t, ãe and o introduced in Section are derived from this 3 schedule.) For Diffusion-X, MM = 8. For Diff we first sample 100 actions, and fit a Gaussian KDE model (width=0.4) from seikit-learn'.
Models were rolled out with 100 random seeds. Each episode lasted 280 timesteps as in Shafiullah et al. (2022)- - 98% humans completed their assigned four tasks within this time.
B.3 CSGO
Game Score: Frags-per-minute, as reported by Pearce and Zhu (2022).
Action-based Wasserstein: Similarly the Kitchen environment above, we approximat- ing the stationary distribution over the action space of the human dataset and our models: W(PAGENT PHUMAN(A))))
For 1xt itimsstess, we compute this in the same way as the state-based Wasserstein above, using POT's emd2 and the Le cost between the dimensional entries.
For 16xt imesteps and 32xt imesteps, for each timestep we concatenate the 16 (and 32 timesteps respectively) following into single vector. If there are not enough timesteps left in the trajectory. then we do not use that timestep. We compute the Wasserstein in the exact same manner as before on these higher dimensional entries.
Hyperparameters. Neural network architectures were unchanged from the Kitchen experiments, except for the observation encoder, for which we tested two options. 1) A lightweight 6-layer CNN
htttps:/scckitt stabbtl//
16

---

Published as a conference paper at ICLR 2023
that operates directly the frames stacked together. 2) A ResNet18 with ImageNet weights the four stacked frames are passed through this independently. The output is then concatenated together and passed through two vanilla convolutional layers. Following average pooling, embedding vector of length 128 is output.
Unlike in the Kitchen environment, found models sensitive overfitting in CSGO. For the smaller observation encoder, we trained for 500 epochs with cosine decay (learning rate of le-4), but tested models both midway through training (250 epochs) and at the end (500 epochs). For the ResNet encoder, we trained with a fixed learning rate of le-4 for a maximum of 100 epochs, but tested models every 20 epochs. All methods were best with either 60 or 80 training epochs.
For diffuusin models, we set T = 20 and standard 3 schedules linearly decaying in [lc-4, 0.02]. For Diffusion-XX M 16. For K-Mecus+Residual, we use 64.
Following Pearce and Zhu (2022), episodes were defined as 10 minutes of playing time. The original data was collected at 16Hz with the game running real-time. We altered the game setting so that it runs at half speed, allowing our models to roll out at 8Hz. It was possible to increase this to 10Hz for diffusion BC.
B.4 ENERGY-BASEE MODEL IMPLEMENTATION
We re-implemented two versions of the Energy-Based Model's (EBM's) described Implicit BC (Florence et al., 2022): I) Derivative-free optimisation ('Deriv-free') and 2) Langevin MCMC ("Langevin'). Whilst we made efforts to run both of these in all three of our environments (Claw, Kitchen, and CSGO), were unable produce rudimentary results in many combinations but doc- ument our experiences here. Ultimately we reported EBM Deriv-free in Claw (Figure 8), Kitchen (Table & and CSGO (Table 3), and EBM Langevin on the Claw only (Figure 8). Other combi- nations failed to produce any reward.
Serivative-Free Optimisation. This follows Algorithm of (Florence al., 2022). At training time, negative examples are sampled uniformly from the action space, and a contrastive loss is optimised. At sampling time, derivative-free optimisation scheme is used (reminiscent of particle- swarm optimisation).
We used the following hyperparameters by default: number counter examples=256, number infer- ence samples=16000, number optimising iterations=3, optimising noise=0.33, noise decay=0.5.
In our experience this method performed reasonably lower dimensional environments. In the Claw (Figure 8), we could recover fair upproximation of the true p(alo) distributions, although there is a segmentation effect. To understand why this occurs, we plotted the energy mesh over the whole action space in Figure 6, where we can observe halo-like effects at the edges of objects, explaining this.
Figure 6: Visualisation of the energy landscapes learnt on the Claw environment by EBM derivative- free optimisation.
17

---

Published as a conference paper at ICLR 2023
Derivative-fret worked poorly in the kitchen environment which has 9 dimensional action space, only occasionally completing two tasks and never more. This was also observed Florence et al. (2022), which showed that Derivative- failed when action dimensionality exceeded six.
Since CSGO's action space is of dimension three, we thought Derivative-free may perform reason- ably, though we were only able to recover rudimentary performance level in Table 3.
were reasonable (Table 6), generally falling between diffusion models and MSE, though it slowed down training significantly since requires drawing many neg- ative samples for each datapoint in batch - this also sometimes reduced the maximum batchsize that could be used (since it effectively gets multiplied).
Sampling times of Derivative-
Langevin MCMC. This follows the description in Appendix B.3 of Florence et al. (2022), along with the gradient penalty in B.3.I. Negative samples are partially optimised at training time, by backpropagating gradients to the action samples. At sampling time, the same gradient-basee op- timisation procedure is run for longer. Florence al. (2022) state that this method should handle higher-ding sebher-immensionll action spaces well.
We used the following hyperparameters by default: number counter examples=32, noise scale 0.5, learning coefficient start = 0.05, learning coefficient end = 0.005, number meme iterations during training=100,, extra meme iterations during sampling=100, M=1.
We were unable to get this method working well in any of our environments. Figure shows the energy landscapes learnt the Claw. We found it produced energy functions that were overly smooth. Also, at sampling time, many samples would get stuck in local maxima at the edges of action space. fundamental issue with the approach is that the negative samples become less 'negative" over the course training the model learns the energy function. This creates distribution drift in the training data and a less stable optimisation objective.
Figure 7: Visualisation of the energy landscapes learnt on the Claw environment by EBM Langevin MCMC optimisation.
In initial experiments, Langevin failed to solve any tasks in the Kitchen or CSGO. Both training and sampling from Langevin are far slower than any other method (Table 6) training requires running an MCMC chain for many iterations for negative samples (set as default to 100). Sampling requires double this number of steps, and additionally cannot be performed within forch.nogras (), since the gradients are required for optimisation.
18

---

Published as a conference paper at ICLR 2023
C
FURTHER RESULTS
Observation o
p(alo)
MSE
Gaussian
Discretised
K-Means
K-Means Res.
EBM Deriv.-Free
EBM Langevin
Diffusion
Diff.X. (M = 4)
DiffXX. (M 16)
Diff.X. (M 32)
Diff. KDE
Figure 8: Full comparison of different distribution modelling choices in the toy claw environment.
19

---

Published as a conference paper at ICLR 2023
Table Robotic control results. Mean + one standard error over three training runs (100 rollouts). These metric capture task completion only. Methods marked with asterisk (*) are our proposed methods.
No.
Tasks completed
Tacks
Taoks
Tasks
Tasks
Tasks
DDiffissio Rosic *Diflisson Basic MLP KDE
0.98
3.44 0.01
0.99
097 0.00
0.36
0.59 0.01
0.04
DDiffission Rosic MLP Extr88
3.37 + 0.03
0.99 + 0.000
0.93 + 0.01
0.83 - 0.02
0.58 + 0.02
0.04 + 0.01
MSE MLP Sieve
3.13 = 0.07
0.99 + 0.0
0.9 0.03
0.72 043
01. 0.02
0,02 = 0.0
Discretised MLP Sieve Disc K-Means
2,26 :: 0.08
0,85 1: 0.00
0.73 1: 0.02
1: 0.03
0.18 1: 0.02
0:0000
0.33 0.03
0.03
0.0
0.00
Means-Besdaall
Deriv-Free, MLP
*Dillision MLP
3,46 0.04
0.99 0.0
0.94
0.82 0.02
0.02
0.02
*Diftusion KDE MLP Sieve *Diftusion Extra MLP Sieve
3.78 = 0.06
:: 0.0
1.0 =:
0:92 =: 0.01
0.79 0.04
0.07 =: 0.01
3.72 1 0.02
1.0 1 0.0
0.97 1 0.0
0.9 0.01
0.77 1 0.02
0.08 1 0.0
MSE Transformer
3,59 =: 0.02
0.99 :: 0.0
0.96 1: 0.01
0.88 0.01
0.69 0.02
0.07 : 0.0
Discretised Transformer
2.71 0.08
0.97
0.02
0.57 0.04
0.34 0.02
0.02 0.0
K-Means Transfonerr
0.05
DDiftision
0312
0.01
*Dillision KDB TransSemmrr
3,93 0.02
0.9
0,89 6.01
*Difiusion Extra Transformer
3.91 2: 0.04
1.0 0.0
0.99 1 0.01
0.94 0.0
0.88
0.01
0.03
Fuummn
3.98
1.0
1.0
1.0
0.98
0.0
Previously repocted (Shadiillah al., 2022)
Behaviour (K-Menas Residunl)
3.09
0.02
Implicii BC EBMi
(Partuh 2021)
PARROT Normalizing Flow (Singh et ol., 2020)
0.04
0.04
0.0
Table 5: Wasserstein of time to completion in the robotic control environment. Methods marked with asterisk (*) are our proposed methods.
Task 1
Task 2
Task 3
Task
DDiffusiom Eettri ML.P Basic
2.36 +: 0.50
3.65 + 0.26
14,40 + 1.05
14.04 +: 0.38 8.86
*Diffusion Basic
0.94
2.66
23.13
"Diffusinn-KT MI.P Basic
3.52 0.45
5.82. 0.40
14.11
0.76
0.74
"Diffusion MLP Sieve
2.69 0.97
3.84 0.51
8.54 = 2.81
9.18 : 1.12
"Diffusion KDE MLP Sieve
3.50 1. 0.68
7.07 1: 1.18
8.51
2.08
7.98 1 0.82
*Diffusion Extra MLP
MSE ML.P Sieve
3.03 0.71
5,70 +: 1.89
9.63
0.94
7.24 :: 0.91
Discrete MLP Sieve
3.74
11.09 1: 1.92
20.98
1.61
9.38
2.12
K-Means MLP Sieve
34.61 7.21
71.20
20.31
K. Meuns+ Residual MLP Sieve
3.57 0.41
10.07
1. 2.96
17.26 J. 4.07
15.51 1L 1.36
*Diffusion Extra Transformer
2.04 0.64
5.12 1 1.02
5.51
0.56
5.92 0.31
*Dilfusion Transformer
1.26 0.01
3,33 1: 0.36
7,89
0.57
3.97 1 0.51
"Diffusion-KD> Transformee
4.21 0.42
6.34 + 1.25
5.22
0.32
5.33
MSE Transfotmerr
3.03 1. 0.74
5.73 1. 0.61
7.99
2 1.39
6.64 1 1.12
Discrete Transforner
Meuns Transformer
55.79
32.64
Means-Residual Transtormer
2.16 0.78
4.55 +: 1.27
10.07 0.23
14.41 + 1.83
20

---

Published as a conference paper at ICLR 2023
Table 6: Timing analysis for various methods, sampling schemes, architectures, environments and hardware. Sampling times were recorded without running the environment looping over the sam- pling process without waiting for the environment step function. Diffusion models do not suffer from any slow down at training time. For sampling time, note that in the Kitchen experiments, where the denoising network forms the majority of the total network, sampling speed is roughly proportional to number of denoising timesteps. However in CSGO, when a large observation encoder is required, the slow down is less impactful.
Training time epoch
Sampling time J
Max sampling rate
Kitchen enviromment, CPU
MSE, MLP Sieve
7.4 seconds
1.2
833 Ilz
Discrete. MLP Sieve
7.6 seconds
2.7 ms
370 Hz
K-Means, MLP Sieve
8.7 seconds
1.5 ms
666 Hz
K-Means-Residune MLP Sieve
10.5 seconds
1.5 ms
666 Hz
DDiffusion BC (I' 50), MLP Sieve
6.4 seconds
44.6 ms
22 Hzz
*Difftsion-X (2' 50, Mf = 8), MLP Sieve
6.4 seconds
55.9 ms
18 Hz
*Diffusio--KDE (T 50. 100 samples), MLP Sieve
6.4 seconds
128.4 ms
Hz
Kitchen environment, V100 GPU MSE, Basic M1.P Discrete, Basic MLP
1.4 seconds
1.1 ms
909 Hx
1.8 seconds
3.9 ms
256 Hz
K-Means, Basic MI.P
2.2 seconds
1.4 ms
714 Hx
K-Means+Residee Basic MLP
2.4 seconds
1.4 ms
714 Hiz
DDiffusion
50), Basic MLP
1.4 seconds
42.1 ms
24 Hx
Diffussion X (T = 50, Mf = 8), Basic MLP
1.4 seconds
47.3 ms
21 Hk
Diffusion-KDD 50, 100 samples), Basic MLP
1.4 seconds
ms
Hz
MSE, MLP Sieve
1.5 seconds
1.5 ms
Hz
Discrete, ML. Sieve
2.2 seconds
4.3 ms
232 1l
K-Means, Sieve
2.9 seconds
1.8 ms
555 Hz.
K. R..MensstRRiine MLP Sieve
3.1 seconds
2.0 ms
500 Hz
EBM Deriv-Frees MLP Sieve
seconds
28.2 ms
35 Hz.
EBM Langevin, MI.P Sieve
142. I seconds
$25.5 ms
*Diflusion
Sieve
2.0 seconds
63.1 mss
16 Hz.
*Diffusio--X (I = 50. Af = 8), MLP Sieve
2.0 seconds
73.8 ms
14 Hz
*Diffusion
50, 100 samples), Sieve
2.0 seconds
ms
12 Hz.
MSE. Transformer
6.8 seconds
5.5 mss
181 Hz.
Discrese, Transformer
7.0 seconds
8.6 ms
116 Hz
K-Means, Transfommer
7.4 seconds
5.9 mss
Hz.
K-Means+Residdn Transformer
7.5 seconds
6.0 ms
167 Hz
EBM Deriv-Free, Transformer
1000 scconds
1390.8
0.7 Hz
EBM Langevin, Transfomer
1000-; seconds
2927.9 ms
0.3 Hz
Diffusion BC (T 50), Transformer
6.8 scconds
244.2
Hz
Difffusion-X (T = 50, Af = 8). Transformer
6.8 seconds
285.3 ms
Hz
Diffusion-K MDaiiiion-DD(((
50, 100 samples), Transformer
6.8 scconds
295.6
Hz
CSGO enviromment, ResNet18 observation encoder, V100 GPU MSF, MI.P Sieve
49 seconds scconds
5.0 ms
200 Hx
Discrete, MLP Sieve
6.0
167 Hz
ERM Deriv-Free, MI.P Sieve
50 seconds
9.3 ms
107 Hz
EBM Langevin, ML. Sieve
240 seconds
537.2
2 11z
K- X-Means, MI. Sieve
51 seconds
5.5 ms
181 Hz
M.MMeM MLP Sieve
51 seconds
5.4
185 Ilz
Diffiusion BC (T 20), MI.P Sieve
49 seconds
31.5 ms
32 Hz.
Difltssion X (T" 20, M 16), MLP Sieve
49 seconds
54.7
18 Ilz
21

---

Published as a conference paper at ICLR 2023
t = 0. Initialised state. Two modes can be made out in action dimensions (e.g. dim 6),
representing the choice between reaching for the kettle and microwave.
action dim
action dim
action dim
action dim S
action dim
action dim 7
action dim
action dim 9
t == 11. The gripper begins moving towards the kettle and distributions become unimodal.
action dim
action dim 2
action dim 3
action dim
action dim S
action dim 6
action dim 8
action dim
t= 27. The gripper initiates fine grasping of the kettle handle, which requires careful coordination between the grippers, which manifests as strong correlation between action dim 8 & 9.
sction
action
action dim
action dim S
action dim 7
action dim 9
63. Strong linear and non-linear correlations between certain action dimensions can be
observed as the gripper pulls away from the kettle.
action dim
action
action
dim
action
action dim
action dim 8
dim
I 87. Kettle task completed. Bimodal distributions can be seen in action dimension 1 & 4, representing the decision point for the selection of the next task.
action dim
action dim
action dim
action dim
action
action dim 6
action dim 8:
action
f = 96. The decision to move towards the bottom burner has been made, so action distributions become unimodal again.
action
action dim
action dim
action dim
action dim 5
action dim 6
action dim
8
action dim 9
Figure 9: Qualitative analysis of demonstrative roll-out for Diffusion BC in the kitchen environ- ment. We visualise the relationships between all 9 action dimensions simultaneously via ma- trix scatter plot, after 200 samples have been drawn from the diffusion model each environment timestep. This allows visualisation of correlations and multimodality between action dimensions that captured by the diffusion model. We have hand-selected six interesting environment timesteps in the roll-out to visualise. Action dimensions correspond to specific joint actuators of the robot, e.g. turn base left/right, or close gripper.
22

---

Published as a conference paper at ICLR 2023
D SAMPLING ALGORITHMS
Algorithm Sampling for Diffusion BC
1: ar (., I)
2:
T. T..... do 3: Z ~~ NV if > 1, else
4: ay-1 -ot
1-0- Vi-u,
co T, 0)) +
var
5: end for
6: return 80
Algorithm 2 Sampling for Diffusion-X
I: ar ~ I.)
2: for i = M do
max(i.
((, if
5: ar-1 vay [a,
Vi-a, colar, T, 0)) + 0,,
I-n
6: end for
return a.
Algorithm 3 Sampling for Diffusion-KII KDE
1: +-
2: for i =1.==...... K do
Use Algorithm to sample action, an.
A .append(ao)
5: end for
6: KDEmodel. (At
7: Likelihoods KDEmodel. A)
arg max, (Likelihoods)
return A
E CLASSIFIER-FREB GUIDANCE ANALYSIS
p(a 0o)
Tumilft
Shright
Right
p(a (O
03
Left
Straightt
Right
0
O2
p(a 02)
Turn tart
Stuigh
Turn tett
p(a 03))
Figure 10: Didactic example of when CFG can lead to generation of less usual trajectories.
23

---

Published as a conference paper at ICLR 2023
CFG can be interpreted as guiding the denoising sampling procedure towards higher values of p (a) (Ho et al., 2022). Given this interpretation, we now provide a grid-world to show concretely why this can lead to sampling of less common trajectories in sequential environment.
Figure 10 shows an environment with four discrete states and a discrete action. The action space allows three ego-centric options; turn left, turn right continue straight forward. Agents always initialised at state 0, giving observation Og, and rolled out for exactly two timesteps, visiting state and ending in state 2 or 3. The figure shows the empirical action distributions demonstration dataset. In states O, 2 & 3, the agent always selects straight. But in state I, the agent makes right turn with 0.1 probability.
Let 01. and denote the observations given by states 1, 2 and 3, respectively. We can apply Bayes rule to find p(oja), which will provide an understanding of what behaviour CFG induces. We are interested in the learnt behaviour at the decision point 01,
p a) = p(a o1 (((11)
(4)
p(a)
Given the agent starts state and is rolled out for two timesteps (sees three states), and P(a Turn right(o;) 0.1, p(a straight = 0.9, find the marginal probability of each observation,
p 0p) = 1/3
(5)
( 01) = -1/3
(6)
p( 02) =1/3 0.1
(7)
p(o 03) = 1/3-0.9.
(8)
Hence, the marginal action distribution is,
3
p(a- = Turn right) = Σ
p(a Turn right|o.)p(o:)
(9)
1=0
= =1/3-00
(10)
p(a Straight) 1/3 1/3 0.9
(11)
We can now compute the quantities p(o) (a),
= p(a Turn right o )(o))) 0.1 1/3
p(o la Turn right)
(12)
p(a Turn right)
0.1
1/3
p(o a Straight) p(a or )p()))
0.9 1/3
0.9
0.31.
(13)
p(a Straight)
1/3 1/3 0.9 1/3
2.9
As such, since CFG favours actions that maximise p(ola), the CFG agent will select the less fre- quently visited right-hand path more often.
24


Reviews:
[{"review_id": "ijv_xxa4aY", "reviewer": "Reviewer_hanC", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes to use diffusion models to imitate human behaviors. To make text-to-image model suitable for observation-to-action prediction, the authors introduce several modifications, including designing different architectures (MLP with embedding + skip connection, and  Transformer-based), removing Classifier-Free Guidance (CFG), and developing reliable sampling schemes. Experiments of task execution in a kitchen environment and playing CSGO show the effectiveness of the proposed method.\n\n----\n_After rebuttal_: The authors have moderately answered the questions and provided additional experiments. Besides, it seems that no critical issues were raised by other reviewers. Given the above, the reviewer would keep the original score (8).", "summary_of_the_review": "- Overall, this is a well-written paper with clear ideas and good designs.\n- More evaluation metrics and comparison methods could be added.", "strength_and_weaknesses": "###### Pros\n- This paper provides (as far as the reviewer knows) one of the first detailed and substantially improved research on applying diffusion models to imitation human demonstrations.\n- The toy experiments and discussions with the arcade claw game (with the figures) provide clear comparisons and convincing motivations.\n- The proposed modifications to standard diffusion models are reasonable and shown to be suitable for sequential environments.\n- The experiments of the kitchen environment with complex joint action space are convincing.\n\n###### Cons\n- The authors compared several methods (i.e., MLP, MLP Sieve, and Transformer), but these are more like ablated models. To demonstrate the superiority of diffusion models for observation-to-action modeling, comparisons to other generative models like GAN and flow models are helpful. Other state-of-the-art methods, like parrot [1] using realNVP for imitation, could also be considered.\n- CSGO experiment is catchy but less convincing than the rest. The results for \"16 x timesteps\" and \"32 x timesteps\" metrics indicate all compared methods are far from satisfactory. More analyses and more baselines should be given for Table 3.\n- Metrics like training time, evaluation time, or GPU consumption during training could better demonstrate the superiority of the proposed method.\n- Does CFG not work due to the lack of data? Compared to text-to-image tasks with massive data, the observation-to-action task only has 566 trajectories with 25 task sequences. The conditions and reasons for CFG not working should be discussed rather than simply declaring that CFG is unsuitable for these tasks.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "RXTU3jj6hY0", "reviewer": "Reviewer_PK6N", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper attempts to use diffusion models for action generation in imitation learning set ups. The paper contributes :\n1. modifications to diffusion models to adapt them to action generation\n2. discussion on effectiveness of CFG\n3. sampling techniques for action generation\nThe paper tests on a few datasets including robot manipulation in kitchen and video games.", "summary_of_the_review": "This paper presents a novel idea and is well written. Discussion on how key insights can be derived from experimental results can be improved. ", "strength_and_weaknesses": "Strengths:\n- Well written, and arguments well supported\n- Detailed results, varying baselines implemented for tasks \n- Core idea is relatively novel and has a lot of potential\n\nWeaknesses:\n- Some implementation details could be improved.\n- Key insight - \" difussion models avoid BC shortcomings by making no model assumptions\", doesn't seem well substantiated from the results as stated.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "MOCLr4-yHI", "reviewer": "Reviewer_zdNu", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the use of diffusion models as observation-to-action models in the context of behavioral cloning. In particular, different version of DMs are explored for three different environments (claw arcade game, robotic control, counter strike). The experiments are technically sound and employing diffusion models for imitations learning shows promising results. ", "summary_of_the_review": "Overall, this is an interesting paper that can be accepted to ICLR 2023. The paper is very well written, the experiments are sound, and the problem the proposed method addressees is important. Using diffusion models for behavioral cloning is an interesting and new direction and the paper also has the potential to stimulate future work in this direction. \n\nHowever, there is also a number of shortcomings that I would like to see addressed prior to acceptance. For one, I am missing a comparison to other generative approaches as baselines. I was somewhat surprised that no other learning-based approach was used as a baseline for the arcade claw game experiment. What is the reason for that? Second, for the claw experiment the paper fails to provide a discussion on how human demonstrations where captured and what exactly the action space is. The appendix states that samples were drawn from the true p(a|o). So here no human demonstrations where used at all? \n\nGiven the paper only discusses environments with rather constraint action spaces, I find that the title 'Imitating Human Behaviour' is quite misleading and should be changed (e.g. to 'Imitating Human Demonstrations'). That the method is able to imitate human behavior arguably is a too strong statement that should be toned down throughout the paper (e.g. in Section 5). \n", "strength_and_weaknesses": "Strengths: \n\n- The paper addresses an important problem with a sound technical solution. \n- The paper is very well written and easy to read. \n- The paper discusses a set of meaningful experiments that show the benefits of employing diffusion models.\n- The comparison of different modeling choices provides an insightful summary.\n\nWeaknesses:\n\n- The discussed baselines (MSE, Gaussian, etc.) for behavior cloning appear rather weak and a comparison to other generative approaches should be provided. \n- The paper makes too strong claims towards modeling 'human behaviour' that should be toned down. \n- The paper does not discuss any limitations. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "2ImR4frU5E", "reviewer": "Reviewer_Ygtk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes to use diffusion models to model human behavior imitation policies learned with behavior cloning. The motivation is that human behaviors can be complex and multi-modal, and simple Gaussian or categorical policies are unable to capture the underlying multi-modal distributions. The paper utilizes DDPM as the diffusion process and proposes simple architectures to denoise the action distribution. A stable sampling scheme is also proposed to avoid bad actions. Experiments are performed on robot control tasks and first-person shooting games, which demonstrate that the method outperforms the baseline action distribution models across various metrics and achieves better coverage of demonstrations. The paper also discussed and showed through experiments that classifier-free guidance is not suitable for modeling the policy, since it may blend in actions from unrelated states.", "summary_of_the_review": "I appreciate the idea of using diffusion models for imitation learning since they can capture multi-modal behaviors better. The technical novelty is not huge, but the experiments are extensive and the method can serve as a good baseline for diffusion-based imitation learning. I’m leaning toward acceptance at this point.", "strength_and_weaknesses": "**Strength:**\n\n- The motivation for using diffusion models for imitation learning is illustrated well with the toy task, where the limitation of standard action models is clear.\n- The proposed sampling methods are effective in improving task completion. The paper also shows that classifier-free guidance is harmful to policy modeling, both through experiments and didactic examples.\n- Experiments are performed on semi-realistic environments such as simulated robotic control and photorealistic video games, which clearly demonstrated the effectiveness of the proposed approach. Metrics are also comprehensive, evaluating both the quality and coverage of the produced actions.\n\n**Weakness:**\n\n- The paper listed the network architecture as the first contribution, but it’s similar to prior work on diffusion models in essence and appears incremental.\n- It would be great to visualize the learned diffusion policies and demonstrate different modes of action produced by the policy (for robotic control and CSGO), which seems to be the main purpose of using diffusion models.\n- Some notations are confusing, especially the use of $t$ since it normally denotes the RL timesteps. I would recommend using another symbol for the diffusion process.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster