Reward Learning as Doubly Nonparametric Bandits:  Optimal Design and Scaling Laws

1 INTRODUCTION . Specifying the reward function accurately for a desired objective , or reward engineering , is challenging to perform by hand , as the consequences of even small errors can be drastic ( Hadfield-Menell et al. , 2017 ) . To address this , reward learning seeks to learn a predictive model of the reward function from data , which is obtained from carefully selected queries to human annotators . The learned reward model is then used as the optimization objective for policy learning . Reward learning has achieved significant empirical success in domains such as text summarization ( Stiennon et al. , 2020 ; Böhm et al. , 2019 ) , robot locomotion ( Daniel et al. , 2014 ) , predicting driving styles ( Kuderer et al. , 2015 ) , and Atari game playing ( Christiano et al. , 2017 ) . Despite their success , reward learning methods still lack theoretical grounding . Moreover , their behavior can be brittle even on simple tasks , due to the difficulty of choosing appropriate queries and due to feedback loops from adaptive querying ( Freire et al. , 2020 ) . Indeed , an ablation study in Christiano et al . ( 2017 ) suggests that random queries can outperform or be competitive with adaptive query procedures . To address these issues , we provide a theoretical framework for analyzing reward learning , framing it as a doubly nonparametric experimental design problem . This framework helps elucidate the role of query selection ( Chaloner & Verdinelli , 1995 ) and also enables us to derive scaling laws—how the sizes of the policy and reward models affect the query complexity—for reward learning ( Kaplan et al. , 2020 ) . Proposed framework . In our framework , we suppose we are given a reward class Cr and policy class Cπ . Our goal is to find a policy π̂ ∈ Cπ that performs well according to an unknown true reward r∗ ∈ Cr . To do this , we query policies π ∈ Cπ , observing noisy estimates of their true reward , and use this information to choose the eventual policy π̂ . To be compatible with modern nonparametric learning methods ( i.e . neural nets ) , we view Cr and Cπ as subsets of Reproducing Kernel Hilbert Spaces ( RKHS ) . The learner therefore optimizes a nonparametric reward function over a nonparametric space of policies , making the task “ doubly ” non- parametric . In contrast , previous work considers a nonparametric function class or reward class , but typically not both . For instance , nonparametric zeroth order or bandit optimization ( Srinivas et al. , 2010 ; Mockus , 2012 ; Wang et al. , 2018 ) considers a nonparametric function on a finite-dimensional input space . Conversely , nonparametric supervised learning ( Wahba , 1990 ; Hofmann et al. , 2008 ) minimizes a known loss function over a nonparametric input space . The doubly nonparametric nature of our task poses new challenges . The ( possibly ) infinitedimensional RKHS requires the learner to select which subspace to explore given a finite number of queries . Furthermore , the unknown reward function makes it challenging for the learner to reason about the information gained from the selected query policies . We address these challenges by deriving a risk upper bound for a family of plug-in estimators based on ridge regression , and then optimizing this bound to solve the optimal design task . In addition to the optimal design problem , our framework allows us to study scaling laws with respect to the reward ( or policy ) class by varying the rate of decay of their corresponding eigenspectrum . This decay rate determines the effective dimensionality of a RKHS ( Zhang , 2002 ) , and provides a natural proxy for varying the the size of the reward or policy class . Qualitatively , our main results show that the excess risk asymptotically vanishes as long as the policy class grows at a slower rate relative to the reward class . Sharpness of analysis . Our risk bounds apply to reward and policy classes of arbitrary or even infinite dimensionality . Despite this generality , we show they provide stronger guarantees than previous known bounds for the specialized settings of compact policy sets and kernel multi-armed bandits . In Section 4.3 , we look at a special case of our problem when the policy set Cπ is a compact subspace and thus has finite rank . For these instances , we show that our learning algorithm obtains a better excess risk O ( n− β β+2 ) versus a rate of O ( n− β−1 2 ( β+1 ) ) obtained by the adaptive GP-UCB algorithm ( Srinivas et al. , 2010 ) , where β > 0 is a power law decay rate . In Section 5 , we specialize our general results to the well-studied problem of Gaussian process bandit optimization ( Williams & Rasmussen , 2006 ) , also known as kernel multi-armed bandit ( MAB ) . Specifically , for the class of Matérn kernels with parameter ν in d dimensions , we show that our algorithm achieves a regret bound of Õ ( T 4ν+d ( 4d+6 ) 6ν+d ( 4d+7 ) ) which is strictly better than those achieved by the GP-UCB and GP-Thompson Sampling ( GP-TS ) ( Chowdhury & Gopalan , 2017 ) algorithms and comparable with π-GP UCB ( Janz et al. , 2020 ) and supKernelUCB ( Valko et al. , 2013 ; Vakili et al. , 2021 ) ; see Table 1 for details . GP-UCB and GP-TS are only yield sub-linear regret bounds when the smoothness of the kernel ν > d2—thus in high dimensions , these bounds essentially become vacuous . The π-GP UCB algorithm was designed specifically to overcome this issue . Our proposed algorithm achieves sublinear regret for all ν > 3/2 . Our Contributions . We propose doubly-nonparametric bandits as a framework for theoretically studying the reward learning problem . Within this framework , we obtain finite sample risk bounds for a ridge regression based plug-in estimator and derive scaling laws for reward learning . From a technical standpoint , we study the optimal design problem for our estimator to select informative query points by showing that the excess risk depends only on the spectral properties of a certain operator of the two RKHSs and the empirical covariance matrix . As a corollary of our risk bounds , we provide sharper regret bounds for a class of kernel MAB problems compared to several existing algorithms , showing that the doubly-nonparametric lens of reward learning is fruitful even for “ singly-nonparametric ” tasks . To obtain these bounds , our reduction carefully constructs two different RKHSs to embed the input space and reward function into a policy and reward class . 2 FRAMEWORK : DOUBLY NONPARAMETRIC BANDITS . Our framework considers non-parametric policy learning with non-parametric reward models . We let π ∈ Hπ denote an arbitrary policy and r ∈ Hr denote an arbitrary reward function , where Hπ and Hr are Reproducing Kernel Hilbert Spaces . For technical reasons , we assume the corresponding kernel functionsKπ andKr both satisfy the Hilbert-Schmidt condition ( see Appendix A for details ) . We let F ( π , r ) ∈ R denote the reward obtained by selecting policy π under reward function r and consider the case where the evaluation functional F is linear in both π and r. In other words , F ( π , r ) = 〈r , Mπ〉Hr where M : Hπ 7→ Hr is a known linear mapping from the policy space to the reward space . Since Hπ and Hr may be infinite-dimensional , linearity is only a weak restriction– e.g . the map f 7→ f ( x ) is linear in f for any RKHS . To incorporate problem structure , we let r∗ denote the true reward function and assume that r∗ ∈ Cr for some known set Cr ⊆ Hr such that ‖r∗‖Hr = 1 . We further assume that policies π are restricted to lie in some Cπ which is a subset of the unit ball in Hπ ( for instance , Cπ might incorporate physical constraints on implementable policies ) . Thus , given the true reward r∗ , the optimal policy ( for a compact Cπ ) is π∗ ∈ argmaxπ∈Cπ F ( π , r ∗ ) . This proposed framework , which allows for infinite-dimensional policy as well as reward classes , allows us to study how both the policy and reward space affect the difficulty of learning . Query access to reward r∗ . The true reward function r∗ is unknown to the learner but is accessible via queries to an oracle ( e.g . a human expert ) , which provide noisy zeroth-order ( or bandit ) evaluations of the reward r∗ . When queried with a policy π ∈ Cπ , the oracle provides a response Oracle Or∗ : π 7→ F ( π , r∗ ) + where ∼ N ( 0 , τ2 ) , ( 1 ) with τ2 denoting the variance of the response . There are two possible query models : passive queries ( Atkinson , 1996 ; Sebastiani & Wynn , 2000 ) , where the learner selects all queries at the same time , and active queries ( Bubeck et al. , 2011 ; Lattimore & Szepesvári , 2020 ) , where the learner is allowed to select queries sequentially . Our focus in this work will be on the passive query model , but in many cases we will outperform existing active query algorithms . Problem statement . Given passive access to the oracleOr∗ , the objective of the learner is to output a policy π̂ ∈ Cπ that has small excess risk ∆ , defined as ∆ ( π̂ ; r∗ ) : = F ( π∗ , r∗ ) − F ( π̂ , r∗ ) . ( 2 ) We think of queries to the oracle as expensive , and are interested in achieving low excess risk with as few queries as possible . This notion of excess risk is also studied by the term simple regret in pure exploration bandit problems ( Lattimore & Szepesvári , 2020 ) . Representations in ` 2 ( N ) . By Mercer ’ s theorem , we can represent any RKHS as a subset of ` 2 ( N ) . Formally , the policy and the reward spaces are isomorphic to the ellipsoids Hπ : = { ∞∑ j=1 κπ , jφπ , j ∣∣∣ ( κπ , j ) ∞j=1 ∈ ` 2 ( N ) with ∞∑ j=1 κ2π , j µ2π , j < ∞ } and Hr : = { ∞∑ j=1 κr , jφr , j ∣∣∣ ( κr , j ) ∞j=1 ∈ ` 2 ( N ) with ∞∑ j=1 κ2r , j µ2r , j < ∞ } , for appropriately chosen eigenfunctions φπ , j and φr , j , and corresponding eigenvalues µπ , j and µr , j ( Wainwright , 2019 ) . These are defined with respect to a base measure P over the input domain ; see Appendix A for details . With a slight abuse of notation , going forward , we will use π and r to denote the corresponding coefficients ( κπ , j ) and ( κr , j ) in the expansion above . 1 With this , the 1While the eigenfunctions φπ and φr can be different , this representation can still be used by modifying the map M appropriately . This is detailed in Appendix A. Algorithm 1 : Policy Learning via Reward Learning Input : Number of queries n , policy set Cπ , oracle Or∗ Select n policies Q = { π1 , . . . , πn } and receive noisy reward evaluations yi = Or∗ ( πi ) . Estimate r̂ using observed responses { ( π1 , y1 ) , . . . , ( πn , yn ) } using ridge regression ( 4 ) . Obtain plug-in policy π̂plug ∈ argmaxπ∈Cπ F ( π , r̂ ) . Output : Policy π̂plug inner products associated with Hπ and Hr simplify 〈π1 , π2〉Hπ : = ∞∑ j=1 π1 , jπ2 , j µπ , j and 〈r1 , r2〉Hr : = ∞∑ j=1 r1 , jr2 , j µr , j . ( 3 ) Also let Sr : = diag ( µ−1r , j ) and Sπ : = diag ( µ −1 π , j ) be diagonal matrices comprising the inverse of the eigenvalues of Hr and Hπ . With this notation , if we view the map M as a ( infinite-dimensional ) matrix , its Hermitian adjoint2 is equal to M∗ = S−1π M > Sr . In order for the evaluation functional F ( π , r∗ ) to be finite for all π ∈ Hπ , the operator norm ‖S 1 2 r MS − 12 π ‖op must be bounded ( see Appendix A ) . We will see later that the decay of this operator ’ s singular values is closely related to the difficulty of learning in our setting .

Summary: This paper is motivated by learning optimal actions in tasks where both the reward function, and policy (actions) are nonparametric. Previous literature has typically only considered one of these two components as being nonparametric. The main focus is on reliably identifying a policy with low instantaneous regret/risk via as few queries as possible, where all queries are specified in advance (i.e. the passive query setting). The proposed approach selects query locations based on an eigenvalue decomposition of the policy space, sampling repeatedly along a set of top eigenvectors to ensure reliable estimation of the reward via a plug-in ridge-regression-based regression. This estimated reward function is then minimised over the policy class to return a suggested policy/action. The accompanying theory shows the decay of the risk of this suggested action, and shows that in the setting of the Gaussian Process bandit, this can enjoy a better rate than existing adaptive approaches such as GP-UCB.