[Paper Review] VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning

Luisa Zintgraf, Kyriacos Shiarlis+5 more|arXiv (Cornell University)|Oct 18, 2019

Domain Adaptation and Few-Shot LearningComputer Science59 references62 citations

TL;DR

VariBAD meta-learns a variational Bayes adaptive deep RL framework that infers task embeddings to perform approximate Bayes-optimal exploration in unknown environments, enabling online adaptation and improved returns.

ABSTRACT

Trading off exploration and exploitation in an unknown environment is key to maximising expected return during learning. A Bayes-optimal policy, which does so optimally, conditions its actions not only on the environment state but on the agent's uncertainty about the environment. Computing a Bayes-optimal policy is however intractable for all but the smallest tasks. In this paper, we introduce variational Bayes-Adaptive Deep RL (variBAD), a way to meta-learn to perform approximate inference in an unknown environment, and incorporate task uncertainty directly during action selection. In a grid-world domain, we illustrate how variBAD performs structured online exploration as a function of task uncertainty. We further evaluate variBAD on MuJoCo domains widely used in meta-RL and show that it achieves higher online return than existing methods.

Motivation & Objective

• Motivate exploration-exploitation trade-offs in unknown environments and achieve near Bayes-optimal behavior through meta-learning.
• Develop a tractable approach to approximate Bayes-optimal exploration for deep RL without privileged task information.
• Learn a task-embedding latent space that captures uncertainty across related MDPs and use it to condition the policy.
• Combine variational inference with meta-learning to perform online inference and action selection under task uncertainty.

Proposed method

• Represent each task with a low-dimensional stochastic latent variable m that captures MDP embeddings.
• Use a variational auto-encoder to infer the posterior q_phi(m|tau_{:t}) online from agent experience.
• Train a decoder to predict past and future states and rewards from current states and actions to learn the embedding.
• Condition the policy pi_psi(a_t|s_t, q_phi(m|tau_{:t})) on the inferred posterior over MDP embeddings.
• Optimize a joint objective that combines RL loss with an ELBO term over online variational inference (Equation 10).
• Avoid backpropagating RL losses through the encoder in practice to stabilize training and improve efficiency.

Experimental results

Research questions

• RQ1 Can a learned latent embedding of tasks enable Bayes-adaptive exploration in deep RL without privileged task information?
• RQ2 Does meta-learned variational inference over MDP embeddings yield near Bayes-optimal exploration and higher online returns compared to existing meta-RL methods?
• RQ3 Can online inference of task uncertainty through a VAE-like encoder improve performance in both simple gridworlds and MuJoCo domains?
• RQ4 Is it feasible to share a learned reward/transition model across tasks while inferring task-specific embeddings during meta-training and testing?

Key findings

• VariBAD achieves exploration behavior closely matching Bayes-optimal performance in a gridworld, outperforming posterior sampling.
• On MuJoCo continuous control tasks, VariBAD adapts within a single rollout and yields higher online returns than competing meta-RL methods like RL^2.
• The approach learns a meaningful latent space over task embeddings that concentrates after task identification (variance down, mean stabilizes).
• The method demonstrates structured online exploration driven by inferred task uncertainty, outperforming baselines that rely on privileged task information during training in some settings.