Lecture 18, Variational Inference, Part 4

What this lecture covers

This final part of the lecture applies the amortized variational inference machinery from earlier parts to concrete architectures. It walks through the variational autoencoder (VAE), explains why its representations tend to disentangle underlying factors of variation, and shows how it can be used for representation learning in RL. It then extends the idea to conditional VAEs for multimodal policies, and finally to sequence VAEs used as state space models for partially observed environments.

After watching, you should be able to describe the VAE training objective and computation graph, explain intuitively why sampling from the unit-variance prior produces valid decoded outputs, and describe how the encoder, decoder, and prior are restructured for conditional and sequential settings.

Key ideas

• Variational autoencoder (VAE): an encoder network q_phi(z|x) outputs a mean and variance for z, a decoder network p_theta(x|z) reconstructs x; both are trained together to maximize the evidence lower bound using the reparameterization trick.
• Why VAE samples look valid: the KL term pushes the encoder's output distribution toward the unit-variance prior, and the encoder is incentivized to use the latent space efficiently, so most points sampled from the prior decode into plausible outputs.
• Representation learning for RL: training a VAE on states (such as Atari frames) and feeding the learned z into a Q-function instead of raw pixels can give a more useful, disentangled state representation.
• Disentanglement: because VAE latent dimensions are encouraged to be independent, they tend to separate underlying factors of variation (position, orientation, lighting) better than raw pixels do.
• Conditional VAE: extends the VAE by giving both encoder and decoder an additional input x, modeling p(y|x,z), useful for representing multimodal policies in imitation learning.
• Sequence VAE (state space model): treats an entire trajectory of latent states z_1...z_T as the latent variable, with a learned, structured prior p(z_{t+1}|z_t,a), a per-timestep decoder p(o_t|z_t), and an encoder that infers z_t from the history of observations.
• Applications: sequence VAEs have been used to learn state space models for planning (for example, Embed to Control) and to provide representations for actor-critic RL on image observations (for example, Stochastic Latent Actor-Critic).

Before you watch

• Watch parts 1 through 3 of this lecture, which derive the evidence lower bound and the reparameterization trick used throughout this part.
• Recall the model-based RL lectures on learning latent state models from images, since the sequence VAE revisits that material with the variational inference framework now available.

Check your understanding

1. Why does the KL-divergence term in the VAE objective encourage the encoder to produce latent codes that resemble samples from the prior?
2. In what sense does a VAE trained on images tend to disentangle the underlying factors of variation, and why might that help a downstream Q-function?
3. What changes in the encoder, decoder, and prior when moving from a standard VAE to a conditional VAE for multimodal policies?
4. Why must the encoder in a sequence VAE depend on the history of observations, while the decoder can be independent per timestep?