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: The paper studies a pure exploration bandit problem in a setting where the reward and policies are assumed to  be in RKHS spaces. The main result is a passive strategy based on experimental design and bounds on the simple regret. The results are extended to the cumulative regret setting via a standard explore-then-commit argument. In the latter setting the results are (depending on the kernel) about competitive with the state-of-the-art (more on this later).