Generative models for robots

Start here: why a robot has to generate, not just predict

Imagine a robot arm that has watched a person reach across a table two hundred times. A cup sits in the way. About half the time the person reached around it on the left, half the time on the right. Now the robot has to reach. The obvious approach, learn to predict the demonstrated motion from the situation, has a failure built in: the best single prediction of a left-or-right motion is the average of the two, and the average goes straight through the cup.

This is the problem I want to work through. Whenever the data has several right answers, a model that outputs one answer learns the wrong one. A generative model instead learns the whole range of plausible answers and lets you draw one at random: sometimes left, sometimes right, never through the cup.

Robotics is full of this. An object can be grasped a dozen ways. A pedestrian might turn left or right. A simulated camera image has many plausible real-world counterparts. In each case the robot needs a model that can propose a good answer and, if asked again, a different good one. I’ll build up these models in the following order:

1. The basics you need: probability distributions, likelihood, the KL divergence, and one inequality. Everything else is built from these.
2. Hidden causes: the idea of a latent variable, and why the obvious way to train such a model is impossible.
3. The ELBO: the bound that makes it possible, derived one step at a time until you can say it out loud.
4. The VAE: the reparameterization trick that lets gradients through a random draw, and a VAE you can train in your browser.
5. Posterior collapse: the way VAEs fail, how to recognize it, and what to do about it.
6. Conditional VAEs: generating given a situation, which is what a robot actually needs.
7. GANs: a completely different way to learn to generate, with a two-player objective, and why the game is hard to keep stable.

Next: The basics you need

1. The basics you need

I’ll start with four ideas that the rest of the discussion builds on. Each gives us a piece we’ll need to understand how these models learn.

A distribution is a landscape over possible outcomes

Point a laser range finder at a wall exactly 2.0 m away and read it 25 times. You will not get 2.000 twenty-five times. You get 1.97, 2.04, 2.01: a cluster around 2.0 with some spread. A probability distribution describes this cluster. For continuous quantities it is drawn as a density, a curve that is tall where readings are common and low where they are rare. The area under the curve between two values is the probability of a reading landing between them; the whole area is 1. Writing p(x) means “the density at value x”.

Drawing a value at random from a distribution is called taking a sample. Tall parts of the curve produce samples often, low parts rarely. The bell-shaped curve you get from most sensors is the Gaussian, written N(μ, σ²): μ is the center, σ the width.

import numpy as np
rng = np.random.default_rng(0)
readings = 2.0 + 0.05 * rng.standard_normal(25)     # mu + sigma * eps, 25 times
# the same thing, spelled out
p = rng.random(25); q = rng.random(25)
eps = np.sqrt(-2 * np.log(1 - p)) * np.cos(2 * np.pi * q)       # Box-Muller
readings2 = 2.0 + 0.05 * eps

# density and log-density of a value under N(mu, sigma^2)
def log_normal(x, mu, sigma):
    return -0.5 * ((x - mu) / sigma) ** 2 - np.log(sigma) - 0.5 * np.log(2 * np.pi)
print(np.exp(log_normal(2.0, 2.0, 0.05)))   # 7.9788
print(log_normal(2.0, 2.0, 0.05))           # 2.0768

Likelihood: how well a model explains the data

Suppose you propose a Gaussian with some μ and σ as a model of the sensor. How good is the proposal? Compute the density the model assigns to each reading, and multiply them all together. That product is the likelihood of the data under the model. It is tiny (25 numbers below 1 multiplied together), so everyone works with its logarithm instead: the log turns the product into a sum, and log p(x) is called the log-likelihood. Bigger is better. Picking the parameters that make the log-likelihood as large as possible is maximum likelihood training, and it is what every model on this page is trying to do, or trying to approximate.

x = np.array([1.97, 2.04, 2.01])
for mu, sigma in [(2.0, 0.05), (1.9, 0.05), (2.0, 0.2)]:
    ll = log_normal(x, mu, sigma)             # one log-density per reading
    print(mu, sigma, ll.round(3), ll.sum().round(2))
# maximum likelihood by hand: mean and (population) standard deviation
mu_hat, sigma_hat = x.mean(), x.std()          # 2.0067, 0.0287

A neural network can play the role of the sliders. Give it whatever the distribution should depend on (the robot's joint angles, a camera image, a goal) and let it output μ and σ. Training by maximum likelihood then nudges the network's weights so that the observed outcomes get high density. This is how a network “outputs a distribution”: it outputs the parameters of one.

import torch, torch.nn as nn

class GaussianHead(nn.Module):
    def __init__(self, n_in, n_out):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(n_in, 64), nn.Tanh(), nn.Linear(64, 2 * n_out))
    def forward(self, c):
        mu, log_sigma = self.net(c).chunk(2, dim=-1)   # split the output in two halves
        return mu, log_sigma.exp()                     # exp() makes sigma positive

model = GaussianHead(n_in=3, n_out=1)                  # e.g. 3 joint angles -> 1 sensor reading
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
for c, x in loader:                                    # batches of (situation, outcome)
    mu, sigma = model(c)
    # sum over dims, mean over batch
    nll = -torch.distributions.Normal(mu, sigma).log_prob(x).sum(-1).mean()
    opt.zero_grad(); nll.backward(); opt.step()

Expectation: the average under a distribution

The expectation of a quantity f(x) under a distribution p, written E_p[f(x)], is its average when x is drawn from p: the values f takes, weighted by how often they occur. You rarely compute it exactly. You estimate it by drawing samples x₁, …, x_n from p and averaging f(x_i). This is a Monte Carlo estimate; it is unbiased, and it gets better slowly, with error shrinking like 1/√n. Chapter 3 will estimate an expectation with a single sample and get away with it.

