Variational Inference

Latent variable models: complex distributions from simple pieces

This lecture covers no new RL algorithm — deliberately. Variational inference “comes up again and again” across model-based RL, inverse RL, and exploration, and it has “a very deep connection to reinforcement learning and learning-based control” that next week’s control-as-inference lecture makes precise. Today is the tooling.

A probabilistic model is anything representing a distribution: $p(x)$ fit to observed points, or a conditional $p(y \mid x)$ — the course has been full of these, since every policy $\pi_\theta(a \mid s)$ is a conditional probabilistic model.

The conditional version already appeared in lecture 2: the mixture density network for multimodal imitation (driving around the tree) is $p(y \mid x) = \sum_z p(y \mid x, z)\, p(z \mid x)$, with the network emitting the mixture means, variances, and weights. Whether $p(z)$ depends on $x$, or $p(y \mid \cdot)$ does, is a design choice.

In general the recipe is: a prior $p(z)$ from some simple class (categorical, or a Gaussian), and a conditional $p(x \mid z)$ that is also simple as a distribution — say $p(x \mid z) = \mathcal{N}(\mu_{\text{nn}}(z), \sigma_{\text{nn}}(z))$ — even though its parameters are arbitrarily complicated neural-net functions of $z$. “Simple” here means explicitly parameterizable: a single Gaussian can capture $p(z)$ perfectly and can represent $p(x \mid z)$, but no single Gaussian captures $p(x)$ itself, which is

$p(x) = \int p(x \mid z)\, p(z)\, dz$

RL supplies the motivating cases: multimodal policies with a noise input $z \sim \mathcal{N}(0, \mathbf{I})$, and the latent state models from the model-based RL lectures, where the “prior” over latents is structured — it models $p(x_{t+1} \mid x_t, u_t)$ and $p(x_1)$, not just an isotropic Gaussian. The lecture returns to those at the end.

Why training them is hard

Given data $x_1, \ldots, x_N$ and a model $p_\theta(x)$, the natural objective is the maximum likelihood fit. Substituting the latent variable decomposition:

$\theta \leftarrow \arg\max_\theta \frac{1}{N} \sum_i \log\left(\int p_\theta(x_i \mid z)\, p(z)\, dz\right)$

That integral, evaluated every gradient step, is “completely intractable” for continuous $z$. Even in the discrete-mixture case where you can sum over components, you often don’t want to — the resulting optimization landscape has “very poor numerical properties.”

The alternative — stated first, justified later — is the expected log likelihood:

$\theta \leftarrow \arg\max_\theta \frac{1}{N} \sum_i \mathbb{E}_{z \sim p(z \mid x_i)}\left[\log p_\theta(x_i, z)\right]$

The intuition: treat the latents as missing halves of the data. “Guess” the $z$ that goes with each $x_i$ and pretend it’s the right one — except that many $z$‘s are plausible, so average over the whole posterior $p(z \mid x_i)$ instead of committing to one. And this objective is tractable in exactly the way the original wasn’t: the log of an integral doesn’t decompose, but an expectation of a log can be estimated by sampling $z \sim p(z \mid x_i)$ and averaging — the same move policy gradient made. Everything now hinges on one question: how do we get $p(z \mid x_i)$?

The variational lower bound

The posterior $p(z \mid x_i)$ is in general complex — a single $x$ can come from many places in $z$-space. The variational move is a crude approximation: fit some $q_i(z)$ from a simple tractable class, say a Gaussian $\mathcal{N}(\mu_i, \sigma_i)$, one per data point. The slides are blunt about it: “this is incorrect but very convenient.” What makes it more than convenient is that any $q_i$ yields a lower bound on $\log p(x_i)$. Multiply and divide by $q_i(z)$ inside the integral, recognize an expectation, and apply Jensen’s inequality ($\log \mathbb{E}[y] \ge \mathbb{E}[\log y]$, since $\log$ is concave):

$\log p(x_i) = \log \mathbb{E}_{z \sim q_i(z)}\!\left[\frac{p(x_i \mid z)\, p(z)}{q_i(z)}\right] \ge \mathbb{E}_{z \sim q_i(z)}\left[\log p(x_i \mid z) + \log p(z)\right] + \mathcal{H}(q_i)$

The right-hand side is the evidence lower bound (ELBO), written $\mathcal{L}_i(p, q_i)$. Every term is tractable: the expectation by sampling from $q_i$, and the entropy $\mathcal{H}(q_i) = -\mathbb{E}_{q_i}[\log q_i(z)]$ in closed form for a Gaussian.

