Skip to the lesson
FixTheRobot Explore Robotics / Robot Learning / M3 Imitation Learning / M3-T6

Robot Learning · Theory · M3

M3-T6 — Multimodal Actions and Beyond Mean Squared Error

Mean squared error predicts the average action. When two different actions are both right, the average can be wrong.

  • Lesson 6 of 6
  • Module 3: Imitation Learning
  • Theory
  • About 13 minutes to read

Cloning as regression (lesson M3-T2) minimises the mean squared error between predicted and demonstrated actions. That is the right loss when the demonstrator is a function of the observation plus a little noise. It is the wrong one when the demonstrations contain two different good ways to act in the same situation, such as going left or right round an obstacle: the best prediction under squared error is their average, which may be a third action that neither would take. This lesson shows the failure on Scout, computes it, and surveys the remedies, from the cheapest (curate the data) to the heaviest (generative policies, among them the diffusion policies and action-chunking methods now common in robot manipulation), together with what each costs.

Learning objectives

After this lesson, you should be able to:

  • Explain why the minimiser of squared error is the conditional mean, and why that fails for multimodal actions.
  • Compute the averaged action at the pillar line for two operators and say what the resulting policy does.
  • Compare four remedies: choose one mode, classify discretised actions, mixture density outputs and conditional generative policies.
  • Choose a remedy for a given task by its cost, and say what each one needs from the data.

The intuition

Two drivers meet a parked car in the road. One passes on the left, the other on the right, and both are right. A third person watching both and asked “what do drivers do here?” who answers with the average of their steering is saying: drive straight into the car. Neither driver ever did that, and it is the one move that is certainly wrong.

Squared error is that third person. For a given observation, the single prediction that minimises the expected squared error over the demonstrators’ actions is their mean. When the actions cluster in one place the mean is a good summary; when they cluster in two places on either side of an obstacle, it lies in the gap between them, which is where the obstacle is.

Two operators, two good paths, one bad averageA top-down view of the hall with a round pillar in the middle. One operator goes round the pillar above it and one below it, both reaching the goal on the right. The average of the two action choices is a straight line into the pillar, shown dashed and ending at the pillar. goal start on the pillar line operator A: passes above operator B: passes below mean of A and Bpillar
Scout at the pillar line. Operator A (blue) passes above the pillar and operator B (green) passes below; both reach the goal. The average of their actions at the line points straight at the pillar (orange).

The mathematics

Why squared error gives the mean

For one observation o, suppose the demonstrated actions a follow a distribution p(a | o). The expected squared error of a prediction â is its variance around the mean plus the squared distance of â from the mean, so it is smallest at the mean.

Squared error is minimised by the conditional mean
E[∥a^−a ∥2 ∣o]=E[∥a−μ(o) ∥2 ∣o]+∥a^−μ(o) ∥2E[\lVert \hat{a}-a\,\rVert ^{2}\,\mid o]=E[\lVert a-\mu (o)\,\rVert ^{2}\,\mid o]+\lVert \hat{a}-\mu (o)\,\rVert ^{2}
μ(o)=E[a ∣o]\mu (o)=E[a\,\mid o]

A flexible model trained to the minimum of the loss on enough data therefore learns μ(o), whatever shape p has. If p has one peak, the mean sits on it. If p has two peaks of weight ½, at a₁ and a₂, then:

Two equal modes
μ(o)=a1+a22\mu (o)=\frac{a_{1}+a_{2}}{2}
loss at μ: ∥a1−a2 ∥24\text{loss at}\,\mu :\,\frac{\lVert a_{1}-a_{2}\,\rVert ^{2}}{4}
loss at a₁ (picking one mode): ‖ a₁ − a₂ ‖² / 2

The loss prefers the mean, at half the expected cost of committing to either mode. The training objective is satisfied by the policy that fails. The same is true of any loss minimised by a mean or a median; the cure is not a smaller step size or more data but a different model of the action, a different loss, or different data.

Why it matters for closed-loop behaviour

A policy that is wrong in a state does not just score badly there; it changes the next state. The mean action at the pillar line is a command to move towards the pillar, which keeps the robot on the line. If the goal pull and the pillar repulsion happen to cancel, the robot stops (a freeze); if they do not, it drives in (a crash). The training loss is the same either way.

