Local policy search with Bayesian optimization

1 Introduction . Reinforcement learning ( RL ) is a notoriously data-hungry machine learning problem , where state-ofthe-art methods easily require tens of thousands of data points to learn a given task [ 1 ] . For every data point , the agent has to carry out potentially complex interactions with its environment , either in simulation or in the physical world . This expensive data collection motivates the development of sample-efficient algorithms . Herein , we consider policy search problems , a type of RL technique where we directly optimize the parameters of a policy with respect to the cumulative reward over a finite episode . The collected data is utilized to estimate the direction of local policy improvement , enabling the use of powerful optimization techniques such as stochastic gradient descent . Policy gradient methods ( e.g. , [ 2–5 ] ) usually rely on random perturbations for data generation , e.g. , in the form of exploration noise in the action space or stochastic policies , and do not reason about uncertainty in their gradient estimation . However , innate in the RL setting is the ability to actively generate data , allowing the agent to decide on informative queries , thereby potentially reducing the amount of data needed to find a ( local ) optimum . Active sampling has the potential to allow those algorithms to improve sample complexity , reducing the number of environment interactions . ∗Equal contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In contrast to random sampling , Bayesian optimization ( BO ) [ 6 ] is a paradigm to optimize expensive-to-evaluate and noisy functions in a sample-efficient manner . At the core of BO is the question of how to query the objective function efficiently to maximize the information contained in each sample . By building a probabilistic model of the objective using past data and , critically , prior knowledge , the algorithm can reason about how to query a noisy oracle to solve the optimization task . Since RL can be framed as a black-box optimization problem , we can use BO to learn policies in a sample-efficient way . However , even though BO has been used to tackle RL , these approaches are often restricted to low-dimensional problems . One reason is that BO aims to find a global optimum ; hence , without further assumptions , BO algorithms need to model and search the entire domain , which needs a lot of data and gets exponentially more difficult as the dimensionality increases . Additionally , as the amount of data grows so does the computational complexity of probabilistic models , which becomes a significant problem . However , the success of RL algorithms using policy gradient methods indicates that for many problems it is sufficient to find a locally optimal policy . Our proposed algorithm combines the strength of gradient-based policy optimization with active sampling of policy parameters using BO . We thereby improve the computational complexity of BO methods on the one hand , and the sample-inefficiency of gradient-based methods on the other hand , especially when proper prior knowledge is available . We achieve these improvements by explicitly learning a probabilistic model of the objective in the form of a Gaussian process ( GP ) [ 7 ] . From this model , we can jointly infer the objective and its gradients with a tractable probabilistic posterior . The resulting Jacobian estimate includes all data points , rendering data usage more efficient . Further , the algorithm infers informative queries from the uncertainty of the probabilistic model to improve the estimate of the local gradient . While in this paper we adapt the setting of Mania et al . [ 1 ] and assume access to zeroth-order information only , the algorithm extends straightforwardly to policy gradient algorithms where additional first-order information is available . In summary , the contribution of this paper is a local BO-like optimizer called Gradient Information with BO ( GIBO ) . The queries of GIBO are chosen optimally to minimize uncertainty about the Jacobian . GIBO uses a local GP model for active sampling and gradient estimation and can be used with existing policy search algorithms . Using only zeroth order information , GIBO is able to • significantly improves sample complexity in extensive within-model comparisons , i.e. , when accurate prior knowledge is available ; • is able to solve RL benchmark problems in a sample efficient manner ; and • reduces variance in the rewards when compared to non-active sampling baselines . 2 Preliminaries . This work presents a local optimizer with active sampling . The objective function and its derivative ’ s joint distribution are modeled using a GP . Since we have developed the optimizer with the RL application in mind , we also introduce the RL problem . For the sake of brevity , we refer the reader to [ 7 ] and [ 8 ] for an introduction into GPs and BO , respectively . 2.1 Problem setting . In the following we phrase policy search as a black-box optimization problem . For a parameterized policy πθ : Θ × S → A that maps states s ∈ S and the static policy parameters θ ∈ Θ to actions a ∈ A , we use the same performance measure as in policy gradient methods for the episodic case . Hence , the objective function J : Rd → R is defined as J ( θ ) = Eπθ [ I∑ i=0 ri ] , where Eπθ is the expectation under policy πθ , ri is the reward at time step i , and I the length of the episode . A BO query is equivalent to the return of one rollout following the policy πθ in the environment . The expected episodic reward is entirely determined by choice of policy parameters ( and the initial conditions ) . Thus , the optimizer explores the reward function in the parameter space rather than in the action space . Since initial conditions might vary and the environment can be non-deterministic , reward evaluations are noisy . Policy search herein is abstracted as a zeroth-order optimization problem of the form θ∗ = arg max θ∈Θ J ( θ ) , ( 1 ) where θ is the variable and Θ ⊂ Rd a bounded set . To solve ( 1 ) , an optimization algorithm can query an oracle for a noisy function evaluation y = J ( θ ) + ω . We assume an i.i.d . noise variable ω ∈ R to follow a normal distribution ω ∼ N ( 0 , σ2 ) with variance σ2 . We do not assume access to gradient information or other higher-order oracles for conciseness . Albeit , GIBO requires that the following critical assumption is fulfilled : Assumption 1 . The objective function J is a sample from a known GP prior J ∼ GP ( m ( θ ) , k ( θ , θ′ ) ) , where the mean function is at least once differentiable and the covariance function k is at least twice differentiable , w.r.t . θ . This is the standard setting for BO with the addition that the mean and kernel need to be differentiable , which is satisfied by some of the most common kernels such as the squared exponential ( SE ) kernel . In the empirical section , we investigate the performance of the developed algorithm with and without Assumption 1 holding true . 2.2 Jacobian GP model . Since GPs are closed under linear operations , the derivative of a GP is again a GP [ 7 ] . This enables us to derive an analytical distribution for the objective ’ s Jacobian , which we can use as a proxy for gradient estimates and enable gradient-based optimization . Following Rasmussen and Williams [ 7 ] , the joint distribution between a GP and its derivative at the point θ∗ is [ ȳ ∇θ∗J∗ ] ∼ N ( [ m ( X ) ∇θ∗m ( θ∗ ) ] , [ K ( X , X ) + σ2I ∇θ∗K ( X , θ∗ ) ∇θ∗K ( θ∗ , X ) ∇2θ∗K ( θ∗ , θ∗ ) ] ) , ( 2 ) where ȳ are the n zeroth-order observations , X ⊂ Θ are the locations of these observations X = [ θ1 , . . . , θn ] , andK the covariance matrix given by the kernel function k : Θ×Θ→ R. The posterior can be derived by conditioning the joint Gaussian prior distribution on the observation [ 7 ] p ( ∇θ∗J∗ ∣∣θ∗ , X , ȳ ) ∼ N ( µ′∗ , Σ′∗ ) µ′∗ = ∇θ∗m ( θ∗ ) ︸ ︷︷ ︸ ∈Rd +∇θ∗K ( θ∗ , X ) ︸ ︷︷ ︸ ∈Rd×n ( K ( X , X ) + σ2I ) −1 ︸ ︷︷ ︸ ∈Rn×n ( ȳ −m ( X ) ) ︸ ︷︷ ︸ ∈Rn ∈ Rd Σ′∗ = ∇2θ∗K ( θ∗ , θ∗ ) ︸ ︷︷ ︸ ∈Rd×d −∇θ∗K ( θ∗ , X ) ︸ ︷︷ ︸ ∈Rd×n ( K ( X , X ) + σ2I ) −1 ︸ ︷︷ ︸ ∈Rn×n ∇θ∗K ( X , θ∗ ) ︸ ︷︷ ︸ ∈Rn×d ∈ Rd×d . ( 3 ) Remark 1 . Note that the term ( K ( X , X ) + σ2I ) −1 with the highest computational cost ( O ( N3 ) ) is the same term that is used to compute the posterior over J . Therefore , calculating the Jacobian does not add to the computational complexity once a GP posterior has been computed . Any twice differentiable kernel is sufficient for the presented framework , but we assume a SE kernel for the remainder of the paper . For the derivatives of the SE kernel function see Appendix A.1 . For a visual example of function- and the Jacobian-posterior , refer to Fig . 1 . The figure indicates that a zeroth-order oracle is enough to form a reasonable belief over the function ’ s gradient . Moreover , Fig . 1 shows that the uncertainty about the Jacobian gets reduced between query points more so than at the query points themselves . To minimize uncertainty about the Jacobian at a specific point , it intuitively makes sense to space out query points in its immediate surrounding . Herein , we formalize this intuition and formulate an optimization problem that sequentially decides on query points that provide the most information about the Jacobian . 2.3 Related work . In the presented contribution we focus on the benefits of active sampling in policy search , specifically on sample efficiency . Therefore , this section focuses on active sampling in model-free RL setting using probabilistic uncertainty estimations . Most of the literature in this setting is based on BO , but generating informative samples is also discussed in literature regarding evolutionary strategies as well as in policy gradient methods . Bayesian optimization as an active sampling method has been used for global policy search , mostly in lower dimensional parameter spaces from 2–15 dimensions [ 9–13 ] . Global BO for RL , exemplified by the mentioned literature and without additional assumptions , is limited to relatively low dimensional problems for two reasons : ( i ) the computational complexity of global probabilistic models does not scale well with the number of data points , ( ii ) global optimization of high-dimensional non-convex objectives is a challenging problem to solve in general . To combat these problems local variants of BO have been proposed and applied to RL problems , see e.g. , [ 14–17 ] . These works rely on restricting the search space of BO by a probabilistic belief over the optimums location [ 14 ] using rectangular trust-regions [ 15 ] , learning a partition [ 16 ] , or by staying close ( as defined by the GP kernel ) to past samples [ 17 ] . Restrictions in the parameter space avoid ’ over-exploration ’ of high-dimensional search spaces and thereby encourage exploitation of ( local ) minima . Our proposed method delegates the exploitation to a gradient-based optimizer after exploring a local property , the function ’ s derivative at the current iterate , for which a local search ( and model ) is sufficient . McLeod et al . [ 18 ] and more recently Shekhar and Javidi [ 19 ] suggest switching from global BO to a local gradient-based method once a locally convex region containing a low-regret solution has been identified , thereby improving convergence properties of BO . In [ 18 ] GIBO can replace the local optimizer of choice and in Shekhar and Javidi [ 19 ] GIBO can be used for optimal uncertainty reduction in the gradient estimate . In general , a GP posterior can incorporate gradient information if the kernel is differentiable and a first-order oracle is available . Bayesian optimization methods that utilize gradient observations are known as first-order BO , and different approaches on how to include the derivative information in the model and acquisition functions have been proposed [ 20–23 ] . Since computing the joint posterior using first- and zeroth-order information is computationally expensive , Ahmed et al . [ 21 ] and Wu et al . [ 22 ] are using a single directional derivative instead of all partial derivatives . A first-order BO approach for RL , where the gradient information is actively used to decide on the following query , is introduced by Prabuchandran et al . [ 23 ] . The method therein actively searches for local optima by querying points where the gradient is expected to be zero . In contrast to this approach , we actively reduce local uncertainty of the Jacobian model and afterwards a gradient-based optimizer decides on the next location . Reinforcement learning problems in the form of ( 1 ) can also be used by evolutionary methods such as [ 1 , 24–26 ] and recently by policy gradient methods Faccio et al . [ 27 ] . These methods typically explore through random perturbations in the parameter space of the policy instead of active sampling . However , generating more informative samples improves evolutionary strategies . Maheswaranathan et al . [ 25 ] shows this by adapting the sampling distribution using surrogate gradient information such as previous estimates , and Choromanski et al . [ 26 ] uses determinantal point processes for informative samples . Policies that generate more informative samples have helped to improve model-free RL algorithms ’ performance during the past decade ; we mention three examples here : Levine and Koltun [ 28 ] propose so-called guiding samples in high reward areas using differential dynamic programming and model knowledge . Soft actor-critic ( SAC ) methods [ 3 ] add the policy ’ s entropy to the reward function to encourage exploration and improve the variance of gradient estimates . Based on SAC an optimistic actor-critic algorithm is introduced in [ 29 ] with a different exploration strategy that samples more informative actions . To reduce variance in the gradient estimate , it is possible to use GIBO as a layer between the policy gradient estimator such as SAC and a gradient-based optimizer , e.g. , stochastic gradient ascent or Adam [ 30 ] . In future work GIBO can be extended to utilize state-of-the art policy gradient methods as an additional oracle for first-order information and help these methods reducing the variance of their gradient estimates through active sampling . Based on the posterior conditioned on all collected rewards , our algorithm can supply posterior gradient estimates and subsequent queries to evaluate . To demonstrate the benefits of GIBO in a simple setup , we adopt the setting proposed by Mania et al . [ 1 ] as a baseline . Augmented Random Search ( ARS ) [ 1 ] assumes a black-box setting without access to gradient samples and estimates the gradient from the finite-difference of random perturbations , effectively solving RL problems . We replace the random sampling strategy of ARS with active sampling and the gradient estimation with a GP model . These changes improve the sample complexity and variance of ARS , especially when prior knowledge about the objective function is available .

Summary: This paper modifies the standard BayesOpt pipeline in 2 ways:  1. Defines a new acquisition function, the "Gradient Information" (GI) which is the reduction in the variance of the Jacobian variance of $J$. 2. Exploration and Exploitation become disjointed: During the BayesOpt process, a batch of $M$ queries are made via argmax of the GI acquisition, purely for exploration (rather than a mix of exploitation and exploration as is standard in BayesOpt), with the goal of producing a final gradient update step $\nabla_{\theta}J$.   This defines a "GIBO" algorithm, which is then experimentally evaluated over the following benchmarks:   1. Sampled functions from GP's 2. LQR Controllers 3. Linear controllers for standard Mujoco tasks.  which show that GIBO outperforms previous ES/ARS and standard BayesOpt methods especially in high dimensional settings, most likely due to the use of "local search"/gradients.