PaperlayerView exact revision on arXiv

Source-linked AI summary

VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning

Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, Shimon Whiteson

arXiv:1910.08348v2cs.LGstat.ML

TL;DR

Unknown-environment RL requires balancing information-seeking exploration against immediate reward, but exact Bayes-optimal planning and posterior inference are often intractable. variBAD meta-learns approximate inference over task uncertainty and conditions its policy on that belief. It closely matches Bayes-optimal behavior in gridworld and outperforms existing methods in online return on MuJoCo tasks.

  • Problem

    Bayes-optimal exploration is generally intractable to compute, while maintaining the required posterior belief can also be computationally infeasible.

  • Method

    variBAD meta-learns approximate inference over an unknown task and conditions action selection on the inferred posterior over low-dimensional MDP embeddings.

  • Results

    Across gridworld and MuJoCo tasks, variBAD closely matches Bayes-optimal behavior and achieves higher online returns during learning than existing methods.

  • Takeaways & Limitations

    The method provides a tractable path toward approximate Bayes-optimal exploration for deep reinforcement learning.

  • Takeaways & Limitations

    Exact BAMDP planning remains intractable for all but the smallest tasks, motivating variBAD as an approximate rather than exact solution.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Unknown-environment RL must balance exploration and exploitation, but exact Bayes-optimal planning is generally intractable. variBAD meta-learns approximate inference and uses task uncertainty for action selection, closely matching Bayes-optimal behavior in gridworld and achieving higher online returns than existing methods.

  • Motivation: Bayes-optimal policies condition actions on both environment state and uncertainty, optimally balancing exploration and exploitation to maximize expected return.This requires maintaining beliefs over possible environments.
  • Motivation: Exact Bayes-adaptive planning is intractable beyond the smallest tasks, while posterior sampling is tractable but can explore inefficiently and far from Bayes-optimal.Posterior sampling follows a policy optimized for a sampled hypothesis MDP.
  • Gridworld illustration: In the gridworld, Bayes-optimal exploration systematically searches possible goal cells, whereas posterior sampling repeatedly routes toward sampled goals and updates its posterior.The figure contrasts these strategies with variBAD’s learned exploration behavior.
  • Proposed method: variBAD combines Bayesian RL, approximate variational inference, and meta-learning to explore unseen but related environments according to task uncertainty.It is designed to maximize expected online return for a distribution of tasks.
  • Proposed method: The method uses a variational auto-encoder for posterior inference and a policy conditioned on posterior belief over MDP embeddings.Together, these components learn how to trade off exploration and exploitation under task uncertainty.
  • Results: VariBAD closely matches Bayes-optimal performance in the gridworld and matches optimal performance from the third rollout.The results are shown against hard-coded optimal, Bayes-optimal, and posterior-sampling strategies.
  • Results: Across gridworld and MuJoCo domains, variBAD exhibits structured exploratory behavior and achieves higher returns during learning than existing methods.The approach is presented as a tractable and flexible method tailored to the training-task distribution.

2 BACKGROUND

The background formulates unknown-task RL through MDPs, beliefs over reward and transition functions, and Bayes-adaptive MDPs. Although this framework gives principled exploration–exploitation behavior, posterior inference and belief-space planning are often intractable, motivating end-to-end meta-learning without test-time planning.

  • MDP and meta-learning setting: An MDP specifies states, actions, reward and transition functions, an initial-state distribution, a discount factor, and a horizon.In the meta-learning setting, tasks are sampled from a distribution of MDPs whose reward and transition functions share structure.
  • MDP and meta-learning setting: Meta-training samples batches of related MDPs, while the task description or task ID remains unknown to the agent.Across tasks, reward and transition functions vary but share some structure.
  • Bayesian reinforcement learning: Bayes-adaptive RL represents uncertainty with a belief over possible reward and transition functions and augments the environment state with that belief.The resulting hyper-state lies in the product of the state space and belief space.
  • Bayesian reinforcement learning: The Bayes-optimal policy takes exploratory actions to reduce task uncertainty only when doing so helps maximize expected return within the horizon.The BAMDP horizon may span multiple MDP episodes, and exploration decisions depend on remaining time.
  • Challenges: Bayesian RL faces three challenges: unknown model parameterization, often-intractable belief updates, and typically intractable planning in belief space.These challenges jointly limit direct BAMDP solution methods.
  • Methodological response: The proposed approach meta-learns models, inference, and belief use jointly, requiring no privileged task information during training and no planning at test time.This is intended to provide a flexible and scalable Bayes-adaptive deep-RL approach.

3 BAYES-ADAPTIVE DEEP RL VIA META-LEARNING

