VariBAD: Bayesian Meta-Learning in Deep RL

• VariBAD is a meta-learning framework that applies variational inference and Bayes-adaptive principles to model latent task dynamics for efficient exploration.
• It jointly optimizes an encoder, decoder, and policy network to update beliefs over unknown environment parameters, enhancing adaptation.
• Empirical evaluations in gridworld and MuJoCo tasks show that variBAD approximates Bayes-optimal performance and outperforms several leading meta-RL methods.

Variational Bayes-Adaptive Deep Reinforcement Learning (variBAD) is a meta-learning framework for performing approximate Bayes-adaptive reinforcement learning in environments with unknown dynamics and rewards. It enables an agent to maintain a belief distribution over latent task parameters, adaptively trading off exploration and exploitation via a structured uncertainty-driven policy. The method achieves this by incorporating variational inference principles into the RL loop, learning both a posterior over task variables and a policy conditioned on the inferred latent state. Empirical results demonstrate that variBAD outperforms previous meta-RL algorithms on both discrete gridworld and continuous control tasks, closely approximating Bayes-optimal performance in key domains (Zintgraf et al., 2019).

1. Bayes-Adaptive MDP Formulation

The Bayes-Adaptive Markov Decision Process (BAMDP) framework addresses optimal exploration-exploitation tradeoff by augmenting the state space with a posterior belief over hidden task parameters. Let $S$ denote the state space and $A$ the action space. Each environment is parameterized by a latent variable $\theta\in\Theta$, influencing both transition dynamics $T_\theta(s'|s,a)$ and reward functions $R_\theta(r|s,a,s')$. The agent maintains a prior belief $b_0(\theta)=p(\theta)$ and, using its trajectory $\tau_{:t}=(s_0,a_0,r_1,s_1,\ldots,s_t)$, updates its posterior $b_t(\theta)=p(\theta|\tau_{:t})$.

In this framework, the "hyper-state" $s_t^+ = (s_t, b_t)$ produces an augmented BAMDP $M^+$ whose transition and reward kernels are given by: $A$0

$A$1

The Bayes-optimal policy $A$2 maximizes the expected discounted return in $A$3 over horizon $A$4: $A$5 Exact inference and planning are intractable for high-dimensional $A$6, motivating the variational-approximate approach of variBAD.

2. Generative Model and Variational Approximation

variBAD leverages a joint generative model $A$7, where: $A$8 Actions $A$9 are sampled from a policy conditioned on current belief.

The method introduces an amortized variational posterior $\theta\in\Theta$0, parameterized as a diagonal Gaussian $\theta\in\Theta$1 output by a recurrent inference network. The training objective employs the evidence lower bound (ELBO) at each time step $\theta\in\Theta$2: $\theta\in\Theta$3 The ELBO enables tractable meta-learning of the underlying task posterior and dynamics decoder.

3. Meta-Training Algorithm

variBAD’s meta-training jointly optimizes three parameter sets: the encoder $\theta\in\Theta$4 for $\theta\in\Theta$5, the decoder $\theta\in\Theta$6 for $\theta\in\Theta$7, and the policy $\theta\in\Theta$8 for $\theta\in\Theta$9, where $T_\theta(s'|s,a)$0 denotes the sampled latent. The total loss at each meta-training iteration is: $T_\theta(s'|s,a)$1 $T_\theta(s'|s,a)$2 is the standard expected RL return; $T_\theta(s'|s,a)$3 controls the balance between RL and ELBO terms. Training proceeds with policy-gradient updates (PPO/A2C) for $T_\theta(s'|s,a)$4 and Adam updates on $T_\theta(s'|s,a)$5 using ELBO gradients.

The meta-training workflow is summarized as follows:

