VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning
Authors: Luisa Zintgraf
Kyriacos Shiarlis
Maximilian Igl
Sebastian Schulze
Yarin Gal
Katja Hofmann
Shimon Whiteson
Venue:  (2020)
Keywords: Meta-Learning
Bayesian Reinforcement Learning
BAMDPs
Deep Reinforcement Learning

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 uncer- tainty 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.

Published as a conference paper at ICLR 2020
VARIBAD: A VERY GOOD METHOD FOR BAYES-ADAPTIVE DEEP RL VIA META-LEARNING
Luisa Zintgraf University of Oxford
Kyriacos Shiarlis Latent Logic
Maximilian Igl University of Oxford
Sebastian Schulze University of Oxford
Yarin Gal University of Oxford
Katja Hofmann Microsoft Research
Shimon Whiteson
University of Oxford
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 an unknown environment, and incorporate task uncer- tainty directly during action selection. In grid-world domain, 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.
INTRODUCTION
Reinforcement learning (RL) is typically concerned with finding an optimal policy that maximises expected return for given Markov decision process (MDP) with an unknown reward and transition function. If these were known, the optimal policy could theory be computed without environment interactions. By contrast, learning in unknown environment usually requires trading off explo- ration (learning about the environment) and exploitation (taking promising actions). Balancing this trade-off is key to maximising expected return during learning, which is desirable in many settings, particularly in high-stakes real-world applications like healthcare and education (Liu et al.) 2014] Yauney & Shah 2018). A Bayes-optimal policy, which does this trade-off optimally, conditions actions not only on the environment state but on the agent's own uncertainty about the current MDP.
In principle, a Bayes-optimal policy can be computed using the framework of Bayes-adaptive Markov decision processes (BAMDPs) (Martin) 1967] Duff & Barto 2002). in which the agent maintains belief distribution over possible environments. Augmenting the state space of the under- lying MDP with this belief yields BAMDP, a special case of a belief MDP (Kaclbling et al.) 1998). Bayes-optimal agent maximises expected return in the BAMDP by systematically seeking out the data needed to quickly reduce uncertainty, but only insofar as doing so helps maximise expected return. performance is bounded from above by the optimal policy for the given MDP, which does not need to take exploratory actions but requires prior knowledge about the MDP to compute.
Unfortunately, planning in a BAMDP, i.e., computing Bayes-optimal policy that conditions on the augmented state, is intractable for all but the smallest tasks. A common shortcut is to rely instead on posterior sampling (Thompson] [933] Strens)] 2000; Osband et al. 2013). Here, the agent periodically samples single hypothesis MDP (c.g., at the beginning of an episode) from its posterior, and the policy that is optimal for the sampled MDP is followed until the next sample is drawn. Planning is far more tractable since it is done regular MDP, not BAMDP. However, posterior sampling's exploration can be highly inefficient and far from Bayes-optimal.
" luisaaiingrra@@essx....uu

---

Published as a conference paper at ICLR 2020
Aestage 1o0m
(a) Environment
(b)
(c) Posterior Sampling
(d)
Bayes- Optimal
VariBAD
Epaadds
(e) Performance
The agent starts the bottom left and has to navigate to an unknown goal, located in the grey area. (b) A Bayes-optimal exploration strategy that systematically searches possible grid cells find the goal, shown in solid (past actions) and dashed (future actions) blue lines. A simplified posterior is shown in the back- ground in grey (p (/uumbe of possible goal positions left) of containing the goal) and white (p 0). (c) Posterior sampling, which repeatedly samples a possible goal position (red squares) from the current posterior, takes the shortest route there, and updates its posterior. (d) Exploration strategy learned by variBAD. The grey background represents the approximate posterior the agent has learned. (e) Average return over all possible environments, over six episodes with 15 steps each (after which the agent is reset to the starting position). VariBAD results are averaged across 20 random seeds. The performance of any exploration strategy is bounded above by the optimal be- haviour (of policy with access the true goal position). The Bayes-optimal agent matches this behaviour from the second episode, whereas posterior sampling needs six rollouts. VariBAD closely approximates Bayes-optimal behaviour in this environment.
Figure 1: Illustration of different exploration strategies. (a) Environment:
Consider the example of a gridworld in Figure D where the agent must navigate to an unknown goal located in the grey area [aa). To maintain a posterior, the agent uniformly assign non-zero probability to cells where the goal could be, and zero to all other cells. A Bayes-optimal strategy strategically searches the set of goal positions that the posterior considers possible, until the goal is found (Tb). Posterior sampling by contrast samples a possible goal position, takes the shortest route there, and then resamples different goal position from the updated posterior (Tc). Doing much less efficient since the agent's 's uncertainty is not reduced optimally (e.g., states are revisited).
As this example illustrates, Bayes-optimal policies explore much more efficiently than posterior sampling. A key challenge is to learn approximately Bayes-optimal policies while retaining the tractability of posterior sampling. In addition, the inference involved in maintaining a posterior belief, needed even for posterior sampling, may itself be intractable.
In this paper, we combine ideas from Bayesian RL, approximate variational inference, and meta- learning to tackle these challenges, and equip an agent with the ability to strategically explore unseen (but related) environments for given distribution, in order maximise its expected online return.
More specifically, we propose variational Bayes-Adaptive Deep RL (variBAD), a way to meta-learn perform approximate inference on an unknown task! and incorporate task uncertainty directly during action selection. Given distribution over MDPs p(M), we represent single MDP M using a learned, low-dimensional stochastic latent variable mn and jointly meta-train:
A variational auto-encoder that can infer the posterior distribution over m new task, given the agent's experience, while interacting with the environment, and
2. A policy that conditions on this posterior belief over MDP embeddings, and thus learns how to trade off exploration and exploitation when selecting actions under task uncertainty.
Figure Te shows the performance of our method versus the hard-coded optimal (with privileged goal information), Bayes-optimal, and posterior sampling exploration strategies. VariBAD's performance closely matches the Bayes-optimal one, matching optimal performance from the third rollout.
We the environment, task, and MDP, interchangeably.
2

---