Four remedies

Choose a mode
Curate the data so each situation has one way of acting: one operator, or one rule for which side to pass.
Classify
Cut the action into bins (directions, or clusters found in the data), predict a probability for each bin with a classifier, then take the most likely bin or sample one. Optionally add a small continuous correction within the bin.
Mixture density
Predict a few Gaussians, each with a weight, a mean and a spread, trained by maximising likelihood. Sample or take the heaviest component (Bishop, 1994).
Conditional generative
Learn to sample actions given the observation: a diffusion model, a conditional variational autoencoder, or an energy-based model. The network need not write down a density; it learns to draw from it.
RemedyHow it handles two modesCost and risks
Choose a modeThere is only one: the data has no conflictNeeds control over the data, or someone to decide the rule; throws away demonstrations; the rule must be applied consistently (lab M3-L4)
Classification over binsOne bin per mode; pick the most likely or sampleDiscretisation error; many bins for a high-dimensional action; picking the argmax each step can flip between modes when they are near equal
Mixture density outputsOne Gaussian per modeThe number of components is chosen by hand; likelihood training is delicate (collapsed or exploding variances); poor for many modes or high dimension
Conditional generative modelsSample from the learned distribution, each draw a coherent actionHeavier to train and to run (several network evaluations per action for diffusion); more hyperparameters; results depend on how the action is represented

Diffusion policies and action chunking

Two ideas dominate modern imitation for manipulation, and both address multimodality, though differently.

Diffusion policies (Chi et al., “Diffusion Policy”, RSS 2023) treat the action sequence as the output of a denoising process. Training adds noise to demonstrated actions and teaches a network to remove it, conditioned on the observation. At run time the policy starts from random noise and removes it step by step, which produces a draw from the learned distribution of actions: different noise gives different modes, and each draw is one coherent action rather than an average. The network is trained with a plain regression loss, but on noise, so it does not collapse the modes. The cost is that producing one action needs several network passes, which is why speed matters on a robot.

Action chunking (Zhao et al., the Action Chunking with Transformers method “ACT”, RSS 2023) predicts a short sequence of future actions at once instead of one at a time, and executes it before predicting again (the original method averages overlapping chunks, which it calls temporal ensembling). Predicting a chunk reduces the effective horizon, which helps with compounding error, and it commits to one coherent plan over the next steps, which prevents flipping between modes from step to step. ACT is trained as a conditional variational autoencoder, so a latent variable can carry the choice between demonstrations’ styles.

The common thread is commitment: each decision is a sample from the distribution, and that sample carries on over several steps. The example below shows why that matters, with two simple stand-ins rather than the real algorithms.

Large pretrained models that act from images and language build on these ideas; see the VLA track and the foundation-models lesson of Module 10 for how.

A worked example

All numbers come from running RLCore in node. Two scripted operators copy the expert, but one always passes above the pillar (y greater than 3) and the other always below. They are the same function with the side fixed, so each succeeds on its own. All results are on 100 starts from seed 500 in the 'near' region (close to the pillar line) unless stated.

The average at the pillar line. Take Scout at (3.5, 3.0), on the line through the pillar and 1.5 m from its centre. Operator A commands (0.533, 0.622) and operator B commands (0.533, −0.622), one going up and one going down. If the data has them equally often at this observation, the minimiser of squared error is the mean:

Mean action at (3.5, 3.0)
μ=12(0.533, 0.622)+12(0.533, −0.622)=(0.533, 0.000)\mu =\tfrac{1}{2}(0.533,\,0.622)+\tfrac{1}{2}(0.533,\,-0.622)=(0.533,\,0.000)

That is a command to move straight east at 0.53 of top speed, towards the pillar. Its squared error against either operator is 0.387, against 0 for the right mode and 1.549 for the wrong one if a mode is picked at random: the expected loss of picking a mode at random is 0.774, double that of the mean. The loss prefers the mean.