import numpy as np
rng = np.random.default_rng(0)

def expectation(f, n):
    z = rng.standard_normal(n)          # 1. draw n samples from p (here N(0, 1))
    fz = f(z)                           # 2. evaluate f at each sample
    est = fz.mean()                     # 3. average: the Monte Carlo estimate
    std_err = fz.std() / np.sqrt(n)     #    the typical error, estimated from the same samples
    return est, std_err

print(expectation(lambda z: z ** 2, 1000))                     # e.g. (0.975, 0.045); truth 1
# almost always (0.0, 0.0); truth 0.0171
print(expectation(lambda z: np.abs(z - 1.3) < 0.05, 10))

KL divergence: how far one distribution is from another

Often two distributions describe the same quantity: your model's and the true one, or a convenient one and the one you actually want. The KL divergence measures how different they are:

KL(q ‖ p) = E_q[ log q(x) − log p(x) ]

• Draw x from q. Compare how surprised you are under q versus under p (log density is a measure of surprise, with the sign flipped). Average over draws.
• It is always ≥ 0, and equals 0 only when q and p are the same distribution. I’ll use that fact throughout the derivation.
• It is not symmetric: KL(q ‖ p) and KL(p ‖ q) differ. The one written here, with the expectation taken under q, is the one that appears in the ELBO.

import numpy as np
from scipy.stats import norm
mu, s = 1.0, 0.6                                        # q = N(1, 0.6^2); p = N(0, 1)

kl_closed = 0.5 * (mu**2 + s**2 - 1 - np.log(s**2))     # 0.6908

z = norm.rvs(mu, s, size=10000, random_state=0)         # samples from q
# E_q[log q - log p] ~ 0.69 +- 0.01
kl_mc = np.mean(norm.logpdf(z, mu, s) - norm.logpdf(z, 0, 1))

# discrete case
q = np.array([0.5, 0.5, 0.0]); p = np.array([0.45, 0.45, 0.10])
mask = q > 0                                            # 0 * log 0 counts as 0
kl_qp = np.sum(q[mask] * np.log(q[mask] / p[mask]))     # 0.1054 nats
# torch: torch.distributions.kl_divergence(Normal(mu, s), Normal(0., 1.)) gives the closed form

One inequality: the log of an average beats the average of logs

The logarithm curve bends downward; it is concave. Take two numbers, 1.5 and 5. Their average is 3.25, and log 3.25 = 1.18. But the average of their logs is (log 1.5 + log 5) / 2 = 1.01, which is smaller. This always happens for a concave function: log(average) ≥ average(log), with equality only when all the numbers are the same. It is called Jensen's inequality, and it holds for averages weighted by any distribution, which is exactly how chapter 3 will use it.

What “learned” and “loss” mean here

A model has weights, random at first. For each training example you compute a single number that says how badly the model did, the loss; then you nudge every weight a little in the direction that shrinks it. That nudging is gradient descent, and the bookkeeping that finds the direction for every weight is called backpropagation. On this page the loss is almost always “minus the log-likelihood” or “minus a bound on it”: making the loss small means making the data probable.

import torch
x = torch.tensor([1.97, 2.04, 2.01]); sigma = 0.05
mu = torch.tensor(1.80, requires_grad=True)             # the one weight, random-ish start

loss = -torch.distributions.Normal(mu, sigma).log_prob(x).sum()
loss.backward()                                         # backpropagation fills mu.grad
print(mu.grad)                                          # tensor(-248.)
with torch.no_grad():
    mu -= 0.0005 * mu.grad                              # one gradient-descent step -> 1.924

Next: Hidden causes

2. Hidden causes: the latent variable

Look at the two hundred demonstrations again. Each one was shaped by a decision the demonstrator made and never wrote down: go left or go right, how wide a berth to give the cup, how much to hurry. The recorded trajectory x is the visible result; the decision is a hidden cause. In modeling, a hidden cause is called a latent variable and written z.

A latent variable model tells a two-step story about how each trajectory came to be. First, draw z from a simple fixed distribution, the prior p(z), usually a standard Gaussian. Second, feed z to a network, the decoder, which outputs the parameters of a distribution over trajectories, p(x | z); draw x from it. Different z values give different trajectories, and because the decoder is a flexible network, a one-dimensional z can already encode “left or right, and how wide”.

One latent number, and the decoder has organized the space: negative values go one way around the cup, positive values the other, and the size of the number sets how wide the berth is. Nobody told it to use z this way. It found this arrangement because it was the arrangement that made the demonstrations most probable, which is what the next two chapters are about.

z = 'L' if rng.random() < 0.5 else 'R'        # step 1: draw the hidden cause from the prior
x = rng.normal({'L': -0.8, 'R': 0.8}[z], 0.2) # step 2: draw the observation given the cause

The problem: training this model asks for an impossible sum

To train by maximum likelihood, we need log p(x) for each demonstration: how probable the model finds it. But the model only defines p(x | z) for a given z. To get p(x) we must average over every hidden cause that could have produced x:

p(x) = ∫ p(x | z) p(z) dz = E_p(z)[ p(x | z) ]

• Read it as: “for every possible z, how well does it explain x, weighted by how likely that z is; add it all up.”
• The integral sign is over all values of z. With one latent number you could approximate it on a grid. With 16 or 32 latent numbers, a grid has more points than there are atoms in the universe.

Worse, most of the sum is wasted. For a given trajectory x, almost every z makes the decoder produce something unlike x, so p(x | z) is essentially zero there. The whole sum comes from a tiny region of z values: a needle in a haystack. Sampling z from the prior and averaging (the Monte Carlo estimate from chapter 1) almost never hits the needle.