Step	Operation	Update Type
Sample task $T_\theta(s'|s,a)$6	Reset environment, encoder hidden state	Initialization
Collect trajectory	Encode $T_\theta(s'|s,a)$7 via GRU to $T_\theta(s'|s,a)$8	Forward Pass
Sample latent	$T_\theta(s'|s,a)$9	Forward Pass
Condition policy	$R_\theta(r|s,a,s')$0s_t$R_\theta(r|s,a,s')$1	Action
Compute ELBO	$R_\theta(r|s,a,s')$2 over batch	Loss Evalu.
Optimizer step	Adam/PPO update on $R_\theta(r|s,a,s')$3	Learning

4. Online Adaptation and Uncertainty-Driven Action Selection

During evaluation, only the encoder $R_\theta(r|s,a,s')$4 and policy $R_\theta(r|s,a,s')$5 are retained. With each new trajectory $R_\theta(r|s,a,s')$6, the encoder maintains and updates the latent posterior $R_\theta(r|s,a,s')$7, yielding: $R_\theta(r|s,a,s')$8 As data accumulates, posterior variance $R_\theta(r|s,a,s')$9 collapses ($b_0(\theta)=p(\theta)$0), smoothly annealing the policy from exploration to exploitation. This enables dynamically structured “uncertainty-driven” exploration, matching Bayes-optimal online adaptation.

5. Architecture and Implementation Specifications

• Encoder $b_0(\theta)=p(\theta)$1: MLP embedding, one layer of size 32 (ReLU); GRU (hidden size 64–128); final linear mapping to $b_0(\theta)=p(\theta)$2 for a $b_0(\theta)=p(\theta)$3-dimensional Gaussian ($b_0(\theta)=p(\theta)$4 typical).
• Decoder $b_0(\theta)=p(\theta)$5: Transition model $b_0(\theta)=p(\theta)$6 – MLP (64,32), ReLU; output Gaussian/categorical for $b_0(\theta)=p(\theta)$7. Reward model $b_0(\theta)=p(\theta)$8 – similar MLP, scalar output.
• Policy network $b_0(\theta)=p(\theta)$9: MLP (32 for grid, 128 for MuJoCo), TanH activation, input $\tau_{:t}=(s_0,a_0,r_1,s_1,\ldots,s_t)$0; critic head of similar dimensions.
• Optimization: PPO/A2C, Adam ($\tau_{:t}=(s_0,a_0,r_1,s_1,\ldots,s_t)$1 grid, $\tau_{:t}=(s_0,a_0,r_1,s_1,\ldots,s_t)$2 MuJoCo), clipping $\tau_{:t}=(s_0,a_0,r_1,s_1,\ldots,s_t)$3, value coefficient $\tau_{:t}=(s_0,a_0,r_1,s_1,\ldots,s_t)$4, entropy $\tau_{:t}=(s_0,a_0,r_1,s_1,\ldots,s_t)$5, $\tau_{:t}=(s_0,a_0,r_1,s_1,\ldots,s_t)$6–$\tau_{:t}=(s_0,a_0,r_1,s_1,\ldots,s_t)$7, GAE $\tau_{:t}=(s_0,a_0,r_1,s_1,\ldots,s_t)$8. VAE: Adam $\tau_{:t}=(s_0,a_0,r_1,s_1,\ldots,s_t)$9; ELBO $b_t(\theta)=p(\theta|\tau_{:t})$0 (grid), $b_t(\theta)=p(\theta|\tau_{:t})$1 (MuJoCo). Max grad norm $b_t(\theta)=p(\theta|\tau_{:t})$2. No extra dropout used; KL regularizes latent $b_t(\theta)=p(\theta|\tau_{:t})$3.

6. Empirical Evaluation

• Gridworld: 5$b_t(\theta)=p(\theta|\tau_{:t})$45 grid, unknown goal in 24 cells. Actions: {up, right, down, left, stay}, horizon $b_t(\theta)=p(\theta|\tau_{:t})$5, BAMDP horizon $b_t(\theta)=p(\theta|\tau_{:t})$6. Sparse reward. variBAD achieves returns matching Bayes-optimal by episode 3, outperforming posterior sampling. Decoder’s $b_t(\theta)=p(\theta|\tau_{:t})$7reward$b_t(\theta)=p(\theta|\tau_{:t})$8cell$b_t(\theta)=p(\theta|\tau_{:t})$9 belief concentrates on true goal, with rapid collapse of latent $s_t^+ = (s_t, b_t)$0 confirming uncertainty-driven exploration.
• MuJoCo Meta-RL: Tasks include AntDir (forward/back, 2), HalfCheetahDir (left/right, 2), HalfCheetahVel (varied speeds, $s_t^+ = (s_t, b_t)$110), Walker (randomized body, $s_t^+ = (s_t, b_t)$220). Evaluation metric is first-episode (online) return in new tasks. In all domains, variBAD’s first-rollout returns exceed those of RL², PEARL (off-policy posterior sampling), E-MAML, and ProMP. For example:

Task	variBAD	RL²	PEARL	E-MAML	ProMP
AntDir	~2150	2000	1500	400	600
HalfCheetahDir	~4000	3500	1000	500	800
HalfCheetahVel	~3200	2800	1100	300	500
Walker	~5000	4500	2000	600	700

PEARL converges in $s_t^+ = (s_t, b_t)$3 frames (off-policy); variBAD/RL² require $s_t^+ = (s_t, b_t)$4 frames (on-policy). At convergence, variBAD matches or exceeds oracle PPO returns (which know the true task). Posterior mean flips sign (e.g., direction tasks) within $s_t^+ = (s_t, b_t)$520 steps, and variance $s_t^+ = (s_t, b_t)$6 declines rapidly, enabling early exploitation.

7. Limitations and Future Directions

variBAD is the first scalable deep-RL algorithm to leverage an explicit approximate Bayesian belief over latent task variables for structured exploration. Its variational inference framework delivers a low-dimensional state to condition policies on, along with a quantifiable uncertainty estimate. However, the approach requires meta-training on a distribution $s_t^+ = (s_t, b_t)$7 representative of test tasks and does not guarantee formal Bayes-optimality due to neural network approximation. Training complexity is substantial due to recurrent inference and on-policy learning, with off-policy methods left for future work. Further research directions include exploiting the decoder $s_t^+ = (s_t, b_t)$8 for model-based planning at test time, learning a faster-adapting prior $s_t^+ = (s_t, b_t)$9, and handling out-of-distribution (OOD) tasks via continual encoder fine-tuning.

Relevant experimental data, exact hyperparameters, and architecture specifications are available in the original codebase (Zintgraf et al., 2019).