How loose is the bound, and what makes a good $q_i$? Intuitively $q_i$ should approximate $p(z \mid x_i)$, compared in KL divergence — and writing out that KL gives the lecture’s central identity.

This immediately suggests an algorithm: alternate maximizing $\mathcal{L}_i$ with respect to $q_i$ (tighten the bound) and with respect to $\theta$ (raise the likelihood). Concretely, per data point: sample $z \sim q_i(z)$, take $\nabla_\theta \mathcal{L}_i \approx \nabla_\theta \log p_\theta(x_i \mid z)$, step $\theta$, then gradient-ascend $\mu_i, \sigma_i$ on $\mathcal{L}_i$.

From per-datapoint posteriors to amortized inference

The catch is the parameter count: a separate $(\mu_i, \sigma_i)$ per data point means $|\theta| + (|\mu_i| + |\sigma_i|) \times N$ parameters — fine for thousands of points, intractable for the millions typical in deep learning. But every $q_i$ is trying to approximate the same function evaluated at a different point: the posterior. So train one inference network $q_\phi(z \mid x) = \mathcal{N}(\mu_\phi(x), \sigma_\phi(x))$ that outputs the posterior parameters for any $x$, amortizing inference across the dataset. Two networks now: the generative model $p_\theta(x \mid z)$ and the inference network $q_\phi(z \mid x)$.

Training $\theta$ works exactly as before. The new question is $\nabla_\phi \mathcal{L}$, where $\phi$ appears both as the sampling distribution of the expectation and in the entropy term. The entropy is closed-form in $\mu_\phi, \sigma_\phi$ — easy. The first term is

$J(\phi) = \mathbb{E}_{z \sim q_\phi(z \mid x_i)}\left[r(x_i, z)\right], \qquad r(x_i, z) = \log p_\theta(x_i \mid z) + \log p(z)$

The reparameterization trick

In RL we’re stuck with the likelihood-ratio estimator because we can’t differentiate through unknown dynamics. Here there are no unknown dynamics — only $q$ — so a better estimator exists. A Gaussian sample can be rewritten as a deterministic function of independent noise:

$z = \mu_\phi(x) + \epsilon\, \sigma_\phi(x), \qquad \epsilon \sim \mathcal{N}(0, 1)$

so that, as a strict equality, $J(\phi) = \mathbb{E}_{\epsilon \sim \mathcal{N}(0,1)}\left[r\big(x_i,\, \mu_\phi(x_i) + \epsilon\, \sigma_\phi(x_i)\big)\right]$. Now $\phi$ parameterizes only deterministic quantities: sample $\epsilon$‘s, and backpropagate $\nabla_\phi r$ directly through $\mu_\phi$ and $\sigma_\phi$. A single sample per data point usually suffices, because this estimator uses the derivatives of $r$ — information the policy-gradient version throws away.

An even cleaner grouping: the $\log p(z)$ and entropy terms together are exactly $-D_{KL}(q_\phi(z \mid x_i) \| p(z))$, which for two Gaussians has an analytic formula. The working form of the ELBO becomes

$\mathcal{L}_i \approx \log p_\theta\big(x_i \mid \mu_\phi(x_i) + \epsilon\, \sigma_\phi(x_i)\big) - D_{KL}\big(q_\phi(z \mid x_i)\, \big\|\, p(z)\big)$

with one Gaussian $\epsilon$ per data point, and autodiff handles everything.

The trade-off between the two estimators: the policy-gradient form handles discrete and continuous latents ($q$ can be any distribution with log probabilities) but has high variance, needing multiple samples and smaller learning rates. Reparameterization needs $r$ differentiable in $z$ — continuous latents only — but is simple, low variance, and works with one sample. Continuous $z$: reparameterize. Discrete $z$: policy gradient.

VAEs, conditional VAEs, and sequence models for RL

The variational autoencoder is the most direct instantiation: an encoder network $q_\phi(z \mid x)$ mapping an image to $(\mu_\phi(x), \sigma_\phi(x))$ over a latent vector (dimension 64 or 128, say), and a decoder $p_\theta(x \mid z)$ mapping $z$ back to a distribution over pixels. To sample: draw $z$ from the prior $\mathcal{N}(0, \mathbf{I})$ and decode.

Why does decoding prior samples give valid images? The KL penalty means the encoder “would pay a very heavy price” for producing $z$‘s unlike prior samples — that keeps encodings inside the prior. The converse — every prior sample decoding to something valid — is an efficiency argument: the encoder “wants to be pretty frugal in its use of the latent space,” expanding its variance toward 1 to fill any unused region, so nearly every $z$ under the prior maps to some valid $x$.