import numpy as np
from scipy.stats import norm
f = lambda z: 1.2 * np.tanh(z)
x0, sigma_x = 0.9, 0.25

z = np.linspace(-6, 6, 1201); dz = z[1] - z[0]                   # a grid over the latent
# log p(x|z) + log p(z), per grid point
log_joint = norm.logpdf(x0, f(z), sigma_x) + norm.logpdf(z)
m = log_joint.max()
log_px = m + np.log(np.sum(np.exp(log_joint - m)) * dz)           # log-sum-exp; -0.878
# p(z|x) on the grid; integrates to 1
posterior = np.exp(log_joint - log_px)

# the Monte Carlo attempt with prior samples
zs = np.random.default_rng(3).standard_normal(10)
print(np.mean(norm.pdf(x0, f(zs), sigma_x)))                     # 0.209, truth 0.416

The needle has a name. The posterior p(z | x) is the distribution over hidden causes given that we saw x: for a left-going trajectory, it sits on the negative z values that decode to that trajectory. By Bayes' rule it equals p(x | z) p(z) / p(x), and the denominator is the very quantity we cannot compute. So the posterior is out of reach too. Chapter 3 gets around both problems with one move.

import torch, torch.nn.functional as F
x = torch.tensor([0.1, 0.3]); fz = torch.tensor([0.0, 0.5]); sigma = 0.1
logp = torch.distributions.Normal(fz, sigma).log_prob(x).sum()           # 0.267: the full thing
# -2.5: what has a gradient
logp_train = -((x - fz) ** 2).sum() / (2 * sigma ** 2)

pix = torch.tensor([1., 0.]); logits = torch.tensor([2.197, 2.197])      # sigmoid(2.197) = 0.9
# log0.9 + log0.1 = -2.41
logp_bern = -F.binary_cross_entropy_with_logits(logits, pix, reduction='sum')

Next: The ELBO

3. The ELBO: a bound you can compute, derived until you can say it

Chapter 2 ended with two things we cannot compute: the likelihood p(x) and the posterior p(z | x). I’ll handle both by introducing a helper distribution q(z | x), our best guess at the posterior: a network, called the encoder, that looks at a trajectory and returns a Gaussian over z centered where it thinks the hidden cause lies. Then, instead of the true log-likelihood, compute a quantity that is always less than or equal to it but computable: the evidence lower bound, or ELBO. Push the bound up and two good things happen at once: the model explains the data better, and the guess gets closer to the true posterior.

The derivation is six lines. The point of this chapter is that you can reproduce them, in words, without notes. Try to predict each line before pressing the button.

Here is the whole thing in one breath, which is the version to memorize: the log-likelihood is the log of an average over hidden causes; rewrite the average under a guess; move the log inside, which can only lower the value; the result splits into a reconstruction term minus a KL to the prior; that is the ELBO, and it misses the truth by exactly the KL from the guess to the true posterior.

Two things to take from the toy. First, the ELBO is never above the truth, no matter how you set the sliders; the inequality is real. Second, the gap depends only on how well the guess matches the posterior, so “climbing the bound” with respect to the guess is literally inference: finding out which hidden causes explain the observation. With the second decoder the posterior has two bumps (two z values produce the same x), a single Gaussian guess cannot match it, and the best the bound can do is settle on one bump and leave a gap. Real encoders make the same compromise.

The two terms as two forces

The reconstruction term wants the encoder to give every trajectory its own precise, narrow spot in latent space, so the decoder can reproduce it exactly. The KL term wants every guess to look like the prior: wide, centered at zero, indistinguishable from the guesses for other trajectories. Training balances the two. The balance is what makes the latent space usable: codes are informative enough to reconstruct from, yet packed closely enough around the prior that a fresh z drawn from N(0, 1) lands somewhere the decoder has seen before. That second property is what lets the robot generate new trajectories at run time by sampling z from the prior.

def elbo_loss(x, encoder, decoder, sigma_x=0.1, beta=1.0):
    # q(z|x) = N(mu, exp(logvar)); shapes [B, d_z]
    mu, logvar = encoder(x).chunk(2, dim=-1)
    # one reparameterized sample per example
    z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)
    # parameters of p(x|z): here the mean
    x_hat = decoder(z)
    # log p(x|z) up to a constant; sum over data dims
    recon = -((x - x_hat) ** 2).sum(-1) / (2 * sigma_x ** 2)
    # KL(q(z|x) || N(0, I)); sum over latent dims
    kl = 0.5 * (mu ** 2 + logvar.exp() - 1 - logvar).sum(-1)
    elbo = recon - beta * kl                                     # one number per example
    # minimize the negative, averaged over the batch
    return -elbo.mean()

Next: The VAE

4. The VAE: gradients through a random draw

Now I’ll address one remaining mechanical problem. The ELBO's first term is an average over z drawn from the encoder's guess, and we estimate it with a single random draw. To train the encoder we need to know how that term changes when the encoder's weights change: its gradient. But a random draw is not a smooth function of anything. If the encoder nudges μ upward, which way does “a random number from N(μ, σ²)” move? Backpropagation, which works by following the chain of operations from the loss back to each weight, hits the draw and stops. The chain is broken.

The reparameterization trick fixes this with a change of bookkeeping. A number drawn from N(μ, σ²) is the same thing as μ + σ·ε where ε is drawn from N(0, 1). So draw ε first, before anything else, and treat it as one more input to the computation, like the data. Then z = μ + σ·ε is an ordinary formula: dz/dμ = 1 and dz/dσ = ε. The dice are rolled outside the machine and the results are fed in; inside the machine everything is differentiable.