What the averaged policy does. Using that conditional mean at every observation (the best that MSE can do with unlimited data and a model flexible enough to fit it), from a start at (1, 3.0) Scout drives east at full speed to x = 3.25, slows as the pillar term of the command turns against the goal pull, and stops at (3.90, 3.00), 1.10 m from the pillar’s centre and 0.15 m outside the crash circle. It stays there for the 80 steps: return −4.00. From 20 starts on the line it succeeds in 0, crashes in 0 and freezes in 20. From the 100 near starts it succeeds in 3, crashes in 34 and times out in 63 (mean return −4.33); on all starts it succeeds in 53, crashes in 21 and times out in 26 (return 2.14).

The remedies on the same task.

PolicySuccessCrashesTimeoutsMean return
Conditional mean of A and B (MSE ideal)33463−4.33
Smoothing regression on 100 demonstrations, bandwidth 0.395507.65
Same, bandwidth 0.6782205.29
Same, bandwidth 1.04440160.78
Choose a mode: all demonstrations from A100008.37
Perfect mixture, one random operator per step31690−1.63
Perfect mixture, one random operator per episode100008.35
Reference: the original expert100008.43

Averaging is a matter of degree. The “smoothing regression” is a Gaussian-kernel average over the 3,272 pairs of 100 demonstrations (49 by A, 51 by B). The wider the kernel the more it blends the two modes: success falls from 95 to 78 to 44 as the bandwidth goes from 0.3 to 1.0 of the observation distance. Models that are very smooth, or that are trained with early stopping or heavy regularisation, average more than models that interpolate their data.

Choosing a mode costs half the data and fixes it. Using only A’s demonstrations gives 100 of 100. It is the fastest remedy and the one lab M3-L4 uses, but it requires agreeing on the rule and discards what the other operator showed.

Sampling is not enough; commitment matters. A perfect mixture model that draws A or B at random at every step succeeds in only 31 of 100 near starts and crashes in 69: the robot zig-zags, each draw undoing the last, and it ends up in the middle of the gap. The same mixture drawn once per episode (a latent choice, as in a conditional VAE) succeeds in 100. Real diffusion policies and chunked policies commit across several steps for this reason. This is a stand-in, not a diffusion model or ACT: I have not implemented either here.

A caution on finite data. A 7-nearest-neighbour clone trained on 40 demonstrations by A and B also succeeds on 100 of 100 near starts, and a classifier that takes the commonest of 8 heading bins among the neighbours succeeds on 99. Neither falls into the trap here, because with finite data the neighbours near a given state are rarely split exactly 50 to 50 and the robot’s first small deviation settles which operator’s data dominates. The trap appears when the data is dense enough that the two modes really overlap, as in the ideal-mean policy, or when the model is smooth enough to blend them. Multimodality is a risk to check for, not a certainty.

What the simulation simplifies. The two modes are exactly symmetric and exactly equal in weight, and the observation says nothing about which operator is acting, so the choice is a pure coin flip. In practice, operators differ by habit, mood or time, and a hidden variable such as “who recorded this” may explain the split and can be added to the input. The real systems named above were evaluated on image-based manipulation tasks, not here.

Engineering implications and limitations

Check for multimodality before choosing a method

  • Look for conflicts. Compute the action variance in nearby states (lesson M3-T5). A few states with a large value may mark a modal decision such as passing left or right.
  • Cluster the actions in those states. Two tight clusters far apart mean modes; one wide cloud means noise.
  • Ask whether a hidden variable explains it. The operator, the day, the goal or the object’s side. If the variable can be observed or labelled, add it to the observation and the conflict goes away without a new model.
  • Watch rollouts, not the loss. A policy that averages has a low loss and a freeze or a crash at a decision point. The failure is at the same place in every episode.

A sensible order of remedies

  1. Remove the conflict from the data if you can: one operator, one protocol, or a labelled mode as an input.
  2. If the task needs both modes, use a classifier over discretised actions or a small mixture density model: they are cheap and easy to debug.
  3. If actions are continuous and high-dimensional, or the distribution is complex, move to a conditional generative policy such as a diffusion policy, and plan for its inference cost.
  4. Use action chunks or another form of commitment so a sampled mode is followed for a few steps, and decide how to blend chunks.
  5. Evaluate over many episodes with confidence intervals (lesson M1-T5), and report success by decision point, not only overall.

