Source-linked AI summary
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
TL;DR
The paper addresses the limited expressiveness and bias of standard behaviour-cloning models for stochastic, multimodal human behaviour. It applies diffusion models as observation-to-action models, adding architectural, guidance, and sampling adaptations for sequential environments. Diffusion models closely match human demonstrations in simulated robotics and modern 3D gameplay, including an 89% task completion rate versus 44% for a recent state-of-the-art method.
Problem
Standard behaviour-cloning choices approximate complex human action distributions, potentially producing average, biased, or uncoordinated cloned policies.
Method
The paper uses diffusion models to learn expressive conditional joint action distributions and adapts them to sequential environments with new architectures, guidance analysis, and sampling schemes.
Results
Diffusion models closely match human demonstrations across simulated robotics and 3D gameplay, achieving an 89% task completion rate versus 44% for a recent state-of-the-art method.
Takeaways & Limitations
Diffusion models improve matching demonstrations in reward and distribution, while Diffusion-X and Diffusion-KDE provide reliable-sampling benefits and CFG should be avoided for observation conditioning.
Takeaways & Limitations
MSE cannot capture variance or multimodality, while independent discretisation can ignore dependencies between action dimensions.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Human behaviour contains diverse, multimodal trajectories that simple behaviour-cloning models may fail to reproduce faithfully. The paper proposes diffusion models to learn full action distributions and reports strong results across robotic control and 3D gameplay.
- Human-AI collaboration benefits from models that respond to plausible human behaviours, especially as complex environments make enumerating all behaviours difficult.
- Behaviour Cloning trains agents to mimic demonstrator actions from offline observation-action datasets and has worked across driving, robotics, and games.
- MSE and independently discretised actions approximate the distribution, encouraging average or uncoordinated policies instead of modelling the full joint action distribution.
- Simple models may be insufficient for human behaviour because demonstrations contain diverse trajectories and multimodality at decision points.
- Diffusion models target complex action distributions, outperform state-of-the-art methods on simulated robotics, and scale to Counter-Strike: Global Offensive gameplay.
- The paper adapts diffusion models to sequential environments through architecture design, guidance analysis, and reliable sampling schemes.
2 MODELLING CHOICES FOR BEHAVIOURAL CLONING
Common behaviour-cloning methods simplify action-distribution modelling through point estimates, independent bins, or finite joint-action clusters. These approximations create averaging, quantisation, dependency, and hyperparameter-selection limitations that motivate diffusion models.
- The claw-machine experiment uses top-down toy images and simultaneous 2D continuous actions, with demonstrations generated from randomly selected valid toy locations.
- MSE: MSE produces a point estimate that cannot capture variance or multimodality and may predict an average or out-of-distribution action between modes.
- Discretised: Independent discretisation introduces quantisation error and learns marginal rather than joint action distributions, causing sampled actions to ignore dependencies between dimensions.
- K-Means: K-Means models joint actions without enumerating all combinations, but its quantisation can be severe and its choice of K critically affects performance.
- K-Means+Residual: K-Means+Residual improves distribution fidelity but still uses at most K point estimates, averages actions within bins, and requires careful K selection.
- Diffusion models: Diffusion models are proposed to learn expressive conditional joint distributions without the approximations used by existing modelling choices.
3 OBSERVATION-TO-ACTION DIFFUSION MODELS
This section adapts diffusion models to observation-to-action sequential environments, covering their denoising process, architectures, guidance, and sampling schemes. The proposed designs address action-vector generation and improve reliability by modifying architecture and sampling.
- 3.1 Diffusion Model Overview: Diffusion models iteratively denoise Gaussian noise into clean action samples, optionally conditioned on an observation.Sampling begins with a_T ∼ N(0, I) and produces progressively denoised actions through T steps.
- 3.1 Diffusion Model Overview: During training, a neural network learns to predict noise added to noisy observation-action inputs across denoising timesteps.The training objective averages over timesteps and observation-action pairs from the demonstration dataset.
- 3.2 Architectural Design: The observation-to-action network takes a noisy action, denoising timestep, and observation as input, then outputs a predicted noise mask.The architecture can incorporate an observation history.
- 3.2 Architectural Design: Diffusion BC generates an action vector conditioned on an observation, unlike text-to-image diffusion models, which generate images conditioned on vectors.The observation may be an image, while the generated output is an action vector.
- 3.2 Architectural Design: Performance improves from Basic MLP to MLP Sieve to Transformer, while separating observation encoding from denoising enables 8Hz rollout in a pixel-based video game.The separation allows one observation-encoder pass and repeated passes through a lighter denoising network.
- 3.3 Why Classifier-Free Guidance Fails: Classifier-Free Guidance can degrade sequential performance by favoring less common observation-specific trajectories rather than higher-likelihood actions.Diffusion-X and Diffusion-KDE instead modify sampling to reduce bad actions while preserving the two distinct modes shown in the predictive distributions.
- 3.4 Reliable Sampling Schemes: Diffusion-X performs extra denoising iterations, whereas Diffusion-KDE samples multiple actions, fits a kernel-density estimator, and selects the highest-likelihood action.Figure 4 distinguishes Diffusion BC at M = 0 from Diffusion-X at M > 0 and identifies Diffusion-KDE separately.
4 EXPERIMENTS
The experiments evaluate diffusion-based behavior cloning against established baselines in robotic kitchen control and CSGO, measuring both task-level outcomes and similarity to human behavior. Diffusion methods generally match demonstrations more closely, while architecture and sampling choices trade performance against inference speed.
- Evaluation: The evaluation compares models with human demonstrations using high-level outcomes and low-level Wasserstein distances over states, actions, task sequences, and completion times.These complementary metrics assess both observable performance and distributional similarity.
- Robotic control: Diffusion BC approaches outperform all baselines on task-completion rate, ordered Diffusion BC<Diffusion X<Diffusion-KDE.The authors associate this ordering with more reliable avoidance of bad actions.
- Efficiency: Diffusion sampling improves behavioral matching but reduces inference speed, with MLP Sieve rates decreasing from 16 Hz for Diffusion BC to 14 Hz for Diffusion-X and 12 Hz for Diffusion-KDE.The paper also reports Diffusion BC sampling rates of 24 Hz for MLP Basic, 16 Hz for MLP Sieve, and 4 Hz for Transformer.
- Robotic control: Across Wasserstein metrics, diffusion models outperform other baselines, while Diffusion-KDE has higher Density but lower Coverage than the other sampling schemes.The authors interpret this as Diffusion-KDE completing more tasks while reproducing less diversity from the demonstration distribution.
- Ablations: Architecture metrics usually improve in the order Basic MLP<MLP Sieve<Transformer.The three architectures were evaluated across the baseline and diffusion variants.
- Classifier-Free Guidance: CFG reduces four-task completion from 0.63 at w = 0 to 0.08 at w = 8 and increases Bottom Burner as the first task from 7% to 25%.Human demonstrators select Bottom Burner first 10% of the time.
- Video-game evaluation: In CSGO, Diffusion-X performs best across both observation encoders for game score and Wasserstein distance to human action distributions, running at 18 Hz versus 200 Hz for MSE.Training time is reported as similar between methods.
5 DISCUSSION & CONCLUSION
The paper argues that diffusion models fit complex human observation-to-action distributions because they avoid coarse approximations, while acknowledging inference-speed, hyperparameter, and temporal-correlation limitations. Experiments show improved demonstration matching and strong task performance across control and image-based settings.
- Diffusion models avoid coarse approximations on the target action distribution, addressing biases and limitations introduced by common behaviour-cloning models.The paper presents this as the central reason diffusion models suit complex human behaviour.
- Diffusion models sample more slowly, introduce additional hyperparameters, and address only single-timestep action-distribution complexity rather than temporal correlations.MSE samples at 666 Hz and 200 Hz, compared with Diffusion BC at 16 Hz and 32 Hz in the kitchen and CSGO environments.
- Diffusion models improved matching of human demonstrations in reward and distribution compared with other methods.
- Diffusion-X and Diffusion-KDE provided benefits over Diffusion BC through reliable sampling schemes.
- Good architecture design enabled a trade-off between performance and sampling speed, while CFG should be avoided for observation conditioning.
- 89% task completion exceeded the recent state-of-the-art result of 44% on a complex control task.
A EXTENDED RELATED WORK
The related work positions this paper as an investigation of diffusion-based behaviour cloning that directly models actions and compares architectures, sampling schemes, and guidance. It distinguishes this approach from discretised, autoregressive, energy-based, and trajectory-diffusion alternatives.
- The paper contrasts diffusion behaviour cloning with GATO and K-Means+Residual, whose discretisation introduces quantisation or finite-policy approximations.
- The paper evaluates diffusion models for behaviour cloning using different neural architectures and sampling schemes, aiming to model policies without implicit approximations.
- Guidance can bias sampling toward low-likelihood trajectories, so the paper introduces two guidance-free methods for high-quality, high-likelihood samples.
- Closed-loop one-step trajectory generation, combined with MLPs or transformers instead of U-Nets, enables learning from high-dimensional image states.
- Unlike sequence-diffusion work that predicts future states, this approach diffuses actions directly and can scale to image observations.
- The work uses diffusion models as a practical behaviour-cloning method within imitation learning and evaluates occupancy similarity with Wasserstein distance.
B.1 CLAW ENVIRONMENT
The claw environment uses image observations and continuous two-dimensional point actions, with demonstrations sampled from a known conditional action distribution. The dataset contains 20,000 demonstrations, and Figure 8 compares model samples across seven observations.
- The dataset contains 20,000 demonstrations generated from seven images with different combinations of valid toys.
- Figure 8 compares methods across seven images, with each point representing one trained-model action sample conditioned on the image observation.
- All methods were trained for 100 epochs with batch size 32; diffusion models used 50 diffusion steps and K-Means used K = 10.
B.2 KITCHEN
The experiments evaluate methods on seven kitchen tasks and an image-based CSGO setting using trajectory rollouts and distributional metrics. They document baseline calibration, Wasserstein-based evaluation, image encoders, and overfitting sensitivity.
- B.2 KITCHEN: The kitchen benchmark comprises seven manipulation tasks, including opening appliances and cabinets, moving a kettle, and activating burners or a light switch.
- B.2 KITCHEN: Time and state-based Wasserstein distances compare model rollouts with human distributions over completion times and states.
- B.2 KITCHEN: Density and Coverage use human states as real samples and model rollouts as generated samples, with 10 nearest neighbours.
- B.2 KITCHEN: K-means+residual achieved a task completion rate of 0.34, below the originally reported 0.44 despite efforts to match the reference settings.
- B.2 KITCHEN: Models were evaluated with 100 trajectories of length 280, while 98% of humans completed their assigned four tasks within that horizon.
- B.2 KITCHEN: Action-based Wasserstein distances compare stationary action distributions at horizons of 1, 16, and 32 timesteps.
- B.2 KITCHEN: CSGO used either a lightweight CNN or ImageNet-weighted ResNet18 observation encoder, and models were generally best after 60 or 80 epochs.
- B.2 KITCHEN: CSGO episodes lasted 10 minutes, with models rolled out at 8 Hz and diffusion behaviour cloning reaching 10 Hz.
B.4 ENERGY-BASED MODEL IMPLEMENTATION
The paper reimplements derivative-free and Langevin EBM variants, finding derivative-free more workable in low-dimensional settings but problematic as action dimensionality increases. Langevin failed broadly and incurred substantial training and sampling costs.
- Derivative-Free Optimisation: Derivative-free EBM recovered a fair approximation of true p(a|o) distributions in Claw, though segmentation and halo-like edge effects remained.The energy mesh suggested halo-like effects around objects as an explanation for the segmentation effect.
- Derivative-Free Optimisation: Derivative-free performed poorly in Kitchen’s 9-dimensional action space, occasionally completing two tasks but never more.The paper also reports rudimentary performance in CSGO despite its three-dimensional action space.
- Derivative-Free Optimisation: Derivative-free sampling times were generally between diffusion models and MSE, but its many negative samples slowed training and sometimes reduced maximum batch size.The default configuration used 256 counterexamples and 16,000 inference samples.
- Langevin MCMC: Langevin failed to work well across the environments, producing overly smooth energy functions and samples stuck at local maxima near action-space edges.The paper attributes training drift partly to negative samples becoming less negative as the energy function is learned.
- Langevin MCMC: Langevin EBM training and sampling were substantially slower because both procedures required many MCMC iterations.Training used 100 MCMC iterations for negative samples, with 100 additional iterations during sampling by default.
C FURTHER RESULTS
Further results examine distribution quality, task completion, timing, and qualitative action structure. Diffusion sampling procedures include standard sampling, Diffusion-X, and Diffusion-KDE variants, while kitchen visualisations show changing multimodality and coordination across action dimensions.
- Further Results: Figure 8 provides a full comparison of distribution-modelling choices in the toy Claw environment.
- Robotic Control: Robotic-control results report task completion means with standard errors over three training runs and 100 rollouts.The table measures task completion only and marks proposed methods with an asterisk.
- Kitchen Qualitative Analysis: At t = 0, some action dimensions show two modes representing the choice between reaching for the kettle and microwave.
- Kitchen Qualitative Analysis: At t = 11, the gripper moves toward the kettle and the action distributions become unimodal.
- Sampling Algorithms: Diffusion BC sampling starts from Gaussian action noise and iteratively denoises it, while Diffusion-X extends the loop and Diffusion-KDE selects among sampled actions using fitted likelihoods.
E CLASSIFIER-FREE GUIDANCE ANALYSIS
The classifier-free guidance analysis interprets guidance as favoring actions with higher p(o|a). In a two-step grid-world, this can make the agent select a less frequently visited right-hand path more often.
- Guidance Interpretation: CFG guides denoising toward actions with higher p(o|a), linking the sampling procedure to observation likelihood.
- Grid-World Analysis: The grid-world contains four states, three discrete actions, and two-step rollouts whose demonstrations define state-dependent action distributions.
- Grid-World Analysis: Bayes-rule analysis evaluates p(o|a) at the decision point where right and straight have probabilities 0.1 and 0.9, respectively.
- Implication of Guidance: CFG favors actions maximizing p(o|a), so the agent selects the less frequently visited right-hand path more often.