mu = torch.tensor(0.5, requires_grad=True); sigma = torch.tensor(0.2, requires_grad=True)
# the dice, rolled outside the graph; no gradient flows into it
eps = torch.randn(())
z = mu + sigma * eps                      # differentiable in mu and sigma
loss = z ** 2
loss.backward()
print(mu.grad, sigma.grad)                # 2z, 2z*eps: one row of the table

# the same thing through the library: rsample() = reparameterized sample, keeps the gradient.
# .sample() would cut the gradient, silently
z = torch.distributions.Normal(mu, sigma).rsample()

The KL term in closed form

The other half of the ELBO needs no sampling at all. When the encoder outputs a Gaussian with independent dimensions, and the prior is the standard Gaussian, the KL divergence is a short formula, summed over latent dimensions:

KL( N(μ, σ²) ‖ N(0, 1) ) = ½ Σ_j ( μ_j² + σ_j² − 1 − log σ_j² )

• j runs over the latent dimensions; each contributes separately.
• μ_j² punishes a guess whose center drifts from 0. σ_j² − 1 − log σ_j² is zero at σ = 1 and positive otherwise: it punishes a guess that is narrower or wider than the prior. A very narrow guess (σ → 0) costs a lot, because −log σ² grows without bound.
• The whole thing is 0 exactly when μ = 0 and σ = 1 in every dimension: when the guess is the prior.

Encoders output log σ² rather than σ, so that any real number is a valid output and σ = exp(½ log σ²) is automatically positive.

mu = torch.tensor([0.5, -0.2]); logvar = torch.tensor([-0.7, 0.1])
kl_per_dim = 0.5 * (mu ** 2 + logvar.exp() - 1 - logvar)     # tensor([0.2233, 0.0226])
kl = kl_per_dim.sum()                                        # 0.2459

The complete recipe

import torch, torch.nn as nn, math
torch.manual_seed(7)

# --- data: where the gripper closed on the rim of a bowl, seen from above
N = 240
angle = torch.rand(N) * 2 * math.pi
radius = 1 + 0.05 * torch.randn(N)
X = torch.stack([radius * angle.cos(), radius * angle.sin()], dim=1)     # [240, 2]

# --- networks
# 4 outputs = mu (2) and logvar (2)
enc = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 4))
dec = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 2))      # z (2) -> mean of x (2)
opt = torch.optim.Adam(list(enc.parameters()) + list(dec.parameters()), lr=0.02)
sigma_x, beta = 0.1, 1.0

for step in range(1, 501):
    mu, logvar = enc(X).chunk(2, dim=-1)                                # step 2 of the recipe
    z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)             # step 3
    x_hat = dec(z)                                                      # step 4
    # step 5: -log p(x|z) + const, per point
    rec = ((x_hat - X) ** 2).sum(-1) / (2 * sigma_x ** 2)
    # KL(q || N(0, I)), per point
    kl = 0.5 * (mu ** 2 + logvar.exp() - 1 - logvar).sum(-1)
    # = -ELBO (beta = 1), batch mean
    loss = (rec + beta * kl).mean()
    opt.zero_grad(); loss.backward(); opt.step()                        # step 6
    if step % 100 == 0:
        print(step, f"rec {rec.mean():.2f}  kl {kl.mean():.2f}  kl/dim {kl.mean()/2:.2f}")

# --- using it
with torch.no_grad():
    # generate: prior -> decoder
    samples = dec(torch.randn(200, 2))
    # ~0.85-0.9; the rim is 1.0
    print("sample radius", samples.norm(dim=1).mean().item())
    # compress: mu for each grasp point
    codes = enc(X)[:, :2]

d loss / d mu     = dZ * dz/dmu     + beta * dKL/dmu
                  = dZ * 1          + beta * mu
d loss / d logvar = dZ * dz/dlogvar + beta * dKL/dlogvar
                  = dZ * eps * (sigma / 2) + beta * (sigma^2 - 1) / 2

With β = 1, after a few hundred steps the violet samples land on the rim: the robot can propose grasp points it never saw, and they are sensible. The latent plot shows why it works. The rim has been laid out around the origin, inside the prior's rings, so that a fresh draw from N(0, I) lands on a code the decoder knows how to read. The colors go around in order: the encoder has learned the rim's geometry without being told about angles. The β menu is for chapter 5; leave it at 1 for now.

Next: Posterior collapse

5. Posterior collapse: when the model stops using its latent

I’ll start with how this failure happens. There is a way to make the KL term exactly zero: have the encoder output μ = 0 and σ = 1 for every input. The guess is then the prior, for every trajectory, so the code z carries no information about which trajectory it came from. The decoder, receiving pure noise, learns to ignore z and output whatever single answer minimizes the reconstruction error on average. That answer is the mean. For the cup, it is the trajectory through the cup.

This failure is called posterior collapse: the approximate posterior collapses onto the prior. It is the VAE's characteristic disease, and for a robot it is exactly the failure this page opened with. A collapsed trajectory model is a regression model in disguise. A collapsed grasp model proposes the same grasp for every draw of z. A collapsed motion-prediction model predicts one future, the average one, for a pedestrian who could go either way.

Why it happens

The objective allows it. Collapse is a legitimate solution of the ELBO. The KL term is at its minimum, zero, and the decoder is doing the best any z-blind model can do. If the decoder is powerful enough to model the data reasonably well on its own (an autoregressive decoder that predicts each waypoint from the previous ones, for instance), the reconstruction term does not suffer much from ignoring z, and every bit of information the code carries costs KL. The trade is not worth it, and the optimizer declines it.

