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

1 INTRODUCTION . Reinforcement learning ( RL ) is typically concerned with finding an optimal policy that maximises expected return for a given Markov decision process ( MDP ) with an unknown reward and transition function . If these were known , the optimal policy could in theory be computed without environment interactions . By contrast , learning in an unknown environment usually requires trading off exploration ( 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 a belief distribution over possible environments . Augmenting the state space of the underlying MDP with this belief yields a BAMDP , a special case of a belief MDP ( Kaelbling et al. , 1998 ) . A 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 . Its 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 a 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 , 1933 ; Strens , 2000 ; Osband et al. , 2013 ) . Here , the agent periodically samples a single hypothesis MDP ( e.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 on a regular MDP , not a BAMDP . However , posterior sampling ’ s exploration can be highly inefficient and far from Bayes-optimal . ∗luisa.zintgraf @ cs.ox.ac.uk Consider the example of a gridworld in Figure 1 , where the agent must navigate to an unknown goal located in the grey area ( 1a ) . To maintain a posterior , the agent can 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 ( 1b ) . Posterior sampling by contrast samples a possible goal position , takes the shortest route there , and then resamples a different goal position from the updated posterior ( 1c ) . Doing so is much less efficient since the agent ’ s uncertainty is not reduced optimally ( e.g. , states are revisited ) . As this example illustrates , Bayes-optimal policies can 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 metalearning to tackle these challenges , and equip an agent with the ability to strategically explore unseen ( but related ) environments for a given distribution , in order to maximise its expected online return . More specifically , we propose variational Bayes-Adaptive Deep RL ( variBAD ) , a way to meta-learn to perform approximate inference on an unknown task,1 and incorporate task uncertainty directly during action selection . Given a distribution over MDPs p ( M ) , we represent a single MDPM using a learned , low-dimensional stochastic latent variable m and jointly meta-train : 1 . A variational auto-encoder that can infer the posterior distribution over m in a 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 1e 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 . 1We use the terms environment , task , and MDP , interchangeably . 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 Bayes-adaptive policies tailored to the training task distribution , with the only assumption that such a distribution is available for meta-training . We evaluate our approach on the gridworld shown above and on MuJoCo domains that are widely used in meta-RL , 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 . 2 BACKGROUND . We define a Markov decision process ( MDP ) as a tuple M = ( S , A , R , T , T0 , γ , H ) with S a set of states , A a set of actions , R ( rt+1|st , at , st+1 ) a reward function , T ( st+1|st , at ) a transition function , T0 ( s0 ) an initial state distribution , γ a discount factor , andH the horizon . In the standard RL setting , we want to learn a policy π that maximises J ( π ) = ET0 , T , π [ ∑H−1 t=0 γ tR ( rt+1|st , at , st+1 ) ] , the expected return . Here , we consider a multi-task meta-learning setting , which we introduce next . 2.1 TRAINING SETUP . We adopt the standard meta-learning setting where we have a distribution p ( M ) over MDPs from which we can sample during meta-training , with an MDP Mi ∼ p ( M ) defined by a tuple Mi = ( S , A , Ri , Ti , Ti,0 , γ , H ) . Across tasks , the reward and transition functions vary but share some structure . The index i represents an unknown task description ( e.g. , a goal position or natural language instruction ) or task ID . Sampling an MDP from p ( M ) is typically done by sampling a reward and transition function from a distribution p ( R , T ) . During meta-training , batches of tasks are 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 4 ) . 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 : ( 1 ) 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 . 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 reinforcement learning formalised as a Bayes-Adaptive MDP ( BAMDP ) , the solution to which is a Bayesoptimal 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 to a prior b0 = p ( R , T ) . Since the agent does not have access to the true reward and transition function , it can maintain a belief bt ( R , T ) = p ( R , T |τ : t ) , which is the posterior over the MDP given the agent ’ s experience τ : t = { s0 , a0 , r1 , s1 , a1 , . . . , st } 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 s+t ∈ S+ = S × B , where B is the belief space . These transition according to T+ ( s+t+1|s + t , at , rt ) = T + ( st+1 , bt+1|st , at , rt , bt ) = T+ ( st+1|st , at , bt ) T+ ( bt+1|st , at , rt , bt , st+1 ) = Ebt [ T ( st+1|st , at ) ] δ ( bt+1 = p ( R , T |τ : t+1 ) ) ( 1 ) i.e. , the new environment state st is the expected new state w.r.t . the current posterior distribution of the transition function , and the belief is updated deterministically according to Bayes rule . The reward function on hyper-states is defined as the expected reward under the current posterior ( after the state transition ) over reward functions , R+ ( s+t , at , s + t+1 ) = R + ( st , bt , at , st+1 , bt+1 ) = Ebt+1 [ R ( st , at , st+1 ) ] . ( 2 ) This results in a BAMDP M+ = ( S+ , A , R+ , T+ , T+0 , γ , H+ ) ( Duff & Barto , 2002 ) , which is a special case of a belief MDP , i.e , the MDP formed by taking the posterior beliefs maintained by an agent in a 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 a given task . The agent ’ s objective is now to maximise the expected return in the BAMDP , J + ( π ) = Eb0 , T+0 , T+ , π H+−1∑ t=0 γtR+ ( rt+1|s+t , at , s+t+1 ) , ( 3 ) i.e. , maximise the expected return in an initially unknown environment , while learning , within the horizonH+ . Note the distinction between the MDP horizonH and BAMDP horizonH+ . Although they often coincide , we might instead want the agent to act Bayes-optimal within the first N MDP episodes , so H+=N × 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 ( 3 ) is maximised by the Bayes-optimal policy , which automatically trades off exploration and exploitation : it takes exploratory actions to reduce its task uncertainty only insofar as it helps to maximise the expected return within the horizon . The BAMDP framework is powerful because it provides a 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 , we propose a method that simultaneously 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 a highly flexible and scalable approach to Bayes-adaptive Deep RL .

Summary: 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.