VariBAD represents unknown tasks with stochastic latent embeddings, infers their posterior online from trajectory experience, and trains a policy to act under task uncertainty. Its objective jointly learns inference, environment models, and policy behavior using meta-training across related MDPs.

  • Task representation: VariBAD represents each unknown MDP with a learned, low-dimensional stochastic task embedding inferred from the agent’s experience.The embedding replaces direct inference over potentially very large reward and transition functions.
  • Approximate inference: An RNN encodes past states, actions, and rewards into the posterior qφ(m|τ:t), while a decoder predicts past and future states and rewards.Future decoding is available during training and helps inference generalize to unseen states from past experience.
  • Uncertainty-aware policy: The policy conditions on both the environment state and posterior over m, enabling action selection that accounts for uncertainty while trading off exploration and exploitation.This learned posterior replaces direct inference over transition and reward functions and can use shared information across tasks.
  • Training objective: The training objective combines the RL policy loss with ELBO terms across context lengths, allowing online inference to improve as more interaction data arrive.A weighting coefficient λ balances supervised model learning against the RL loss, and ELBO terms may be subsampled for efficiency.
  • Meta-test procedure: At meta-test time, the encoder and policy act directly in sampled tasks without decoder use or gradient adaptation.The policy is trained during meta-training to behave approximately Bayes-optimal at deployment.

4 RELATED WORK

VariBAD differs from related meta-RL and Bayesian RL methods by meta-learning unsupervised inference over task uncertainty and conditioning its policy on that posterior. The approach targets tractable approximate Bayes-adaptive exploration without privileged task information, while acknowledging that deep neural networks lack some formal guarantees.

  • Meta Reinforcement Learning: Unlike recurrent model-free meta-RL, variBAD adds a stochastic latent variable, decoder, and variational objective to represent uncertainty about tasks.Removing the decoder and VAE objective reduces variBAD to the recurrent adaptation setting described for RL2.
  • Meta Reinforcement Learning: Unlike MAML and ProMP, variBAD directly accounts for exploration during adaptation by conditioning action selection on a posterior over MDPs.MAML and ProMP instead learn an initialization from which a few gradient steps can produce good performance.
  • Distinction from embedding methods: VariBAD’s embedding represents task uncertainty, and its posterior-conditioned policy learns to trade off exploration and exploitation online.Its objective explicitly optimizes for Bayes-optimal behavior rather than using the model for test-time planning.
  • Bayesian Reinforcement Learning: VariBAD offers tractable approximate Bayes-optimal exploration assuming meta-training on related tasks, but deep neural networks do not provide the formal guarantees of some Bayesian RL methods.Posterior sampling remains more tractable but is described as less efficient than Bayes-optimal behavior.
  • Variational inference and meta-learning: VariBAD meta-learns an unsupervised inference procedure and does not require privileged task descriptions during training.This contrasts with approaches that meta-train posterior representations using task descriptions or other privileged information.

5 EXPERIMENTS

The experiments examine variBAD’s structured online exploration in gridworld and its ability to adapt during the first rollout on MuJoCo tasks. Across these settings, variBAD approximates Bayes-optimal behaviour and achieves strong single-episode performance.

  • Experimental settings: The experiments cover a didactic gridworld and four MuJoCo continuous-control tasks commonly used in meta-RL.The MuJoCo tasks include AntDir, HalfCheetahDir, HalfCheetahVel, and Walker.
  • Gridworld: In gridworld, variBAD uses deterministic actions to infer the hidden goal while exploring unvisited cells.Its learned reward predictions exclude visited cells and guide exploration until the goal is found.
  • Gridworld: VariBAD closely matches Bayes-optimal behaviour and reaches optimal performance from the third rollout.The comparison includes a privileged optimal policy, Bayes-optimal control, and posterior sampling.
  • MuJoCo continuous control: Only variBAD and RL2 adapt within a single MuJoCo episode, while variBAD outperforms RL2 on HalfCheetahDir.The other evaluated methods require substantially more interactions in each new task to perform well.
  • MuJoCo continuous control: VariBAD’s first rollout matches the optimal oracle policy up to a small margin and supports maximising expected reward within one episode.The first rollout includes exploratory steps, but the reported performance remains close to the oracle conditioned on the true task.

6 CONCLUSION & FUTURE WORK

The paper presents variBAD as a meta-learning approach to approximate Bayes-optimal exploration, with results spanning gridworld and more challenging MuJoCo tasks. It concludes that the method offers a path toward tractable approximate Bayes-optimal exploration while identifying out-of-distribution generalisation as future work.

  • Conclusion: VariBAD uses meta-learning to exploit knowledge from related tasks and perform approximate inference in unknown environments.The method is presented as a deep RL approach to approximate Bayes-optimal behaviour.
  • Conclusion: In gridworld, variBAD closely matches Bayes-optimal behaviour, while in MuJoCo it outperforms existing methods during a single episode.These results support the paper’s central empirical claim across simple and more challenging environments.
  • Conclusion: The authors identify tractable approximate Bayes-optimal exploration for deep RL as the broader direction opened by variBAD.This is stated as the paper’s concluding perspective rather than as a demonstrated guarantee.
  • Future work: Out-of-distribution task generalisation may require retraining inference components and updating the policy or adding explicit planning.The stated concerns are incorrect prior or posterior updates and difficulty interpreting a changed posterior.