The dynamics push toward it. At the start of training the encoder is random, so z is noise to the decoder. The decoder's fastest route to a lower loss is to ignore this noise, which it learns within a few hundred steps. Once ignored, the encoder receives no gradient rewarding informative codes, only the KL gradient pulling toward the prior. The two networks lock each other into the collapsed state; it is a local optimum that is easy to enter and hard to leave.

The balance can be mis-set. A large β, or equivalently a large decoder spread σ_x, makes the KL cost dominate and produces collapse even with a weak decoder. This is what the lab shows. It happens by accident more often than people expect, because “MSE plus KL” silently fixes σ_x, and the right value depends on the scale of the data.

# 1. KL annealing (warm-up): beta climbs from 0 to 1 over the first T steps
def beta_at(step, T=5000):
    return min(1.0, step / T)

# 2. Free bits: per-dimension KL below a floor lam is not penalized
kl_dim = 0.5 * (mu ** 2 + logvar.exp() - 1 - logvar)       # [B, d_z]
# floor applied to the batch-average per dimension
kl_free = torch.clamp(kl_dim.mean(0), min=lam).sum()
# example: KL per dim (2.1, 0.03, 0.5, 0.0), lam = 0.1 -> penalized (2.1, 0.1, 0.5, 0.1):
# dims 2 and 4 pay the floor whether they carry information or not, so they may as well carry
# some.

# 3. beta < 1: just use beta = 0.1 .. 0.5 in elbo_loss; report the beta = 1 ELBO separately for
# honesty

# Diagnostic: active units. Dead dimensions have KL ~ 0 and mu that does not vary across the
# data.
with torch.no_grad():
    mu, logvar = enc(X).chunk(2, dim=-1)
    # per dimension, averaged over data
    kl_dim = (0.5 * (mu ** 2 + logvar.exp() - 1 - logvar)).mean(0)
    # both tests, per dimension
    active = (mu.var(0) > 0.01) & (kl_dim > 0.02)
    print(f"{active.sum().item()} of {mu.shape[1]} latent dimensions are active")
    # the two-z test: decode two different codes and look
    z1, z2 = torch.randn(1, mu.shape[1]), torch.randn(1, mu.shape[1])
    print((dec(z1) - dec(z2)).abs().max().item())                     # ~0 means z is ignored

Next: Conditional VAEs

6. Conditional VAEs: generating for the situation at hand

Next, I’ll make the model depend on the situation. A robot never needs “a trajectory”. It needs a trajectory for this cup position, a grasp for this object, a prediction of where this pedestrian goes next. The quantity it wants is a conditional distribution p(x | c), where c is whatever describes the situation: the cup's position, a point cloud of the object, the last second of a person's motion.

The conditional VAE (CVAE) is the VAE with c handed to both networks. The encoder looks at x and c and guesses z; the decoder produces x from z and c. The ELBO is unchanged except that c appears everywhere. Something useful happens to the meaning of z: since the decoder already knows the situation, the latent only has to carry what the situation does not determine. For the cup, that is the one thing the demonstrations disagreed about: left or right. The latent becomes the model's word for “the remaining choice”.

Sample a few times at different cup positions. The model has learned the demonstrators' habit: with the cup on the right, most samples go left, and the proportions track the data. The regression network, trained on the same demonstrations, goes through the cup whenever the two sides were about equally common. Now drag the z slider: one continuous number selects the side and the berth, exactly as in chapter 2, and it does so for every cup position, because the decoder reads the cup position separately.

You will also see a few red samples. They come from z values near the boundary between “left” and “right”, where the decoder has to interpolate between two behaviors and produces a trajectory that is neither. This is a genuine weakness of a continuous latent with a discrete choice inside it, discussed under Go deeper. In practice robots handle it the way the demo invites you to: draw several samples and let a collision check or a cost function discard the bad ones. Generative models propose; planners check.

class CVAE(nn.Module):
    def __init__(self, d_x=8, d_c=1, d_z=1, h=48):
        super().__init__()
        # q(z | x, c)
        self.enc = nn.Sequential(nn.Linear(d_x + d_c, h), nn.Tanh(), nn.Linear(h, 2 * d_z))
        # p(x | z, c)
        self.dec = nn.Sequential(nn.Linear(d_z + d_c, h), nn.Tanh(), nn.Linear(h, d_x))
        self.d_z = d_z
    def loss(self, x, c, sigma_x=0.1, beta=1.0):
        mu, logvar = self.enc(torch.cat([x, c], -1)).chunk(2, -1)
        z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)
        # the decoder sees c directly ...
        x_hat = self.dec(torch.cat([z, c], -1))
        # ... so z only has to carry what c does not fix
        rec = ((x - x_hat) ** 2).sum(-1) / (2 * sigma_x ** 2)
        kl = 0.5 * (mu ** 2 + logvar.exp() - 1 - logvar).sum(-1)
        return (rec + beta * kl).mean()
    @torch.no_grad()
    def sample(self, c, n=10, tau=1.0):                             # run time: no encoder, no x
        # tau < 1: typical; tau > 1: adventurous
        z = tau * torch.randn(n, self.d_z)
        # n candidate trajectories for this c
        return self.dec(torch.cat([z, c.expand(n, -1)], -1))

# deployment: propose, then check
cands = model.sample(cup_position, n=10)                            # [10, 8] waypoints
# your collision checker / cost function
ok = ~collides(cands, cup_position)
# none passed: draw 50 more and check again
best = cands[ok][0] if ok.any() else None

# what the model believes about the split, for a given cup position
many = model.sample(cup_position, n=1000)
# fraction passing the cup on the left
frac_left = (many[:, 3] < 0).float().mean()

