Bayesian Policy Optimization for Model Uncertainty

1 Introduction . At its core , real-world robotics focuses on operating under uncertainty . An autonomous car must drive alongside unpredictable human drivers under road conditions that change from day to day . An assistive home robot must simultaneously infer users ’ intended goals as it helps them . A robot arm must recognize and manipulate varied objects . These examples share common themes : ( 1 ) an underlying dynamical system with unknown latent parameters ( road conditions , human goals , object identities ) , ( 2 ) an agent that can probe the system via exploration , while ultimately ( 3 ) maximizing an expected long-term reward via exploitation . The Bayes-Adaptive Markov Decision Process ( BAMDP ) framework ( Ghavamzadeh et al. , 2015 ) elegantly captures the exploration-exploitation dilemma that the agent faces . Here , the agent maintains a belief , which is a posterior distribution over the latent parameters φ given a history of observations . A BAMDP can be cast as a Partially Observable Markov Decision Process ( POMDP ) ( Duff & Barto , 2002 ) whose state is ( s , φ ) , where s corresponds to the observable world state . By planning in the belief space of this POMDP , the agent balances explorative and exploitative actions . In this paper , we focus on BAMDP problems in which the latent parameter space is either a discrete finite set or a bounded continuous set that can be approximated via discretization . For this class of BAMDPs , the belief is a categorical distribution , allowing us to represent it using a vector of weights . The core problem for BAMDPs with continuous state-action spaces is how to explore the reachable belief space . In particular , discretizing the latent space can result in an arbitrarily large belief vector , which causes the belief space to grow exponentially . Approximating the value function over the reachable belief space can be challenging : although point-based value approximations ( Kurniawati et al. , 2008 ; Pineau et al. , 2003 ) have been largely successful for approximating value functions of discrete POMDP problems , these approaches do not easily extend to continuous state-action spaces . Monte-Carlo Tree Search approaches ( Silver & Veness , 2010 ; Guez et al. , 2012 ) are also prohibitively expensive in continuous state-action spaces : the width of the search tree after a single iteration is too large , preventing an adequate search depth from being reached . Our key insight is that we can bypass learning the value function and directly learn a policy that maps beliefs to actions by leveraging the latest advancements in batch policy optimization algorithms ( Schulman et al. , 2015 ; 2017 ) . Inspired by previous approaches that train learning algorithms with an ensemble of models ( Rajeswaran et al. , 2017 ; Yu et al. , 2017 ) , we examine model uncertainty through a BAMDP lens . Although our approach provides only locally optimal policies , we believe that it offers a practical and scalable solution for continuous BAMDPs . Our method , Bayesian Policy Optimization ( BPO ) , is a batch policy optimization method which utilizes a black-box Bayesian filter and augmented state-belief representation . During offline training , BPO simulates the policy on multiple latent models sampled from the source distribution ( Figure 1a ) . At each simulation timestep , it computes the posterior belief using a Bayes filter and inputs the state-belief pair ( s , b ) to the policy . Our algorithm only needs to update the posterior along the simulated trajectory in each sampled model , rather than branching at each possible action and observation as in MCTS-based approaches . Our key contribution is the following . We introduce a Bayesian policy optimization algorithm to learn policies that directly reason about model uncertainty while maximizing the expected long-term reward ( Section 4 ) . To address the challenge of large belief representations , we introduce two encoder networks that balance the size of belief and state embeddings in the policy network ( Figure 1b ) . In addition , we show that our method , while designed for BAMDPs , can be applied to continuous POMDPs when a compact belief representation is available ( Section 4.2 ) . Through experiments on classical POMDP problems and BAMDP variants of OpenAI Gym benchmarks , we show that BPO significantly outperforms algorithms that address model uncertainty without explicitly reasoning about beliefs and is competitive with state-of-the-art POMDP algorithms ( Section 5 ) . 2 Preliminaries : Bayesian Reinforcement Learning . The Bayes-Adaptive Markov Decision Process framework ( Duff & Barto , 2002 ; Ross et al. , 2008 ; Kolter & Ng , 2009 ) was originally proposed to address uncertainty in the transition function of an MDP . The uncertainty is captured by a latent variable , φ ∈ Φ , which is either directly the transition function , e.g . φsas′ = T ( s , a , s′ ) , or is a parameter of the transition , e.g . physical properties of the system . The latent variable is either fixed or has a known transition function . We extend the previous formulation of φ to address uncertainty in the reward function as well . Formally , a BAMDP is defined by a tuple 〈S , Φ , A , T , R , P0 , γ〉 , where S is the observable state space of the underlying MDP , Φ is the latent space , and A is the action space . T and R are the parameterized transition and reward functions , respectively . The transition function is defined as : T ( s , φ , a′ , s′ , φ′ ) = P ( s′ , φ′|s , φ , a′ ) = P ( s′|s , φ , a′ ) P ( φ′|s , φ , a′ , s′ ) . The initial distribution over ( s , φ ) is given by P0 : S × Φ→ R+ , and γ is the discount . Bayesian Reinforcement Learning ( BRL ) considers the long-term expected reward with respect to the uncertainty over φ rather than the true ( unknown ) value of φ . The uncertainty is represented as a belief distribution b ∈ B over latent variables φ. BRL maximizes the following Bayesian value function , which is the expected value given the uncertainty : Vπ ( s , b ) = R ( s , b , a ′ ) + γ ∑ s′∈S , b′∈B P ( s′ , b′|s , b , a′ ) Vπ ( s′ , b′ ) = R ( s , b , a′ ) + γ ∑ s′∈S , b′∈B P ( s′|s , b , a′ ) P ( b′|s , b , a′ , s′ ) Vπ ( s′ , b′ ) ( 1 ) where the action is a′ = π ( s , b ) .1 The Bayesian reward and transition functions are defined in expectation with respect to φ : R ( s , b , a′ ) = ∑ φ∈Φ b ( φ ) R ( s , φ , a ′ ) , P ( s′|s , b , a′ ) = ∑ φ∈Φ b ( φ ) P ( s ′|s , φ , a′ ) . The belief distribution can be maintained recursively , with a black-box Bayes filter performing posterior updates given observations . We describe how to implement such a Bayes filter in Section 4.1 . The use of ( s , b ) casts the partially observable BAMDP as a fully observable MDP in belief space , which permits the use of any policy gradient method . We highlight that a reactive Bayesian policy in belief space is equivalent to a policy with memory in observable space ( Kaelbling et al. , 1998 ) . In our work , the complexity of memory is delegated to a Bayes filter that computes a sufficient statistic of the history . In partially observable MDPs ( POMDPs ) , the states can be observed only via a noisy observation function . Mixed-observability MDPs ( MOMDPs ) ( Ong et al. , 2010 ) are similar to BAMDPs : their states are ( s , φ ) , where s is observable and φ is latent . Although any BAMDP problem can be cast as a POMDP or a MOMDP problem ( Duff & Barto , 2002 ) , the source of uncertainty in a BAMDP usually comes from the transition function , not the unobservability of the state as it does with POMDPs and MOMDPs . 3 Related Work . A long history of research addresses belief-space reinforcement learning and robust reinforcement learning . Here , we highlight the most relevant work and refer the reader to Ghavamzadeh et al . ( 2015 ) , Shani et al . ( 2013 ) , and Aberdeen ( 2003 ) for more comprehensive reviews of the Bayes-Adaptive and Partially Observable MDP literatures . Belief-Space Reinforcement Learning . Planning in belief space , where part of the state representation is a belief distribution , is intractable ( Papadimitriou & Tsitsiklis , 1987 ) . This is a consequence of the curse of dimensionality : the dimensionality of belief space over a finite set of variables equals the size of that set , so the size of belief space grows exponentially . Many approximate solvers focus on one or more of the following : 1 ) value function approximation , 2 ) compact , approximate belief representation , or 3 ) direct mapping of belief to an action . QMDP ( Littman et al. , 1995 ) assumes full observability after one step to approximate Q-value . Point-based solvers , like SARSOP ( Kurniawati et al. , 2008 ) and PBVI ( Pineau et al. , 2003 ) , exploit the piecewise-linear-convex structure of POMDP value functions ( under mild assumptions ) to approximate the value of a belief state . Samplingbased approaches , such as BAMCP ( Guez et al. , 2012 ) and POMCP ( Silver & Veness , 2010 ) , combine Monte Carlo sampling and simple rollout policies to approximate Q-values at the root node in a search tree . Except for QMDP , these approaches target discrete POMDPs and can not be easily extended to continuous spaces . Sunberg & Kochenderfer ( 2018 ) extend POMCP to continuous spaces using double progressive widening . Model-based trajectory 1 The state space S can be either discrete or continuous . The belief space B is always continuous , but we use ∑ notation for simplicity . optimization methods ( Platt et al. , 2010 ; van den Berg et al. , 2012 ) have also been successful for navigation on systems like unmanned aerial vehicles and other mobile robots . Neural network variants of POMDP algorithms are well suited for compressing highdimensional belief states into compact representations . For example , QMDP-Net ( Karkus et al. , 2017 ) jointly trains a Bayes-filter network and a policy network to approximate Q-value . Deep Variational Reinforcement Learning ( Igl et al. , 2018 ) learns to approximate the belief using variational inference and a particle filter , and it uses the belief to generate actions . Our method is closely related to Exp-GPOMDP ( Aberdeen & Baxter , 2002 ) , a model-free policy gradient method for POMDPs , but we leverage model knowledge from the BAMDP and revisit the underlying policy optimization method with recent advancements . Peng et al . ( 2018 ) use Long Short-Term Memory ( LSTM ) ( Hochreiter & Schmidhuber , 1997 ) to encode a history of observations to generate an action . The key difference between our method and Peng et al . ( 2018 ) is that BPO explicitly utilizes the belief distribution , while in Peng et al . ( 2018 ) the LSTM must implicitly learn an embedding for the distribution . We believe that explicitly using a Bayes filter improves data efficiency and interpretability . Robust ( Adversarial ) Reinforcement Learning . One can bypass the burden of maintaining belief and still find a robust policy by maximizing the return for worst-case scenarios . Commonly referred to as Robust Reinforcement Learning ( Morimoto & Doya , 2001 ) , this approach uses a min-max objective and is conceptually equivalent to H-infinity control ( Başar & Bernhard , 2008 ) from classical robust control theory . Recent works have adapted this objective to train agents against various external disturbances and adversarial scenarios ( Pinto et al. , 2017 ; Bansal et al. , 2018 ; Pattanaik et al. , 2018 ) . Interestingly , instead of training against an adversary , an agent can also train to be robust against model uncertainty with an ensemble of models . For example , Ensemble Policy Optimization ( EPOpt ) ( Rajeswaran et al. , 2017 ) trains an agent on multiple MDPs and strives to improve worst-case performance by concentrating rollouts on MDPs where the current policy performs poorly . Ensemble-CIO ( Mordatch et al. , 2015 ) optimizes trajectories across a finite set of MDPs . While adversarial and ensemble model approaches have proven to be robust even to unmodeled effects , they may result in overly conservative behavior when the worst-case scenario is extreme . In addition , since these methods do not infer or utilize uncertainty , they perform poorly when explicit information-gathering actions are required . Our approach is fundamentally different from them because it internally maintains a belief distribution . As a result , its policies outperform robust policies in many scenarios . Adaptive Policy Methods . Some approaches can adapt to changing model estimates without operating in belief space . Adaptive-EPOpt ( Rajeswaran et al. , 2017 ) retrains an agent with an updated source distribution after real-world interactions . PSRL ( Osband et al. , 2013 ) samples from a source distribution , executes an optimal policy for the sample for a fixed horizon , and then re-samples from an updated source distribution . These approaches can work well for scenarios in which the latent MDP is fixed throughout multiple episodes . Universal Policy with Online System Identification ( UP-OSI ) ( Yu et al. , 2017 ) learns to predict the maximum likelihood estimate φMLE and trains a universal policy that maps ( s , φMLE ) to an action . However , without a notion of belief , both PSRL and UP-OSI can over-confidently execute policies that are optimal for the single estimate , causing poor performance in expectation over different MDPs .

Summary: In this paper, the authors propose a policy gradient algorithm for solving a Bayes-Adaptive MDP (BAMDP). At each iteration, the algorithm samples several MDPs from the prior distribution and simulates a trajectory for each sampled MDP. During the simulation, the algorithm uses a Bayes filter to update the posterior belief distribution at each time step. Finally, the algorithm uses the sampled trajectories and update the policy using the TRPO algorithm.