Published as a conference paper at ICLR 2020
Previous approaches to BAMDPs are only tractable in environments with small action and state spaces or rely on privileged information about the task during training. VariBAD offers a tractable and flexible approach for learning Eayes-adaptie with the only assumption that such a distribution is available for meta-training. We evaluate our approach the gridworld shown above and on MuloCo domains that are widely used in meta-RI., and show that variBAD exhibits superior exploratory behaviour at test time compared to existing meta-learning methods, achieving higher returns during learning. As such, variBAD opens a path to tractable approximate Bayes-optimal exploration for deep reinforcement learning.
policies tailored to the training task distribution,
2
BACKGROUND
with S a set of states, A a set of actions, R(r++1| Se, de St+1) 1 a reward function, T( Stc, ar) transition function, (s0) an initial state distribution, discount factor, and H the horizon. In the standard RL setting, we want to learn a policy æ that maximises (T) EIo.T.= E 7 ((R(r+i1ss ar s4+1) the expected return. Here, we consider a multi-task meta-learning setting, which we introduce next.
We define a Markov decision process (MDP) as tuple = (S,,,,,,,,777,
2.1 TRAINING SETUP
We adopt the standard meta-learning setting where we have distribution p(M) over MDPs from which we can sample during meta-training, with an MDP M; ~ p(M) defined by a tuple M, A, A,,,R R4, T;, Tho, 7, H).. Across tasks, the reward and transition functions vary but share some structure. The index represents an unknown task description (e.g., a goal position or natural language instruction) reward and transition function from a distribution pl R. T). During meta- teaitttriiinn...n batches of tasks repeatedly sampled, and a small training procedure is performed on each of them, with the goal of learning to learn (for an overview of existing methods see Sec At meta-test time, the agent is evaluated based on the average return it achieves during learning, for tasks drawn from p. Doing this well requires at least two things: (I) incorporating prior knowledge obtained in related tasks, and (2) reasoning about task uncertainty when selecting actions to trade off exploration and exploitation. In the following, we combine ideas from meta-learning and Bayesian RL to tackle these challenges.
task ID. Sampling MDP from p(M) is typically done by sampling
2.2 BAYESIAN REINFORCEMENT LEARNING
When the MDP is unknown, optimal decision making has to trade off exploration and exploitation when selecting actions. In principle, this can be done by taking a Bayesian approach to reinforce- ment learning formalised as a Anyes-Adeptivee MDP (BAMDP), the solution to which is a Bayes- optimal policy (Bellman 1956 Duff & Barto 2002 Ghavamzadeh et al. (2015).
In the Bayesian formulation of RL, we assume that the transition and reward functions are distributed according prior do p( R, T). Since the agent does not have access the true reward and transition function, it can maintain a belief by R, T) = p(R. R, T(r:)), which is the posterior over the MDP given the agent's experience T. = {so, ao, r1, $1, 01 ..., 3t3 up until the current timestep. This is often done by maintaining a distribution over the model parameters.
To allow the agent to incorporate the task uncertainty into its decision-making, this belief can be augmented to the state, resulting in hyper-states st E S = S x B, where B is the belief space. These transition according
( 1ss Tt) = T be41lst, ae,
=T (st+ 1 ar bi) T (br+1152, at, St41) = =R,, [8(4+119), 4:) 5(b41 = p(R. T(T()1))
ie., the new environment state 5, is the expected new state w.r.t. the current posterior distribution of the transition function, and the belief updated deterministically according to Bayes rule. The reward function hyper-states is defined as the expected reward under the current posterior (after the state transition) over reward functions,
Rt (st at, st.) R+ be, oe, 81+1, (441,b141) Es,-1 [R(s at 8t+1))
3

---

Published as a conference paper at ICLR 2020
This results in BAMDP M+ (S+ "A, ,7, h) Duff & Barto 2002 which is special case of a belief MDP, i.. the MDP formed by taking the postenior beliets maintained by an agent in partially observable MDP and reinterpreting them as Markov states (Cassandra et al. 1994). In an arbitrary belief MDP, the belief is over a hidden state that can change over time. In a BAMDP, the belief is over the transition and reward functions, which are constant for given task. The agent's objective is now to maximise the expected return in the BAMDP.
H+-1 Σ
J+ E by Tt.
(re4als77 ae st+1)
i... maximise the expected return in an initially unknown environment, while learning, within the horizon H*. Note the distinction between the MDP horizon H and BAMDP horizon H A.t.Alohggh they often coincide, we might instead want the agent to act Bayes-optimal within the first N MDP episodes, so H*=N x H. Trading off exploration and exploitation optimally depends heavily on how much time the agent has left (e.g., to decide whether information-seeking actions are worth it).
The objective in is maximised by the Bayes-optimal policy, which automatically trades off ex- ploration and exploitation: takes exploratory actions to reduce task uncertainty only insofar as it helps to maximise the expected return within the horizon. The BAMDP framework is powerful because it provides principled way of formulating Bayes-optimal behaviour. However, solving the BAMDP is hopelessly intractable for most interesting problems.
The main challenges are as follows.
We typically do not know the parameterisation of the true reward and/or transition model,
The belief update (computing the posterior p R, T T.)) is often intractable, and
Even with the correct posterior, planning in belief space is typically intractable.
In the following, propose method that simultancously meta-learns the reward and transition functions, how to perform inference in an unknown MDP, and how to use the belief to maximise expected online return. Since the Bayes-adaptive policy is learned end-to-end with the inference framework, no planning is necessary at test time. We make minimal assumptions (no privileged task information is required during training), resulting in highly flexible and scalable approach to Bayes-adaptive Deep RL.
3 BAYES- GAYES-ADAPIVEE DEEP RL VIA META-LEARNING
In this section, we present variBAD, and describe how we tackle the challenges outlined above. We start by describing how to represent reward and transition functions, and (posterior) distributions over these. We then consider how to meta-learn to perform approximate variational inference in a given task and finally put all the pieces together to form our training objective.
In the typical meta-learning setting, the reward and transition functions that are unique to each MDP unknown, but also share some structure across the MDPs M; in p(M). We know that there exists a true i which represents either a task description or task ID, but we do not have access to this information. We therefore represent this value using learned stochastic latent variable my. For given MDP M, we can then write
R (r(441/5 ar 5t-1) Re R(r+41|st ar St-1; m4);
(4)
T(($++1/5..a) ar) xe T(s ax; m;),
(5)
where R and T are shared across tasks. Since we do not have access to the true task description or ID, we need to infer m; given the agent's experience up to time step collected in M,.
(so, ag, T1, 81, 01 "i,,,,,.. 8t- at-1, Tt, 8t);
(6)
ie., we want to infer the posterior distribution p(m, ((())(()) over m, given T (from now on, we drop the sub- and superscript for ease of notation).
(9)
Recall that our goal is to learn a distribution over the MDPs, and given a posteriori knowledge of the environment compute the optimal action. Given the above reformulation, it is sufficient to
4

---

Published as a conference paper at ICLR 2020
Encoder
Decoder
0
h/-1
8:
a1
R
T++1
Si+1
Po
8t-
at
bb = (m/
m
r't
pg
84+1
8;
Policy
Txp
Oz
81
Figure 2: VariBAD architecture: A trajectory of states, actions and rewards processed online using an RNN to produce the posterior over task embeddings, 2s (m|T:)) The posterior is trained using decoder which attempts to predict past and future states and rewards from current states and actions. The policy conditions on the posterior in order to act in the environment and is trained using RL.
reason about the embedding m, instead of the transition and reward dynamics. This is particularly useful when deploying deep learning strategies, where the reward and transition function can consist of millions of parameters, but the embedding can be small vector.
3.1 APPROXIMATE INFERENCE.
Computing the exact posterior typically not possible: we do have access the MDP (and hence the transition and reward function), and marginalising over tasks is computationally infeasible. Consequently, we need lear model the environment Po(T.H+ |a. (0-. -1); parameterised by 0, together with an amortised inference network (m) parameterised by ø, which allows fast inference runtime at each timestep t. The action-selectio policy is not part of the MDP, so an environmental model can only give rise to distribution of trajectories when conditioned on actions, which we typically draw from our current policy, learning objective is thus to maximise
л. any given time step our model
Ep [log po (T (o
[0.H+-1)],
where T is the trajectory distribution induced by our policy and we slightly abuse notation by denoting by 7 the state-reward trajectories, excluding the actions. In the following, we drop the conditioning -1 simplify notation.
Instead of optimising 77. which is intractable, we can optimise tractable lower bound, defined with learned approximate posterior go (m)(T.) which can be estimated by Monte Carlo sampling (for the full derivation see AppendixA
E (p(M. lo > EE (mm(ae) [logp((T.==m)) - K [[(gg(m/7:a)lpommmmmm (8)
EELBO,,
The term Elo H+ [m]] often referred as the reconstruction loss, and P(T:e (m) as the decoder. The term K (g(m([::)) (mom)))) is the KL-divergence between our variational posterior 4s and the prior over the embeddings pomm). We the prior to our previous posterior, qo (m/ with initial prior (m)m) = (0.
As can be seen Equation ( and Figure 2 when the agent is at timestep f, we encode the past trajectory T. to get the current posterior qmm Fir) since this is all the information available to perform inference about the current task. We then decode the entire trajectory T.u+ including the future, i.e., model [p(,,rr mm). This is different than the conventional VAE setup (and possible since have access to this information during training). Decoding not only the past but also the future is important because this way, variBAD learns perform inference about unseen states given the past.
5

---

Published as a conference paper at ICLR 2020
The reconstruction term 1og (m) factorises
log p(T (m, a. log (s0, ro..... -0
(9)
H
= log pog mm) +
Σ [logp(s:1lss,.0 a,, m) + log as 8441, m)]
Here, p(solm) is the initial state distribution TV p(s;+1|s:; a;; m) the transition function T", and p(re41/8 ag, 8141; m) the reward function R'. From now, we include Tý in T" for ease of notation.
3.2 TRAINING OBJECTIVE
We now formulate training objective for learning the approximate posterior distribution task embeddings, the policy, and the generalised reward and transition functions Ht' and T", We use deep neural networks to represent the individual components. These are:
1. The encoder (o(m|T:)) (T)), parameterised by Φ);
2. An approximate transition function T p (5 (&, ag; m) and an approximate reward function R' p (r:+1/ss at 8141; m)) which are jointly parameterised by 0; and
3. A policy (a, (e (o(mm Tx)) parameterised by và and dependent ø.
The policy is conditioned on both the environment state and the posterior over m, (a, St. q(m/7:))) ))
This is similar to the formulation of BAMDPs introduced in 2.2, with the difference that we learn a unifying distribution over MDP embeddings, instead of the transition/reward function directly. This makes learning easier since there are fewer parameters to perform inference over, and we can use data from tasks to learn the shared reward and transition function. The posterior can be represented by the distribution's 's parameters (e.g., mean and standard deviation if q is Gaussian).
Our overall objective is to maximise
N+ 0) + ELBO
((0,0,3) = Ep(M) J(W,
(10)
Expectations are approximated by Monte Carlo samples, and the ELBO can be optimised using the reparameterisation trick (Kingma & Welling 2014). For 0, We encode past trajectories using recurrent network as in Duan et al. (2016); Wang et al. (0166) but other types of encoders could be considered like the ones used in Zaheer al. (211 Garnelo
the prior (m)) = A 7).
et al. (0188; Rakelly et al. (2019). The network architecture is shown in Figure
we that the ELBO appears for all possible context lengths 4. This way, variBAD can learn how to perform inference online (while the agent is interacting with an envi- ronment), and decrease its uncertainty over time given more data. In practice, we may subsample fixed number of ELBO terms (for random time steps t) for computational efficiency if H+ + is large.
In Equation 10
Equation (10) is trained end-to-end, and À weights the supervised model learning objective against the RL loss. This is necessary since parameters ó are shared between the model and the policy. However, we found that backpropagating the RL loss through the encoder typically unnecessary in practice. Not doing also speeds up training considerably, avoids the need to trade off these losses, and prevents interference between gradients of opposing losses. In our experiments, we therefore optimise the policy and the VAE using different optimisers and learning rates. We train the RL agent and the VAE using different data buffers: the policy is only trained with the most recent data since larger buffer of observed trajectories.
on-policy algorithms in experiments; and for the VAE we maintain separate,
At meta-test time, roll out the policy in randomly sampled test tasks (via forward passes through the encoder and policy) to evaluate performance. The decoder not used at test time, and no gradient adaptation is done: the policy has learned to act approximately Bayes-optimal during meta-training.

---

Published as a conference paper at ICLR 2020
4 RELATED WORK
Meta Reinforcement Learning. A prominent model-free meta-RL, approach is to utilise the dy- namics of recurrent networks for fast adaptation (RI.2, 2 Wang et al. (2016); Duan et al. (2016)) At every time step, the network gets an auxiliary comprised of the preceding action and reward. This allows learning within a task to happen online, entirely in the dynamics of the recurrent network, If we remove the decoder (Fig 2 and the VAE objective (Eg 7 variBAD reduces to this setting, i.e., the main differences are that we use a stochastic latent variable (an inductive bias for representing uncertainty) together with decoder to reconstruct previous and future transitions / rewards (which acts as an auxiliary loss (Jaderberg et al. 2017) to encode the task in latent space and deduce infor- mation about unseen states). Ortega et al. (2019) provide an in-depth discussion of meta-learning sequential strategies and how to recast memory- based meta-learning within a Bayesian framework.
Another popular approach to meta RL is to learn initialisation of the model, such that at test time, only a few gradient steps are necessary to achieve good performance (Finn et al. 2017 7. Nichol & Schulman 2018). These methods do not directly account for the fact that the initial policy needs to explore, a problem addressed, a.o., by [Stadie et al. (2018) (E-MAML) and Rothfuss et al. (2019) (ProMP). In terms of model complexity. MAML and ProMP are relatively lightweight, since they typically consist of a feedforward policy. RL2 and variBAD use recurrent modules, which increases model complexity but allows online adaptation. Other methods that perform gradient adaptation at test time are, e.g., Hoothhooof et al. (2018) who meta-learn a loss function conditioned on the agent's s experience that is used at lest time so learn a policy (from scratch); and [Snn et al. (2017) who learn meta-critic that can criticise any actor for any task, and is used at test time to train policy. Compared to variBAD, these methods usually separate exploration (before gradient adaptation) and exploitation (after gradient adaptation) at test time by design, making them less sample efficient.
Skill / Task Embeddings. Learning (variational) task or skill embeddings for meta / transfer rein- forcement learning is used in variety of approaches. Hausman et al. (2018) use approximate vari- ational inference learn an embedding space of skills (with a different lower bound than variBAD). At test time the policy is fixed, and new embedder is learned that interpolates between already learned skills. Arnekvist et al.] (2019) learn a stochastic embedding of optimal Q-functions for dif- ferent skills, and condition the policy on (samples of) this embedding. Adaptation at test time is done in latent space. [Co-Reyes et al. (2018) learn a latent space of low-level skills that can be controlled by higher-level policy, framed within the setting of hierarchical RL. This embedding is learned using a VAF to encode state trajectories and decode states and actions. Zintgraf et al.) (019)) learn a deterministic task embedding trained similarly to MAMI. (Finn et al. 2017). Similar to variBAD, Zhang et al. (2018) use learned dynamics and reward modules to learn a latent representation which the policy conditions on and show that transferring the (fixed) encoder to new environments helps learning. Perez et al. (2018) learn dynamic models with auxiliary latent variables, and use them for mrodel-reedcctvve control. Lan et al. (2019) learn a task embedding with an optimisation procedure similar to MAML, where the encoder is updated at test time, and the policy is fixed. [Sæemndssson eet a (2018) explicitly learn an embedding of the environment model, which is subsequently used for model predictive control (and not, like in variBAD, for exploration). In the field of imitation learn- ing, some approaches embed expert demonstrations to represent the task; e.g., Wang et al. (2017) use variational methods and Duan et al. (2017) learn deterministic embeddings.
VariBAD differs from the above methods mainly in what the embedding represents (i.c., task uncer- tainty) and how it used: the policy conditions on the posterior distribution over MDPs, allowing it to reason about task uncertainty and trade off exploration and exploitation online. Our objective (8) explicitly optimises for Bayes-optimal behaviour. Unlike some of the above methods, we do not use the model at test time, but modell-aseed planning is a natural extension for future work.
Bayesian Reinforcement Learning. Bayesian methods for RL can be used to quantify uncertainty to support cection-sllection and provide a way to incorporate prior knowledge into the algorithms (see Ghavamzadeh et al. (2015) for a review). A Bayes-optimal policy is one that optimally trades off exploration and exploitation, and thus maximises expected return during learning. While such policy can in principle be computed using the BAMDP framework, it is hopelessly intractable for all but the smallest tasks. Existing methods are therefore restricted to small and discrete state / action spaces (Asmuth & Littman| 201 G Guez et al.] 2012 2013). or a discrete set of tasks (Brunskill 2012) Poupart al. 2006). VariBAD opens a path to tractable approximate Bayes-optimal exploration for
7

---

Published as a conference paper at ICLR 2020
deep RL by leveraging ideas from meta-leaming and approximate variational inference, with the only assumption that we can meta-train on a set of related tasks. Existing approximate Bayesian RL methods often require us to define prior belief update the reward transition function, and rely on (possibly expensive) sample-based planning procedures. Due to the use of deep neural networks however, variBAI lacks the formal guarantees enjoyed by some of the methods mentioned above.
Closely related to our approach is the recent work of Humplik et al. 2019). Like variBAD, they condition the policy on posterior distribution over the MDP, which is meta-trained using privileged information such as a task description. In comparison, variBAD meta-learns to represent the belief in an unsupervised way, and does not rely on privileged task information during training.
Posterior sampling (Strens 2000:Osband Osband et al. 2013 which extends Thompson sampling TTomm- son 1933) from bandits MDPs, estimates a posterior distribution over MDPs (i.c., model and re- ward functions), in the same spirit as variBAD. This posterior is used to periodically sample single hypothesis MDP (e.g., at the beginning of an episode), and the policy that is optimal for the sampled MDP is followed subsequently. This approach is less efficient than Bayes-optimal behaviour and therefore typically has lower expected return during learning.
A related approach for inter-task transfer of abstract knowledge is to pose policy search with pri- ors as Markov Chain Monte Carlo inference (Wingate et al. 201T). Similarly Guez et al.] (2013) propose a Monte Carlo Tree Search based method for Bayesian planning to get a tractable, sample- based method for obtaining approximate Bayes-optimal behaviour. [Osand et al. (018) note that non-Bayesian treatment for decision making can be arbitrarily suboptimal and propose a simple randomised prior based approach for structured exploration. Some recent deep RL methods use stochastic latent variables for structured exploration (Gupta et al. 2018 Rakelly et al. 2019). which gives rise to behaviour similar to posterior sampling. Other ways to use the posterior for explo- ration are, e.g., certain reward bonuses Kolter Ng (200 et al. (2012) and methods based on optimism in the face of uncertainty (Kearns & Singh [2022] Brafman & Tennenholtz (2022). Non- Bayesian methods for exploration are often used in practice, such as other exploration bonuses (c.g., via state- visitaiiitttinn counts) or using uninformed sampling of actions (e.g., <-greeyy action selection). Such methods are prone to wasteful exploration that does not help maximise expected reward.
Related to BAMDPs are contextual MDPs, where the task description is referred to as context, on which the environment dynamics and rewards depend (Hallak et al. 2015)] Diang et al., 2017 7 Dann et al. 2018 Modi & Tewari 2019). Research in this area is often concerned with developing light bounds by putting assumptions on the context, such as having a small known number of contexts, or that there is a linear relationship between the contexts and dynamics/rewards. Similarly, the framework of hidden parameter (HiP-) MDPs assumes that there is a set of low-dimensional latent factors which define family of related dynamical systems (with shared reward structure), similar to the assumption we make in Equation (5) (Doshi- Velez & Konidaris) 2016; Killian et al.] [017)] Yao
al. 2018). These methods however don't directly learn Bayes-optimal behaviour but allow for longer training period in new environments to infer the latents and train the policy.
Variational Inference and Meta-Learning. A main difference of variBAD to many existing Bayesian RL methods is that we meta-learn the inference procedure, i.c., how to do a posterior update. Apart from (RL) methods mentioned above, related work in this direction can be found, a.o., in Garneoo et al. (2018): Gordon et al. (2019); Choi et al. (2019). By comparison, variBAD has an inference procedure tailored to the setting of Bayes-optimal RL.
POMDPs. Several deep learning approaches to model free reinforcement learning (gg et al. 2019 and model learning for planning (Tsshiatsshhe et al., 2018) partially observable Markov deci- sion processes have recently been proposed and utilise approximate variational inference methods. VariBAD by contrast focuses on BAMDPs (Martin| 1967, Duff & Barto 2002 Ghavamzadeh et al,) 2015). special case of POMDPs where the transition and reward functions constitute the hidden state and the agent must maintain a belief over them. While in general the hidden state in a POMDP can change at each time-step, in BAMDP the underlying task, and therefore the hidden state, is fixed per task. We exploit this property by learning an embedding that is fixed over time, unlike approaches like [gl al.] (2019) which use filtering to track the changing hidden state. While we utilise the power of deep approximate variational inference, other approaches for BAMDPs often use more accurate but less scalable methods, e.g.. Lee et al. (2019) discretise the latent distribution and use Bayesian filtering for the posterior update.
8

---

Published as a conference paper at ICLR 2020
t- 15
Epleade
cal
Episode 3 Episode
EpPice
Latent variaooe
environment
(a) Example Rollout
(b) Reward Predictions
(c) Latent Space
Figure 3: Behaviour of variBAD in the gridworld environment. (a) Hand-picked but representative example test rollout. The blue background indicates the posterior probability of receiving a reward that cell. (b) Probability of receiving reward for each cell, predicted by the decoder, over the course of interacting with the environment (average black, goal state in green). (c) Visualisation of the latent space; each line is one latent dimension, the black line is the average.
5 EXPERIMENTS
In this section first investigate the properties of variBAD didactic gridworld domain. We show that variBAD performs structured and online exploration as it infers the task at hand. Then we consider more complex meta-learning settings by employing tasks commonly used the meta-RL. literature. We show that variBAD learns to adapt to the task during the first rollout, unlike many existing meta-learning methods. Details and hyperparameters can be found in the appendix, and at htttts://ggiiii com/ Imzint iirtinngggnn /
four MuloCo continuous control
5.1 GRIDWORLD
To gain insight into variBAD's properties, we start with didactic gridworld environment. The task to reach goal (selected uniformly at random) in the agent, inducing task uncertainty and necessitating exploration. The goal can be anywhere except around the starting cell, which is at the bottom left. Actions are: up, right, down, left, stay (executed deterministically), and after 15 steps the agent is reset. The horizon within the MDP is H 15, but we choose a horizon of
gridworld. The goal unobserved by
H = 45 for the BAMDP. Le., we train our agent to maximise performance for 4 MDP episodes. The agent gets sparse reward signal: -0.1 non-goal cells, and +1 on the goal cell. The best strategy is to explore until the goal is found, and stay at the goal or return when reset the initial position. We
latent dimensionality of
Figure Billustrates how variBAD behaves at test time with deterministic actions (i.e., all exploration how the agent interacts with the environment, with the blue background visualising the posterior belief by using the learned reward function. VariBAD learns the correct prior and adjusts its belief correctly over time. It predicts no reward for cells it has visited, and explores the remaining cells until it finds the goal.
done by the policy). In 3a
nice property of variBAD is that we can gain insight into the agent's belief about the environment by analysing what the decoder predicts, and how the latent space changes while the agent interacts with the environment. Figure (3b show the reward predictions: each line represents grid cell and its value the probability of receiving reward at that cell. As the agent gathers more data, more and more cells are excluded (p(rew = 1) = 0), until eventually the agent finds the goal. Figure 3c visualise the 5-dimensional latent space. We see that once the agent finds the goal, the posterior concentrates: the variance drops close to zero, and the mean settles on a value.
As we showed in Figure le the behaviour of variBAD closely matches that of the Bayes-optimal policy. Recall that the Bayes-optimal policy is the one which optimally trades off exploration and exploitation in unknown environment, and outperforms posterior sampling. Our results on this gridworld indicate that variBAD is an effective way to approximate Bayes-optimal control, and has the additional benefit of giving insight into the task belief of the policy.

---

Published as a conference paper at ICLR 2020
Ant-Dir
Cheetah Dir
Cheetah- Vel
Walker
1400
300
Avg Return
1000
1903
110
00
1923
120
ver BAD
600
RL2
033
PreMP
300
-06
E-MAM.
Poarl
200
200
Episodes
Episodes
Episodes
Episodes
Figure 4: Average test performance for the first rollouts of MuJoCo environments (using 5 seeds).
5.2 MUJOCO CONTINUOUS CONTROL META-LEARNING TASKS
We show that variBAD can scale to more complex meta learning settings by emploving it on MuJoCo Todorov et al. 2012) locomotion tasks commonly used in the meta-RI. literature [ We consider the AntDir and HalfCheetahDDrr environment where the agent has to run either forwards or backwards (i.e., there are only two tasks), the HalfCheetah Vel environment where the agent has to run differ- ent velocities, and the Walker environment where the system parameters are randomised.
Figure 4 shows the performance at test time compared existing methods. While we show per- formance for multiple rollouts for the sake of completeness, anything beyond the first rollout is not directly relevant to our goal. which single episode. Only variBAD and RL are able to adapt to the task at hand within a single episode. RI.² underperforms variBAD on the HalfCheetahDirr environment, and learning slower and less stable (see learning curves and runtime comparisons in Appendix C.. Even though the first rollout includes exploratory steps, this matches the optimal oracle policy (which conditioned on the true task description) up to small margin. The other methods (PEARL Rakelly et al. (2019), E-MAML Stadie et al.] (2018) and ProMP Rothfuss et al.] 2019) are not designed to maximise reward during single rollout, and perform poorly in this case. They all require substantially more environment interactions each new task to achieve good performance. PEARL, which is akin to posterior sam- pling, only starts performing well starting from the third episode (Note: PEARL outperforms our oracle slightly, likely since our oracle is based on PPO, and PEARL is based on SAC).
maximise performance on new task, while learning, within
Overall, our empirical results confirm that variBAD can scale up current benchmarks and max- imise expected reward within a single episode.
6 CONCLUSION & FUTURE WORK
We presented variBAD, a novel deep RL method to approximate Bayes-optimal behaviour, which uses meta-learning utilise knowledge obtained in related tasks and perform approximate inference in unknown environments. In a didactic gridworld environment, our agent closely matches Bayes- optimal behaviour, and in more challenging MuJoCo tasks, variBAD outperforms existing methods in terms of achieved reward during a single episode. In summary. we believe variBAD opens path tractable approximate Bayes-optimal exploration for deep reinforcement learning.
There are several interesting directions of future work based on variBAD. For example, we currently do not use the decoder at test time. One could instead use the decoder for model- predictive planning, or to get sense for how wrong the predictions are (which might indicate we are out of distribution, and further training is necessary). Another exciting direction for future research is considering settings where the training and test distribution of environments are not the same. Generalising to dut-of-distribution tasks poses additional challenges and in particular for variBAD two problems are likely to arise: the inference procedure will be wrong (the prior and/or posterior update) and the policy will not be able to interpret a changed posterior. In this case, further training of both the encoder/decoder might be necessary, together with updates the policy and/or explicit planning.
Ennirromments taken from http:://ggihhbb.cc//aaeeeaalll
10

---

Published as a conference paper at ICLR 2020
ACKNOWLEDGMENTS
We thank Anuj Mahajan who contributed early work on this topic. We thank Joost van Amers- foort, Andrei Rusu and Dushyant Rao for useful discussions and feedback. Luisa Zintgraf is sup- ported by the Microsoft Research PhD Scholarship Program. Maximilian Igl supported by the UK EPSRC CDT in Autonomous Intelligent Machines and Systems. Sebastian Schulze is supported by Dyson. This work was supported by a generous equipment grant and donated DGX-1 from NVIDIA. This project has received funding from the European Research Council under the Euro- pean Union's 's Horizon 2020 research and innovation programme (grant agreement number 637713).
REFERENCES
Isac Arnekvist, Danica Kragic, and Johannes Stork. Vpe:
Variational policy embedding for transfer reinforcement learning. In International Conference on Robotics and Automation, 2019.
John Asmuth and Michael L Littman. Learning is planning: near bayes-optimal reinforcement learning via monte-carlo tree search. In Conf on Uncertainty in Artificial Intelligence, 2011.
Richard Bellman. problem in the sequential design of experiments. Sankhya: The Indian Journal of Statistics (1933-1960), 16(3/4):21-2299 1956.
Ronen Brafman and Moshe Tennenholtz. R-max-a general polynomial time algorithm for near- optimal reinforcement learning. Journal of Machine Learning Research, pp. 3:213-231, 2002.
Emma Brunskill. Bayes-optimal reinforcement learning for discrete uncertainty domains. In Inter- national Conference on Autonomous Agents and Mulliagent Systems, 2012.
Anthony R. Cassandra, Leslie Pack Kaelbling, and Michael L. Littman. Acting optimally in partially observable stochastic domains. In Twelfth National Conference on Artificial Intelligence, 1994. AAAI Classic Paper Award, 2013.
Kristy Choi, Mike Wu, Noah Goodman, and Stefano Ermon. Meta-amortized variational inference and learning. In International Conference on Learning Representation, 2019.
John Co-Reyes, YuXuan Liu, Abhishek Gupta, Benjamin Eysenbach, Pieter Abbeel, and Sergey Levine. Self-consistent trajectory autoencoder: Hierarchical reinforcement learning with trajec- tory embeddings. International Conference on Machine Learning, 2018.
Christoph Dann, Lihong Li, Wei Wei, and Emma Brunskill. Policy certificates: Towards accountable reinforcement learning. arXiv preprint vrXiv:1811.03055 2018.
Finale Velee-Vlez and George Konidaris. Hidden parameter markov decision processes: A semi- parametric regression approach for discovering latent task parametrizations. In International Joint Conference Artificial Intelligence, 2016.
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbecl. RL3: Fast reinforcement learning via slow reinforcement learning. arXiv:1611.0277) 2016.
Yan Duan, Marcin Andrychowicz, Bradly Stadie, OpenAl Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojeiech Zaremba. One-shot imitation learning. In Advances in Neural Information Processing Systems, 2017.
Michael O'Gordon Duff and Andrew Barto. Optimal Learning: Computational procedures for Bayes-adaptive Markov decision processes. PhD thesis, Univ of Massachusetts at Amherst, 2002.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learnng for fast adaptation of deep networks. In International Conference on Machine Learning, 2017.
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. In ICML 2018 Workshop Thheoretccll Foundations and Applications of Deep Generative Models, 2018.
Mohammad Ghavamzadeh, Shie Mannor, Joelle Pineau, Aviy Tamar, et al. Bayesian reinforcement learning: A survey. Foundations and Trends in Machine Learning, 8(5-6)3359-43 2015.
11

---

Published as a conference paper at ICLR 2020
Jonathan Gordon, John Bronskill, Matthias Bauer, Sebastian Nowozin, and Richard E Turner. Meta- learning probabilistic inference for prediction. In Internationaa Conference on Learning Repre- sentation, 2019.
Arthur Guez, David Silver, and Peter Dayan. Efficient bayes-adaptive reinforcement learning using sample-based search. In Advances in Neural Processing Systems, 2012.
Arthur Guez, David Silver, and Peter Dayan. Scalable and efficient bayes-adaptive reinforcement learning based on monte-carlo tree search. Journal of Artificial Intelligence Research, 48:841- 883. 2013.
Abhishek Gupta, Russell Mendonca, YuXuan Liu, Pieter Abbeel, and Sergey Levine. Meta- reinforcement learning of structured exploration strategies. In Advances in Neural Processing Systems, 2018.
Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes. arXiv preprint 1.X.::1522.02253 2015.
Karol Hausman, Jost Tobias Springenberg, Ziyu Wang, Nicolas Heess, and Martin Riedmiller. Learning an embedding space for transferable robot skills. International Conference on Learn- ing Representation, 2018
Rein Houthooft, Yuhua Chen, Phillip Isola, Bradly Stadie, Filip Wolski, OpenAl Jonathan Ho, and Pieter Abbeel. Evolved policy gradients. In Advances in Neural Information Processing Systems, 2018.
Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A Ortega, Yee Whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv:1905.06442 2019.
Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for pomdps. In International Conference on Machine Learning, 2019.
Max Jaderberg, Volodymyr Mnih, Wojeiech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representation, 2017.
Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire Con- textual decision processes with low bellman rank are pac-leamable. In International Conference on Machine Searning-Volume 2017.
Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-34, 1998.
Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. Ma- chine learning, 49(2-3):209-233 2002.
Taylor W Killian, Samuel Daulton, George Konidaris, and Finale Deshi-Velez. Robust and effi- cient transfer learning with hidden parameter markov decision processes. In Advances in neural information processing systems, 2017.
Diederik P Kingma and Max Welling, Auto-encoding variational bayes. In International Conference on Learning Representation, 2014.
J Zico Kolter and Andrew Y Ng. Near-bayesian exploration in polynomial time. In International Conference on Machine Learning, 2009.
Lin Lan, Zhenguo Li, Xiaohong Guan, and Pinghui Wang. Meta reinforcement learning with task embedding and shared policy. In INTERNATIONAL JOINT CONFERENCE ON ARTIFICIAL IN- TELLIGENCE, 2019.
Gilwoo Lee, Brian Hou, Aditya Mandalika, Jeongseok Lee, and Siddhartha S Srinivasa. Bayesian policy optimization for model uncertainty. In International Conference on Learning Representa- fion, 2019.
12

---

Published as a conference paper at ICLR 2020
Yun-En Liu, Travis Mandel, Emma Brunskill, and Zoran Popovic. Trading off scientific knowledge and user learning with multi-anee anmeed bandits. EDM, pp. 161-168, 2014.
James John Martin, Bayesian decision problems and Markov chains. Wiley, 1967.
Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Picter Abbeel. A simple neural attentive meta- learner. arXiv: /707.03 141, 2017.
Aditya Modi and Ambuj Tewari. Contextual markov decision processes using generalized linear models. In Reinforcement Learning for Real Life (RLARealLife) Workshop at the International Conference on Machine Learning, 2019.
Alex Nichol and John Schulman. Reptile: a scalable metalearning algorithm. arXiv:1803.0299 2018.
Pedro A Ortega, Jane X Wang, Mark Rowland, Tim Genewein, Zeb Kurth-Nelson, Razvan Pascanu, Nicolas Heess, Joel Veness, Alex Pritzel, Pablo Sprechmann, al. Meta-learning of sequential strategies. arXiv: 03030, 2019.
Ian Osband, Daniel Russo, and Benjamin Van Roy. (more) efficient reinforcement learning via posterior sampling. In Advances Neural Information Processing Systems, 2013.
lan Osband, John Aslanides, and Albin Cassirer. Randomized prior functions for deep reinforcement learning. In Advances in Neural Information Processing Systems, 2018.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Christian F Perez, Felipe Petroski Such, and Theofanis Karaletsos. Efficient transfer learning and online adaptation with latent variable models for continuous control. In Continual Learning Work- shop, NeurlPS 2018, 2018.
Pascal Poupart, Nikos Vlassis, Jesse Hoey, and Kevin Regan. An analytic solution to discrete bayesian reinforcement learning. International Conference on Machine Learning. 2006.
Kate Rakelly, Aurick Zhou, Deirdre Quillen, Chelsea Finn, and Sergey Levine. Efficient off-policy meta-reiiforreement learning via probabilistic context variables. In International Conference Machine Learning, 2019.
Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. Promp: Proximal meta-policy search. In International Conference on Learning Representation, 2019.
Steindór Sæmundsson, Katja Hofmann, and Mare Peter Deisenroth. Meta reinforcement learning with latent variable gaussian processes. In Conference on Uncertainty in Artificial Intelligence, 2018.
Jonathan Sorg. Satinder Singh, and Richard Lewis. Variance-based rewards for approximate bayesian reinforcement learning. In Conference on Uncertainty in Artificial Intelligence, 2012.
Bradly C Stadie, Ge Yang, Rein Houthooft, Xi Chen, Yan Duan, Yuhuai Wu, Pieter Abbeel, and Ilya Sutskever. Some considerations on learning to explore via meta-reinforcemment learning. In Advances in Neural Processing Systems, 2018.
Malcolm Strens. A bayesian framework for reinforcement learning. In International Conference Machine Learning. 2000.
Flood Sung, Li Zhang, Tao Xiang, Timothy Hospedales, and Yongxin Yang. Learning to learn: Meta-critic networks for sample efficient learning. a7Xin: 2017.
William R Thompson. On the likelihood that unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3/4):285-29) 1933.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In IROS, pp. 5026-5033. IEEE, 2012. ISBN 978-1-4663-1137--5
13

---

Published as a conference paper at ICLR 2020
Sebastian Tschiatschek, Kai Arulkumaran, Jan Stühmer, and Katja Hofmann. Variational inference for data-efficient model learning in pomdps. arXiv:1805.0928// 2018.
Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning reinforcement learn. In Annual Meeting of the Cognitive Science Community (CogSci), 2016.
Ziyu Wang, Josh S Merel, Scott E Reed, Nando de Freitas, Gregory Wayne, and Nicolas Heess. Robust imitation of diverse behaviors. In Advances Neural Information Processing Systems, 2017.
David Wingate, Noah D Goodman, Daniel M Roy, Leslie P Kaelbling, and Joshua Tenenbaum. Bayesian policy search with policy priors. In International Joint Conference on Artificial Intelli- gence, 2011.
Jiayu Yao, Taylor Killian, George Konidaris, and Finale Doshi-Velez. Direct policy transfer via hidden parameter markov decision processes. In LLARLA Workshop, FAIM, 2018.
Gregory Yauney and Pratik Shah. Reinforcement learning with action-derived rewards for chemotherapy and clinical trial dosing regimen selection. In Machine Learning for Healthcare Conference, pp. 161-226, 2018.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan R Salakhutdinov, and Alexander J Smola. Deep sets. In Advances in Neural Processing Systems, 2017.
Amy Zhang, Harsh Satija, and Joelle Pineau. Decoupling dynamics and reward for transfer learning. In ICLR workshop track, 2018.
Luisa M Zintgraf, Kyriacos Shiarlis, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast context adaptation via meta-learning, In International Conference on Machine Learning, 2019.
14

---

Published as a conference paper at ICLR 2020
Bayes-Adaptiiv Deep Reinforcement Learning via Meta-Learning
Supplementary Material
A FULL ELBO DERIVATION
Equation 8 can be derived as follows.
(m(m)) dm (o(mlFa) Po po(T.H,
E (( [log (T.:() = Ep log po(T. m)
Ep
log (mlmma)))
log Po(T.Hrmm) m)
Ep, Cm mi))
(ma
Ep, (m ri
(T. + log - log (m)
Ep (matre))) (logpo(TTmmm) [mm] (l(ga(mm(r:))bmmmmm)
(11)
ELBO,,
B EXPERIMENTS: GRIDWORLD
B.1 ADDITIONAL. REMARKS
Figure 3e visualises how the latent space changes the agent interacts with the environment. As we can see, the value of the latent dimensions starts around mean 1 and variance 0, which the prior we chose for the beginning of an episode. Given that the variance increases for a little bit before the agent finds the goal, this prior might not be optimal, natural extension of variBAD therefore to also learn the prior to match the task at hand.
B.2 HYPERPARAMETERS
We used the PyTorch framework for our experiments. Hyperparameters are listed below, and the source code can be found at thttts:://githut httpe//b....com..iigggaa//araaaa
Hyperparameters for variBAD are:
RL Algorithm
A2C
Number of policy steps
60
Number of parallel processes
16
Epsilon
le-5
Discount factor
0.95
Max grad norm
0.5
Value loss coefficient
0.5
Entropy coefficient
0.01
GAE parameter tau
0.95
ELBO loss coefficient
1.0
Policy LR
0.001
Policy VAE
0.001
Task embedding size
Policy architecture
hidden layers, 32 nodes each, TanH activations
Encoder architecture
FC layer with 40 nodes, GRU with hidden size 64, output layer with 10 outputs (μ and σ)) ReLu activations
Reward decoder architecture
hidden layers, 32 nodes each,
25 outputs heads, ReLu activations
Decoder loss function
Binary cross entropy
15

---

Published as a conference paper at ICLR 2020
Avarrgg
Auerage
Orande
vanHAD
Bayes Optial
RL2
ver WrDD
Opyyel-Opümel
RL2
Posterior Sampling
Posterior Sampping
up
Frames
Episocies
(a) Learning
(b) Average return per episade.
Figure 5: Results for the gridworld toy environment. Results are averages over 20 seeds (with 95% confidence intervals for the learning curve).
Hyperparameters for RL2 the same as above, with the following changes:
Policy architecture
States embedded using fc linear layer, output size 32.
Rewards are embedding using a fc layer, output size 8. Results are concatenate and passed output size 32. After an additional fe layer with hidden size 32, the network outputs the actions. We used TanH activations throughout.
GRU, hidden size 128,
B.3 COMPARISON TO RL2
Figure 5a shows the learning curves for variBAD and RL2, comparison an oracle policy (which has access to the goal position). We trained these policies horizon of H+ 4 x H = 60. i.e., on a BAMDP the values of hard-coded Bayes-optimal policy, and hard-coded posterior sampling policy using dashed lines.
which the agent has to maximise online return within four episodes. We indicate
Figure 5b shows the end- pedfoorrormecc of variBAD and RL2, compared to the hard-coded optimal policy (which has access to the goal position). Bayes-optimal policy, and posterior sampling policy. VariBAD and RL2 both closely approximate the Bayes-optimal solution. By inspecting the individ- ual runs, we found that VariBAD learned the Bayes-optimal solution for 4 out of 20 seeds, RL2 zero times. Both otherwise find solutions that are very close to Bayes-optimal, with the difference that during the second rollout, the cells left to search are not all on the shortest path from the starting point.
Note that both variBAD and RL2 were trained on only four episodes, but we evaluate them on six episodes here. After the fourth rollout, we do not fix the latent hidden state, but continue rolling out the policy as before. As we can see, the performance of RL2 drops again after the fourth episode: this likely due to instabilities in 128-dimensional hidden state. VariBAD's latent representation, the approximate task posterior, is concentrated and does not change with more data.
C EXPERIMENTS: MuJoCo
C.1 LEARNING CURVES
Figure 6 shows the learning curves for the MuJoCo environments for all approaches. The ora- cle policy was trained using PPO. PEARL (Rakelly et al. 2019) was trained using the reference used are also taken from this imple-
implementation provided by the authors. The environments mentation. E-MAML (Stadie et al., 2018) and ProMP (Rothfuss et al. 2019) were trained using the reference implementation provided by |oothfuss et al. (019))
16

---

Published as a conference paper at ICLR 2020
Ant-Dir
Cheslah-Di
Cheetah-Vel
Walker
1502
2500
1090
Avg return
1250
2000
-00
180
1500
-10
-150
002
peiat
1000
282
pertrip
230
280
350
1802
1000
Avg return
1232
3000
-20
1000
1500
080
-10
482
-060
294
300
Frames
168
Framea
100
Framea
Frames
Figure 6: Learning curves for the MuJoCo results presented in Section 5.2 The top row shows performance evaluated at the first rollout, and the second row shows the performance at the N-th rollout. For variBAD and RL2, N = 2. For ProMP and E-MAML, N = 20. For PEARL, N = 10.
As we can see, PEARL is much more sample efficient in terms number of frames than the other methods (Fig 6). which is because it is an off-policy method. On-policy off-policy training is an orthogonal issue to our contribution, but an extension of variBAD to off-policy methods is an interesting direction for future work. Doing posterior sampling using off-policy methods also requires PEARL to use different encoder (to maintain order invariance of the sampled trajectories) which is non-recurrent (and hence faster train, see next section) but restrictive since assumes independence between individual transitions.
Note than in Figure 4 for the Walker environment evaluation, we used the models obtained after half the training time (5e + frames) for variBAD and the Oracle, since performance declined again after that.
For all MuloCo environments, we trained variBAD with reward decoder only (even for Walker, where the dynamics change, we found that this has superior performance).
C.2 TRAINING DETAILS AND COMPARISON TO RL2
We are interested in maximising performance within single rollout (H 200). However in order to compare better to existing methods, we trained variBAD and the R1.2 baseline to maximise performance within flag to the states, so that the agent knows when it gets reset in-between episodes. This allows us to evaluate on multiple rollouts (without re-setting the hidden states of the RNN) because the agents have learned to handle re-sets to the starting position.
rollouts (/1 + 400). We implemented task resets by adding a 'done'
We observe that RL2 is sometimes unstable when comes to maintaining its performance over multiple rollouts, e.g., in the Cheetah Vel task (Figure (). We hypothesise that the drop of RL2's performance in Vheetah occurs because it has not properly learned to deal with environment re- sets. The sudden change in state space (with includes joint positions and velocities) could lead to dramatic shift in the hidden state, which then might not represent the task at hand properly. addi- tion, once the Cheetah is running at the correct velocity, it can infer which task it is in from its own velocity (which is part of the environment state) and stop doing inference, which might be another reason observe this drop when the environment resets and the state suddenly has different (very low) velocity. For variBAD this is less of a problem, since we train the latent embedding to represent the task, and only the task. Therefore, the agent does not have do the inference procedure again when reset to the starting position, but can rely on the latent task description that is given by the approximate posterior. It might also just be due implementation details, and, e.g., Mishra et al. (017) do not observe this problem (see their Fig 4).
17

---

Published as a conference paper at ICLR 2020
Episode
Episode 2
Episode
206
120
Steps
50
X-position
x-position
x-positon
Figure 7:
Behaviour at test time for the for the task "walk left" in HaliCheetahDir. The x-axis reflects the position of the agent; the y-axis the steps in the environment (to be read from bottom to top). Rows are separate examples, columns the number of rollouts.
C.3 CHEETAHDIR TEST TIME BEHAVIOUR
To get a sense for where these differences between the different approaches might stem from, con- sider Figure 7 which shows example behaviour of the policies during the first three rollouts in the HalfCheetahDir environment, when the task is "go left". Both variBAD and RL adapt to the task online, whereas PEARL acts according the current sample, which the first two rollouts can mean walking in the wrong direction. For visualisation of the variBAD latent space at test time for this environment see Appendix C.5 While we outperform at meta-test time, PEARL is more sample efficient during meta-training (see Fig 6). since it is an off-policy method. Extending variBAD to off off-policy methods is an interesting but orthogonal direction for future work.
C.4 RUNTIME COMPARISON
The following are rough estimates of average run-times for the HalfCheetah-Dii environment (from what we have experienced; we often ran multiple experiments per machine, so some of these might be overestimated and should be mostly understood as giving relative sense of ordering).
ProMP, E-MAML: 5-8 hours
variBAD: 48 hours
RL2: 60 hours
PEARL: hours
E-MAMI. and ProMP have the advantage that they not have recurrent part such as variBAD or RL2. Forward and backward passes through recurrent networks can be slow, especially with large horizons.
Even though both variBAD and RL2 use recurrent modules, we observed that variBAD is faster when training the policy with PPO. This is because we do not backpropagate the RL-loss through the recurrent part, which allows us to make the PPO mini-batch updates without having re-compute the embeddings (so it saves us a lot of forward/backwwad passes through the recurrent model). This difference is less pronounced with other RL methods that do not rely this many forward/backwwad passes per policy update.
18

---

Published as a conference paper at ICLR 2020
Task: walk right
Insent vertancee
positio
150
-3.0
-258 position
4550 38
4577 50agss
Figure 8: Visualisation of the latent space at meta-test time, for the HalfCheetahDir environment and the tasks "go right" (top) and the task "go left" (bottom). Left: value of the posterior mean during single rollout (200 environment steps). The black line the average value. Middle: value of the posterior log-variance during single rollout. Right: Behaviour of the policy during single rollout. The x-axis show the position of the Cheetah, and the y-axis the step (should be read from bottom top).
C.5 LATENT SPACE VISUALISATION
nice feature of variBAD that can give us insight into the uncertainty the agent about what task it is in. Figure 8 shows the latent space for the HalfChettahDir tasks "go right" (top row) and "go left" (bottom row). We observe that the latent mean and log- vor-vnriance adapt rapidly, within just few environment steps (left and middle figures). This is also how fast the agent adapts the current task (right figures). As expected, the variance decreases over time as the agent gets more certain. interesting to note that the values of the latent dimensions swap signs between the two tasks.
Visualising the belief the reward/state space directly, as we have done in the gridworld example, is more difficult for MuJoCo tasks, since do instead, is to additionally train a model that predicts a ground-truth task description (separate from the main objective and just for further analysis, since we do want to this privileged information for meta-training). This would give us a more direct sense of what task it thinks it is in.
have continuous states and actions. What could
19

---

Published as a conference paper at ICLR 2020
C.6 HYPERPARAMETERS
We used the PyTorch framework (Paszke et al. 2017) for our experiments. The default arguments for MuJoCo experiments be found below, for details reference implementation
http::///ithbb.. wtt.:/ooommmmmmiiieee http:///ww.ngggg////rrbaaaaa
RL Algorithm
PPO
Batch size
3200
Epochs
Minibatches
Max grad norm
0.5
Clip parameter
Value loss coefficient
0.5
Entropy coefficient
0.01
Notes
We use Huber loss in the RL loss
Weight of KL term in ELBO
0.1
Policy LR
0.0007
Policy VAE
0.001
Task embedding size
Policy architecture
hidden layers, 128 nodes each, TanH activations States, actions, rewards encoder: FC layer (32/16/16-dim), GRU with hidden size 128, output layer with 5 outputs, ReLu activations
Hncoder architecture
Reward decoder architecture
hidden layers, and 32 nodes,
ReLu activations
Reward decoder loss function
Mean squared error
20


Reviews:
[{"review_id": "SyxKPjfroH", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "*EDIT: Score increased after discussion with authors clarified many concerns raised below*\n\nSummary:\nThis paper presents an algorithmic approach toward learning Bayes Optimal policies under the uncertainty of the environment. Leveraging meta-learning, the proposed variBAD approximate inference procedure is capable of adapting within the first episode at, what the authors term, meta-test time. \n\nComments:\nIt is my estimation that this paper is well positioned to further current state-of-the-art adaptive RL frameworks or methodologies, whether they are meta-learned, transferred or directly inferred through probabilistic mechanisms. The primary contribution of this paper is in how variBAD learns the variational inference procedure. As noted in the related work section, many contemporary policy learning approaches via variational inference are limited by their construction, selection of prior distributions, etc. The advantage of the proposed methodology is that it is capable of efficiently inferring the current environment and adapting the policy learning procedure accordingly. The experiments successfully compare with relevant baselines and prior approaches. The discussion is well framed in highlighting the benefits and limitations of the proposed methodology in relief to prior approaches. One possible weakness in the experimentation, given how closely RL^2 matches the performance of variBAD, is that the specific contributions of each architectural choice or optimization protocol are unclear. There are a few areas in the paper where the authors suggest that ablating their model in specific ways would recover the core approaches present in RL^2. It would be instructive to see how/if performance degrades or converges toward that of RL^2 as the variBAD methodology is ablated.\n\nThe paper is well grounded in the literature, albeit skewed perhaps a bit too far toward recent meta-learning results. This is understandable given the focus of this paper, however there are other approaches that might deserve a mention as they similarly parameterize variation over possible MDPs with some latent variables. Namely, I have in mind a few lines of research such as Contextual MDPs (Hallak, et al, 2015; Jiang, et al, 2017; Dann, et al, 2018; etc.), Successor Features (Barretto, et al, 2016,2017,2019; Lehnert, et al, 2017; etc.)  and HiP-MDPs (Doshi-Velez and Konidaris, 2016; Killian, et al, 2017). In particular use of the HiP-MDP framework, Yao, et al (2018) also use the inferred latent variable used to identify the task to condition the policy. While it's always easy to dig into the rabbit holes of related research and distract from the overall objective of a paper, I thought that there was sufficient overlap with these other lines of research that the authors may find interesting. I do not claim that any one of these additional sources of prior work have been overlooked to the detriment of the current paper, they are offered as merely a suggestion to broaden the author's anchoring in the literature.\n\nNow, some more specific questions about the paper and proposed approach. Further clarity along any of these questions would greatly improve the presentation of the paper as well as further convince me of the paper's suitability for publication.\n1) What is the advantage of decoding the entire trajectory? It is well understood that this is advantageous in training as that data is available and allows for better inference of the variational parameters. However, under test conditions where the framework may be operating in environments that lie outside the distribution of MDPs it was trained on, I can imagine that errors in trajectory prediction may compound and throw off the entire inference procedure. The experimental set-up did not allay these concerns as there was no mention for holding out-of-distribution tasks/environments aside and the variation in environments is pretty narrow. \n2) How is the proposed trajectory decoding more stable than model predictive control? Is stability a large consideration for variBAD when exploring? How much can one trust the exploratory actions under variBAD? \n3) The visualization and careful explanation in Section 5.1 of how variBAD executes inference and learning was greatly appreciated. However, are these intuitions valid when extending beyond discrete state and action spaces? Can one make the same claims about the overall approach or procedure in the MuJoCo domains? It was mildly disappointing that a similar explanatory effort was not made in more complex environments. Even an acknowledgement of this being unreasonable would help round out the discussion in Section 5.2.\n4) It is not clear what the connection is between the horizon H and the number of rollouts used for evaluation/inference/training. I spent a bit more time than necessary going over and over these items in the paper to where I think that I may understand but I'm still not 100% confident about what is impacted by the number of rollouts used.\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "SkeDSGonKS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a new deep reinforcement learning method that can efficiently trade-off exploration and exploitation. An optimal policy for this trade-off can be solved under the Bayesian-adaptive MDP framework, but in practice, the computation is often intractable. To solve the challenge and approximate a Bayesian-optimal policy, the proposed method VariBAD combines meta-learning, variational inference, and bayesian RL. Specifically, the algorithm learns latent representations of task embeddings and performs tractable approximate inference by optimizing a tractable lower bound of the objective.\n\nThe paper is well-written and easy to follow. The combination of meta-learning, variational inference and BAMDP is a clear and neat way to approximate Bayes-optimal policy. The idea also sounds practical for RL as it can approximately solve larger tasks with unknown priors. Experiments on Gridworld and Mujoco show the effectiveness of the proposed method. On Gridworld the performance of the proposed algorithm is close to the performance of the Bayes-optimal policy.\n\nOne concern for this paper is the level of novelty, as each major component of the proposed solution has been explored quite extensively in the existing literature (as mentioned in the related work section). \n\nIn addition, since comparing many existing Bayesian RL methods, VariBAD meta-learns the inference procedure. This can add additional computation complexity to Bayesian RL, which is not explained or mentioned in neither the method part nor the experiment. I hope the authors can add some discussions on this aspects", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "HkgOjdCb5r", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \nThis paper considers a version of reinforcement learning problem where an unknown prior distribution over Markov decision processes are assumed and the learner can sample from it. After sampling a MDP, a standard reinforcement learning is done. Then the paper investigates the Bayes-optimal strategy for such meta-learning setting. The experiments are done for an artificial maze solving tasks. \n\nComments: \nConsidering a Bayesian setting of reinforcement learning is sound and well-motivated in a mathematical or statistical sense. On the other hand, I wonder what kind of practical applications motivate such formulation. Unfortunately, I don’t have any examples in mind and the paper also shows some artificial experiments. So, the formulation seems, so far, not to be convincing in a practical sense. \n\nAnother concern in my mind is that the proposed methods are not supported by any theoretical analyses. I think mathematical papers without practical applications are acceptable if they contain strong mathematical analyses. The present paper, however, does not contain such analyses. \n\nAs a summary, I feel that the paper is not strong for theoretical analyses nor practical usefulness, and thus further investigation for either side is necessary.  \n\n\nComments after Rebuttal:\nI modified my score according to authors' comments. ", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "B1lTqBu_KB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The proposed method represents a single MDP using a learned, low-dimensional stochastic latent variable. On these grounds, given a set of tasks sampled from a distribution, the method jointly trains: (1) a variational auto-encoder that can infer the posterior distribution over the postulated latent variable when it encounters a new task while interacting with the environment, and (2) a policy that conditions on this posterior distribution over the MDP embeddings, and thus learns how to trade off exploration and exploitation when selecting actions.\n\nSuch variational inference arguments for transfer learning in the context of MDPs are not new. The authors have not made a good job reviewing the related literature. Most importantly, their experimental evaluations lack substantial comparison to such related methods. This is totally disappointing.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper considers the problem of transfer learning among families of MDP, and proposes a variational Bayesian approach to learn a probabilistic model of a new problem drawn from the same distribution as previous tasks, which is then leveraged during action selection. 

After discussion, the three respondent reviewers converged to the opinion that the paper is novel and interesting, and well evaluated. (Reviewer 1 never responded to any questions the authors or me, so I have disregarded their review.) I am therefore recommending an accept.
Decision: Accept (Poster)