m, log_s2 = prior_net(c).chunk(2, -1)                                # conditional prior p(z|c)
kl = (0.5 * (log_s2 - logvar)
      + (logvar.exp() + (mu - m) ** 2) / (2 * log_s2.exp()) - 0.5).sum(-1)
# at run time: z = m + torch.exp(0.5 * log_s2) * torch.randn_like(m)

Next: GANs

7. GANs: learning to generate by playing a game

So far, I’ve described generators that ask “how probable is the data under this model?” and push a bound on that probability up. A generative adversarial network (GAN, 2014) never asks that question. It has no encoder, no likelihood and no ELBO. Instead it sets up a contest between two networks.

The generator G takes noise z and produces a sample: a grasp pose, a trajectory, a camera image. The discriminator D receives a sample and outputs the probability that it is real. D is trained to say “real” for recorded data and “fake” for the generator's output. G is trained to make D say “real”. Neither is told what a good sample looks like; that knowledge emerges from the contest, the same way a forger and an inspector sharpen each other.

Robots meet GANs in three places. Sim-to-real: a generator turns simulator images into images the discriminator cannot tell from the real camera, so a policy trained in simulation sees realistic inputs. Imitation: in GAIL, the discriminator tells expert behavior from the robot's, and its verdict is used as the reward the robot learns from. And data: generators propose grasps or trajectories from a small set of recordings.

The objective

min_G max_D   E_x~data[ log D(x) ] + E_z~N(0,I)[ log ( 1 − D(G(z)) ) ]

• The first term is large when D assigns high probability to real samples. The second is large when D assigns low probability to fakes. D wants both large: this is the log-likelihood of a classifier with labels “real” and “fake”.
• G controls only the second term, through G(z), and wants it small: it wants D(G(z)) near 1.
• “min max” means: D is chosen to make the expression as large as possible for the current G; G is chosen to make that best-case value as small as possible. In practice the two take turns with gradient steps, one or a few for D, then one for G.

What is G really minimizing? If D were perfect for the current generator, D(x) would equal p_data(x) / (p_data(x) + p_G(x)): the fraction of samples near x that are real. Plug that in and the expression becomes a divergence between the real and generated distributions (the Jensen–Shannon divergence, a symmetric relative of the KL). So a GAN, like a VAE, is minimizing a mismatch between two distributions. The difference is that it never writes either density down; the discriminator estimates the mismatch from samples alone. That is why GANs can produce sharp, realistic images where VAEs produce blurry averages: nothing in the objective rewards hedging.

With balanced settings, the four clusters are usually all covered within one to two thousand steps, though not always, and not permanently: keep training and you may see the coverage drift. Switch to the eager preset, reset, and train: the generator typically piles all of its proposals on one side, the discriminator learns that side is fake, and the proposals jump to another side. That is mode collapse, and the jumping is mode hopping. Watch the loss traces while it happens: they oscillate and tell you almost nothing about whether the samples are good.

Mode collapse

Mode collapse is the generator producing one kind of sample, or a few, for every z: all grasps from the east, all pedestrians turning left, all simulated images rendered with the same lighting. The objective allows it because the discriminator is only ever asked “is this sample real?”, never “has the generator covered everything?”. A generator that makes one perfectly realistic kind of output fools D as well as one that makes all kinds. It is often easier to make one kind well, so that is where gradient descent goes first.

Collapse rarely stays still. Once every fake approaches from the east, D learns that eastern approaches are suspicious, so G's easiest improvement is to move everything north. D follows. The generator chases a discriminator that chases the generator, and the pair can circle indefinitely. For a robot the consequence is a proposal model that is confident and narrow: it will offer only eastern grasps when the east side is blocked, and offer them with conviction.

Instability: why the game is hard to keep still

Training a VAE is ordinary minimization: one loss, one landscape, downhill until it flattens. Training a GAN is not. Each player's landscape is shaped by the other player and moves every step, so there is no fixed quantity being minimized and no fixed “bottom” to reach. The equilibrium the game is looking for (a generator D cannot beat, and a D that is the best response to it) is a saddle, not a valley, and gradient steps taken by both players at once do not naturally settle into a saddle. They can circle it forever, or spiral away from it.

x, y, lr = 1.2, 0.6, 0.2
def simultaneous(x, y):  return x - lr * y,        y + lr * x
def alternating(x, y):   x2 = x - lr * y;          return x2, y + lr * x2
def extragradient(x, y):
    xh, yh = x - lr * y, y + lr * x                 # look ahead with a plain step
    return x - lr * yh, y + lr * xh                 # step using the gradients from there

Three more things make it worse in a real GAN. If the discriminator becomes too good, D(G(z)) is essentially 0 for every fake, log(1 − D) is flat there, and the generator receives no gradient at all; the original paper already recommended the non-saturating replacement −log D(G(z)), which keeps a gradient alive but does not remove the imbalance. When the real and generated distributions do not overlap, the Jensen–Shannon divergence is constant and its gradient says nothing about which way to move. And the losses themselves are useless as progress bars: D's loss going up may mean G is improving or that D has been outpaced, and there is no way to tell from the numbers alone.

import torch, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(3)
# four approach directions
centers = torch.tensor([[0.8, 0.], [0., 0.8], [-0.8, 0.], [0., -0.8]])
# a side at random, plus 0.1 noise
def real_batch(B):
    return centers[torch.randint(0, 4, (B,))] + 0.1 * torch.randn(B, 2)

G = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 2))
D = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
optG = torch.optim.Adam(G.parameters(), lr=1e-3, betas=(0.5, 0.999))
optD = torch.optim.Adam(D.parameters(), lr=1e-3, betas=(0.5, 0.999))
B = 64; ones, zeros = torch.ones(B, 1), torch.zeros(B, 1)