Three RL applications close the lecture:

Representation learning. Train a VAE on the states in your replay buffer (Atari images, say) and feed the Q-function $z$ instead of raw pixels. Because the prior factorizes into independent dimensions, the VAE is pushed toward disentangling the underlying factors of variation — the position and velocity of the player character and the skull in Montezuma’s Revenge, not per-pixel colors. The Higgins et al. 2017 interpolations (chairs rotating, faces relighting) show latent traversals tracking factors rather than pixel blends — though Levine notes it’s “debatable in practice” how fully VAEs achieve this. The loop also pre-trains cleanly from prior data.

Conditional VAEs for multimodal policies. Condition encoder and decoder on $x$ (the prior usually stays unconditional): the decoder $p_\theta(y \mid x, z)$ is a policy taking a noise sample as extra input. These appear mostly in imitation learning — fully observed MDPs admit deterministic optimal policies, but human demonstrations are multimodal and non-Markovian (left or right around the tree, never the middle). Examples from the lecture: latent “plans” learned from free-form human play data, and a Transformer-based conditional VAE driving a real bimanual robot putting a shoe on a foot.

Sequence VAEs for partially observed systems. The latent variable is the entire sequence $z_1, \ldots, z_T$; the observed variable is the trajectory of observations; and the whole thing is conditioned on the action sequence. The prior carries the dynamics — $p(z_1) \prod_t p(z_{t+1} \mid z_t, a_t)$, learned, so latents correlate across time. The decoder is deliberately independent per step, $\prod_t p(o_t \mid z_t)$, forcing the $z$‘s to constitute a Markovian state space. The encoder cannot be per-step — a single $o_t$ underdetermines $z_t$ — so it is the most complex piece, a sequence model (LSTM or Transformer) inferring $q(z_t \mid o_{1:t}, a_{1:t})$, optionally conditioned on $z_{t-1}$; any choice yields a valid bound, better encoders just tighten it. This is the machinery behind embed-to-control’s recovered pendulum state spaces, latent-space planning from pixels, and stochastic latent actor-critic running RL in the learned state space.

Check yourself

Derivation and modeling exercises in the lecture’s spirit — work them on paper.

Problem 18.1 ELBO from the KL side

Without using Jensen’s inequality, derive the evidence lower bound: expand $D_{KL}(q_i(z) \| p(z \mid x_i))$, arrive at $\log p(x_i) = D_{KL}(q_i \| p(z \mid x_i)) + \mathcal{L}_i(p, q_i)$, and state the two optimization consequences this equality delivers.

Show solution

Start from the definition and substitute $p(z \mid x_i) = p(x_i, z)/p(x_i)$: $D_{KL} = \mathbb{E}_{q_i}\!\left[\log \frac{q_i(z)\, p(x_i)}{p(x_i, z)}\right] = -\mathbb{E}_{q_i}[\log p(x_i \mid z) + \log p(z)] + \mathbb{E}_{q_i}[\log q_i(z)] + \log p(x_i)$, where $\log p(x_i)$ exits the expectation because it doesn’t depend on $z$. The middle term is $-\mathcal{H}(q_i)$, so $D_{KL} = -\mathcal{L}_i(p, q_i) + \log p(x_i)$; rearranging gives the identity. Consequence one: $D_{KL} \ge 0$ always, so $\log p(x_i) \ge \mathcal{L}_i$ — the bound, with no Jensen step. Consequence two: the left side is independent of $q_i$, so maximizing $\mathcal{L}_i$ with respect to $q_i$ must minimize the KL — the same objective simultaneously tightens the bound (over $q$) and raises the likelihood (over $\theta$), which is what licenses alternating ascent on a single quantity.

Problem 18.2 bound tightness by hand

A model has one binary latent: $p(z = A) = p(z = B) = \tfrac{1}{2}$, and for a particular data point, $p(x_i \mid A) = 0.8$ and $p(x_i \mid B) = 0.2$. Compute $\log p(x_i)$, the true posterior, and the ELBO for (a) $q$ equal to the posterior and (b) $q$ a point mass on $A$. Verify that each gap equals $D_{KL}(q \| p(z \mid x_i))$.

Show solution

