---
title: 'VariBAD: Bayesian Meta-Learning in Deep RL'
url: https://www.emergentmind.com/topics/variational-bayes-adaptive-deep-reinforcement-learning-varibad
type: topic
---

# VariBAD: Bayesian Meta-Learning in Deep RL

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 [1910.08348].

## 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 [1910.08348].

Source: https://www.emergentmind.com/topics/variational-bayes-adaptive-deep-reinforcement-learning-varibad