for step in range(2000):
    # --- D step: a binary classifier, labels 1 = real, 0 = fake
    x_real = real_batch(B)
    x_fake = G(torch.randn(B, 2)).detach()          # detach: this step must not train G
    # lossD = -[ log D(x) + log(1 - D(G(z))) ], the classifier's cross-entropy
    lossD = (F.binary_cross_entropy_with_logits(D(x_real), ones)
             + F.binary_cross_entropy_with_logits(D(x_fake), zeros))
    optD.zero_grad(); lossD.backward(); optD.step()

    # --- G step: fresh noise, fool D. D's weights get gradients but are never stepped here.
    x_fake = G(torch.randn(B, 2))
    lossG = F.binary_cross_entropy_with_logits(D(x_fake), ones)   # non-saturating: -log D(G(z))
    # original minimax would be: lossG = -F.binary_cross_entropy_with_logits(D(x_fake), zeros)
    optG.zero_grad(); lossG.backward(); optG.step()

# coverage: nearest side for 200 proposals
with torch.no_grad():
    prop = G(torch.randn(200, 2))
    side = torch.cdist(prop, centers).argmin(1)
    # healthy: about 0.25 each; collapsed: one ~1.0
    print(torch.bincount(side, minlength=4) / 200.)

# W: [out, in]; u: running estimate of the top left-singular vector
def spectral_norm_step(W, u):
    v = F.normalize(W.t() @ u, dim=0)           # one power-iteration step
    u = F.normalize(W @ v, dim=0)
    sigma = u @ W @ v                           # ~ largest singular value of W
    # the normalized weight actually used in the forward pass
    return W / sigma, u

# the library version, applied per layer
SN = nn.utils.parametrizations.spectral_norm
D = nn.Sequential(SN(nn.Linear(2, 32)), nn.Tanh(),
                  SN(nn.Linear(32, 32)), nn.Tanh(), SN(nn.Linear(32, 1)))

Next: Putting it together

Putting it together

I’ll bring the two families back together here. Both answer the same need: a robot facing several right answers should be able to produce any of them, not their average. A VAE reaches that by explaining the data with hidden causes, training a decoder and an encoder together on a bound on the log-likelihood that you can now derive in six lines. Its failure mode is to stop using the hidden cause. A GAN reaches it by a contest between a generator and a judge; its failure modes are to satisfy the judge with one kind of answer and to never stop circling.

The field's current default for multimodal robot actions, diffusion policies, borrows from both: a likelihood-style bound like the ELBO for training, and sample quality that GANs used to be needed for. If you meet one, you will recognize the parts.

The originals

• Kingma and Welling, 2013. Auto-Encoding Variational Bayes. The VAE and the reparameterization trick.
• Rezende, Mohamed and Wierstra, 2014. Stochastic Backpropagation and Approximate Inference in Deep Generative Models. The same idea, independently.
• Sohn, Lee and Yan, 2015. Learning Structured Output Representation using Deep Conditional Generative Models. The CVAE.
• Bowman et al., 2016. Generating Sentences from a Continuous Space. Posterior collapse named, and KL annealing.
• Higgins et al., 2017. β-VAE. The β knob.
• Goodfellow et al., 2014. Generative Adversarial Nets.
• Arjovsky, Chintala and Bottou, 2017. Wasserstein GAN; Gulrajani et al., 2017, Improved Training of Wasserstein GANs.
• Mescheder, Geiger and Nowozin, 2018. Which Training Methods for GANs do actually Converge?
• Ho and Ermon, 2016. Generative Adversarial Imitation Learning.
• Mousavian, Eppner and Fox, 2019. 6-DOF GraspNet. A CVAE for grasps.
• Rao et al., 2020. RL-CycleGAN. Sim-to-real for robot grasping.
• Chi et al., 2023. Diffusion Policy. Where the story goes next.

Next: Background and related work

Glossary

The terms I’ve used, in roughly the order they first appear.

Generative model

A model of the distribution of the data that can produce new, plausible samples from it.

Probability distribution, density

A rule assigning probability to outcomes; for continuous outcomes, a curve whose area over a range is the probability of that range.

Sample

One value drawn at random from a distribution.

Gaussian

The bell-shaped distribution N(μ, σ²), with center μ and width σ.

Likelihood, log-likelihood

The probability a model assigns to the observed data, as a function of the model's parameters; and its logarithm, which sums over data points.

Maximum likelihood

Choosing parameters to make the observed data as probable as possible.

Expectation

The average of a quantity under a distribution; estimated by averaging over samples (Monte Carlo).

KL divergence

KL(q ‖ p) = E_q[log q − log p]: how far q is from p. Non-negative, zero only when equal, not symmetric.

Jensen's inequality

For a concave function like log, the function of an average is at least the average of the function.

Latent variable

A hidden cause in the model's story of the data, written z; never observed.

Prior

The distribution of z before anything is observed; in a VAE, N(0, I).

Decoder

The network mapping z (and, in a CVAE, the situation c) to the parameters of p(x | z).

Posterior

p(z | x): which hidden causes could have produced this observation. Intractable in a VAE.

Encoder

The network that outputs the guess q(z | x), a Gaussian over z, from an observation.

ELBO

Evidence lower bound: E_q[log p(x | z)] − KL(q(z | x) ‖ p(z)) ≤ log p(x). The training objective of a VAE.

Variational inference

Approximating a posterior by the closest member of a simple family, measured by KL; maximizing the ELBO over q.

Amortized inference

Using one network to produce q for any input, instead of optimizing q separately per data point.