A FULL ELBO DERIVATION

The supplied passage only states that Equation (8) can be derived; it does not provide the derivation or explain the equation’s role.

  • Equation (8) is introduced as derivable, but the supplied passage gives no derivation steps or mathematical details.

B EXPERIMENTS: GRIDWORLD

The gridworld experiments compare variBAD with oracle, Bayes-optimal, posterior-sampling, and RL2 policies across multiple rollouts. They show near-Bayes-optimal behaviour and a concentrated latent task representation.

  • Additional remarks: The latent dimensions begin near mean 1 and variance 0, then variance increases briefly before the goal is found.This behaviour motivates learning the prior to better match the task distribution.
  • Gridworld evaluation: Figure 5 reports gridworld results averaged over 20 seeds, with 95% confidence intervals for the learning curve.
  • Gridworld evaluation: The learning curves compare variBAD and RL2 with an oracle policy, alongside hard-coded Bayes-optimal and posterior-sampling policies.Training used H+ = 4 × H = 60 across four episodes.
  • Gridworld evaluation: VariBAD and RL2 both closely approximate the Bayes-optimal solution, but variBAD learned it in 4 of 20 seeds compared with zero RL2 seeds.Their remaining solutions were reported to be very close to Bayes-optimal.
  • Additional remarks: After the fourth rollout, RL2 performance drops again, whereas variBAD’s concentrated latent representation remains unchanged with additional data.The reported RL2 decline is attributed to likely instability in its 128-dimensional hidden state.

C.1 LEARNING CURVES

The MuJoCo learning-curve comparison evaluates performance at the first and later rollouts, while accounting for differences in rollout counts and training setups. PEARL is more frame-efficient because it is off-policy, whereas variBAD uses a reward decoder across all MuJoCo environments.

  • Learning-curve setup: Figure 6 compares first-rollout performance with performance at the N-th rollout across the MuJoCo approaches.N = 2 for variBAD and RL2, 20 for ProMP and E-MAML, and 10 for PEARL.
  • Learning-curve comparison: PEARL is much more sample efficient in number of frames than the other methods because it is off-policy.The authors describe on-policy versus off-policy training as orthogonal to their contribution.
  • Evaluation details: The Walker evaluation uses models obtained after 5e+7 frames for variBAD and the Oracle because performance declined after half the training time.This is a specific evaluation choice rather than a general training rule.
  • Training details: VariBAD was trained with a reward decoder only for all MuJoCo environments, including Walker despite its changing dynamics.The authors report that this configuration had superior performance for Walker.

C.2 TRAINING DETAILS AND COMPARISON TO RL2

The comparison examines multi-rollout evaluation, test-time behavior, runtime, and latent-space adaptation. VariBAD avoids redoing task inference after resets, adapts online in HalfCheetahDir, and trains faster than RL2 with PPO, while direct belief visualization remains difficult in MuJoCo.

  • Training details: VariBAD and RL2 are trained for two rollouts, with H = 200 per rollout and H+ = 400, while hidden states persist across resets.A done flag tells the agent when it is reset between episodes.
  • Comparison to RL2: RL2 can be unstable across multiple rollouts, whereas variBAD can rely on its task-only latent embedding after a reset.The authors hypothesize that RL2’s drop in CheetahVel may involve reset-induced hidden-state changes, task inference from velocity, or implementation details.
  • Test-time behavior: Figure 7 organizes HalfCheetahDir behavior by position on the x-axis, environment steps on the y-axis, example rows, and rollout-count columns.The y-axis is read from bottom to top.
  • Test-time behavior: VariBAD and RL2 adapt online to “go left,” whereas PEARL can initially walk in the wrong direction according to its current sample.The comparison concerns the first three test-time rollouts in HalfCheetahDir.
  • Runtime comparison: Recurrent networks make forward and backward passes slow for long horizons, but variBAD trains faster than RL2 with PPO by avoiding RL-loss backpropagation through its recurrent part.This runtime advantage is less pronounced for other RL methods requiring fewer forward and backward passes per policy update.
  • Latent-space visualization: VariBAD’s latent posterior mean and log-variance adapt within a few environment steps, with variance decreasing as task certainty increases.The latent dimensions also swap signs between the “go right” and “go left” tasks.
  • Latent-space visualization: Direct visualization of belief in reward or state space is more difficult for MuJoCo because its states and actions are continuous.The authors suggest a separate ground-truth task-description predictor for analysis without using privileged information during meta-training.
  • Implementation: The experiments use PyTorch, with MuJoCo default arguments documented in the reference implementation.The implementation is identified by the paper’s linked variBAD repository.
Loading 1910.08348v2…