Lecture 22, Part 5: Meta-RL as Partially Observed MDPs

What this lecture covers

This closing segment on meta-reinforcement learning reframes everything covered so far as a single unifying problem: solving a partially observed MDP (POMDP) in which the hidden state includes an unknown task variable Z. Once you accept that the policy needs to act on s and Z, and that Z must be inferred from a history of states, actions and rewards, meta-learning becomes a special case of POMDP solving, solvable by explicit state estimation or memory-based policies.

The lecture introduces PEARL, a method that trains a variational encoder to infer a distribution over Z from experience and explores by sampling from that posterior, connecting to the earlier discussion of posterior sampling in exploration. It then closes the meta-RL unit by comparing the three perspectives covered across this lecture series (RNN-based, gradient-based, and inference-based) and discusses emergent phenomena linking meta-RL to cognitive science.

Key ideas

• Meta-RL as a POMDP: augmenting the state with an unobserved task variable Z turns solving new tasks into standard POMDP inference: infer Z from history, then act.
• Two ways to handle partial observability: explicit state estimation (estimating a distribution over Z) or memory-based policies (the RNN approach from earlier parts).
• PEARL algorithm: trains a policy conditioned on state and Z, plus an inference network that predicts Z from a history of transitions, using variational inference similar to a variational autoencoder.
• Exploration via posterior sampling: PEARL explores by sampling a Z from its current belief, acting as if that Z were correct, and updating its belief as more data arrives, echoing posterior sampling from the exploration lectures.
• Simple permutation-invariant encoder: PEARL's encoder featurizes each transition independently and averages the features, which works because value-based learning does not depend on transition order.
• Sub-optimality of posterior sampling: this exploration strategy is good but not optimal, since a single well-planned sweep can sometimes find the goal faster than repeated posterior sampling.
• Unifying the three meta-RL perspectives: the RNN approach, gradient-based (MAML) approach, and inference-based approach are all instances of the same F_theta(M_i) template, differing mainly in architecture and whether adaptation is a forward pass, a gradient step, or posterior inference.
• Trade-offs between approaches: RNNs are simple but prone to meta-overfitting with no recourse at test time; gradient-based methods extrapolate well with more test-time gradient steps but are harder to apply beyond policy gradients; inference-based methods give elegant exploration but share the overfitting risk.

Walkthrough

Meta-RL as a partially observed MDP (0:00)

The lecture recaps POMDPs and argues that conditioning a policy on state and an unknown context Z is exactly the POMDP problem: the augmented state includes Z, but only s and the reward are observed. Solving this modified POMDP is shown to be equivalent to meta-learning.

PEARL: variational inference for task context (5:00)

PEARL trains a policy conditioned on s and Z alongside an inference network that predicts a posterior over Z from a history of transitions, using variational inference to maximize expected post-adaptation reward while keeping the posterior close to a prior. Because the encoder is stochastic, sampling Z from it provides a natural, if sometimes sub-optimal, exploration strategy, illustrated with a 2D point-mass task where the agent searches a semicircle of possible goals.

Encoder design and training (8:00)

A simple permutation-invariant encoder, averaging featurized transitions, is shown to work well because off-policy value-based learning treats transitions independently. PEARL is meta-trained with an off-policy actor-critic algorithm such as soft actor-critic, updating the encoder alongside the policy using sampled histories from a replay buffer.

Comparing the three perspectives on meta-RL (9:00)

The lecture closes by comparing the RNN-based, gradient-based, and inference-based approaches as different architectural instantiations of the same F_theta(M_i) framework, discussing their relative strengths around meta-overfitting, extrapolation at test time, and ease of optimization, and noting that a stochastic version of MAML starts to resemble the inference-based approach.

Emergent phenomena and cognitive science connections (12:02)

The lecture ends by connecting meta-RL to cognitive science, noting research showing that meta-RL training can give rise to emergent behaviors resembling episodic recall, model-based adaptation, and even causal reasoning, even when the meta-learning algorithm itself was not designed to produce those specific behaviors.

Before you watch

• Watch Parts 2 through 4 of this lecture, which introduce the RNN-based and gradient-based meta-RL approaches that this segment unifies.
• Review the posterior sampling exploration strategy from the exploration lectures, since PEARL's exploration mechanism builds directly on it.
• Familiarity with variational inference and variational autoencoders helps with the PEARL training objective.

Check your understanding

1. How does framing meta-RL as a POMDP unify the RNN-based and inference-based approaches?
2. Why does PEARL's encoder average featurized transitions rather than processing them sequentially like an RNN?
3. In what sense is posterior sampling exploration sub-optimal, as illustrated by the point-mass example?
4. What are the relative trade-offs between the RNN-based, gradient-based, and inference-based meta-RL approaches discussed at the end of the lecture?