An Experimental Design Perspective on Exploration in Reinforcement Learning

1 INTRODUCTION . Reinforcement learning ( RL ) has suffered for years from a curse of poor sample complexity . State-ofthe-art model-free reinforcement learning algorithms routinely take tens of thousands of sampled transitions to solve very simple tasks and millions to solve moderately complex ones ( Haarnoja et al. , 2018 ; Lillicrap et al. , 2015 ) . The current best model-based reinforcement learning ( MBRL ) algorithms are better , requiring thousands of samples for simple problems and hundreds of thousands of samples for harder ones ( Chua et al. , 2018 ) . In settings where each sample is expensive , even this smaller cost can be prohibitive for the practical application of RL . For example , in the physical sciences many simulators require the solution of computationally demanding spatial PDEs in plasma control ( Breslau et al. , 2018 ) or aerodynamics applications ( Jameson & Fatica , 2006 ) . In robotics , due to the difficulty of simulating more complicated objects ( Heiden et al. , 2021 ) , typical RL methods operate on fast but limited rigid-body simulators ( Todorov et al. , 2012 ) . These costly transition functions prompt the question : “ If we were to collect one additional datapoint from anywhere in the state-action space to best improve our solution to the task , which one would it be ? ” An answer to this question can be used to guide data collection in RL . Across the fields of black-box optimization and experimental design , techniques have been developed which choose data to collect that are particularly useful in improving the value of the objective function of the problem . For example , Bayesian optimization ( BO ) focuses on maximizing a function for which only a black-box model is available and queries are expensive ( Frazier , 2018 ; Shahriari et al. , 2015 ) . More generally , Bayesian optimal experimental design ( BOED ) aims to choose data to collect which are maximally informative about the value of some derived quantity ( Chaloner & Verdinelli , 1995 ) . We aim to leverage these ideas for data-efficiency in reinforcement learning . Along these lines , several works in the realm of Bayesian RL address this problem in the sequential setting . Bayes-adaptive MDP ( Ross et al. , 2007 ) constructs a modified MDP by augmenting it using the posterior of the MDP , leading to a policy that can optimally trade off between acquiring more information and exploiting the knowledge it already has . However such an MDP is intractable to exactly solve in large spaces so approximations and heuristics have been developed for the solution ( Smith , 2007 ; Guez et al. , 2012 ) . A particularly relevant heuristic is the value of perfect information ( VPI ) from Dearden et al . ( 1998 ) , which attempts to capture the potential change in value of a state if the value of a particular action at that state was perfectly known , specifically in a tabular setting . However , VPI doesn ’ t attempt to distinguish between states that are visited during the execution of the optimal policy and those that aren ’ t . This is a critical distinction when evaluating data to collect in continuous spaces as otherwise queries would be wasted learning the optimal policy in irrelevant parts of the state space . Motivated by our opening question , in this paper we study the setting where the agent collects data by sequentially making queries to the transition function with free choice of both the initial state and the action . We refer to this setting as transition-query reinforcement learning ( TQRL ) and formally define it in Section 3.1 . An effective algorithm for this setting would need to choose informative initial states as well as actions , which to the best of our knowledge has not been studied in the continuous MDP literature . In this work , we draw a connection between MBRL and the world of BOED by deriving an acquisition function that quantifies how much information a state-action pair would provide about the optimal solution to a MDP . Like the techniques in Bayesian RL , our acquisition function is able to determine which state-action pairs are worth acquiring in a way which takes into account the reward function and the uncertainty in the dynamics . However , like the full Bayes-Adaptive MDP and unlike the VPI heuristic , this function takes into account the current estimates of which states the optimal policy will actually visit and values potential queries accordingly . It is also scalable enough to apply to multidimensional continuous control problems . In particular , our acquisition function is the expected information gain ( EIG ) about the trajectory taken by an optimal policy in the MDP that would be achieved if we were to query the transition function at a given state-action pair . Finally , we assess the performance of our acquisition function as a data selection strategy in the TQRL setting . Using this method we are able to solve several continuous reinforcement learning tasks ( including a nuclear fusion example ) using orders of magnitude less data than a variety of competitor methods . In summary , the contributions of our paper are : • We construct a novel acquisition function that quantifies how much information a stateaction pair would provide about the optimal solution to a continuous MDP if the next state were observed from the ground truth transition function . Our function is able to select relevant datapoints for control purposes leading to much-improved data efficiency . • We propose a practical algorithm for computing this acquisition function and use it to solve continuous MDPs in the TQRL setting . • We evaluate the algorithm on five widely varying control tasks , where it is often orders of magnitude more sample-efficient than competitor methods and reaches similar asymptotic performance . 2 RELATED WORK . Transition Query Reinforcement Learning In many RL algorithms , data is collected by initializing a policy at a start state and executing actions in the environment in an episodic manner . Kearns et al . ( 2002 ) introduced the setting where the agent collects data by sampling transitions in a sequential manner from the ground truth transition model by querying at a state and action of its choice , which they refer to as RL with access to a generative model . We refer to this setting for brevity as TQRL . This setting is relevant in a variety of real-world applications where there is a simulator of the transition model available . We in particular see the setting in nuclear fusion research , where plasma dynamics are modeled by solving large partial differential equations where 200ms of plasma time can take up to an hour in simulation ( Breslau et al. , 2018 ) . There is substantial theoretical work on TQRL for finite MDPs . In particular , Azar et al . ( 2013 ) give matching log-linear upper and lower PAC sample complexity bounds , a substantial speedup to the upper bound for the standard problem which is quadratic in state size ( Kakade , 2003 ) . This is achieved simply by the naive algorithm of learning a transition model by uniformly sampling the space and then performing value iteration on the estimate of the MDP for an optimal policy . More recently , the bound for this setting was tightened to hold for smaller numbers of samples by Li et al . ( 2020 ) , meaning that for any dataset size in a continuous problem , the PAC performance can be quantified . Finally , Agarwal et al . ( 2020 ) show that the naive ‘ plug-in ’ estimator used in the previous works is minimax optimal for this setting . In summary , this setting is thoroughly understood for finite MDPs and it gives a sample complexity reduction from quadratic to linear in the state space size . To our knowledge there do not exist works specifically solving the TQRL setting , though every standard RL algorithm can be applied in this setting by continuing to collect data in an online manner . In this work , we give an algorithm specifically designed for this setting in continuous MDPs . Our work empirically shows sample complexity benefits reminiscent of those theoretically shown in the tabular setting . Exploration in Reinforcement Learning To encourage exploration in RL , agents often use an - greedy approach ( Mnih et al. , 2013 ) , upper confidence bounds ( UCB ) ( Chen et al. , 2017 ) , Thompson sampling ( Osband et al. , 2016 ) , added Ornstein-Uhlenbeck action noise ( Lillicrap et al. , 2015 ) , or entropy bonuses ( Haarnoja et al. , 2018 ) to add noise to a policy which is otherwise optimizing the RL objective . Though UCB , Thompson Sampling , and an entropy bonus all try to explore in a way which can adapt to the problem , they all tackle which action to take from a predetermined state and don ’ t explicitly consider which states would be good to acquire data from . As we mentioned , an ideal method of exploration would be to solve the intractable Bayes-adaptive MDP ( Ross et al. , 2007 ) , giving an optimal tradeoff between exploration and exploitation . Kolter & Ng ( 2009 ) ; Guez et al . ( 2012 ) show that even approximating these techniques in the sequential setting can result in substantial theoretical reductions in sample complexity compared to frequentist PACMDP bounds as in Kakade ( 2003 ) . Other methods stemming from Dearden et al . ( 1998 ; 1999 ) address this by using the myopic Value of Perfect Information as a heuristic for similar Bayesian exploration . These methods don ’ t scale to continuous problems and don ’ t give any method of choosing at which states to query , only actions given a current state . These methods were further extended with the development of Knowledge Gradient policies ( Ryzhov et al. , 2019 ; Ryzhov & Powell , 2011 ) , which approximate the value function of the Bayes-adaptive MDP , and Information-Directed Sampling ( Russo & Van Roy , 2014 ) , which takes actions based on minimizing the ratio between squared regret and information gain over dynamics . This was extended to continuous-state finite-action settings in Nikolov et al . ( 2019 ) . However , this work doesn ’ t solve fully continuous problems , operates in the rollout setting rather than TQRL , and computes the information gain with respect to the dynamics rather than some notion of the optimal policy . A tutorial on Bayesian methods can be found in Ghavamzadeh et al . ( 2016 ) for further reference . Separate from the techniques used in RL for a particular task , several methods tackle the problem of unsupervised exploration ( Schmidhuber , 1991 ) , where the goal is to learn as much as possible about the transition model without a task or reward function . One approach synthesizes a reward from modeling errors ( Pathak et al. , 2017 ) . Another estimates learning progress by estimating model accuracy ( Lopes et al. , 2012 ) . Others use an information gain-motivated formulation of model disagreement ( Pathak et al. , 2019 ; Shyam et al. , 2019 ) as a reward . Other methods incentivize the policy to explore regions it hasn ’ t been before using hash-based counts ( Tang et al. , 2017 ) , predictions mimicking a randomly initialized network ( Burda et al. , 2019 ) , or a density estimate ( Bellemare et al. , 2016 ) . However , these methods all assume that there is no reward function and would be inefficient if used in that setting as they spend time exploring areas of state space which can be quickly determined to be bad for maximizing reward on a task . Bayesian Algorithm Execution and BOED Recently , a flexible framework known as Bayesian algorithm execution ( BAX ) ( Neiswanger et al. , 2021 ) has been proposed for efficiently estimating properties of expensive black-box functions , which builds off of a large literature from Bayesian Optimal Experiment Design ( Chaloner & Verdinelli , 1995 ) . The BAX framework gives a general procedure for sampling points which are informative about the future execution of an algorithm . In this paper , we extend this framework to the setting of model-predictive control , when we have expensive dynamics ( i.e . transition function ) which we treat as a black-box function in the BAX framework . Via this strategy , we are able to use similar techniques to develop acquisition functions for data collection in reinforcement learning . Gaussian Processes ( GPs ) in Reinforcement Learning There has been substantial prior work using GPs in reinforcement learning . Most well-known is PILCO ( Deisenroth & Rasmussen , 2011 ) , which computes approximate analytic gradients of policy parameters through the GP dynamics model while accounting for uncertainty . Most related to our eventual MPC method is ( Kamthe & Deisenroth , 2018 ) , which gives a principled probabilistic model-predictive control algorithm for GPs .

Summary: This paper uses insight from Bayesian optimal experimental design and define an *acquisition function* which quantifies information gain about an MDP's optimal solution given a state-action pair. Such an acquisition function would provide a mechanism for selecting which transitions to query from a ground truth model in model-based RL, in hopes to greatly improve sample efficiency of a learning algorithm. They propose a method for computing this acquisition function, and evaluate it empirically.