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: $$T^+(s_{t+1},b_{t+1}|s_t,b_t,a_t) = \mathbb{E}_{\theta\sim b_t}\left[ T_\theta(s_{t+1}|s_t,a_t) \right]\cdot \delta[b_{t+1}=p(\theta|\tau_{:t+1})]$$

$$R^+(s_t,b_t,a_t,s_{t+1},b_{t+1}) = \mathbb{E}_{\theta\sim b_{t+1}}\left[ R_\theta(r|s_t,a_t,s_{t+1}) \right]$$

The Bayes-optimal policy $\pi^{+*}$ maximizes the expected discounted return in $M^+$ over horizon $H^+$: $$J^+(\pi) = \mathbb{E}_{b_0,T^+,\pi}\left[ \sum_{t=0}^{H^+-1} \gamma^t R^+(\cdot) \right]$$ Exact inference and planning are intractable for high-dimensional $\theta$, motivating the variational-approximate approach of variBAD.

2. Generative Model and Variational Approximation

variBAD leverages a joint generative model $$p_\theta(\theta,\tau_{:H^+}) = p(\theta)\,p_\theta(\tau_{:H^+}|\theta)$$, where: $$p_\theta(\tau_{:H^+}|\theta) = p(s_0)\prod_{t=0}^{H^+-1} p_\theta(s_{t+1}|s_t,a_t,\theta)\,p_\theta(r_{t+1}|s_t,a_t,s_{t+1},\theta)$$ Actions $a_t$ are sampled from a policy conditioned on current belief.

The method introduces an amortized variational posterior $q_\phi(\theta|\tau_{:t})$, parameterized as a diagonal Gaussian $(\mu_t,\sigma_t)$ output by a recurrent inference network. The training objective employs the evidence lower bound (ELBO) at each time step $t$: $$\log p_\theta(\tau_{:H^+}) \ge \mathbb{E}_{\theta\sim q_\phi(\cdot|\tau_{:t})}[\log p_\theta(\tau_{:H^+}|\theta)] - \text{KL}[q_\phi(\theta|\tau_{:t})\,||\,p(\theta)]$$ 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 $\phi$ for $q_\phi$, the decoder $\theta$ for $p_\theta$, and the policy $\psi$ for $\pi_\psi(a|s,z)$, where $z$ denotes the sampled latent. The total loss at each meta-training iteration is: $$L(\phi,\theta,\psi) = -\mathbb{E}_{M\sim p(M)}[J_{RL}(\psi,\phi)] + \lambda\,\mathbb{E}_{M,\tau}\Bigg[ \sum_{t=0}^{H^+} \text{KL}[q_\phi(\theta|\tau_{:t})||q_\phi(\theta|\tau_{:t-1})] - \mathbb{E}_{\theta\sim q_\phi}[\log p_\theta(\tau_{:H^+}|\theta)] \Bigg]$$ $J_{RL}$ is the standard expected RL return; $\lambda \in [0,1]$ controls the balance between RL and ELBO terms. Training proceeds with policy-gradient updates (PPO/A2C) for $\psi$ and Adam updates on $(\phi,\theta)$ using ELBO gradients.

The meta-training workflow is summarized as follows:

Step	Operation	Update Type
Sample task $M_i$	Reset environment, encoder hidden state	Initialization
Collect trajectory	Encode $(s_t,a_t,r_{t+1})$ via GRU to $(\mu_t,\sigma_t)$	Forward Pass
Sample latent	$\theta_t = \mu_t + \sigma_t \odot \epsilon$	Forward Pass
Condition policy	$s_t^z = [$s_t$; \theta_t],\quad a_t\sim\pi_\psi(a|s_t^z)$	Action
Compute ELBO	$ELBO_t$ over batch	Loss Evalu.
Optimizer step	Adam/PPO update on $(\phi, \theta, \psi)$	Learning

4. Online Adaptation and Uncertainty-Driven Action Selection

During evaluation, only the encoder $q_\phi$ and policy $\pi_\psi$ are retained. With each new trajectory $\tau_{:t}$, the encoder maintains and updates the latent posterior $(\mu_t,\sigma_t)$, yielding: $$\theta_t \sim \mathcal{N}(\mu_t, \sigma_t),\quad s_t^z = [s_t;\theta_t],\quad a_t \sim \pi_\psi(a|s_t^z)$$ As data accumulates, posterior variance $\sigma_t$ collapses ($\to 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 $q_\phi$: MLP embedding, one layer of size 32 (ReLU); GRU (hidden size 64–128); final linear mapping to $(\mu,\log\sigma)$ for a $d$-dimensional Gaussian ($d=5$ typical).
• Decoder $p_\theta$: Transition model $T_\theta$ – MLP (64,32), ReLU; output Gaussian/categorical for $s'$. Reward model $R_\theta$ – similar MLP, scalar output.
• Policy network $\pi_\psi$: MLP (32 for grid, 128 for MuJoCo), TanH activation, input $[s,\theta]$; critic head of similar dimensions.
• Optimization: PPO/A2C, Adam ($lr=1e^{-3}$ grid, $7e^{-4}$ MuJoCo), clipping $0.1$, value coefficient $0.5$, entropy $0.01$, $\gamma=0.95$–$0.99$, GAE $\tau=0.95$. VAE: Adam $1e^{-3}$; ELBO $\lambda=1.0$ (grid), $0.1$ (MuJoCo). Max grad norm $0.5$. No extra dropout used; KL regularizes latent $\theta$.

6. Empirical Evaluation

• Gridworld: 5$\times$5 grid, unknown goal in 24 cells. Actions: {up, right, down, left, stay}, horizon $H=15$, BAMDP horizon $H^+=60$. Sparse reward. variBAD achieves returns matching Bayes-optimal by episode 3, outperforming posterior sampling. Decoder’s $P($reward$|$cell$)$ belief concentrates on true goal, with rapid collapse of latent $\sigma$ confirming uncertainty-driven exploration.
• MuJoCo Meta-RL: Tasks include AntDir (forward/back, 2), HalfCheetahDir (left/right, 2), HalfCheetahVel (varied speeds, $\sim$10), Walker (randomized body, $\sim$20). 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 $<2\times 10^6$ frames (off-policy); variBAD/RL² require $\sim5\times 10^7$ 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 $\sim$20 steps, and variance $\sigma$ 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 $p(M)$ 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 $p_\theta$ for model-based planning at test time, learning a faster-adapting prior $p(\theta)$, 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).