$p(x_i) = \tfrac{1}{2}(0.8) + \tfrac{1}{2}(0.2) = 0.5$, so $\log p(x_i) = \log 0.5 \approx -0.693$. Posterior: $p(A \mid x_i) = 0.4/0.5 = 0.8$, $p(B \mid x_i) = 0.2$. (a) $\mathcal{L} = 0.8(\log 0.8 + \log 0.5) + 0.2(\log 0.2 + \log 0.5) + \mathcal{H}(q)$ with $\mathcal{H}(q) = -(0.8 \log 0.8 + 0.2 \log 0.2) \approx 0.500$. The expectation term is $0.8(-0.916) + 0.2(-2.303) \approx -1.194$, so $\mathcal{L} \approx -0.693 = \log p(x_i)$: the bound is tight, and indeed $D_{KL}(q \| p(z \mid x_i)) = 0$ since $q$ is the posterior. (b) Point mass on $A$: $\mathcal{H}(q) = 0$ and $\mathcal{L} = \log 0.8 + \log 0.5 \approx -0.916$. Gap: $-0.693 - (-0.916) = 0.223$. And $D_{KL} = 1 \cdot \log(1/0.8) \approx 0.223$ — exactly the gap, as the decomposition demands. Committing to the single most likely $z$ costs precisely the KL to the true posterior.

Problem 18.3 two gradient estimators

For $J(\phi) = \mathbb{E}_{z \sim q_\phi(z \mid x)}[r(x, z)]$ with $q_\phi = \mathcal{N}(\mu_\phi(x), \sigma_\phi(x))$: (a) write both the policy-gradient estimator and the reparameterized estimator of $\nabla_\phi J$; (b) explain why the reparameterized one has lower variance; (c) explain why the same trick is generally unavailable for the RL objective it resembles.

Show solution

(a) Policy gradient: sample $z_j \sim q_\phi$, use $\nabla_\phi J \approx \frac{1}{M} \sum_j \nabla_\phi \log q_\phi(z_j \mid x)\, r(x, z_j)$. Reparameterized: sample $\epsilon_j \sim \mathcal{N}(0, 1)$, use $\nabla_\phi J \approx \frac{1}{M} \sum_j \nabla_\phi\, r\big(x,\, \mu_\phi(x) + \epsilon_j\, \sigma_\phi(x)\big)$, valid because $z = \mu_\phi + \epsilon \sigma_\phi$ makes the expectation over $\epsilon$, which is independent of $\phi$, a strict equality. (b) The likelihood-ratio form treats $r$ as a black-box scalar weight — it uses only $r$‘s value, so noise in the samples multiplies large score terms. The reparameterized form differentiates through $r$, injecting first-order information about how $r$ changes with $z$; empirically a single sample per data point suffices where the policy gradient needs many. (c) In RL the role of $r$ is played by the return, which depends on the environment dynamics: writing the trajectory as a deterministic differentiable function of policy parameters and external noise would require differentiating through $p(s_{t+1} \mid s_t, a_t)$, which is unknown and not assumed differentiable. In amortized VI, $r = \log p_\theta(x \mid z) + \log p(z)$ is just another neural network — its derivatives are available to autodiff. (The trick also fails for discrete $z$, where $\nabla_z r$ is undefined; there the policy-gradient estimator remains the tool.)

Problem 18.4 wiring a sequence VAE

You have trajectories of image observations $o_{1:T}$ and actions $a_{1:T}$ from a partially observed system and want a latent state-space model as a variational autoencoder. Identify: the latent variable, the observed variable, the conditioning information, and the factorized forms of the prior and decoder. Then explain why the decoder should be independent per time step while the encoder must not be.

Show solution

The latent variable is the whole sequence $z_{1:T}$ (not one latent per step in separate VAEs); the observed variable is the trajectory $o_{1:T}$; and the model is a conditional sequence VAE conditioned on the actions $a_{1:T}$ — the one part of the data not being modeled. The prior carries the structure: $p(z_{1:T} \mid a_{1:T}) = p(z_1) \prod_{t} p(z_{t+1} \mid z_t, a_t)$, with $p(z_1)$ a unit Gaussian but the dynamics factors learned — so $z_2, z_3, \ldots$ are not unit Gaussians, and the latent dimensions deliberately correlate across time. The decoder factorizes per step: $\prod_t p(o_t \mid z_t)$. This is a modeling choice that forces each $z_t$ to summarize everything needed to render its own observation — i.e., forces the latents to form a Markovian state space; letting the decoder peek at other steps would let $z_t$ off the hook. The encoder cannot be per-step precisely because the system is partially observed: $o_t$ alone underdetermines $z_t$, so the encoder must aggregate history, $q(z_t \mid o_{1:t}, a_{1:t})$, typically with an LSTM or Transformer, optionally also conditioning on $z_{t-1}$. Any such encoder yields a valid lower bound; richer ones give more accurate posteriors and hence tighter bounds.