Reparameterization trick

Writing z = μ + σ·ε with ε ~ N(0, 1), so that a random draw becomes differentiable in μ and σ.

β

A weight on the KL term. β = 1 is the ELBO; larger values push toward the prior; equivalent to changing the decoder's spread.

Posterior collapse

The VAE failure where q(z | x) equals the prior for every x, the KL is zero, and the decoder ignores z.

KL annealing, free bits

Two remedies for collapse: raise β from 0 gradually; or exempt a floor of KL per dimension from the penalty.

Conditional VAE (CVAE)

A VAE whose encoder and decoder also receive a situation c, modeling p(x | c).

Aggregate posterior

The average of the encoder's guesses over the dataset; where it differs from the prior, prior samples fall into holes.

GAN

Generative adversarial network: a generator trained to fool a discriminator that is trained to tell generated from real.

Generator, discriminator

The two players: noise → sample, and sample → probability real.

Minimax objective

E[log D(x)] + E[log(1 − D(G(z)))], maximized by D and minimized by G.

Non-saturating loss

Training G on −log D(G(z)) instead of log(1 − D(G(z))), to keep gradients alive when fakes are bad.

Jensen–Shannon divergence

A symmetric relative of KL; what the original GAN minimizes with an optimal discriminator.

Mode collapse, mode hopping

The generator producing one kind of output; and moving that one kind around as the discriminator adapts.

Spectral normalization, gradient penalty

Constraints that keep the discriminator smooth so its gradients stay useful.

Wasserstein GAN

A GAN whose critic estimates the Earth mover's distance, with informative gradients even when distributions do not overlap.

GAIL

Adversarial imitation learning: a discriminator between expert and policy behavior serves as the reward.

Sim-to-real

Making a policy trained in simulation work on a real robot; GANs are used to translate simulated images toward real ones.

Marginal, joint, conditional

p(x), the probability of an observation summed over all causes; p(x, z), of observation and cause together; p(x | z) or p(z | x), of one given the other. Bayes’ rule relates them (crash course, A1).

Monte Carlo estimate, unbiased, standard error

Estimating an expectation by averaging f over samples. Unbiased: correct on average over repeats. Standard error: spread(f)/√n, the typical error of the estimate.

Entropy, cross-entropy, nat, bit

H(p) = −Σ p log p, the average surprise; H(p, q) = −Σ p log q; KL is their difference. A nat uses the natural log, a bit uses log base 2; 1 nat = 1.44 bits.

Log-sum-exp

Computing log Σ e^ℓ_i as m + log Σ e^{ℓ_i − m} with m the largest term, to avoid underflow.

Backpropagation

The chain rule applied backwards from the loss through every operation, giving the gradient for every weight in one pass. Cannot pass through a random draw; the reparameterization trick moves the draw out of its way.

Adam

An optimizer that scales each weight’s step by a running estimate of its gradient’s size, so steps are about the learning rate regardless of gradient scale; β₁ sets the momentum’s memory (0.9 ≈ 10 steps, 0.5 ≈ 2).

Standardization

Subtracting the mean and dividing by the standard deviation of each data coordinate before training, so that a decoder spread such as σ_x = 0.1 means the same thing on every dataset.

Reduction (sum vs mean)

Whether a loss adds or averages over data dimensions. A log-likelihood sums; averaging silently divides the reconstruction term by the data dimension and multiplies the effective β by it.

Mixture density network (MDN)

A network whose output is the weights, means and spreads of a mixture of Gaussians; trained by exact negative log-likelihood with log-sum-exp. The other classic answer to multimodal outputs.

Straight-through estimator

Using a non-differentiable operation (argmax, nearest code) in the forward pass and pretending it was the identity in the backward pass. Used by VQ-VAE and hard Gumbel-softmax.

VQ-VAE, codebook collapse

A VAE whose latent is the nearest of K learned code vectors; no KL term. Codebook collapse: most codes are never selected.

Gumbel-softmax

A reparameterization for categorical variables: argmax(log π + Gumbel noise) is an exact sample, and softmax((log π + noise)/τ) is its differentiable relaxation.

Action chunking, temporal ensembling

Predicting a sequence of future actions at once, and averaging the overlapping predictions made at successive steps with exponentially decaying weights. Used by ACT.

ACT

Action Chunking with Transformers: a CVAE policy whose decoder emits 100-step action chunks; L1 reconstruction, β = 10, latent set to zero at run time.

Covariate shift, DAgger

A cloned policy visits states the demonstrator never did, and errs there; errors compound with the horizon. DAgger fixes it by having the expert label the states the policy actually reaches.

minADE_K, minFDE_K

Best-of-K metrics for multimodal predictors: the average (or final) displacement error of the closest of K sampled trajectories to the one that happened.

Importance-weighted bound (IWAE)

logsumexp_k[log p(x | z_k) + log p(z_k) − log q(z_k | x)] − log K: a lower bound on log p(x) that tightens as K grows; the honest evaluation number for a VAE.

Change of variables

If x = g(ε) with g invertible, log p(x) = log p(ε) − log |det ∂g/∂ε|. Explains why μ + σε has distribution N(μ, σ²), and is the basis of normalizing flows.

Lipschitz, spectral normalization, gradient penalty

A function is 1-Lipschitz if it never stretches distances; the Wasserstein critic must be. Spectral normalization divides each weight matrix by its largest singular value; the gradient penalty penalizes ‖∇D‖ ≠ 1 at interpolated points.

Earth mover’s (Wasserstein-1) distance

The cost of moving one distribution’s mass onto another’s. Unlike JS, it grows with the distance between non-overlapping distributions, so it has a useful gradient.