Costs to plan for

  • Compute. Several network passes per action for diffusion; a fast controller must run on its own clock, with the chunk covering the gap.
  • Data. Generative models learn distributions, which needs the data to represent the variation you want reproduced.
  • Debugging. A stochastic policy gives different results per run; fix seeds for evaluation and use enough episodes.
  • Scale. Large models that act from images and language build on these methods; their treatment belongs to the VLA track, and Module 10 discusses what is proven about them.

Common misconceptions

“MSE learns the demonstrated behaviour.”

It learns the conditional average of the demonstrated actions, which is a demonstrated behaviour only when the actions in a state agree. For two sides of a pillar it is a command into the pillar.

“More data or a bigger network will fix the averaging.”

A more flexible model fits the mean better. The failure is in the loss, not in the model’s capacity.

“Sampling from the right distribution at every step is enough.”

On Scout, drawing A or B at random each step succeeded in 31 of 100 starts; drawing once per episode, 100. The sample must be held for a stretch of steps.

“Multimodality always breaks a cloned policy.”

Not always: a nearest-neighbour clone on the same two operators succeeded in 100 of 100. It depends on the model and whether the modes overlap in the data. Test for it.

Engineering takeaways

  1. Squared error is minimised by the conditional mean; with two good actions in one state the mean may be a bad one.
  2. On Scout the mean of the two operators at (3.5, 3.0) is (0.533, 0.000), straight at the pillar; the policy built from it froze or crashed in 97 of 100 near starts.
  3. Remedies: choose one mode, classify discretised actions, predict a mixture, or sample from a conditional generative model; each trades simplicity for expressiveness.
  4. Commitment is as important as sampling: a mode chosen per step zig-zagged (31 of 100) and chosen per episode worked (100 of 100). Diffusion policies and action chunking carry this idea.
  5. Check for multimodality with nearby-state action variance and rollouts before reaching for a heavier model; the VLA track continues with large pretrained policies.

References and further reading

The primary sources for the diffusion policy, action-chunking and mixture density methods named in the lesson.

  • C. M. Bishop, “Mixture Density Networks”, Technical Report NCRG/94/004, Aston University, 1994.
  • C. Chi, S. Feng, Y. Du, Z. Xu, E. Cousineau, B. Burchfiel and S. Song, “Diffusion Policy: Visuomotor Policy Learning via Action Diffusion”, Robotics: Science and Systems (RSS), 2023.
  • T. Z. Zhao, V. Kumar, S. Levine and C. Finn, “Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware”, Robotics: Science and Systems (RSS), 2023.
  • J. Ho, A. Jain and P. Abbeel, “Denoising Diffusion Probabilistic Models”, NeurIPS, 2020.

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

Question 1

Two operators pass a pillar on opposite sides and a regression network is trained with squared error on both. What does it predict on the line through the pillar, and why?

Explanation

The average of their actions, which is a command straight at the pillar. Squared error is minimised by the conditional mean, and the mean of two opposite lateral actions has no lateral part.

Question 2

Why is sampling a mode at every step worse than sampling once per episode on Scout?

Explanation

The draws alternate between the operators, so the robot zig-zags and does not leave the line; on Scout it succeeded in 31 of 100 near starts against 100 when the mode was chosen once. Chunking and latent variables commit to a mode for several steps.

Question 3

Name two remedies for multimodal actions and one cost of each.

Explanation

Classification over discretised actions: discretisation error and flips between nearly equal bins. Diffusion policy: several network passes per action and heavier training. Choosing a mode: it discards demonstrations and needs an agreed rule.

What’s next

Next: the lab M3-L4 puts this on Scout: two operators disagree around the pillar, the clone freezes, and you find the conflict and rebuild the dataset.

Practice this lesson

Clean the Demonstrations

Two operators disagree around the pillar and the clone freezes: find the conflict and rebuild the dataset.

Next lesson

M4-T1 — Markov Decision Processes

Write a task as states, actions, transitions, rewards and a discount, and say what the Markov property assumes.

Nothing below is required to finish this lesson.