Sharp Analysis for KL-Regularized Contextual Bandits and RLHF

## Abstract

Abstract Reverse-Kullback-Leibler (KL) regularization has emerged to be a predominant technique used to enhance policy optimization in reinforcement learning (RL) and reinforcement learning from human feedback (RLHF), which forces the learned policy to stay close to a reference policy. While the effectiveness and necessity of KL-regularization have been empirically demonstrated in various practical scenarios, current theoretical analysis of KL-regularized RLHF still obtains the same 𝒪 ⁢ ( 1 / ϵ 2 ) 𝒪 1 superscript italic-ϵ 2 \mathcal{O}(1/\epsilon^{2}) caligraphic_O ( 1 / italic_ϵ start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT ) sample complexity as problems without KL-regularization. To understand the fundamental distinction between policy learning objectives with KL-regularization and ones without KL-regularization, we are the first to theoretically demonstrate the power of KL-regularization by providing a sharp analysis for KL-regularized contextual bandits and RLHF, revealing an 𝒪 ⁢ ( 1 / ϵ ) 𝒪 1 italic-ϵ \mathcal{O}(1/\epsilon) caligraphic_O ( 1 / italic_ϵ ) sample complexity when ϵ italic-ϵ \epsilon italic_ϵ is sufficiently small. We further explore the role of data coverage in contextual bandits and RLHF. While the coverage assumption is commonly employed in offline RLHF to link the samples from the reference policy to the optimal policy, often at the cost of a multiplicative dependence on the coverage coefficient, its impact on the sample complexity of online RLHF remains unclear. Previous theoretical analyses of online RLHF typically require explicit exploration and additional structural assumptions on the reward function class. In contrast, we show that with sufficient coverage from the reference policy, a simple two-stage mixed sampling strategy can achieve a sample complexity with only an additive dependence on the coverage coefficient. Our results provide a comprehensive understanding of the roles of KL-regularization and data coverage in RLHF, shedding light on the design of more efficient RLHF algorithms.

## 1 Introduction

Recently, *Reinforcement Learning from Human Feedback* (RLHF) has emerged as a central tool for aligning large language models (LLMs) and diffusion models with human values and preferences (Christiano et al., ; Ziegler et al., ; Ouyang et al., ; Bai et al., ; Rafailov et al., ), exhibiting impressive capabilities in applications, such as Chatgpt (Achiam et al., ), Claude (Anthropic, ), Gemini (Team et al., ), and LLaMA-3 (Meta, ).

RLHF methods treat the language model as a policy that takes a prompt $x$ and produces a response $a$ conditioned on $x$, and they optimize the policy by aligning it with human feedback. There are mainly two kinds of feedback: absolute rating and preference comparison. For absolute rating, the collection typically involves human annotators to provide rating scores like 1 to 5 (Wang et al., , ) for responses or hard 0-1 scores for math reasoning tasks since the reasoning tasks often have gold standard answers (Cobbe et al., ; Hendrycks et al., ; Xiong et al., ).

On the other hand, preference comparison is frequently applied in chat tasks when making comparisons is much easier for human labeler (Achiam et al., ). Although preference feedback is believed to be more intuitive for human users and easier to collect, it also poses more challenges for the RLHF algorithms to effectively leverage the feedback signals since the reward signals are not directly observed. In practice, the learning process typically involves (a) constructing a reward model based on the maximum likelihood estimation (MLE) of *Bradley-Terry* (BT) (Bradley and Terry, ) model from the preference feedback; (b) applying RL algorithms like PPO (Schulman et al., ) to train the language model so that it maximizes the reward signals with KL regularization (Ouyang et al., ; Bai et al., ; Touvron et al., ). There is also a series of work on direct preference learning algorithms, including Slic (Zhao et al., ), DPO (Rafailov et al., ), IPO (Azar et al., ), which directly optimize the preference signals to train the language model without constructing the reward model. However, despite their efficiency, these types of algorithms often suffer from overfitting, and DPO suffers from a drop of chosen probability (Song et al., ; Liu et al., ; Amini et al., ; Huang et al., ).

Since human value and preference are so complicated that they are unlikely to be encompassed by the considered preference model classes, the learned model is easy to be hacked and become biased. Practically, the policy may generate disproportionate bold words or emoji to please the learned reward (Zhang et al., ). Hence, the KL-regularization between the learned policy and a reference policy (the pre-trained model after supervised fine-tuning) plays a fundamental role in RLHF to avoid overfitting. There is a line of RLHF work that realizes the significance of KL-regularization and regards the problem as a reverse-KL regularized contextual bandit (Ziegler et al., ; Wu et al., ; Ouyang et al., ; Rafailov et al., ; Xiong et al., ; Ye et al., ). However, they adopt the techniques from bandit framework and neglect the characteristic of reverse-KL-regularization, thus obtaining almost the same sample complexity with problems without KL-regularization. Therefore, the question of *whether there exists a fundamental distinction between policy learning objectives with and without KL-regularization* is still largely under-explored.

Compared to the offline RLHF algorithms (Rafailov et al., ; Azar et al., ; Chen et al., ) that can only use planning to approximate the solution to the relative entropy minimization problem (Ziebart et al., ; Song et al., ), online RLHF has been demonstrated to outperform offline methods empirically and theoretically (Bai et al., ; Meta, ; Xiong et al., ; Tajwar et al., ; Song et al., ; Wu et al., ), because it has further interactions with human or the preference oracle. Most standard theoretical online RL techniques apply optimism to balance exploration and exploitation (Abbasi-Yadkori et al., ; Wang et al., ; Jin et al., ). However, it is inefficient to implement exploration for practical RLHF algorithms.
Meanwhile, an emerging line of offline RLHF literature highlights the coverage of the reference policy $\pi_{0}$. The coverage of $\pi_{0}$ refers to the ability of the model to generate diverse responses for a wide range of prompts. A model with good coverage can generalize well to unseen contexts and actions, which is essential for the learned reward function to also generalize well. In practice, this is evidenced by the fact that the simple best-of-$n$ sampling based on $\pi_{0}$ is competitive with the well-tuned PPO algorithm for general open-ended conversation tasks (Dong et al., ), and the fact that the $\pi_{0}$ can solve a majority of the math problems with multiple responses (Shao et al., ; Nakano et al., ). However, the theoretical understanding of the role of coverage in online RLHF is still largely understudied. Thus, it is natural to ask *if explicit exploration is necessary for online RLHF with good coverage of $\pi_{0}$* and *how the coverage of $\pi_{0}$ affects the sample complexity of online RLHF*.

In this paper, we answer the above questions by

- •
providing a novel fine-grained analysis for KL-regularized contextual bandits and RLHF, which adapts to the optimization landscape of the reverse-KL regularization and reveals a sharper sample complexity than the existing results,
- •
proposing an efficient 2-stage mixed sampling strategy for online RLHF with good coverage of $\pi_{0}$, which achieves sample complexity with only an additive dependence on the coverage coefficient.

### 1.1 Our Contributions

In this work, we make a first attempt to illustrate the statistical benefits of KL-regularization for policy optimization in contextual bandits and reinforcement learning (RL) from preference feedback. Our main contributions are summarized as follows:

- •
In Section 3, we formulate RLHF with absolute-rating feedback as a contextual bandit problem with KL-regularization. First, we provide a lower bound for the KL-regularized contextual bandit problem, which indicates that the sample complexity of the problem is $\Omega(\eta\log N_{\mathcal{R}}(\epsilon)/\epsilon)$ when $\epsilon$ is sufficiently small, where $N_{\mathcal{R}}(\epsilon)$ is the covering number of the reward function class and $\eta$ is the KL-regularization coefficient.
- •
Then we showcase a novel analysis to upper bound the suboptimality gap of the KL-regularized objective in contextual bandits, and propose a simple two-stage mixed sampling strategy for online RLHF which achieves a sample complexity of $\mathcal{O}(\max(\eta^{2}D^{2},\eta/\epsilon)\log N_{\mathcal{R}}(\epsilon/
\delta))$ when the reward scale is a constant, where $D$ is the coverage coefficient of the reference policy $\pi_{0}$ and $\delta$ is the confidence parameter. To the best of our knowledge, this is the first work to provide a sharp sample complexity for KL-regularized contextual bandits.
- •
In Section 4, we extend our analysis to reinforcement learning from preference feedback. We rigorously demonstrate that KL-regularization is essential for more efficient policy learning in RLHF with preference data. We further propose a two-stage mixed sampling strategy for online preference learning setting with good coverage of $\pi_{0}$, which achieves a sample complexity of $\mathcal{O}(\max(\eta^{2}D^{2},\eta/\epsilon)\log N_{\mathcal{R}}(\epsilon/
\delta))$ when the reward scale is a constant.

### 1.2 Previous Understanding of KL Regularization in RL

While we mainly focus on the theoretical understanding of KL-regularization in RLHF, it is also worth mentioning that our analysis for KL-regularized contextual bandits also contributes to the theoretical understanding of the impact of KL-regularization in RL since contextual bandits can be viewed as a simplified version of Markov decision processes (MDPs).

In RL, KL-regularization has been widely used to stabilize the learning process and prevent the policy from deviating too far from the reference policy. Here, we provide a brief overview of the existing understanding of KL-regularization in decision-making problems.
From the perspective of policy optimization, KL-regularization captures entropy regularization as a special case (^1^11We can regard the entropy regularization as a special case of KL-regularization by setting the reference policy as the uniform distribution.), which is also an extensively used technique in RL literature (Sutton, ; Szepesvári, ).
There is a large body of literature that has explored the benefits of entropy regularization or KL-regularization in RL (Schulman et al., ; Fox et al., ; Schulman et al., ; Haarnoja et al., , ; Ahmed et al., ). Most related to our work, Ahmed et al. () provided a comprehensive understanding of the role of entropy regularization in RL, showing that entropy regularization can improve the training efficiency and stability of the policy optimization process by changing the optimization landscape through experiments on continuous control tasks (Brockman, ).

Theoretically, Neu et al. () provided a unified view of entropy regularization as approximate variants of Mirror Descent or Dual Averaging, and left the statistical justification for using entropy regularization in RL as an open question. Geist et al. () provided a framework for analyzing the error propagation in regularized MDPs, which also focused on the proof of the convergence for the policy optimization methods with regularization and lacked a sharp sample complexity analysis.

### 1.3 Other Related Work

#### Analyses for Policy Optimization with Regularization

While it is previously unknown whether regularization can improve the sample complexity of policy optimization without additional assumptions, there are some works that provided a sharp convergence rate in the presence of regularization (Mei et al., ; Shani et al., ; Agarwal et al., , ). However, these works either assumed the access of exact or unbiased policy gradient or required uniform value function approximation error, which are not the standard case in sample-based RL setting.

#### RLHF Algorithms

There are mainly three types of RLHF algorithms: offline, online and hyrbid. The most well-known offline algorithms are Slic (Zhao et al., ), Direct Preference Optimization (DPO) (Rafailov et al., ), Identity-PO (IPO) (Azar et al., ) and (SPIN) (Chen et al., ). They aim to approximate the closed-form solution of the optimization problem on a fixed offline dataset. For the online algorithms, the most representative one is Proximal Policy Optimization (PPO) (Schulman et al., ). PPO has been used in the Chat-GPT (OpenAI, ), Gemini (Team et al., ), and Claude (Bai et al., ). However, the deep RL method PPO is known to be sample inefficient and unstable, making its success hard to reproduce for the open-source community. In response to this, there have been many efforts to propose alternative algorithms to the PPO algorithm. The Reward ranked fine-tuning (RAFT) (also known as rejection sampling finetuning) (Dong et al., ; Touvron et al., ; Gulcehre et al., ; Gui et al., ) is a stable framework requiring minimal hyper-parameter tuning, which iteratively learns from the best-of-n policy (Nakano et al., ). This framework proves to be particularly effective in the reasoning task such as (Gou et al., ; Tong et al., ). However, the RAFT-like algorithms only use the positive signal by imitating the best-of-n sampling. To further improve the efficiency, there is an emerging body of literature that proposes online direct preference optimization by extending DPO or IPO to an online iterative framework (Xiong et al., ; Guo et al., ; Wu et al., ; Calandriello et al., ; Xiong et al., ). Finally, for the third type, the common point of hybrid and online algorithms is that they both require further interaction with the preference oracle and on-policy data collection. The difference is that hybrid algorithms start with a pre-collected dataset (Xiong et al., ; Song et al., ; Touvron et al., ), while the online algorithms learn from scratch.

#### RLHF Theory

The theoretical study of RLHF can date back to the dueling bandits (Yue et al., ; Dudík et al., ; Saha, ; Di et al., ; Li et al., ; Verma et al., ; Wu et al., ) and follow-up work on MDPs (Wang et al., ; Zhu et al., ). However, these works deviate from the practice because they do not realize the significance of KL-regularization and still choose the greedy policy that simply maximizes the reward. After this line of work, Xiong et al. (); Ye et al. (); Song et al. () highlight the KL-regularization theoretically and incorporate the KL term into the learning objective. However, they circumvent the special advantages of KL-regularization and still follow the techniques in bandit analysis, thus obtaining loose bounds. In our paper, we establish a new lower bound and a sharper upper bound for the KL-regularized framework, thus validating the empirical advantage of KL-regularization. There are also some works extending KL-regularized RLHF from bandit problems to the Markov decision process (MDP) problems (Zhong et al., ; Xiong et al., ). We expect that our techniques can also be extended to the MDP setting, which we leave for future work.

## 2 Preliminaries

In this section, we formally state the problem settings of RL from human feedback (RLHF), where we consider two types of feedback: absolute rating and preference.

### 2.1 Contextual Bandits with KL Regularization

The first setting is the absolute-rating feedback, where we can query the ground-truth reward function to measure the quality of the responses by providing absolute reward value. For instance, in the NVIDIA Helpsteer project (Wang et al., , ), human labelers are required to provide absolute score in five attributes: helpfulness, correctness, coherence, complexity, and verbosity. The dataset leads to many high-ranking open-source reward models, including the ArmoRM-Llama3-8B-v0.1 (Wang et al., , ), URM-LLaMa-3.1-8B(^2^22[https://huggingface.co/LxzGordon/URM-LLaMa-3.1-8B](https://huggingface.co/LxzGordon/URM-LLaMa-3.1-8B)), and Llama-3.1-Nemotron-70B-Reward(^3^33[https://huggingface.co/nvidia/Llama-3.1-Nemotron-70B-Reward](https://huggingface.co/nvidia/Llama-3.1-Nemotron-70B-Reward)). We also notice that recently this feedback framework is extended to other task such as video generation (He et al., ).

The absolute-rating feedback is directly modeled as reward functions (Wang et al., ; Xiong et al., ), and is regarded as contextual bandits with KL regularization.
In the contextual bandit setting, at each round $t$, the agent observes a context $x_{t}\in\mathcal{X}$ generated from a distribution $d_{0}$ and chooses an action $a_{t}\in\mathcal{A}$. The agent receives a stochastic reward $r_{t}\in\mathbb{R}$ depending on the context $x_{t}$ and the action $a_{t}$. The goal is to maximize the expected cumulative reward over $T$ rounds.

The learner has access to a family of reward functions $R(\theta,x,a)$ parameterized by $\theta\in\Theta$, such that there exists $\theta_{*}\in\Theta$ satisfying $\mathbb{E}[r_{t}|x_{1:t},a_{1:t}]=R(\theta_{*},x_{t},a_{t})$. WLOG, we assume that the reward feedback $r_{t}$ at all rounds is a non-negative real number bounded by $B$. We consider a KL-regularized objective as follows:

$$ $\displaystyle Q(\pi)=\mathbb{E}_{x\sim d_{0}}\mathbb{E}_{a\sim\pi(\cdot|x)} \bigg{[}R(\theta_{*},x,a)-\eta^{-1}\ln\frac{\pi(a|x)}{\pi_{0}(a|x)}\bigg{]},$ (2.1) $$

where $\pi_{0}$ is a known fixed policy, and $\eta>0$ is a hyperparameter that controls the trade-off between maximizing rewards and staying close to the reference policy $\pi_{0}$.

###### Remark 2.1 .

It is worth noting that entropy or Kullback-Leibler (KL) regularization is also widely used in contextual bandits (Berthet and Perchet, ; Wu et al., ) and deep RL algorithms (Schulman et al., ; Fox et al., ; Schulman et al., ; Haarnoja et al., , ), where KL-divergence regularization is a popular technique for preventing drastic updates to the policy. Algorithms such as Trust Region Policy Optimization (TRPO) (Schulman et al., ) explicitly incorporate KL-regularization to limit the policy updates during optimization, ensuring that the updated policy does not deviate too much from the current policy. This constraint promotes stable and reliable learning, particularly in high-dimensional state-action spaces. Additionally, KL-regularization is central to Proximal Policy Optimization (PPO) (Schulman et al., ), where a penalty term involving KL-divergence ensures updates remain within “trust region”.

###### Remark 2.1 .

### 2.2 Reinforcement Learning from Preference Feedback

The second framework we consider is preference feedback, which is widely applied in projects such as Chat-GPT (OpenAI, ) and Claude (Bai et al., ). Specifically, when receiving a prompt $x\in\mathcal{X}$, and two actions (responses) $a^{1},a^{2}\in\mathcal{A}$ from some LLM policy $\pi(\cdot|x)$, a preference oracle will give feedback $y$ defined as follows:

###### Definition 2.2 (Preference Oracle) .

A Preference Oracle is a function $\mathbb{P}:\mathcal{X}\times\mathcal{A}\times\mathcal{A}\to\{0,1\}$. Given a context $x\in\mathcal{X}$ and two actions $a_{1},a_{2}\in\mathcal{A}$, the oracle can be queried to obtain a preference signal $y\sim Bernoulli(\mathbb{P}(x,a_{1},a_{2}))$,
where $y=1$ indicates that $a_{1}$ is preferred to $a_{2}$ in the context $x$, and $y=0$ indicates the opposite.

To learn the preference, we follow Ouyang et al. (); Zhu et al. (); Rafailov et al. (); Liu et al. (); Xiong et al. () and assume that the preference oracle is measured by the difference of ground-truth reward functions $R(\theta_{*},x,a)$, which is named the Bradley-Terry (BT) model (Bradley and Terry, ).

###### Definition 2.3 (Bradley-Terry Model) .

Given a context $x\in\mathcal{X}$ and two actions $a_{1},a_{2}\in\mathcal{A}$, the probability of $a_{1}$ being preferred to $a_{2}$ is modeled as

$$ $\displaystyle\mathbb{P}(x,a_{1},a_{2})=\frac{\exp(R(\theta_{*},x,a_{1}))}{\exp (R(\theta_{*},x,a_{1}))+\exp(R(\theta_{*},x,a_{2}))}=\sigma(R(\theta_{*},x,a_{ 1})-R(\theta_{*},x,a_{2})),$ (2.2) $$

where $\sigma(\cdot)$ is the sigmoid function.

The RLHF training always follows the fine-tuning process, which yields a reference policy $\pi_{0}$. When performing RLHF on specific tasks, to avoid overfitting, we impose KL-regularization to the learned reward model when optimizing the policy. Hence, our objective function is also (2.1).

###### Definition 2.2 (Preference Oracle) .

###### Definition 2.3 (Bradley-Terry Model) .

### 2.3 Additional Notations and Definitions

In this subsection, we introduce the definitions shared by both settings.

#### Reward function class.

We consider a function class $\mathcal{R}=\{R(\theta,\cdot,\cdot)|\theta\in\Theta\}$ and for the realizability, we assume that the ground truth reward function $R(\theta_{*},x,a)$ is in the function class $\mathcal{R}$. Then, we define the covering number of $\mathcal{R}$ as follows.

###### Definition 2.4 ( $\epsilon$ -cover and covering number) .

Given a function class $\mathcal{F}$, for each $\epsilon>0$, an $\epsilon$-cover of $\mathcal{F}$ with respect to $||\cdot||_{\infty}$, denoted by $\mathcal{C}(\mathcal{F},\epsilon)$, satisfies that for any $f\in\mathcal{F}$, we can find $f^{\prime}\in\mathcal{C}(\mathcal{F},\epsilon)$ such that $||f-f^{\prime}||_{\infty}\leq\epsilon$. The $\epsilon$-covering number, denoted as $N_{\mathcal{F}}(\epsilon)$, is the smallest cardinality of such $\mathcal{C}(\mathcal{F},\epsilon)$.

###### Definition 2.4 ( ϵ italic-ϵ \epsilon italic_ϵ -cover and covering number) .

#### Planning oracle.

Given a reward model, we can learn the policy by optimizing the KL-regularized objective in (2.1). To simplify the analysis, we assume that there exists a planning oracle, which in empirical can be efficiently approximated by rejection sampling (Liu et al., ), Gibbs sampling (Xiong et al., ), and iterative preference learning with a known reward (Dong et al., ).

###### Definition 2.5 (Policy Improvement Oracle) .

For a reward function $R(\theta,\cdot,\cdot)\in\mathcal{R}$ and a reference policy $\pi_{0}$, for any prompt $x\sim d_{0}$, we can compute:

$$ $\displaystyle\pi_{\theta}^{\eta}(\cdot|x):=\mathop{\mathrm{argmax}}_{\pi(\cdot |x)\in\Delta(\mathcal{A})}\mathbb{E}_{a\sim\pi(\cdot|x)}\Big{[}R(\theta,x,a)- \eta^{-1}\log\frac{\pi(a|x)}{\pi_{0}(a|x)}\Big{]}\propto\pi_{0}(\cdot|x)\cdot \exp\big{(}\eta R(\theta,x,\cdot)\big{)}.$ $$

Hence, the comparator policy is the solution to the oracle given the true reward function $R(\theta^{*},\cdot,\cdot)$: $\pi^{*}(\cdot|\cdot)\propto\pi_{0}(\cdot|\cdot)\cdot\exp(\eta R(\theta^{*},
\cdot,\cdot))$. The goal is to minimize the sub-optimality of our learned policy $\widehat{\pi}$ with respect to $\pi^{*}$: $Q(\pi^{*})-Q(\widehat{\pi})$.

###### Definition 2.5 (Policy Improvement Oracle) .

#### Coverage conditions.

It is crucial to assume that our data-collector policy $\pi_{0}$ possesses good coverage, which can ensure that the learned reward function can generalize well to unseen contexts (prompts) and actions (responses), and thus can enable us to approximate the optimal policy. The global coverage is the uniform cover over all the policies in the considered class $\Pi$, which is standard in offline RL (Munos and Szepesvári, ; Song et al., ) and online RL (Xie et al., ; Rosset et al., ). Essentially, Song et al. () demonstrated that global coverage is necessary for offline framework and Direct Preference Optimization (DPO) fails without global coverage. Hence, we introduce two types of global-coverage conditions.

###### Definition 2.6 (Data Coverage) .

Given a reference policy $\pi_{0}$, $D^{2}$ is the minimum positive real number satisfying $\forall\ {(x,a)\in\mathcal{X}\times\mathcal{A},s.t.\ \pi(a|x)>0}$ and $\forall\ b:\mathcal{X}\to[-B,B]$, we have

$$ $\displaystyle\sup_{\theta,\theta^{\prime}\in\Theta}\frac{|R(\theta^{\prime},x, a)-R(\theta,x,a)-b(x)|^{2}}{\mathbb{E}_{x^{\prime}\sim d_{0}}\mathbb{E}_{a^{ \prime}\sim\pi_{0}(\cdot|x^{\prime})}|R(\theta^{\prime},x^{\prime},a^{\prime}) -R(\theta,x^{\prime},a^{\prime})-b(x^{\prime})|^{2}}\leq D^{2}.$ $$

The coverage coefficient $D$ measures how well the in-sample error induced by distribution $d_{0}\times\pi_{0}$ can cover the out-of-sample error, i.e., it depicts the ability of $\pi_{0}$ to cover the action space. This concept is adapted from the F-design for online RL under general function approximation (Agarwal et al., ), and resembles the coverage coefficient for offline RL (Ye et al., , ), and the eluder dimension (Wang et al., ; Ye et al., ; Agarwal et al., ; Zhao et al., ) for online RL.

###### Definition 2.7 (Global-Policy Coverage) .

Given a reference policy $\pi_{0}$, $C_{\mathrm{GL}}$ is the minimum positive real number satisfying that for any $\pi:\mathcal{X}\rightarrow\mathcal{A}$

$$ $\displaystyle\sup_{x\sim d_{0},a\in\mathcal{A}}\frac{\pi(a|x)}{\pi_{0}(a|x)} \leq C_{\mathrm{GL}}.$ $$

The two conditions require the reference policy to cover all possible policy distributions, which is standard and common in RL literature. Additionally, although the two conditions defined above are both global, it is obvious that $D^{2}\leq C_{\mathrm{GL}}$, indicating that it is more general to assume a finite $D$.

Because of the KL-regularization for RLHF, the learned policy will not move too far from the reference policy. Hence, it is natural to relax the global coverage to local coverage inside the KL-ball (Song et al., ).

###### Definition 2.8 (Local KL-ball Coverage) .

Given a reference policy $\pi_{0}$, for a positive constant $\rho_{\mathrm{KL}}<\infty$, and all policy satisfying that $\mathbb{E}_{x\sim d_{0}}[\mathrm{KL}(\pi,\pi_{0})]\leq\rho_{\mathrm{KL}}$, we define

$$ $\displaystyle\sup_{x\sim d_{0},a\in\mathcal{A}}\frac{\pi(a|x)}{\pi_{0}(a|x)}:= C_{\rho_{\mathrm{KL}}}.$ $$

###### Remark 2.9 (Relation between Local and Global Coverage Conditions) .

The local-coverage condition (Definition 2.8) is more precise because compared to the global conditions targeting all possible policies, it only constrains the coverage to a KL-ball. In Song et al. (), because of the specific form of the oracle (Definition 2.5), the considered policy class is
$\Pi=\{\pi(\cdot|\cdot)\propto\pi_{0}(\cdot|\cdot)\exp(\eta R(\theta,\cdot,
\cdot)):R(\theta,\cdot,\cdot)\in\mathcal{R}\}$. Thus, they only need to assume that the condition hold for $\rho=2\eta B$, indicating that $C_{\rho_{\mathrm{KL}}}\leq C_{\mathrm{GL}}$. On the other hand, the data coverage condition (Definition 2.6) is measured on the level of reward functions instead of policies. In this sense, the data coverage condition and local-coverage condition do not encompass each other.

###### Definition 2.6 (Data Coverage) .

###### Definition 2.7 (Global-Policy Coverage) .

###### Definition 2.8 (Local KL-ball Coverage) .

###### Remark 2.9 (Relation between Local and Global Coverage Conditions) .

## 3 KL-Regularized Contextual Bandits

### 3.1 Lower Bound

In this section, we provide a lower bound for the KL-regularized contextual bandit problem.

###### Theorem 3.1 .

For any $\epsilon\in(0,1),\eta>0$, and any algorithm $A$, there exists a KL-regularized contextual bandit problem with $O(1)$ coverage coefficient and reward function class $\mathcal{R}$ such that $A$ requires at least $\Omega\big{(}\min(\frac{\eta\log N_{\mathcal{R}}(\epsilon)}{\epsilon},\frac{
\log N_{\mathcal{R}}(\epsilon)}{\epsilon^{2}})\big{)}$ rounds to achieve a suboptimality gap of $\epsilon$.

###### Remark 3.2 .

The lower bound in Theorem 3.1 indicates that the sample complexity of the KL-regularized contextual bandit problem is $\Omega(\eta\log N_{\mathcal{R}}(\epsilon)/\epsilon)$ when $\epsilon$ is sufficiently small. In our proof, the KL-regularization term shifts the local landscape of the objective function, which prevents us to directly apply the standard bandit analysis, and thus requires a novel analysis to derive the new lower bound. This $\Omega(\eta\log N_{\mathcal{R}}(\epsilon)/\epsilon)$ lower bound suggests that the KL-regularized contextual bandit problem enjoys a lower sample complexity compared to the standard contextual bandit.

###### Theorem 3.1 .

###### Remark 3.2 .

### 3.2 The Proposed Algorithm

Figure: Algorithm 1 Two-Stage mixed-policy sampling

We present the algorithmic framework in Algorithm 1 for the KL-regularized contextual bandit problem, which serves as a theoretical model for online RLHF with absolute-rating feedback. The algorithm consists of two states:

- •
In the first stage, we sample $m$ contexts (prompts) and actions (answers) from the foundation model $\pi_{0}$ and observe the corresponding rewards (absolute ratings). These ratings can be regarded as noisy observations of the underlying reward function $R(\theta_{*},x,a)$. In line 7, we compute an estimate of the reward function $\widehat{\theta}_{0}$ using least squares regression based on the collected data. In line 8, we apply the planning oracle to obtain the policy $\pi_{\widehat{\theta}_{0}}^{\eta}$ which maximizes the following KL-regularized estimated objective in Definition 2.5 with reward function $R(\theta,\cdot,\cdot)=R(\widehat{\theta}_{0},\cdot,\cdot)$.
- •
In the second stage, we utilize the trained policy $\pi_{\widehat{\theta}_{0}}^{\eta}$ to sample $n$ contexts (prompts) and actions (responses). With the intermediate policy $\pi_{\widehat{\theta}_{0}}^{\eta}$, we can collect new data $\{(x_{i},a_{i},r_{i})\}_{i=1}^{n}$ which is more aligned with the data distribution induced by the optimal policy $\pi_{*}$.
In line 14, the algorithm combines data from both stages $\{(x_{i},a_{i},r_{i})\}_{i=1}^{n}$ and $\{(x_{i}^{0},a_{i}^{0},r_{i}^{0})\}_{i=1}^{m}$ to compute a refined least squares estimate $\widehat{\theta}$ of the reward function, minimizing the sum of squared errors across both datasets. By aggregating the two datasets together, there is an overlap between the data to compute $\widehat{\theta}$ and $\widehat{\theta}_{0}$, so that the output policy $\pi_{\widehat{\theta}}^{\eta}$ is well covered by the intermediate policy $\pi_{\widehat{\theta}_{0}}^{\eta}$.

### 3.3 Theoretical Guarantees

#### Revisiting Previous Analysis.

The previous analysis (e.g. Xiong et al., ) is not tight since they basically follow the techniques of bandits and neglect the significance of KL-regularization. For simplicity, We use short-hand notation $R(\theta,x,\pi)=\mathbb{E}_{a\sim\pi(\cdot|x)}R(\theta,x,a)$ and denote $\mathrm{KL}(\pi(\cdot|x)\|\pi^{\prime}(\cdot|x))$ by $\mathrm{KL}(\pi\|\pi^{\prime})$ when there is no confusion. We make the estimation on a dataset $\{(x_{i},a_{i},r_{i}):x_{i}\sim d_{0},s_{i}\sim\pi_{0}\}_{i=1}^{n}$: $\pi_{\widehat{\theta}}^{\eta}=\mathop{\mathrm{argmax}}_{\pi\in\Pi}\mathbb{E}_{
x\sim d_{0}}[R(\widehat{\theta},x,\pi)-\eta^{-1}\mathrm{KL}(\pi\|\pi_{0})]$, and has a small in-sample-error:
$\mathbb{E}_{x\sim d_{0}}\mathbb{E}_{a\sim\pi_{o}(\cdot|x)}\big{[}(R(\widehat{
\theta},x,a)-R(\theta^{*},x,a))^{2}\big{]}=O(1/n).$
The sub-optimality is decomposed as:

$$ $\displaystyle Q(\pi^{*})-Q(\pi_{\widehat{\theta}}^{\eta})=$ $\displaystyle\mathbb{E}_{x\sim d_{0}}\big{[}R(\theta^{*},x,\pi^{*})-R(\widehat {\theta},x,\pi^{*})\big{]}+\mathbb{E}_{x\sim d_{0}}\big{[}R(\widehat{\theta},x ,\pi_{\widehat{\theta}}^{\eta})-R(\theta^{*},x,\pi_{\widehat{\theta}}^{\eta}) \big{]}$ $\displaystyle+\mathbb{E}_{x\sim d_{0}}\big{[}R(\widehat{\theta},x,\pi^{*})- \eta^{-1}\mathrm{KL}(\pi^{*}\|\pi_{0})\big{]}-\mathbb{E}_{x\sim d_{0}}\big{[}R (\widehat{\theta},x,\pi_{\widehat{\theta}}^{\eta})-\eta^{-1}\mathrm{KL}(\pi_{ \widehat{\theta}}^{\eta}\|\pi_{0}))\big{]}$ $\displaystyle\leq$ $\displaystyle\mathbb{E}_{x\sim d_{0}}\big{[}R(\theta^{*},x,\pi^{*})-R(\widehat {\theta},x,\pi^{*})+R(\widehat{\theta},x,\pi_{\widehat{\theta}}^{\eta})-R( \theta^{*},x,\pi_{\widehat{\theta}}^{\eta})\big{]},$ $$

where the inequality holds since $\pi_{\widehat{\theta}}^{\eta}$ is the maximum.

Then, the suboptimality can be further bounded by using the coverage condition (Definition 2.7) and concentration inequalities:

$$ $\displaystyle Q(\pi^{*})-Q(\pi_{\widehat{\theta}}^{\eta})\leq$ $\displaystyle 2C_{\mathrm{GL}}\mathbb{E}_{x\sim d_{0}}\mathbb{E}_{a\sim\pi_{0} (\cdot|x)}\big{[}|R(\theta^{*},x,a)-R(\widehat{\theta},x,a)|\big{]}$ $\displaystyle\leq$ $\displaystyle 2C_{\mathrm{GL}}\sqrt{\mathbb{E}_{a\sim\pi_{0}(\cdot|x)}\big{[}( R(\theta^{*},x,a)-R(\widehat{\theta},x,a))^{2}\big{]}}=O(C_{\mathrm{GL}}/\sqrt {n}).$ $$

Hence, they need $\Theta(C_{\mathrm{GL}}^{2}/\epsilon^{2})$ sample complexity to ensure $O(\epsilon)$ sub-optimility.

#### Power of KL-regularization

The crucial point of the sharper result is utilizing the strong convexity of the objective $Q$ because of the KL-regularization. Specifically, we take the first-order Taylor expansion of sub-optimality with respect to $\{\Delta(x,a)=R(\widehat{\theta},x,a)-R(\theta^{*},x,a):a\in\mathcal{A}\}$

$$ $\displaystyle Q(\pi^{*})-Q(\pi_{\widehat{\theta}}^{\eta})=$ $\displaystyle\eta\mathbb{E}_{x\sim d_{0}}\Big{[}\sum_{a\in\mathcal{A}}\pi_{f}^ {\eta}(a|x)\Delta^{2}(x,a)-\sum_{a_{1},a_{2}\in\mathcal{A}}\pi_{f}^{\eta}(a_{1 }|x)\pi_{f}^{\eta}(a_{2}|x)\Delta(x,a_{1})\Delta(x,a_{2})\Big{]}$ $\displaystyle\leq$ $\displaystyle\eta\mathbb{E}_{x\sim d_{0}}\Big{[}\sum_{a\in\mathcal{A}}\pi_{f}^ {\eta}(a|x)\Delta^{2}(x,a)\Big{]},$ $$

where $f(\cdot,\cdot)=\gamma R(\widehat{\theta},\cdot,\cdot)+(1-\gamma)R(\theta_{*},
\cdot,\cdot)~{}(\gamma\in(0,1))$ the inequality uses the fact that second term on the right-hand side of the equality is $(\sum_{a\in\mathcal{A}}\pi_{f}^{\eta}(a|x)\Delta(x,a))^{2}\geq 0$.

Now, under Algorithm 1, the coverage condition (Definition 2.6) and with concentration inequalities, if the datasize $m=\Theta(\eta^{2}D^{2})$, we can prove that for $\|R(\widehat{\theta},\cdot,\cdot)-R(\theta^{*},\cdot,\cdot)\|_{\infty}\leq\eta
^{-1}$ and $\|R(\widehat{\theta}_{0},\cdot,\cdot)-R(\theta^{*},\cdot,\cdot)\|_{\infty}\leq
\eta^{-1}$, which implies the whole-policy coverage condition for all contexts: $\|\pi_{f}^{\eta}(\cdot|\cdot)/\pi_{\widehat{\theta}_{0}}^{\eta}(\cdot|\cdot)\|
_{\infty}\leq e^{4}.$ Therefore, by setting $n=\Theta(\eta/\epsilon)$, we obtain that $\pi_{\widehat{\theta}}^{\eta}$ is $O(\epsilon)$ optimal. The conclusion is presented in the following theorem.

###### Theorem 3.3 .

Suppose that Assumption 2.6 holds. For any $\delta\in(0,1/5)$, $\epsilon>0$ and constant $c_{m,n}>0$, if we set $m=\Theta(\eta^{2}D^{2}\cdot B^{2}\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta))$ and $n=\Theta(\eta/\epsilon\cdot B^{2}\log(N_{\mathcal{R}}(\epsilon_{c})/\delta))$ and $\epsilon_{c}=\min\{\frac{\epsilon}{2(1+c_{m,n}^{-1})B},\frac{1}{8(1+c_{m,n})B
\eta^{2}D^{2}}\}$, then with probability at least $1-5\delta$ the output policy of Algorithm 1 $\pi_{\widehat{\theta}}^{\eta}$ is $\mathcal{O}(\epsilon)$ optimal.

###### Remark 3.4 .

Theorem 3.3 shows that the sample complexity of Algorithm 1 is $\mathcal{O}(\eta/\epsilon\log N_{\mathcal{R}}(\epsilon/\delta))$ when the reward scale is a constant and $\epsilon$ is sufficiently small. The result indicates that the proposed two-stage mixed sampling strategy can achieve a suboptimality gap of $\epsilon$ with only an additive dependence on the coverage coefficient $D^{2}$.

###### Theorem 3.3 .

###### Remark 3.4 .

### 3.4 Discussion: Result for Local-Coverage

In this subsection, we consider a more general assumption as described in Definition 2.8.

###### Corollary 3.5 .

Let $C_{\rho_{\mathrm{KL}}}$ be in Definition 2.8 where $\rho_{\mathrm{KL}}=2\eta B$. For any $\delta\in(0,1/6)$ and $\epsilon>0$, if we set $n=c_{m,n}m=\Theta(C_{\rho_{\mathrm{KL}}}\eta/\epsilon\cdot B\log(N_{\mathcal{R
}}(\epsilon_{c})/\delta))$ (where constant $c_{m,n}>0$, $\epsilon_{c}=\epsilon/(2(1+c_{m,n}^{-1})B)$) then with probability at least $1-6\delta$ the output policy of Algorithm 2 $\pi_{\widehat{\theta}}^{\eta}$ is $O(\epsilon)$ optimal.

In comparison with the sample complexity $\Theta(\eta^{2}D^{2}+\eta/\epsilon)$ under global-coverage in Theorem 3.3, the order $\Theta(C_{\rho_{\mathrm{KL}}}\eta/\epsilon)$ depends on a weaker coverage $C_{\rho_{\mathrm{KL}}}$, but only has a multiplicative dependence on the coverage coefficient instead of additive dependence. whether the additive dependence can be achieved under the local-coverage condition is left as future work. Moreover, we compare this result with Theorem 4.2 in Song et al. () and suppose that the in-sample-error $\epsilon_{\mathrm{reward}}$ of Song et al. () is $O(1/n)$, their sample complexity is $\Theta(C_{\rho_{\mathrm{KL}}}^{2}/\epsilon^{2})$, which is looser than ours $\Theta(C_{\rho_{\mathrm{KL}}}\eta/\epsilon)$ when $\eta=o(C_{\rho_{\mathrm{KL}}}/\epsilon)$.

###### Corollary 3.5 .

## 4 Reinforcement Learning from Preference Feedback

In this section, we consider the problem of aligning the language model with preference feedback. As discussed in Section 2.2, at each round, we can sample a pair of actions (responses) $a_{1},a_{2}$ and call a preference oracle to get the preference label $y\in\{0,1\}$, where $y=1$ means that the user prefers $a_{1}$ over $a_{2}$ (Definition 2.2). To learn the reward function, we introduce the following assumption for step 1 to ensure the existence of an MLE estimation oracle that can globally maximize the likelihood of the BT model over all possible reward functions.

###### Definition 4.1 (MLE estimation oracle) .

Given a set of context-action pairs $\{(x_{i},a_{i}^{1},a_{i}^{2},y_{i})\}_{i=1}^{n}$ generated from the BT model, can output the parameter $\widehat{\theta}$ such that

$$ $\displaystyle\widehat{\theta}=\mathop{\mathrm{argmax}}_{\theta\in\Theta}\sum_{ i=1}^{n}\underbrace{y_{i}\cdot\log\sigma(R(\theta,x_{i},a_{i}^{1})-R(\theta,x_ {i},a_{i}^{2}))+(1-y_{i})\cdot\log\sigma(R(\theta,x_{i},a_{i}^{2})-R(\theta,x_ {i},a_{i}^{1}))}_{\mathcal{L}(\theta|x_{i},a_{i}^{1},a_{i}^{2},y_{i})}.$ $$

Following the previous analysis for RLHF (Xiong et al., ), we assume the existence of a policy improvement oracle (Definition 2.5, corresponding to step 2) that can compute the optimal policy $\pi_{\widehat{\theta}}^{\eta}$ based on the reward function $\widehat{\theta}$.

###### Remark 4.2 .

We learn the reward function since we can always control the reward (like clipping and normalization) to ensure that the reward function is always bounded by $B$. The bounded assumption does not apply for direct preference learning like DPO (Rafailov et al., ) since there is no intrinsic policy function class encompassing the soundness (Song et al., ), thus increasing the cases of overfitting.

###### Definition 4.1 (MLE estimation oracle) .

###### Remark 4.2 .

### 4.1 Theoretical guarantees

#### Lower Bound

We provide a lower bound for the RLHF problem with preference feedback. The lower bound is derived by constructing a hard instance where the reward function is difficult to estimate from the preference feedback.

###### Theorem 4.3 .

For any $\epsilon\in(0,1),\eta>0$, and any algorithm $A$, there exists a KL-regularized preference learning problem as defined in Section 2.2 with $O(1)$ coverage coefficient and reward function class $\mathcal{R}$ such that $A$ requires at least $\Omega\big{(}\min(\frac{\eta\log N_{\mathcal{R}}(\epsilon)}{\epsilon},\frac{
\log N_{\mathcal{R}}(\epsilon)}{\epsilon^{2}})\big{)}$ samples to achieve a suboptimality gap of $\epsilon$.

We defer Algorithm 2, a 2-stage mixed-policy sampling algorithm for RLHF with preference feedback, to Appendix A for conciseness because of its similarity to Algorithm 1.

###### Theorem 4.3 .

#### Upper bound for global coverage

We provide the theoretical guarantees for Algorithm 2 in the following theorem.

###### Theorem 4.4 .

Suppose that Assumption 2.6 holds. For any $\delta\in(0,1/6)$ and $\epsilon>0$, if we set

$$ $m=\Theta(\eta^{2}D^{2}\cdot e^{B}\log(N_{\mathcal{R}}(\epsilon_{c})/\delta))\ \text{and}\ n=\Theta(\eta/\epsilon\cdot e^{B}\log(N_{\mathcal{R}}(\epsilon_{c} )/\delta))$ $$

where $\epsilon_{c}=\min\{\frac{\epsilon}{2(1+c_{m,n}^{-1})e^{B}},\frac{1}{8(1+c_{m,n
})e^{B}\eta^{2}D^{2}}\}$ then with probability at least $1-6\delta$ the output policy of Algorithm 2 $\pi_{\widehat{\theta}}^{\eta}$ is $O(\epsilon)$ optimal.

###### Remark 4.5 (Comparison with Hybrid Framework) .

We compare our two-stage mixed sampling method with hybrid frameworks. From the algorithmic perspective, a hybrid algorithm first learns from an offline dataset and then requires sufficient online iterations to ensure the performance (Xiong et al., ). For example, for a finite action space with $A$ actions, the number of online iterations should be $\Theta(A)$. In contrast, our method only requires two iterations of sampling from mixed policy and interacting with the environment. Moreover, the results of hybrid literature depend on both the coverage coefficient and the structure complexity of the function class (like the dimension for a linear function class or eluder dimension (Russo and Van Roy, )). Our result only needs the coverage condition of the reference policy. More importantly, we obtain a sharper bound on the sample complexity and derive the additive dependence on the coverage coefficient.

###### Remark 4.6 .

Although the coefficient $e^{B}$ appearing in sample size $m,n$ can be exponentially large, this term is caused by the non-linearity of the link function for the preference model, and is common in RLHF literature (Zhu et al., ; Xiong et al., ; Ye et al., ; Song et al., ).

Theorem 4.4 shows that the sample complexity of Algorithm 2 is $\mathcal{O}(\eta/\epsilon\log N_{\mathcal{R}}(\epsilon/\delta))$ when the reward scale is a constant and $\epsilon$ is sufficiently small. The result indicates that the proposed two-stage mixed sampling strategy can achieve a suboptimality gap of $\epsilon$ with only an additive dependence on the coverage coefficient $D^{2}$.

Besides, the algorithm only requires sampling from the reference policy $\pi_{0}$ and the intermediate policy $\pi_{\widehat{\theta}_{0}}^{\eta}$, which is more aligned with the practical scenarios where the preference feedback is collected from the human users and it is expensive to collect the data while the language model is being updated. Our result implies that we may achieve a near-optimal sample complexity by simply leveraging an intermediate policy to collect more data, and the training process of the reward model and the policy (language model) can be highly decoupled.

###### Theorem 4.4 .

###### Remark 4.5 (Comparison with Hybrid Framework) .

###### Remark 4.6 .

#### Upper bound for local coverage

We also show the result under the local-coverage assumption (Definition 2.8) as follows.

###### Corollary 4.7 .

Let $C_{\rho_{\mathrm{KL}}}$ be in Definition 2.8 where $\rho=2\eta B$. For any $\delta\in(0,1/6)$ and $\epsilon>0$, if we set $n=c_{m,n}m=\Theta(C_{\rho_{\mathrm{KL}}}\eta/\epsilon\cdot e^{B}\log(N_{
\mathcal{R}}(\epsilon_{c})/\delta))$ (where constant $c_{m,n}>0$, $\epsilon_{c}=\frac{\epsilon}{2(1+c_{m,n}^{-1})e^{B}}$) then with probability at least $1-6\delta$ the output policy of Algorithm 2 $\pi_{\widehat{\theta}}^{\eta}$ is $O(\epsilon)$ optimal.

###### Corollary 4.7 .

## 5 Experimental Results

Figure: (a) Benefit of mixed-policy sampling
Refer to caption: x1.png

In this section, we conduct experiments with synthetic data to investigate the benefit of mixed-policy sampling and the effect of KL-regularization coefficient on the sample complexity of the problem. We plot the experimental results for RL from preference feedback in Figure 1 and defer the results for KL-regularized contextual bandits in Appendix B. All the trials are repeated for 10 times and plotted with the standard variation.

We consider the case where context distribution $d_{0}$ is a projected Gaussian distribution over the unit sphere and $\mathcal{A}$ is a discrete set with $|\mathcal{A}|=5$.
We construct the reward functions as $R(\phi,x,a)=\langle x,\phi(a)\rangle$, parameterized by a mapping $\phi$ from $\mathcal{A}$ to $\mathbb{R}^{10}$. To generate $\phi_{*}$, we sample $\phi_{*}(a)$ independently for each $a\in\mathcal{A}$ according to another projected gaussian distribution over the sphere with radius equal to 5. In Figure 1(a), we compare the suboptimality gaps of mixed-policy sampling with $m=n$ to those of offline sampling using $\pi_{0}$ under the same sample sizes. The result indicates that the usage of mixed-policy sampling reduces the suboptimality gap by a large margin. In Figure 1(b), it is shown that the sample complexity is remarkably affected by the KL-regularization term, corroborating our sharp analysis for regularized RLHF.

## 6 Conclusion

We have presented a comprehensive theoretical analysis of the role of reverse-KL regularization in decision-making models including contextual bandits and reinforcement learning from preference feedback, highlighting its significance in terms of sample complexity. Our results provide new insights into the power of regularization extending beyond its traditional role of mitigating errors from the current critic (or reward) model.

Additionally, we examined the role of data coverage in both contextual bandits and RLHF. Our analysis shows that with sufficient coverage from the reference policy, a mixed sampling strategy can achieve a sample complexity that exhibits only an additive dependence on the coverage coefficient without the need for explicit exploration or additional structural assumptions.

## Acknowledgements

We are grateful to Wei Xiong for the helpful discussions and insightful feedback in revising this work.

## Appendix A Algorithm for Preference Feedback

Figure: Algorithm 2 2-Stage mixed-policy sampling for preference feedback

In the first stage, we sample $m$ context-action pairs $\{(\widetilde{x}_{i},\widetilde{a}_{i}^{1},\widetilde{a}_{i}^{2},\widetilde{y}
_{i})\}_{i=1}^{m}$ from the BT model and call the preference oracle to get the preference labels. We then compute the MLE estimator of the reward function $\widehat{\theta}_{0}$ based on the preference feedback in line 7. Afterwards, we apply the planning oracle to compute the optimal policy $\pi_{\widehat{\theta}_{0}}^{\eta}$ based on the reward function $\widehat{\theta}_{0}$ in line 8. Line 7 and line 8 correspond to the practical implementation of RLHF(Ouyang et al., ; Bai et al., ; Touvron et al., ) given a dataset of preference feedback.

In the second stage, we sample $n$ context-action pairs $\{(x_{i},a_{i}^{1},a_{i}^{2},y_{i})\}_{i=1}^{n}$ using the intermediate policy $\pi_{\widehat{\theta}_{0}}^{\eta}$ and call the preference oracle to get the preference labels. We then compute the MLE estimator of the reward function $\widehat{\theta}$ based on the preference feedback from both stages. Finally, we apply the planning oracle to compute the optimal policy $\pi_{\widehat{\theta}}^{\eta}$ based on the reward function $\widehat{\theta}$.

## Appendix B Additional Experimental Results

Figure: (a) Benefit of mixed-policy sampling
Refer to caption: x3.png

We show our experimental results for KL-regularized contextual bandits in the following figure 2, which also corroborate our theory in Section 3.

## Appendix C Proofs from Section 3

### C.1 Proof of Theorem 3.1

###### Proof of Theorem 3.1 .

Consider a simple case when $|\mathcal{X}|=M$ and $|\mathcal{A}|=2$. We suppose that the context $x$ is drawn uniformly from $\mathcal{X}$ at the beginning of each round.
Let $\Theta$ be the set consisting of mappings from $\mathcal{X}$ to $\mathcal{A}=\{0,1\}$. For each $\theta\in\Theta$, we have $R(\theta,x,a)=\begin{cases}1/2+c&\text{if }a=\theta(x),\\
1/2&\text{if }a\neq\theta(x),\end{cases}$ where $c>0$ is a constant, and $\theta(x)$ is the optimal action under context $x$ when the model is $\theta$.

For any $(\theta,x,a)\in\Theta\times\mathcal{X}\times\mathcal{A}$, we assume the reward feedback $r\sim Bernoulli(R(\theta,x,a))$ when the model is $\theta$ and $a$ is chosen under context $x$.

We pick a pair of model $\theta_{1},\theta_{2}$ in $\Theta$, such that $\theta_{1}(x)=\begin{cases}\theta_{2}(x)&\text{if }x\neq x_{0},\\
1-\theta_{2}(x)&\text{if }x=x_{0}.\end{cases}$

We denote by $\mathbb{P}_{\theta}$, $\mathbb{E}_{\theta}$ the probability measure and expectation under the model $\theta$. Let $N(x)$ be the number of times the context $x$ is observed in the first $T$ rounds for an $x\in\mathcal{X}$.

For two Bernoulli random variables $X$ and $Y$ with parameters $1/2-c$ and $1/2+c$, we have

$$ $\displaystyle\mathrm{KL}(X\|Y)$ $\displaystyle=(1/2-c)\log\frac{1/2-c}{1/2+c}+(1/2+c)\log\frac{1/2+c}{1/2-c}$ $\displaystyle\leq 2\cdot\Bigl{(}\frac{1}{2}\log\frac{1}{1+2c}+\frac{1}{2}\log \frac{1}{1-2c}\Bigr{)}=-\log(1-4c^{2}).$ $$

Applying Pinsker’s inequality (Lemma E.3), we have for all event $A$ measurable with respect to the filtration generated by the observations,

$$ $\displaystyle|\mathbb{P}_{\theta_{1}}(A)-\mathbb{P}_{\theta_{2}}(A)|\leq\sqrt{ -\frac{1}{2}\log(1-4c^{2})\mathbb{E}_{\theta_{1}}[N(x_{0})]}\leq\sqrt{2c^{2} \mathbb{E}_{\theta_{1}}[N(x_{0})]}=\sqrt{2c^{2}T/M},$ $$

where the first inequality follows from the chain rule of KL divergence, and the fact that $\mathbb{E}_{\theta_{1}}[N(x_{0})]=T/M$.

Set $A$ to be the event that $\pi_{out}(\theta_{1}(x_{0})|x_{0})>1/2$. Then we have

$$ $\displaystyle\mathbb{P}_{\theta_{1}}(\pi_{out}(\theta_{1}(x_{0})|x_{0})\leq 1/ 2)+\mathbb{P}_{\theta_{2}}(\pi_{out}(\theta_{2}(x_{0})|x_{0})\leq 1/2)\geq 1-| \mathbb{P}_{\theta_{1}}(A)-\mathbb{P}_{\theta_{2}}(A)|\geq 1-\sqrt{2c^{2}T/M}.$ $$

If the model $\theta$ is uniformly drawn from $\Theta$, then we have

$$ $\displaystyle\mathbb{E}_{\theta\sim Unif(\Theta)}\mathbb{P}_{\theta}(\pi_{out} (\theta(x_{0}))\leq 1/2)\geq\frac{1}{2}-\sqrt{c^{2}T/2M}$ $$

for an arbitrary $x_{0}$.

Then we consider the following suboptimality gap:

$$ $\displaystyle\mathbb{E}_{\pi_{\theta_{*}}^{\eta}}\bigg{[}R(\theta_{*},x,a)- \frac{1}{\eta}\ln\frac{\pi_{\theta_{*}}^{\eta}(a|x)}{\pi_{0}(a|x)}\bigg{]}- \mathbb{E}_{\pi_{out}}\bigg{[}R(\theta_{*},x,a)-\frac{1}{\eta}\ln\frac{\pi_{ out}(a|x)}{\pi_{0}(a|x)}\bigg{]}$ $\displaystyle=\frac{1}{\eta}\mathbb{E}_{\pi_{\theta_{*}}^{\eta}}\bigg{[}\ln \frac{\pi_{0}(a|x)\cdot\exp\bigl{(}\eta R(\theta_{*},x,a)\bigr{)}}{\pi_{\theta _{*}}^{\eta}(a|x)}\bigg{]}-\frac{1}{\eta}\mathbb{E}_{\pi_{out}}\bigg{[}\ln \frac{\pi_{0}(a|x)\cdot\exp\bigl{(}\eta R(\theta_{*},x,a)\bigr{)}}{\pi_{out}(a |x)}\bigg{]}$ $\displaystyle=\frac{1}{\eta}\mathbb{E}_{\pi_{out}}\Bigl{[}\ln\frac{\pi_{out}(a |x)}{\pi^{*}(a|x)}\Bigr{]},$ $$

where the last equality follows from the fact that $\pi_{\theta_{*}}^{\eta}\propto\pi_{0}(a|x)\cdot\exp(\eta R(\theta_{*},x,a))$. Note that we handle the difference in reward and the KL-divergence term together, which is distinct from the standard analysis of the lower bound for contextual bandits.

To bound the suboptimality gap, we further have

$$ $\displaystyle\mathbb{E}_{\theta\sim Unif(\Theta)}\mathbb{E}_{\pi_{out}}\Bigl{[ }\ln\frac{\pi_{out}(a|x)}{\pi^{*}(a|x)}\Bigr{]}$ $\displaystyle=\mathbb{E}_{\theta\sim Unif(\Theta)}\frac{1}{M}\sum_{x\in \mathcal{X}}\mathbb{E}_{a\sim\pi_{out}(\cdot|x)}\Bigl{[}\ln\frac{\pi_{out}(a|x )}{\pi^{*}(a|x)}\Bigr{]}$ $\displaystyle\geq\mathbb{E}_{\theta\sim Unif(\Theta)}\frac{1}{M}\sum_{x\in \mathcal{X}}\mathbb{P}_{\theta}(\pi_{out}(\theta(x))\leq 1/2)\cdot\biggl{[} \frac{1}{2}\ln\frac{1+\exp(-\eta c)}{2}+\frac{1}{2}\ln\frac{1+\exp(\eta c)}{2} \biggr{]}$ $\displaystyle\geq\Bigl{(}\frac{1}{2}-\sqrt{c^{2}T/2M}\Bigr{)}\Bigl{[}\frac{1}{ 2}\ln\frac{1+\exp(-\eta c)}{2}+\frac{1}{2}\ln\frac{1+\exp(\eta c)}{2}\Bigr{]}$ (C.1) $$

Note that

$$ $\displaystyle\frac{\text{d}}{\text{d}u}\Bigl{[}\frac{1}{2}\ln\frac{1+e^{-u}}{2 }+\frac{1}{2}\ln\frac{1+e^{u}}{2}\Bigr{]}\Big{|}_{u=0}$ $\displaystyle=\frac{1}{2}\Bigl{[}\frac{1}{1+\exp(-u)}-\frac{1}{1+\exp(u)}\Bigr {]}\Big{|}_{u=0}=0,$ $\displaystyle\frac{\text{d}^{2}}{\text{d}u^{2}}\Bigl{[}\frac{1}{2}\ln\frac{1+e ^{-u}}{2}+\frac{1}{2}\ln\frac{1+e^{u}}{2}\Bigr{]}$ $\displaystyle=\frac{\exp(u)}{[1+\exp(u)]^{2}}.$ $$

Thus, applying Taylor’s expansion on the right-hand side of (C.1), we have

$$ $\displaystyle\mathbb{E}_{\theta\sim Unif(\Theta)}\mathbb{E}_{\pi_{out}}\Bigl{[ }\ln\frac{\pi_{out}(a|x)}{\pi^{*}(a|x)}\Bigr{]}\geq\frac{1}{2}\cdot\Bigl{(} \frac{1}{2}-\sqrt{c^{2}T/2M}\Bigr{)}\eta^{2}c^{2}\cdot\frac{1}{3+\exp(\eta c)},$ $$

which follows from the Taylor’s expansion where $f(x)=f(0)+f^{\prime}(0)x+\frac{1}{2}f^{\prime\prime}(z)x^{2}$ where $z\in[0,\eta c]$, and the fact that $\frac{1}{2}\frac{e^{z}}{(1+e^{z})^{2}}=\frac{1}{2}\frac{1}{e^{-z}+e^{z}+2}\leq
\frac{1}{2}\frac{1}{3+e^{\eta c}}$.

When $\epsilon<1/64\eta$, we can set $c=8\sqrt{\epsilon/\eta}$. To achieve a suboptimality gap of $\epsilon$, we need to satisfy:

$$ $\displaystyle\frac{1}{2}\cdot\Bigl{(}\frac{1}{2}-\sqrt{c^{2}T/2M}\Bigr{)}\eta^ {2}c^{2}\cdot\frac{1}{3+\exp(\eta c)}\leq\eta\epsilon,$ $$

indicating that $T\geq\frac{\eta M}{512\epsilon}=\Omega(\frac{\eta M}{\epsilon})$.

When $\epsilon\geq 1/64\eta$, or equivalently, $\eta\geq 1/64\epsilon$, we employ a different lower bound for (C.1) as follows:

$$ $\displaystyle\frac{1}{2}\ln\frac{1+\exp(-\eta c)}{2}+\frac{1}{2}\ln\frac{1+ \exp(\eta c)}{2}$ $\displaystyle=\frac{1}{2}\ln\frac{2+\exp(\eta c)+\exp(-\eta c)}{4}$ $\displaystyle\geq\frac{1}{2}\cdot\frac{1}{2}\Bigl{(}\ln\frac{\exp(\eta c)+\exp (-\eta c)}{2}\Bigr{)}$ $\displaystyle\geq\frac{1}{4}(\eta c-\ln 2),$ (C.2) $$

where the first inequality follows from Jensen’s inequality.

Substituting (C.2) into (C.1), we have

$$ $\displaystyle\epsilon\geq\frac{1}{\eta}\mathbb{E}_{\theta\sim Unif(\Theta)} \mathbb{E}_{\pi_{out}}\Bigl{[}\ln\frac{\pi_{out}(a|x)}{\pi^{*}(a|x)}\Bigr{]} \geq\frac{1}{4}\cdot\Bigl{(}\frac{1}{2}-\sqrt{c^{2}T/2M}\Bigr{)}(\eta c-\ln 2) \cdot\frac{1}{\eta}.$ $$

Set $c=64\epsilon$. Then we have $T=\Omega(M/\epsilon^{2})$.

∎

###### Proof of Theorem 3.1 .

### C.2 Proof of Theorem 3.3

We start with the following lemma, which provides an on-policy generalization bound for the reward function. Due to the on-policy nature of the algorithm (i.e., the usage of intermediate $\pi_{\widehat{\theta}_{0}}^{\eta}$), we can leverage the covering number of the reward function class $\mathcal{R}$ to derive the generalization error. Since we are using a fixed policy $\pi_{\widehat{\theta}_{0}}^{\eta}$ to sample in the second stage, we can derive the generalization error of the reward function as follows:

###### Lemma C.1 (Generalization error of reward function) .

For an arbitrary policy $\pi$, a set of context-action pairs $\{(x_{i},a_{i})\}_{i=1}^{n}$ generated i.i.d. from $\pi$, and a distance threshold $0<\epsilon_{c}\leq B$, we have with probability at least $1-\delta$, for any pair of parameters $\theta_{1}$ and $\theta_{2}$,

$$ $\displaystyle\mathbb{E}_{\pi}|R(\theta_{1},x,a)-R(\theta_{2},x,a)|^{2}$ $\displaystyle\leq\frac{2}{n}\sum_{i=1}^{n}|R(\theta_{1},x_{i},a_{i})-R(\theta_ {2},x_{i},a_{i})|^{2}+\frac{32B^{2}}{3n}\log(2N_{\mathcal{R}}(\epsilon_{c})/ \delta)+10\epsilon_{c}B.$ $$

###### Proof.

We first consider an $\epsilon_{c}$-net $\mathcal{R}^{c}$ of the reward function class $\mathcal{R}$ where $\mathcal{R}^{c}=\{R(\theta,\cdot,\cdot)|\theta\in\Theta^{c}\}$ with size $N_{\mathcal{R}}(\epsilon_{c})$. For any $R(\theta,\cdot,\cdot)\in\mathcal{R}$, there exists $\theta^{c}$ such that $\|R(\theta,\cdot,\cdot)-R(\theta^{c},\cdot,\cdot)\|_{\infty}\leq\epsilon_{c}$.

By Lemma E.1, for each pair of $\theta_{1}^{c},\theta_{2}^{c}\in\Theta^{c}$ (corresponding to $\theta_{1},\theta_{2}$), we have with probability at least $1-\delta$,

$$ $\displaystyle\bigg{|}\frac{1}{n}\sum_{i=1}^{n}(R(\theta_{1}^{c},x_{i},a_{i})-R (\theta_{2}^{c},x_{i},a_{i}))^{2}-\mathbb{E}_{\pi}|R(\theta_{1}^{c},x,a)-R( \theta_{2}^{c},x,a)|^{2}\bigg{|}$ $\displaystyle\leq\sqrt{\frac{2\text{Var}_{\pi}|R(\theta_{1}^{c},x,a)-R(\theta_ {2}^{c},x,a)|^{2}}{n}\log(2/\delta)}+\frac{2}{3n}B^{2}\log(2/\delta)$ $\displaystyle\leq\sqrt{\frac{2B^{2}\mathbb{E}_{\pi}|R(\theta_{1}^{c},x,a)-R( \theta_{2}^{c},x,a)|^{2}}{n}\log(2/\delta)}+\frac{2}{3n}B^{2}\log(2/\delta)$ $$

where the second inequality follows from the fact that $R(\theta_{1}^{c},x,a),R(\theta_{2}^{c},x,a)\leq B$.

Using union bound over all $\theta_{1}^{c},\theta_{2}^{c}\in\Theta^{c}$, we have with probability at least $1-\delta$, for all $\theta_{1}^{c},\theta_{2}^{c}\in\Theta^{c}$,

$$ $\displaystyle\mathbb{E}_{\pi}|R(\theta_{1}^{c},x,a)-R(\theta_{2}^{c},x,a)|^{2} -\frac{1}{n}\sum_{i=1}^{n}(R(\theta_{1}^{c},x_{i},a_{i})-R(\theta_{2}^{c},x_{i },a_{i}))^{2}$ $\displaystyle\leq\sqrt{\frac{4B^{2}\mathbb{E}_{\pi}|R(\theta_{1}^{c},x,a)-R( \theta_{2}^{c},x,a)|^{2}}{n}\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta)}+\frac {4B^{2}}{3n}\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta),$ $$

from which we further obtain the following inequality by Lemma E.2,

$$ $\displaystyle\mathbb{E}_{\pi}|R(\theta_{1}^{c},x,a)-R(\theta_{2}^{c},x,a)|^{2} \leq\frac{2}{n}\sum_{i=1}^{n}(R(\theta_{1}^{c},x_{i},a_{i})-R(\theta_{2}^{c},x _{i},a_{i}))^{2}+\frac{32B^{2}}{3n}\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta).$ (C.3) $$

Then we can complete the proof by the definition of $\epsilon$-net.
∎

Next, we provide the following lemma, which gives an upper bound on the cumulative square error of the learned reward function.

###### Lemma C.2 (Confidence bound for reward function) .

For an arbitrary policy $\pi$, and a set of data $\{(x_{i},a_{i},r_{i})\}_{i=1}^{n}$ generated i.i.d. from $\pi$, suppose that $\widehat{\theta}$ is the least squares estimator of $\theta_{*}$, i.e., $\widehat{\theta}=\arg\min_{\theta\in\Theta}\sum_{i=1}^{n}(R(\theta,x_{i},a_{i}
)-r_{i})^{2}$. Then for any threshold $\epsilon_{c}>0$, with probability at least $1-\delta$, it holds that

$$ $\displaystyle\sum_{i=1}^{n}(R(\widehat{\theta},x_{i},a_{i})-R(\theta_{*},x_{i} ,a_{i}))^{2}\leq 16B^{2}\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta)+4\epsilon_ {c}nB.$ $$

###### Proof.

We have the following inequality for $\sum_{i=1}^{n}(R(\widehat{\theta},x_{i},a_{i})-R(\theta_{*},x_{i},a_{i}))^{2}$,

$$ $\displaystyle\sum_{i=1}^{n}(R(\widehat{\theta},x_{i},a_{i})-R(\theta_{*},x_{i} ,a_{i}))^{2}$ $\displaystyle=\sum_{i=1}^{n}(R(\widehat{\theta},x_{i},a_{i})-r_{i})^{2}-\sum_{ i=1}^{n}(R(\theta_{*},x_{i},a_{i})-r_{i})^{2}$ $\displaystyle\quad+2\sum_{i=1}^{n}(R(\widehat{\theta},x_{i},a_{i})-R(\theta_{* },x_{i},a_{i})(r_{i}-R(\theta_{*},x_{i},a_{i}))$ $\displaystyle\leq 2\sum_{i=1}^{n}(R(\widehat{\theta},x_{i},a_{i})-R(\theta_{*} ,x_{i},a_{i}))(r_{i}-R(\theta_{*},x_{i},a_{i})),$ $$

where the last inequality follows from the fact that $\sum_{i=1}^{n}(R(\widehat{\theta},x_{i},a_{i})-r_{i})^{2}\leq\sum_{i=1}^{n}(R(
\theta_{*},x_{i},a_{i})-r_{i})^{2}$.

We then consider an $\epsilon_{c}$-net $\mathcal{R}^{c}$ of the reward function class $\mathcal{R}$ where $\mathcal{R}^{c}=\{R(\theta,\cdot,\cdot)|\theta\in\Theta^{c}\}$ with size $N_{\mathcal{R}}(\epsilon_{c})$. For any $R(\theta,\cdot,\cdot)\in\mathcal{R}$, there exists $\theta^{c}$ such that $\|R(\theta,x,a)-R(\theta^{c},x,a)\|_{\infty}\leq\epsilon_{c}$.
From Azuma-Hoeffding inequality, with probability at least $1-\delta$, it holds for all $\theta\in\Theta^{c}$ that

$$ $\displaystyle\sum_{i=1}^{n}(R(\theta,x_{i},a_{i})-R(\theta_{*},x_{i},a_{i}))(r _{i}-R(\theta_{*},x_{i},a_{i}))$ $\displaystyle\leq\sqrt{2B^{2}\sum_{i=1}^{n}(R(\theta,x_{i},a_{i})-R(\theta_{*} ,x_{i},a_{i}))^{2}\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta)}.$ $$

Then we further have with probability at least $1-\delta$, there exists $\|R(\theta^{c},\cdot,\cdot)-R(\widehat{\theta},\cdot,\cdot)\|\leq\epsilon_{c}$ such that

$$ $\displaystyle\sum_{i=1}^{n}(R(\widehat{\theta},x_{i},a_{i})-R(\theta_{*},x_{i} ,a_{i}))(r_{i}-R(\theta_{*},x_{i},a_{i}))$ $\displaystyle\leq\sqrt{2B^{2}\sum_{i=1}^{n}(R(\theta,x_{i},a_{i})-R(\theta_{*} ,x_{i},a_{i}))^{2}\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta)}+2\epsilon_{c}nB,$ $$

which implies that

$$ $\displaystyle\sum_{i=1}^{n}(R(\widehat{\theta},x_{i},a_{i})-R(\theta_{*},x_{i} ,a_{i}))^{2}\leq 16B^{2}\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta)+4\epsilon_ {c}nB$ (C.4) $$

from Lemma E.2.
∎

With the above lemmas, we are now ready to prove the following lemma that bounds the estimation error of the reward function $R(\widehat{\theta},\cdot,\cdot)$ under the sampled policy $\pi_{\widehat{\theta}_{0}}^{\eta}$.

###### Lemma C.3 .

Let $\widehat{\theta}_{0}$ be the least squares estimator of the reward function based on the data $\{(x_{i}^{0},a_{i}^{0},r_{i}^{0})\}_{i=1}^{m}$ generated from $\pi_{0}$ as defined in Algorithm 1. Then for any threshold $\epsilon_{c}>0$,
with probability at least $1-2\delta$, we have

$$ $\displaystyle\mathbb{E}_{\pi_{\widehat{\theta}_{0}}^{\eta}}|R(\widehat{\theta} ,x,a)-R(\theta_{*},x,a)|^{2}\leq\frac{43B^{2}}{n}\log(2N_{\mathcal{R}}( \epsilon_{c})/\delta)+10\epsilon_{c}(1+m/n)B.$ $$

###### Proof.

By Lemma C.1, we have with probability at least $1-\delta$, the following upper bound holds for $\mathbb{E}_{\pi_{\widehat{\theta}_{0}}^{\eta}}|R(\theta_{1},x,a)-R(\theta_{2},
x,a)|^{2}$,

$$ $\displaystyle\mathbb{E}_{\pi_{\widehat{\theta}_{0}}^{\eta}}|R(\theta_{1},x,a)- R(\theta_{2},x,a)|^{2}$ $\displaystyle\leq\frac{2}{n}\sum_{i=1}^{n}|R(\theta_{1},x_{i},a_{i})-R(\theta_ {2},x_{i},a_{i})|^{2}+\frac{32B^{2}}{3n}\log(2N_{\mathcal{R}}(\epsilon_{c})/ \delta)+10\epsilon_{c}B.$ (C.5) $$

By Lemma C.2, with probability at least $1-\delta$

$$ $\displaystyle\sum_{i=1}^{n}|R(\theta_{*},x_{i},a_{i})-R(\widehat{\theta},x_{i} ,a_{i})|^{2}\leq 16B^{2}\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta)+4\epsilon_ {c}(n+m)B.$ (C.6) $$

Then we can complete the proof using a union bound and substituting (C.6) into (C.5).
∎

###### Lemma C.4 .

If $m\geq 128\eta^{2}D^{2}B^{2}\cdot\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta))$, and there exists a positive constant $c_{m,n}>0$ such that $n=c_{m,n}m$ in Algorithm 1 and Assumption 2.6 holds, then by taking $\epsilon_{c}\leq\min\{B,(8(1+c_{m,n})B\eta^{2}D^{2})^{-1}\}$, with probability at least $1-3\delta$, we have

$$ $\displaystyle\eta|R(\widehat{\theta}_{0},x,a)-R(\theta_{*},x,a)|\leq 1,\quad \eta|R(\widehat{\theta},x,a)-R(\theta_{*},x,a)|\leq 1$ $$

for any pair $(x,a)\in\mathcal{X}\times\mathcal{A}$ such that $\pi_{0}(a|x)>0$.

###### Proof.

By Lemma C.1, with probability at least $1-\delta$, for all $\theta_{1},\theta_{2}\in\Theta$, we have

$$ $\displaystyle\mathbb{E}_{\pi_{0}}|R(\theta_{1},x,a)-R(\theta_{2},x,a)|^{2}\leq \frac{2}{m}\sum_{i=1}^{m}|R(\theta_{1},x_{i}^{0},a_{i}^{0})-R(\theta_{2},x_{i} ^{0},a_{i}^{0})|^{2}+\frac{32B^{2}}{3m}\log(2N_{\mathcal{R}}(\epsilon_{c})/ \delta).$ $$

By Lemma C.2, with probability at least $1-\delta$, we have

$$ $\displaystyle\sum_{i=1}^{m}|R(\widehat{\theta}_{0},x_{i}^{0},a_{i}^{0})-R( \theta_{*},x_{i}^{0},a_{i}^{0})|^{2}\leq 16B^{2}\log(2N_{\mathcal{R}}(\epsilon _{c})/\delta)+4\epsilon_{c}m.$ $$

Also, with probability at least $1-\delta$, we have

$$ $\displaystyle\sum_{i=1}^{m}|R(\theta_{*},x_{i}^{0},a_{i}^{0})-R(\widehat{ \theta},x_{i}^{0},a_{i}^{0})|^{2}\leq 16B^{2}\log(2N_{\mathcal{R}}(\epsilon_{c })/\delta)+4\epsilon_{c}(m+n)B.$ $$

Similar to the proof of Lemma C.3, we have if $m\geq 128\eta^{2}D^{2}B^{2}\cdot\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta),~{
}n=c_{m,n}n$, then with probability at least $1-3\delta$,

$$ $\displaystyle\ \mathbb{E}_{\pi_{0}}|R(\theta_{*},x,a)-R(\widehat{\theta}_{0},x ,a)|^{2}\leq 1/\eta^{2}D^{2},\quad\mathbb{E}_{\pi_{0}}|R(\theta_{*},x,a)-R( \widehat{\theta},x,a)|^{2}\leq 1/\eta^{2}D^{2}.$ $$

which implies that $\eta|R(\widehat{\theta}_{0},x,a)-R(\theta_{*},x,a)|\leq 1$ and $\eta|R(\widehat{\theta},x,a)-R(\theta_{*},x,a)|\leq 1$
for all $(x,a)\in\mathcal{X}\times\mathcal{A}$ such that $\pi_{0}(a|x)>0$.
∎

###### Proof of Theorem 3.3 .

We have

$$ $\displaystyle\mathbb{E}_{\pi_{\theta_{*}}^{\eta}}\bigg{[}R(\theta_{*},x,a)- \frac{1}{\eta}\ln\frac{\pi_{\theta_{*}}^{\eta}(a|x)}{\pi_{0}(a|x)}\bigg{]}- \mathbb{E}_{\pi_{\widehat{\theta}}^{\eta}}\bigg{[}R(\theta_{*},x,a)-\frac{1}{ \eta}\ln\frac{\pi_{\widehat{\theta}}^{\eta}(a|x)}{\pi_{0}(a|x)}\bigg{]}$ $\displaystyle=\frac{1}{\eta}\mathbb{E}_{\pi_{\theta_{*}}^{\eta}}\bigg{[}\ln \frac{\pi_{0}(a|x)\cdot\exp\bigl{(}\eta R(\theta_{*},x,a)\bigr{)}}{\pi_{\theta _{*}}^{\eta}(a|x)}\bigg{]}-\frac{1}{\eta}\mathbb{E}_{\pi_{\widehat{\theta}}^{ \eta}}\bigg{[}\ln\frac{\pi_{0}(a|x)\cdot\exp\bigl{(}\eta R(\theta_{*},x,a) \bigr{)}}{\pi_{\widehat{\theta}}^{\eta}(a|x)}\bigg{]}$ $\displaystyle=\frac{1}{\eta}\mathbb{E}_{x\sim d_{0}}\big{[}\ln Z_{\theta_{*}}^ {\eta}(x)\big{]}-\frac{1}{\eta}\mathbb{E}_{x\sim d_{0}}\big{[}\ln Z_{\widehat{ \theta}}^{\eta}(x)\big{]}-\mathbb{E}_{x\sim d_{0}}\bigg{[}\sum_{a\in\mathcal{A }}\pi_{\widehat{\theta}}^{\eta}(a|x)\cdot\big{(}R(\theta_{*},x,a)-R(\widehat{ \theta},x,a)\big{)}\bigg{]}$ $$

For an arbitrary reward function $f:\mathcal{X}\times\mathcal{A}\to\mathbb{R}$, let $\Delta(x,a)=f(x,a)-R(\theta_{*},x,a)$. Consider the following first derivative of $J(f)=\ln Z_{f}^{\eta}(x)-\eta\sum_{a\in\mathcal{A}}\pi_{f}^{\eta}(a|x)\cdot
\Delta(x,a)$, where $Z_{f}^{\eta}(x)=\sum_{a\in\mathcal{A}}\pi_{0}(a|x)\cdot\exp(\eta\cdot f(x,a))$ and $\pi_{f}^{\eta}(a|x)\propto\pi_{0}(a|x)\cdot\exp(\eta\cdot f(x,a))$.

$$ $\displaystyle\quad\frac{\partial}{\partial\Delta(x,a)}\bigg{[}\ln Z_{f}^{\eta} (x)-\eta\sum_{a\in\mathcal{A}}\pi_{f}^{\eta}(a|x)\cdot\Delta(x,a)\bigg{]}$ $\displaystyle=\frac{1}{Z_{f}^{\eta}(x)}\cdot\pi_{0}(a|x)\exp\bigl{(}\eta\cdot f (x,a)\bigr{)}\cdot\eta-\eta\cdot\pi_{f}^{\eta}(a|x)$ $\displaystyle\quad-\eta\cdot\Delta(x,a)\cdot\frac{\pi_{0}(a|x)\cdot\exp\bigl{( }\eta\cdot f(x,a)\bigr{)}}{Z_{f}^{\eta}(x)}\cdot\eta+\eta\cdot\Delta(x,a)\cdot \frac{\bigl{[}\pi_{0}(a|x)\cdot\exp\bigl{(}\eta\cdot f(x,a)\bigr{)}\bigr{]}^{2 }}{[Z_{f}^{\eta}(x)]^{2}}\cdot\eta$ $\displaystyle\quad+\eta\sum_{a^{\prime}\in\mathcal{A}\backslash\{a\}}\frac{\pi _{0}(a^{\prime}|x)\cdot\exp\bigl{(}\eta\cdot f(x,a^{\prime})\bigr{)}}{Z_{f}^{ \eta}(x)}\cdot\eta\cdot\Delta(x,a^{\prime})\cdot\frac{\pi_{0}(a|x)\cdot\exp \bigl{(}\eta\cdot f(x,a)\bigr{)}}{Z_{f}^{\eta}(x)}$ $\displaystyle=-\eta^{2}\pi_{f}^{\eta}(a|x)\Delta(x,a)+\eta^{2}[\pi_{f}^{\eta}( a|x)]^{2}\cdot\Delta(x,a)+\eta^{2}\sum_{a^{\prime}\in\mathcal{A}\backslash\{a \}}\pi_{f}^{\eta}(a^{\prime}|x)\pi_{f}^{\eta}(a|x)\Delta(x,a^{\prime}).$ $$

Therefore, there exists $f(\cdot,\cdot)=\gamma R(\widehat{\theta},\cdot,\cdot)+(1-\gamma)R(\theta_{*},
\cdot,\cdot)$ such that  ($\gamma\in(0,1)$)

$$ $\displaystyle\mathbb{E}_{x\sim d_{0}}[J(R(\widehat{\theta},\cdot,\cdot))-J(R( \theta_{*},\cdot,\cdot))]=\frac{1}{\eta}\mathbb{E}_{x\sim d_{0}}\biggl{[}-\eta ^{2}\sum_{a\in\mathcal{A}}\pi_{f}^{\eta}(a|x)\cdot\gamma\cdot\bigl{(}R( \widehat{\theta},x,a)-R(\theta_{*},x,a)\bigr{)}^{2}\biggr{]}$ $\displaystyle\quad+\frac{1}{\eta}\mathbb{E}_{x\sim d_{0}}\biggl{[}\gamma\eta^{ 2}\sum_{a_{1}\in\mathcal{A}}\sum_{a_{2}\in\mathcal{A}}\pi_{f}^{\eta}(a_{1}|x) \pi_{f}^{\eta}(a_{2}|x)\bigl{(}R(\widehat{\theta},x,a_{1})-R(\theta_{*},x,a_{1 })\bigr{)}\bigl{(}R(\widehat{\theta},x,a_{2})-R(\theta_{*},x,a_{2})\bigr{)} \biggr{]}$ $\displaystyle\geq-\eta\cdot\mathbb{E}_{\pi_{f}^{\eta}}\bigl{[}\bigl{(}R( \widehat{\theta},x,a)-R(\theta_{*},x,a)\bigr{)}^{2}\bigr{]}$ $$

From Lemma C.4, if $m\geq 128\eta^{2}D^{2}B^{2}\cdot\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta)$, for any $(x,a)\in\mathcal{X}\times\mathcal{A}$ such that $\pi_{0}(a|x)>0$, it holds that

$$ $\displaystyle\eta|R(\widehat{\theta}_{0},x,a)-R(\theta_{*},x,a)|\leq 1,\quad \eta|R(\widehat{\theta},x,a)-R(\theta_{*},x,a)|\leq 1,$ $$

which means that for any $(x,a)\in\mathcal{X}\times\mathcal{A}$

$$ $\frac{\pi_{f}^{\eta}(a|x)}{\pi_{\widehat{\theta}_{0}}^{\eta}(a|x)}\leq e^{4}.$ $$

Let $\epsilon_{c}=\min\{\frac{\epsilon}{(1+c_{m,n}^{-1})B},\frac{1}{8(1+c_{m,n})B
\eta^{2}D^{2}},B\}$. From Lemma C.3, if $m\geq 128\eta^{2}D^{2}B^{2}\cdot\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta)$ and $n\geq\eta/\epsilon\cdot B^{2}\log(N_{\mathcal{R}}(\epsilon_{c})/\delta))$ and $n=c_{m,n}m$ then with high probability the output policy $\pi_{\widehat{\theta}}^{\eta}$ is $O(\epsilon)$ optimal.
∎

###### Lemma C.1 (Generalization error of reward function) .

###### Proof.

###### Lemma C.2 (Confidence bound for reward function) .

###### Proof.

###### Lemma C.3 .

###### Proof.

###### Lemma C.4 .

###### Proof.

###### Proof of Theorem 3.3 .

### C.3 Proof of Corollary 3.5

###### Proof of Corollary 3.5 .

The proof follows the same lines as Theorem 4.4 by replacing the global-coverage condition with the local-coverage condition.
It still holds that

$$ $\displaystyle Q(\pi^{*})-Q(\pi_{\widehat{\theta}_{0}}^{\eta})\leq\eta\cdot \mathbb{E}_{\pi_{f}^{\eta}}\bigl{[}\bigl{(}R(\widehat{\theta}_{0},x,a)-R( \theta_{*},x,a)\bigr{)}^{2}\bigr{]}$ $$

where $\pi_{f}^{\eta}(a|x)\propto\pi_{0}(a|x)\cdot\exp(\eta\cdot f(x,a))$ and $f(\cdot,\cdot)=\gamma R(\widehat{\theta}_{0},\cdot,\cdot)+(1-\gamma)R(\theta_{
*},\cdot,\cdot)$ for some $\gamma\in(0,1)$.
Thus, We have $\mathrm{KL}(\pi_{f}^{\eta}(a|x)\|\pi_{0})\leq 2\eta B$, which further implies that

$$ $\displaystyle Q(\pi^{*})-Q(\pi_{\widehat{\theta}}^{\eta})\leq\eta\cdot C_{\rho _{\mathrm{KL}}}\cdot O\big{(}\frac{1}{n}B\log(N_{\mathcal{R}}(\epsilon_{c})/ \delta)+B(1+c_{m,n}^{-1})\epsilon_{c}\big{)}$ $$

by Lemma D.4.
Then we can conclude by substituting the value of $m$ into the suboptimality gap.
∎

###### Proof of Corollary 3.5 .

## Appendix D Proofs from Section 4

### D.1 Proof of Theorem 4.3

###### Proof of Theorem 4.3 .

The proof follows a similar construction as the one for Theorem 3.1.
Consider a simple case when $|\mathcal{X}|=M$ and $|\mathcal{A}|=2$. We suppose that the context $x$ is drawn uniformly from $\mathcal{X}$ at the beginning of each round.
Let $\Theta$ be the set consisting of mappings from $\mathcal{X}$ to $\mathcal{A}=\{0,1\}$. For each $\theta\in\Theta$, we have $R(\theta,x,a)=\begin{cases}c&\text{if }a=\theta(x),\\
0&\text{if }a\neq\theta(x),\end{cases}$ where $c>0$ is a constant, and $\theta(x)$ is the optimal action under context $x$ when the model is $\theta$.

We pick a pair of model $\theta_{1},\theta_{2}$ in $\Theta$, such that $\theta_{1}(x)=\begin{cases}\theta_{2}(x)&\text{if }x\neq x_{0},\\
1-\theta_{2}(x)&\text{if }x=x_{0}.\end{cases}$

We denote by $\mathbb{P}_{\theta}$, $\mathbb{E}_{\theta}$ the probability measure and expectation under the model $\theta$.

Applying Pinsker’s inequality (Lemma E.3), we have for all event $A$ measurable with respect to the filtration generated by the observations,

$$ $\displaystyle|\mathbb{P}_{\theta_{1}}(A)-\mathbb{P}_{\theta_{2}}(A)|\leq\sqrt{ \log(1/2+e^{c}/4+e^{-c}/4)\mathbb{E}_{\theta_{1}}[N(x_{0})]}\leq\sqrt{c^{2} \mathbb{E}_{\theta_{1}}[N(x_{0})]}=\sqrt{c^{2}T/M},$ $$

where the first inequality follows from the chain rule of KL divergence, and the fact that $\mathbb{E}_{\theta_{1}}[N(x_{0})]=T/M$.

Set $A$ to be the event that $\pi_{out}(\theta_{1}(x_{0})|x_{0})>1/2$. Then we have

$$ $\displaystyle\mathbb{P}_{\theta_{1}}(\pi_{out}(\theta_{1}(x_{0})|x_{0})\leq 1/ 2)+\mathbb{P}_{\theta_{2}}(\pi_{out}(\theta_{2}(x_{0})|x_{0})\leq 1/2)\geq 1-| \mathbb{P}_{\theta_{1}}(A)-\mathbb{P}_{\theta_{2}}(A)|\geq 1-\sqrt{c^{2}T/M}.$ $$

If the model $\theta$ is uniformly drawn from $\Theta$, then we have

$$ $\displaystyle\mathbb{E}_{\theta\sim Unif(\Theta)}\mathbb{P}_{\theta}(\pi_{out} (\theta(x_{0}))\leq 1/2)\geq\frac{1}{2}-\sqrt{c^{2}T/4M}$ $$

for an arbitrary $x_{0}$.

Then we consider the following suboptimality gap:

$$ $\displaystyle\mathbb{E}_{\pi_{\theta_{*}}^{\eta}}\bigg{[}R(\theta_{*},x,a)- \frac{1}{\eta}\ln\frac{\pi_{\theta_{*}}^{\eta}(a|x)}{\pi_{0}(a|x)}\bigg{]}- \mathbb{E}_{\pi_{out}}\bigg{[}R(\theta_{*},x,a)-\frac{1}{\eta}\ln\frac{\pi_{ out}(a|x)}{\pi_{0}(a|x)}\bigg{]}$ $\displaystyle=\frac{1}{\eta}\mathbb{E}_{\pi_{\theta_{*}}^{\eta}}\bigg{[}\ln \frac{\pi_{0}(a|x)\cdot\exp\bigl{(}\eta R(\theta_{*},x,a)\bigr{)}}{\pi_{\theta _{*}}^{\eta}(a|x)}\bigg{]}-\frac{1}{\eta}\mathbb{E}_{\pi_{out}}\bigg{[}\ln \frac{\pi_{0}(a|x)\cdot\exp\bigl{(}\eta R(\theta_{*},x,a)\bigr{)}}{\pi_{out}(a |x)}\bigg{]}$ $\displaystyle=\frac{1}{\eta}\mathbb{E}_{\pi_{out}}\Bigl{[}\ln\frac{\pi_{out}(a |x)}{\pi^{*}(a|x)}\Bigr{]},$ $$

where the last equality follows from the fact that $\pi_{\theta_{*}}^{\eta}\propto\pi_{0}(a|x)\cdot\exp(\eta R(\theta_{*},x,a))$.

To bound the suboptimality gap, we further have

$$ $\displaystyle\mathbb{E}_{\theta\sim Unif(\Theta)}\mathbb{E}_{\pi_{out}}\Bigl{[ }\ln\frac{\pi_{out}(a|x)}{\pi^{*}(a|x)}\Bigr{]}$ $\displaystyle=\mathbb{E}_{\theta\sim Unif(\Theta)}\frac{1}{M}\sum_{x\in \mathcal{X}}\mathbb{E}_{a\sim\pi_{out}(\cdot|x)}\Bigl{[}\ln\frac{\pi_{out}(a|x )}{\pi^{*}(a|x)}\Bigr{]}$ $\displaystyle\geq\mathbb{E}_{\theta\sim Unif(\Theta)}\frac{1}{M}\sum_{x\in \mathcal{X}}\mathbb{P}_{\theta}(\pi_{out}(\theta(x))\leq 1/2)\cdot\biggl{[} \frac{1}{2}\ln\frac{1+\exp(-\eta c)}{2}+\frac{1}{2}\ln\frac{1+\exp(\eta c)}{2} \biggr{]}$ $\displaystyle\geq\Bigl{(}\frac{1}{2}-\sqrt{c^{2}T/4M}\Bigr{)}\Bigl{[}\frac{1}{ 2}\ln\frac{1+\exp(-\eta c)}{2}+\frac{1}{2}\ln\frac{1+\exp(\eta c)}{2}\Bigr{]}$ (D.1) $$

Note that

$$ $\displaystyle\frac{\text{d}}{\text{d}u}\Bigl{[}\frac{1}{2}\ln\frac{1+e^{-u}}{2 }+\frac{1}{2}\ln\frac{1+e^{u}}{2}\Bigr{]}\Big{|}_{u=0}$ $\displaystyle=\frac{1}{2}\Bigl{[}\frac{1}{1+\exp(-u)}-\frac{1}{1+\exp(u)}\Bigr {]}\Big{|}_{u=0}=0,$ $\displaystyle\frac{\text{d}^{2}}{\text{d}u^{2}}\Bigl{[}\frac{1}{2}\ln\frac{1+e ^{-u}}{2}+\frac{1}{2}\ln\frac{1+e^{u}}{2}\Bigr{]}$ $\displaystyle=\frac{\exp(u)}{[1+\exp(u)]^{2}}.$ $$

Thus, applying Taylor’s expansion on the right-hand side of (D.1), we have

$$ $\displaystyle\mathbb{E}_{\theta\sim Unif(\Theta)}\mathbb{E}_{\pi_{out}}\Bigl{[ }\ln\frac{\pi_{out}(a|x)}{\pi^{*}(a|x)}\Bigr{]}\geq\frac{1}{2}\cdot\Bigl{(} \frac{1}{2}-\sqrt{c^{2}T/4M}\Bigr{)}\eta^{2}c^{2}\cdot\frac{1}{3+\exp(\eta c)}$ $$

When $\epsilon<1/64\eta$, we can set $c=8\sqrt{\epsilon/\eta}$. To achieve a suboptimality gap of $\epsilon$, we need to satisfy:

$$ $\displaystyle\frac{1}{2}\cdot\Bigl{(}\frac{1}{2}-\sqrt{c^{2}T/4M}\Bigr{)}\eta^ {2}c^{2}\cdot\frac{1}{3+\exp(\eta c)}\leq\eta\epsilon,$ $$

indicating that $T\geq\frac{\eta M}{512\epsilon}=\Omega(\frac{\eta M}{\epsilon})$.

When $\epsilon\geq 1/64\eta$, or equivalently, $\eta\geq 1/64\epsilon$, we employ a different lower bound for (C.1) as follows:

$$ $\displaystyle\frac{1}{2}\ln\frac{1+\exp(-\eta c)}{2}+\frac{1}{2}\ln\frac{1+ \exp(\eta c)}{2}$ $\displaystyle=\frac{1}{2}\ln\frac{2+\exp(\eta c)+\exp(-\eta c)}{4}$ $\displaystyle\geq\frac{1}{2}\cdot\frac{1}{2}\Bigl{(}\ln\frac{\exp(\eta c)+\exp (-\eta c)}{2}\Bigr{)}$ $\displaystyle\geq\frac{1}{4}(\eta c-\ln 2),$ (D.2) $$

where the first inequality follows from Jensen’s inequality.

Substituting (D.2) into (D.1), we have

$$ $\displaystyle\epsilon\geq\frac{1}{\eta}\mathbb{E}_{\theta\sim Unif(\Theta)} \mathbb{E}_{\pi_{out}}\Bigl{[}\ln\frac{\pi_{out}(a|x)}{\pi^{*}(a|x)}\Bigr{]} \geq\frac{1}{4}\cdot\Bigl{(}\frac{1}{2}-\sqrt{c^{2}T/4M}\Bigr{)}(\eta c-\ln 2) \cdot\frac{1}{\eta}.$ $$

Set $c=64\epsilon$. Then we have $T=\Omega(M/\epsilon^{2})$.

∎

###### Proof of Theorem 4.3 .

### D.2 Proof of Theorem 4.4

First, we provide the following lemma for the connection between the likelihood loss and the reward difference, which is a key step to upper bound the reward difference between $\widehat{\theta}$ and $\theta_{*}$.

###### Lemma D.1 .

For an arbitrary policy $\pi$, and a set of context-action pairs $\{(x_{i},a_{i}^{1},a_{i}^{2},y_{i})\}_{i=1}^{n}$ generated i.i.d. from the BT model and $\pi$, we have with probability at least $1-\delta$, for any $s\leq n$,

$$ $\displaystyle\frac{1}{2}\sum_{i=1}^{s}\mathcal{L}(\theta|x_{i},a_{i}^{1},a_{i} ^{2},y_{i})-\mathcal{L}(\theta_{*}|x_{i},a_{i}^{1},a_{i}^{2},y_{i})$ $\displaystyle\leq\log(1/\delta)-\frac{1}{8}e^{-B}\sum_{i=1}^{s}\big{(}[R( \theta,x_{i},a_{i}^{2})-R(\theta,x_{i},a_{i}^{1})]-[R(\theta_{*},x_{i},a_{i}^{ 2})-R(\theta_{*},x_{i},a_{i}^{1})]\big{)}^{2}$ $$

###### Proof.

Applying Lemma E.4 to the sequence

$$ $\bigg{\{}-\frac{1}{2}y_{i}\cdot\log\frac{\sigma(R(\theta_{*},x_{i},a_{i}^{1})- R(\theta_{*},x_{i},a_{i}^{2}))}{\sigma(R(\theta,x_{i},a_{i}^{1})-R(\theta,x_{i },a_{i}^{2}))}-\frac{1}{2}(1-y_{i})\cdot\log\frac{\sigma(R(\theta_{*},x_{i},a_ {i}^{2})-R(\theta_{*},x_{i},a_{i}^{1}))}{\sigma(R(\theta,x_{i},a_{i}^{2})-R( \theta,x_{i},a_{i}^{1}))}\bigg{\}}_{i=1}^{n},$ $$

We have with probability at least $1-\delta$, for all $s\leq n$,

$$ $\displaystyle\frac{1}{2}\sum_{i=1}^{s}\mathcal{L}(\theta|x_{i},a_{i}^{1},a_{i} ^{2},y_{i})-\mathcal{L}(\theta_{*}|x_{i},a_{i}^{1},a_{i}^{2},y_{i})$ $\displaystyle\leq\log(1/\delta)+\sum_{i=1}^{s}\log\bigg{(}\sqrt{\sigma(R( \theta_{*},x_{i},a_{i}^{2})-R(\theta_{*},x_{i},a_{i}^{1}))\cdot\sigma(R(\theta ,x_{i},a_{i}^{2})-R(\theta,x_{i},a_{i}^{1}))}$ $\displaystyle\quad\quad+\sqrt{\sigma(R(\theta_{*},x_{i},a_{i}^{1})-R(\theta_{* },x_{i},a_{i}^{2}))\cdot\sigma(R(\theta,x_{i},a_{i}^{1})-R(\theta,x_{i},a_{i}^ {2}))}\bigg{)}$ $\displaystyle=\log(1/\delta)-\frac{1}{2}\sum_{i=1}^{s}\bigg{(}\sqrt{\sigma(R( \theta_{*},x_{i},a_{i}^{2})-R(\theta_{*},x_{i},a_{i}^{1}))}-\sqrt{\sigma(R( \theta,x_{i},a_{i}^{2})-R(\theta,x_{i},a_{i}^{1}))}\bigg{)}^{2}$ $\displaystyle\leq\log(1/\delta)-\frac{1}{8}\sum_{i=1}^{s}\bigl{(}\sigma(R( \theta_{*},x_{i},a_{i}^{2})-R(\theta_{*},x_{i},a_{i}^{1}))-\sigma(R(\theta,x_{ i},a_{i}^{2})-R(\theta,x_{i},a_{i}^{1}))\bigr{)}^{2}$ $\displaystyle\leq\log(1/\delta)-\frac{1}{8}e^{-B}\sum_{i=1}^{s}\big{(}[R( \theta,x_{i},a_{i}^{2})-R(\theta,x_{i},a_{i}^{1})]-[R(\theta_{*},x_{i},a_{i}^{ 2})-R(\theta_{*},x_{i},a_{i}^{1})]\big{)}^{2},$ $$

where the equality follows from the fact that $\sigma(r)+\sigma(-r)=1$ and the last inequality holds since $\sigma^{\prime}(r)=\sigma(r)\cdot(1-\sigma(r))\geq e^{-B}$ for all $r\in[-B,B]$.
∎

To further control the error bound for the reward function with the help of Lemma D.1, we introduce the following lemma to show that the likelihood function class $\mathcal{L}$ can be well-covered by the $\epsilon$-net of the reward function class $\mathcal{R}$.

###### Lemma D.2 (Covering number of $\mathcal{L}$ ) .

For any $\epsilon_{c}>0$, consider an $\epsilon_{c}$-net $\mathcal{R}^{c}=\{R(\theta,\cdot,\cdot)|\theta\in\Theta^{c}\}$ for the reward function class $\mathcal{R}$ with size $N_{\mathcal{R}}(\epsilon_{c})$. Then for any $\theta\in\Theta$, there exists $\theta^{c}\in\Theta^{c}$ such that for any $s\in[n]$,

$$ $\displaystyle\sum_{i=1}^{s}\mathcal{L}(\theta|x_{i},a_{i}^{1},a_{i}^{2},y_{i}) \leq\sum_{i=1}^{s}\mathcal{L}(\theta^{c}|x_{i},a_{i}^{1},a_{i}^{2},y_{i})+2s \epsilon_{c}.$ $$

###### Proof.

For any $r\in\mathbb{R}$, we have

$$ $\displaystyle\frac{\text{d}\log(\sigma(r))}{\text{d}r}=\frac{1}{\sigma(r)} \cdot\sigma(r)\cdot(1-\sigma(r))=1-\sigma(r)\in(0,1).$ $$

Thus, for any $\theta\in\Theta$, $x\in\mathcal{X},a^{1},a^{2}\in\mathcal{A}$ and $y\in\{0,1\}$, there exists $\theta^{c}\in\Theta^{c}$ such that

$$ $\displaystyle\big{|}\mathcal{L}(\theta|x,a^{1},a^{2},y)-\mathcal{L}(\theta^{c} |x,a^{1},a^{2},y)\big{|}$ $\displaystyle\leq\big{|}[R(\theta,x,a^{1})-R(\theta,x,a^{2})]-[R(\theta^{c},x, a^{1})-R(\theta^{c},x,a^{2})]\big{|}=2\epsilon_{c}.$ $$

∎

With the above two lemmas, we are now ready to provide the confidence bound for the MLE estimator of the reward function.

###### Lemma D.3 .

Consider a set of context-action pairs $\{(x_{i},a_{i}^{1},a_{i}^{2},y_{i})\}_{i=1}^{n}$ where labels $\{y_{i}\}_{i=1}^{n}$ are generated independently from the BT model. Suppose that $\widehat{\theta}$ is the MLE estimator as defined in Definition 4.1. We have with probability at least $1-\delta$,

$$ $\displaystyle\sum_{i=1}^{n}\big{(}[R(\widehat{\theta},x_{i},a_{i}^{2})-R( \widehat{\theta},x_{i},a_{i}^{1})]-[R(\theta_{*},x_{i},a_{i}^{2})-R(\theta_{*} ,x_{i},a_{i}^{1})]\big{)}^{2}\leq O(e^{B}\log(N_{\mathcal{R}}(\epsilon_{c})/ \delta)+e^{B}n\epsilon_{c}).$ $$

###### Proof.

By Lemma D.1 and Lemma D.2, we have with probability at least $1-\delta$, for any $\theta\in\Theta$,

$$ $\displaystyle\frac{1}{2}\sum_{i=1}^{n}\mathcal{L}(\theta|x_{i},a_{i}^{1},a_{i} ^{2},y_{i})-\mathcal{L}(\theta_{*}|x_{i},a_{i}^{1},a_{i}^{2},y_{i})$ $\displaystyle\leq\log(N_{\mathcal{R}}(\epsilon_{c})/\delta)-\frac{1}{8}e^{-B} \sum_{i=1}^{n}\big{(}[R(\theta,x_{i},a_{i}^{2})-R(\theta,x_{i},a_{i}^{1})]-[R( \theta_{*},x_{i},a_{i}^{2})-R(\theta_{*},x_{i},a_{i}^{1})]\big{)}^{2}+O(n \epsilon_{c}).$ $$

Since $\widehat{\theta}$ is the MLE estimator, we have $\sum_{i=1}^{n}\mathcal{L}(\theta|x_{i},a_{i}^{1},a_{i}^{2},y_{i})-\mathcal{L}(
\theta_{*}|x_{i},a_{i}^{1},a_{i}^{2},y_{i})\geq 0$, which further implies

$$ $\displaystyle 0\leq\log(N_{\mathcal{R}}(\epsilon_{c})/\delta)-\frac{1}{8}e^{-B }\sum_{i=1}^{n}\big{(}[R(\theta,x_{i},a_{i}^{2})-R(\theta,x_{i},a_{i}^{1})]-[R (\theta_{*},x_{i},a_{i}^{2})-R(\theta_{*},x_{i},a_{i}^{1})]\big{)}^{2}+O(n \epsilon_{c}).$ $$

Then we have

$$ $\displaystyle\sum_{i=1}^{n}\big{(}[R(\widehat{\theta},x_{i},a_{i}^{2})-R( \widehat{\theta},x_{i},a_{i}^{1})]-[R(\theta_{*},x_{i},a_{i}^{2})-R(\theta_{*} ,x_{i},a_{i}^{1})]\big{)}^{2}\leq O(e^{B}\log(N_{\mathcal{R}}(\epsilon_{c})/ \delta)+e^{B}n\epsilon_{c}).$ $$

∎

Finally, we provide the on-policy confidence bound for the squared reward difference between the MLE estimator $\widehat{\theta}$ and the optimal reward function $\theta_{*}$.

###### Lemma D.4 .

Consider an arbitrary policy $\pi$, and a set of context-action pairs $\{(x_{i},a_{i}^{1},a_{i}^{2},y_{i})\}_{i=1}^{n}$ generated i.i.d. from the BT model and $\pi$. Suppose that $\widehat{\theta}$ is the MLE estimator. We have with probability at least $1-2\delta$, there exists a mapping $b:\mathcal{X}\to\mathbb{R}$ such that

$$ $\displaystyle\mathbb{E}_{\pi}\big{[}\big{(}R(\widehat{\theta},x,a)-R(\theta_{* },x,a)-b(x)\big{)}^{2}\big{]}\leq O\bigg{(}\frac{1}{n}e^{B}\log(N_{\mathcal{R} }(\epsilon_{c})/\delta)+e^{B}\epsilon_{c}\bigg{)}.$ $$

###### Proof.

By Lemma D.3, we have with probability at least $1-\delta$,

$$ $\displaystyle\sum_{i=1}^{n}\big{(}[R(\widehat{\theta},x_{i},a_{i}^{2})-R( \widehat{\theta},x_{i},a_{i}^{1})]-[R(\theta_{*},x_{i},a_{i}^{2})-R(\theta_{*} ,x_{i},a_{i}^{1})]\big{)}^{2}\leq O(e^{B}\log(N_{\mathcal{R}}(\epsilon_{c})/ \delta)+e^{B}n\epsilon_{c}).$ $$

We consider an $\epsilon_{c}$-net $\mathcal{R}^{c}=\{R(\theta,\cdot,\cdot)|\theta\in\Theta^{c}\}$ for the reward function class $\mathcal{R}$ with size $N_{\mathcal{R}}(\epsilon_{c})$. For any $R(\theta,\cdot,\cdot)$, there exists $R(\theta^{c},\cdot,\cdot)$ such that

$$ $\displaystyle\big{|}R(\theta,x,a)-R(\theta^{c},x,a)\big{|}\leq O(\epsilon_{c})$ $$

for all $x\in\mathcal{X},a\in\mathcal{A}$.

Applying Lemma E.1, with probability at least $1-\delta$, we have

$$ $\displaystyle\sum_{i=1}^{n}\big{(}[R(\theta^{c},x_{i},a_{i}^{2})-R(\theta^{c}, x_{i},a_{i}^{1})]-[R(\theta_{*},x_{i},a_{i}^{2})-R(\theta_{*},x_{i},a_{i}^{1}) ]\big{)}^{2}$ $\displaystyle\quad-n\mathbb{E}_{x\sim d_{0}}\mathbb{E}_{a^{1},a^{2}\sim\pi} \big{[}\big{(}R(\theta^{c},x,a^{1})-R(\theta_{*},x,a^{1})-R(\theta^{c},x,a^{2} )+R(\theta_{*},x,a^{2})\big{)}^{2}\big{]}$ $\displaystyle\leq\sqrt{\sum_{i=1}^{n}4B^{2}\mathbb{E}_{x\sim d_{0}}\mathbb{E}_ {a^{1},a^{2}\sim\pi}\big{[}\big{(}R(\theta^{c},x,a^{1})-R(\theta_{*},x,a^{1})- R(\theta^{c},x,a^{2})+R(\theta_{*},x,a^{2})\big{)}^{2}\big{]}\log(N_{\mathcal{ R}}(\epsilon_{c})/\delta)}$ $\displaystyle\quad+\frac{8}{3}B^{2}\log(N_{\mathcal{R}}(\epsilon_{c})/\delta)$ $$

for all $\theta^{c}\in\Theta^{c}$.

From Lemma E.2 and the definition of $\Theta^{c}$, we further have

$$ $\displaystyle\mathbb{E}_{x\sim d_{0}}\mathbb{E}_{a^{1},a^{2}\sim\pi}\big{[} \big{(}R(\widehat{\theta},x,a^{1})-R(\theta_{*},x,a^{1})-R(\widehat{\theta},x, a^{2})+R(\theta_{*},x,a^{2})\big{)}^{2}\big{]}$ $\displaystyle\leq O(\frac{1}{n}B^{2}\log(N_{\mathcal{R}}(\epsilon_{c})/\delta) +\frac{1}{n}\sum_{i=1}^{n}\big{(}[R(\widehat{\theta},x_{i},a_{i}^{2})-R( \widehat{\theta},x_{i},a_{i}^{1})]-[R(\theta_{*},x_{i},a_{i}^{2})-R(\theta_{*} ,x_{i},a_{i}^{1})]\big{)}^{2}+B\epsilon_{c}),$ (D.3) $$

from which we can further derive that

$$ $\displaystyle\mathbb{E}_{x\sim d_{0}}\mathbb{E}_{a^{1},a^{2}\sim\pi}\big{[} \big{(}R(\widehat{\theta},x,a^{1})-R(\theta_{*},x,a^{1})-R(\widehat{\theta},x, a^{2})+R(\theta_{*},x,a^{2})\big{)}^{2}\big{]}$ $\displaystyle\leq O\big{(}\frac{1}{n}e^{B}\log(N_{\mathcal{R}}(\epsilon_{c})/ \delta)+e^{B}\epsilon_{c}\big{)}$ $$

with probability at least $1-2\delta$
from Lemma D.3 and the union bound.

We can then complete the proof by setting

$$ $\displaystyle b(x)=\mathbb{E}_{a^{2}\sim\pi(\cdot|x)}\big{[}R(\widehat{\theta} ,x,a^{2})-R(\theta_{*},x,a^{2})\big{]}.$ $$

∎

###### Lemma D.5 (Coverage of $\pi_{*}$ and $\pi_{\widehat{\theta}}$ by $\pi_{\widehat{\theta}_{0}}$ ) .

If $m\geq 32\eta^{2}D^{2}e^{B}\log(N_{\mathcal{R}}(\epsilon_{c}))$, $n=c_{m,n}m$ and $\epsilon_{c}\leq\frac{1}{(1+c_{m,n})e^{B}\eta^{2}D^{2}}$ in Algorithm 2 and Assumption 2.6 holds, then with probability at least $1-4\delta$, there exists $b_{1}:\mathcal{X}\to\mathbb{R}$ and $b_{2}:\mathcal{X}\to\mathbb{R}$ such that

$$ $\displaystyle\eta|R(\widehat{\theta}_{0},x,a)-R(\theta_{*},x,a)-b_{1}(x)|\leq 1 ,\quad\eta|R(\widehat{\theta},x,a)-R(\theta_{*},x,a)-b_{2}(x)|\leq 1$ $$

for all $x\in\mathcal{X},a\in\mathcal{A}$ such that $\pi_{0}(a|x)>0$.

###### Proof.

By Lemma D.3 and the union bound, we have with probability at least $1-\delta$, it holds that

$$ $\displaystyle\sum_{i=1}^{m}\big{(}[R(\widehat{\theta},\widetilde{x}_{i}, \widetilde{a}_{i}^{2})-R(\widehat{\theta},\widetilde{x}_{i},\widetilde{a}_{i}^ {1})]-[R(\theta_{*},\widetilde{x}_{i},\widetilde{a}_{i}^{2})-R(\theta_{*}, \widetilde{x}_{i},\widetilde{a}_{i}^{1})]\big{)}^{2}$ $\displaystyle\quad+\sum_{i=1}^{n}\big{(}[R(\widehat{\theta},x_{i},a_{i}^{2})-R (\widehat{\theta},x_{i},a_{i}^{1})]-[R(\theta_{*},x_{i},a_{i}^{2})-R(\theta_{* },x_{i},a_{i}^{1})]\big{)}^{2}$ $\displaystyle\leq O(e^{B}\log(N_{\mathcal{R}}(\epsilon_{c})/\delta)+e^{B}(n+m) \epsilon_{c}).$ (D.4) $$

Consider an $\epsilon_{c}$-net $\mathcal{R}^{c}=\{R(\theta,\cdot,\cdot)|\theta\in\Theta^{c}\}$ for the reward function class $\mathcal{R}$ with size $N_{\mathcal{R}}(\epsilon_{c})$. For any $R(\theta,\cdot,\cdot)$, there exists $R(\theta^{c},\cdot,\cdot)$ such that

$$ $\displaystyle\big{|}R(\theta,x,a)-R(\theta^{c},x,a)\big{|}\leq O(\epsilon_{c})$ $$

for all $x\in\mathcal{X},a\in\mathcal{A}$.

Applying Lemma E.1, with probability at least $1-\delta$, we have

$$ $\displaystyle\sum_{i=1}^{m}\big{(}[R(\theta^{c},\widetilde{x}_{i},\widetilde{a }_{i}^{2})-R(\theta^{c},\widetilde{x}_{i},\widetilde{a}_{i}^{1})]-[R(\theta_{* },x_{i},a_{i}^{2})-R(\theta_{*},x_{i},a_{i}^{1})]\big{)}^{2}$ $\displaystyle\quad-m\mathbb{E}_{x\sim d_{0}}\mathbb{E}_{a^{1},a^{2}\sim\pi_{0} }\big{[}\big{(}R(\theta^{c},x,a^{1})-R(\theta_{*},x,a^{1})-R(\theta^{c},x,a^{2 })+R(\theta_{*},x,a^{2})\big{)}^{2}\big{]}$ $\displaystyle\leq\sqrt{\sum_{i=1}^{m}4B^{2}\mathbb{E}_{x\sim d_{0}}\mathbb{E}_ {a^{1},a^{2}\sim\pi_{0}}\big{[}\big{(}R(\theta^{c},x,a^{1})-R(\theta_{*},x,a^{ 1})-R(\theta^{c},x,a^{2})+R(\theta_{*},x,a^{2})\big{)}^{2}\big{]}\log(N_{ \mathcal{R}}(\epsilon_{c})/\delta)}$ $\displaystyle\quad+\frac{8}{3}B^{2}\log(N_{\mathcal{R}}(\epsilon_{c})/\delta)$ $$

for all $\theta^{c}\in\Theta^{c}$.

From Lemma E.2 and the definition of $\Theta^{c}$, we further have

$$ $\displaystyle\mathbb{E}_{x\sim d_{0}}\mathbb{E}_{a^{1},a^{2}\sim\pi}\big{[} \big{(}R(\widehat{\theta},x,a^{1})-R(\theta_{*},x,a^{1})-R(\widehat{\theta},x, a^{2})+R(\theta_{*},x,a^{2})\big{)}^{2}\big{]}$ $\displaystyle\leq O(\frac{1}{m}B^{2}\log(N_{\mathcal{R}}(\epsilon_{c})/\delta) +\frac{1}{m}\sum_{i=1}^{n}\big{(}[R(\widehat{\theta},\widetilde{x}_{i}, \widetilde{a}_{i}^{2})-R(\widehat{\theta},\widetilde{x}_{i},\widetilde{a}_{i}^ {1})]-[R(\theta_{*},\widetilde{x}_{i},\widetilde{a}_{i}^{2})-R(\theta_{*}, \widetilde{x}_{i},\widetilde{a}_{i}^{1})]\big{)}^{2}+B\epsilon_{c}).$ (D.5) $$

Substituting (D.4) into (D.5), we have with probability at least $1-2\delta$,

$$ $\displaystyle\mathbb{E}_{x\sim d_{0}}\mathbb{E}_{a^{1},a^{2}\sim\pi_{0}}\big{[ }\big{(}R(\widehat{\theta},x,a^{1})-R(\theta_{*},x,a^{1})-R(\widehat{\theta},x ,a^{2})+R(\theta_{*},x,a^{2})\big{)}^{2}\big{]}$ $\displaystyle\leq O\big{(}\frac{1}{m}e^{B}\log(N_{\mathcal{R}}(\epsilon_{c})/ \delta)+e^{B}\cdot\frac{n+m}{m}\cdot\epsilon_{c}\big{)}.$ $$

Therefore, there exists a mapping $b_{2}:\mathcal{X}\to\mathbb{R}$ such that

$$ $\displaystyle\mathbb{E}_{\pi_{0}}\big{[}\big{(}R(\widehat{\theta},x,a)-R( \theta_{*},x,a)-b_{2}(x)\big{)}^{2}\big{]}\leq O\big{(}\frac{1}{m}e^{B}\log(N_ {\mathcal{R}}(\epsilon_{c})/\delta)+e^{B}\cdot\frac{n+m}{m}\cdot\epsilon_{c} \big{)}.$ $$

From Lemma D.4, we have with probability at least $1-2\delta$, there exists a mapping $b_{1}:\mathcal{X}\to\mathbb{R}$ such that

$$ $\displaystyle\mathbb{E}_{\pi_{0}}\big{[}\big{(}R(\widehat{\theta}_{0},x,a)-R( \theta_{*},x,a)-b_{1}(x)\big{)}^{2}\big{]}\leq O\big{(}\frac{1}{m}e^{B}\log(N_ {\mathcal{R}}(\epsilon_{c})/\delta)+e^{B}(1+c_{m,n})\epsilon_{c}\big{)}.$ $$

Hence, we can complete the proof by a union bound over the two events and Assumption 2.6.
∎

###### Proof of Theorem 4.4 .

Let $b$ be the mapping defined in Lemma D.4 for $\widehat{\theta}$
We have

$$ $\displaystyle\mathbb{E}_{\pi_{\theta_{*}}^{\eta}}\bigg{[}R(\theta_{*},x,a)- \frac{1}{\eta}\ln\frac{\pi_{\theta_{*}}^{\eta}(a|x)}{\pi_{0}(a|x)}\bigg{]}- \mathbb{E}_{\pi_{\widehat{\theta}}^{\eta}}\bigg{[}R(\theta_{*},x,a)-\frac{1}{ \eta}\ln\frac{\pi_{\widehat{\theta}}^{\eta}(a|x)}{\pi_{0}(a|x)}\bigg{]}$ $\displaystyle=\frac{1}{\eta}\mathbb{E}_{\pi_{\theta_{*}}^{\eta}}\bigg{[}\ln \frac{\pi_{0}(a|x)\cdot\exp\bigl{(}\eta R(\theta_{*},x,a)\bigr{)}}{\pi_{\theta _{*}}^{\eta}(a|x)}\bigg{]}-\frac{1}{\eta}\mathbb{E}_{\pi_{\widehat{\theta}}^{ \eta}}\bigg{[}\ln\frac{\pi_{0}(a|x)\cdot\exp\bigl{(}\eta R(\theta_{*},x,a) \bigr{)}}{\pi_{\widehat{\theta}}^{\eta}(a|x)}\bigg{]}$ $\displaystyle=\frac{1}{\eta}\mathbb{E}_{x\sim d_{0}}\big{[}\ln Z_{\theta_{*}}^ {\eta}(x)\big{]}-\frac{1}{\eta}\mathbb{E}_{x\sim d_{0}}\big{[}\ln Z_{\widehat{ \theta}}^{\eta}(x)\big{]}-\mathbb{E}_{x\sim d_{0}}\bigg{[}\sum_{a\in\mathcal{A }}\pi_{\widehat{\theta}}^{\eta}(a|x)\cdot\big{(}R(\theta_{*},x,a)-R(\widehat{ \theta},x,a)\big{)}\bigg{]}.$ $$

For an arbitrary reward function $f:\mathcal{X}\times\mathcal{A}\to\mathbb{R}$, let $\Delta(x,a)=f(x,a)-R(\theta_{*},x,a)$. Consider the following first derivative of $J(f)=\ln Z_{f}^{\eta}(x)-\eta\sum_{a\in\mathcal{A}}\pi_{f}^{\eta}(a|x)\cdot
\Delta(x,a)$, where $Z_{f}^{\eta}(x)=\sum_{a\in\mathcal{A}}\pi_{0}(a|x)\cdot\exp(\eta\cdot f(x,a))$ and $\pi_{f}^{\eta}(a|x)\propto\pi_{0}(a|x)\cdot\exp(\eta\cdot f(x,a))$.

Similar to the proof of Theorem 3.3, we still have

$$ $\displaystyle\quad\frac{\partial}{\partial\Delta(x,a)}\bigg{[}\ln Z_{f}^{\eta} (x)-\eta\sum_{a\in\mathcal{A}}\pi_{f}^{\eta}(a|x)\cdot\Delta(x,a)\bigg{]}$ $\displaystyle=\frac{1}{Z_{f}^{\eta}(x)}\cdot\pi_{0}(a|x)\exp\bigl{(}\eta\cdot f (x,a)\bigr{)}\cdot\eta-\eta\cdot\pi_{f}^{\eta}(a|x)$ $\displaystyle\quad-\eta\cdot\Delta(x,a)\cdot\frac{\pi_{0}(a|x)\cdot\exp\bigl{( }\eta\cdot f(x,a)\bigr{)}}{Z_{f}^{\eta}(x)}\cdot\eta+\eta\cdot\Delta(x,a)\cdot \frac{\bigl{[}\pi_{0}(a|x)\cdot\exp\bigl{(}\eta\cdot f(x,a)\bigr{)}\bigr{]}^{2 }}{[Z_{f}^{\eta}(x)]^{2}}\cdot\eta$ $\displaystyle\quad+\eta\sum_{a^{\prime}\in\mathcal{A}\backslash\{a\}}\frac{\pi _{0}(a^{\prime}|x)\cdot\exp\bigl{(}\eta\cdot f(x,a^{\prime})\bigr{)}}{Z_{f}^{ \eta}(x)}\cdot\eta\cdot\Delta(x,a^{\prime})\cdot\frac{\pi_{0}(a|x)\cdot\exp \bigl{(}\eta\cdot f(x,a)\bigr{)}}{Z_{f}^{\eta}(x)}$ $\displaystyle=-\eta^{2}\pi_{f}^{\eta}(a|x)\Delta(x,a)+\eta^{2}[\pi_{f}^{\eta}( a|x)]^{2}\cdot\Delta(x,a)+\eta^{2}\sum_{a^{\prime}\in\mathcal{A}\backslash\{a \}}\pi_{f}^{\eta}(a^{\prime}|x)\pi_{f}^{\eta}(a|x)\Delta(x,a^{\prime}).$ $$

Note that

$$ $\displaystyle J(R(\widehat{\theta},x,\cdot))=\ln Z_{\widehat{\theta}}^{\eta}(x )-\eta\sum_{a\in\mathcal{A}}\pi_{\widehat{\theta}}^{\eta}(a|x)\cdot\bigl{(}R( \widehat{\theta},x,a)-R(\theta_{*},x,a)\bigr{)}$ $\displaystyle=\ln\sum_{a\in\mathcal{A}}\pi_{0}(a|x)\cdot\exp(\eta(R(\widehat{ \theta},x,a)-b(x)))-\eta\sum_{a\in\mathcal{A}}\pi_{\widehat{\theta}}^{\eta}(a| x)\cdot\bigl{(}R(\widehat{\theta},x,a)-R(\theta_{*},x,a)-b(x)\bigr{)}$ $\displaystyle=J(R(\widehat{\theta},x,\cdot)-b(x)).$ $$

Therefore, there exists $f(\cdot,\cdot)=\gamma[R(\widehat{\theta},\cdot,\cdot)-b(\cdot)]+(1-\gamma)R(
\theta_{*},\cdot,\cdot)$ such that  ($\gamma\in(0,1)$)

$$ $\displaystyle\mathbb{E}_{x\sim d_{0}}[J(R(\widehat{\theta},\cdot,\cdot))-J(R( \theta_{*},\cdot,\cdot))]$ $\displaystyle=\frac{1}{\eta}\mathbb{E}_{x\sim d_{0}}\biggl{[}-\eta^{2}\sum_{a \in\mathcal{A}}\pi_{f}^{\eta}(a|x)\cdot\gamma\cdot\bigl{(}R(\widehat{\theta},x ,a)-R(\theta_{*},x,a)-b(x)\bigr{)}^{2}\biggr{]}$ $\displaystyle\quad+\frac{1}{\eta}\mathbb{E}_{x\sim d_{0}}\bigg{[}\gamma\eta^{2 }\sum_{a_{1}\in\mathcal{A}}\sum_{a_{2}\in\mathcal{A}}\pi_{f}^{\eta}(a_{1}|x) \pi_{f}^{\eta}(a_{2}|x)\bigl{(}R(\widehat{\theta},x,a_{1})-R(\theta_{*},x,a_{1 })-b(x)\bigr{)}$ $\displaystyle\quad\bigl{(}R(\widehat{\theta},x,a_{2})-R(\theta_{*},x,a_{2})-b( x)\bigr{)}\bigg{]}$ $\displaystyle\geq-\eta\cdot\mathbb{E}_{\pi_{f}^{\eta}}\bigl{[}\bigl{(}R( \widehat{\theta},x,a)-R(\theta_{*},x,a)-b(x)\bigr{)}^{2}\bigr{]}$ $$

From Lemma D.2, if $m\geq 32\eta^{2}D^{2}e^{B}\cdot\log(2N_{\mathcal{R}}(\epsilon_{c})/\delta)$, for any $(x,a)\in\mathcal{X}\times\mathcal{A}$ such that $\pi_{0}(a|x)>0$, it holds that

$$ $\displaystyle\eta|R(\widehat{\theta}_{0},x,a)-R(\theta_{*},x,a)-b_{1}(x)|\leq 1 ,\quad\eta|R(\widehat{\theta},x,a)-R(\theta_{*},x,a)-b_{2}(x)|\leq 1,$ $$

which means that

$$ $\frac{\pi_{f}^{\eta}}{\pi_{\widehat{\theta}_{0}}^{\eta}}\leq e^{4}.$ $$

Let $\epsilon_{c}=\min\{\frac{\epsilon}{2(1+c_{m,n}^{-1})e^{B}},\frac{1}{(1+c_{m,n}
)e^{B}\eta^{2}D^{2}}\}$. From Lemma D.4, under the condition of the theorem, with high probability the output policy $\pi_{\widehat{\theta}}^{\eta}$ is $O(\epsilon)$ optimal.
∎

###### Lemma D.1 .

###### Proof.

###### Lemma D.2 (Covering number of ℒ ℒ \mathcal{L} caligraphic_L ) .

###### Proof.

###### Lemma D.3 .

###### Proof.

###### Lemma D.4 .

###### Proof.

###### Lemma D.5 (Coverage of π ∗ subscript 𝜋 \pi_{*} italic_π start_POSTSUBSCRIPT ∗ end_POSTSUBSCRIPT and π θ ^ subscript 𝜋 ^ 𝜃 \pi_{\widehat{\theta}} italic_π start_POSTSUBSCRIPT over^ start_ARG italic_θ end_ARG end_POSTSUBSCRIPT by π θ ^ 0 subscript 𝜋 subscript ^ 𝜃 0 \pi_{\widehat{\theta}_{0}} italic_π start_POSTSUBSCRIPT over^ start_ARG italic_θ end_ARG start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT end_POSTSUBSCRIPT ) .

###### Proof.

###### Proof of Theorem 4.4 .

### D.3 Proof of Corollary 4.7

In this subsection, we also discuss our result under the local-coverage condition (Definition 2.8).

###### Proof of Corollary 4.7 .

The proof follows the same lines as Theorem 4.4 by replacing the global-coverage condition with the local-coverage condition.
It still holds that

$$ $\displaystyle Q(\pi^{*})-Q(\pi_{\widehat{\theta}_{0}}^{\eta})\leq\eta\cdot \mathbb{E}_{\pi_{f}^{\eta}}\bigl{[}\bigl{(}R(\widehat{\theta}_{0},x,a)-R( \theta_{*},x,a)-b(x)\bigr{)}^{2}\bigr{]}$ $$

where $\pi_{f}^{\eta}(a|x)\propto\pi_{0}(a|x)\cdot\exp(\eta\cdot f(x,a))$ and $f(\cdot,\cdot)=\gamma[R(\widehat{\theta}_{0},\cdot,\cdot)-b(\cdot)]+(1-\gamma)
R(\theta_{*},\cdot,\cdot)$ for some $\gamma\in(0,1)$.
Thus, We have $\mathrm{KL}(\pi_{f}^{\eta}(a|x)\|\pi_{0})\leq 2\eta B$, which further implies that

$$ $\displaystyle Q(\pi^{*})-Q(\pi_{\widehat{\theta}}^{\eta})\leq\eta\cdot C_{\rho _{\mathrm{KL}}}\cdot O\big{(}\frac{1}{n}e^{B}\log(N_{\mathcal{R}}(\epsilon_{c} )/\delta)+e^{B}(1+c_{m,n}^{-1})\epsilon_{c}\big{)}$ $$

by Lemma D.4.
Then we can conclude by substituting the value of $m$ into the suboptimality gap.
∎

###### Proof of Corollary 4.7 .

## Appendix E Auxiliary Lemmas

###### Lemma E.1 (Freedman inequality) .

Let $M,v>0$ be fixed constants. Let $\{X_{i}\}_{i=1}^{n}$ be a stochastic process, $\{\mathcal{G}_{i}\}_{i}$ be a sequence of $\sigma$-fields, and $X_{i}$ be $\mathcal{G}_{i}$-measurable, while almost surely

$$ $\displaystyle\mathbb{E}[X_{i}|\mathcal{G}_{i}]=0,|X_{i}|\leq M,\text{ and } \sum_{i=1}^{n}\mathbb{E}[X_{i}^{2}|\mathcal{G}_{i-1}]\leq v.$ $$

Then for any $\delta>0$, with probability at least $1-\delta$, it holds that

$$ $\displaystyle\sum_{i=1}^{n}X_{i}\leq\sqrt{2v\log(1/\delta)}+\frac{2}{3}M\log(1 /\delta).$ $$

###### Lemma E.2 .

Suppose $a,b\geq 0$. If $x^{2}\leq a+b\cdot x$, then $x^{2}\leq 2b^{2}+2a$.

###### Proof.

By solving the root of quadratic polynomial $q(x):=x^{2}-b\cdot x-a$, we obtain $\max\{x_{1},x_{2}\}=(b+\sqrt{b^{2}+4a})/2$. Hence, we have $x\leq(b+\sqrt{b^{2}+4a})/2$ provided that $q(x)\leq 0$. Then we further have

$$ $\displaystyle x^{2}\leq\frac{1}{4}\left(b+\sqrt{b^{2}+4a}\right)^{2}\leq\frac{ 1}{4}\cdot 2\left(b^{2}+b^{2}+4a\right)\leq 2b^{2}+2a.$ (E.1) $$

∎

###### Lemma E.3 (Pinsker’s inequality) .

If $\mathbb{P}_{1}$, $\mathbb{P}_{2}$ are two probability measures on a common measurable space $(\Omega,\mathcal{F})$, then it holds that

$$ $\displaystyle\delta(\mathbb{P}_{1},\mathbb{P}_{2})\leq\sqrt{\frac{1}{2}\mathrm {KL}(\mathbb{P}_{1}\|\mathbb{P}_{2})},$ $$

where $\delta(\cdot,\cdot)$ is the total variation distance and $\mathrm{KL}(\cdot\|\cdot)$ is the Kullback-Leibler divergence.

###### Lemma E.4 (Lemma A.4, Foster et al. 2021 ) .

For any sequence of real-valued random variables $(X_{t})_{t\leq T}$ adapted to a filtration $(\mathcal{F}_{t})_{t\leq T}$, it holds that with probability at least $1-\delta$, for all $T^{\prime}\leq T$,

$$ $\displaystyle\sum_{t=1}^{T^{\prime}}X_{t}\leq\sum_{t=1}^{T^{\prime}}\log\bigl{ (}\mathbb{E}_{t-1}[e^{X_{t}}]\bigr{)}+\log(1/\delta).$ $$

###### Lemma E.1 (Freedman inequality) .

###### Lemma E.2 .

###### Proof.

###### Lemma E.3 (Pinsker’s inequality) .

###### Lemma E.4 (Lemma A.4, Foster et al. 2021 ) .

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1) In line 294 it says that "it is obvious that $D^2 \leq C_{GL}$. Could you please provide a proof for that?

2) In line 374 is $s_i \sim \pi_0$ a typo, where does $s_i$ come from?

3) Can you provide a specific example where you compute the coverage of the reference policy and it is O(1)?
soundness: 3
strengths: 1) Improved sample complexity for contextual bandits under KL regularization with a novel analysis using the property of strong convexity due to the KL regularizer

2) Lower bound for the contextual bandit problem under KL regularization that is tight with the upper for sufficiently small $\epsilon$

3) Lower bound for the RLHF problem with preference feedback

4) Design of an algorithm for the RLHF problem with guarantees that match the lower bound.
summary: In this paper, the authors provide a new analysis of contextual bandits under KL regularization that achieves an improved sample complexity guarantee. Then, they study the RLHF problem where under coverage assumptions on the reference policy and they provide a tight algorithm under this assumption. In the end, they provide experimental results.
weaknesses: 1) Cannot see the tradeoff between $\eta$ and the number of samples needed. Even from the experimental results, the lower $\eta$ the better the performance.

2) The O(1) coverage assumption is unclear if it is a reasonable one or not a provided way to check so.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. In Definition 2.7 (and 2.8), is the sup taken over $x\in\mathrm{supp}(d_0)$? The current notation is a bit confusing.
soundness: 3
strengths: 1. The paper is written clearly and the main results are well delivered. 

2. The paper establishes an integrated theory that provides both lower and matching upper bounds for CB/RLHF sample complexity. The derivations are solid. 

3. RLHF with KL regularization is predominant in LLM alignment. Studying this problem through a theoretical lens has sufficient significance in helping gather insights into designing more efficient RLHF methods.
summary: The paper provides sharp theoretical analyses for KL-regularized contextual bandit (CB) and RLHF problems. It first studies the theoretical benefits of KL regularization in CB and RLHF and shows that KL regularization improves the sample complexity to $O(1/\epsilon)$, while for unregularized problems $O(1/\epsilon^2)$ samples are required. The paper then studies the role of data coverage. In particular, a two-stage mixed sampling strategy is proposed to achieve sample complexity with only additive dependence on the policy coverage coefficient if the data coverage is sufficient, while previous results often depend on the coverage coefficient multiplicatively. The paper also provides a local policy coverage coefficient and derives sample complexity which has multiplicative dependence on this weaker notion than global policy coverage. Numerical experiments are provided to support the theories.
weaknesses: 1. I have concerns about the coverage assumptions in the paper and their relation to the practical use case of RLHF in LLM alignment. The policy coverage (Definition 2.7 and 2.8) is assumed for reference policy $\pi_0$. In RLHF, such reference policy is typically a fine-tuned LLM that can have extremely low if not zero probability for some actions (e.g. nonsense responses). In this case, the global policy coverage coefficient can blow up to infinity, and the local policy coverage coefficient can be large (as the KL constraint is in expectation). The current paper only derives additive dependence results for global policy coverage, which can be vacuous if the global coefficient is infinity. On the other hand, the dependence on the local coefficient is still multiplicative (as discussed in Section 3.4), which can be extremely large. 

2. I also have concerns regarding the sample complexity upper bounds. The paper claims to first study the effect of KL regularization in improving the sample complexity for policy optimization from $O(1/\epsilon^2)$ to $O(1/\epsilon)$. However, such $O(1/\epsilon)$ sample complexity result already exists for general strongly convex regularizers [1], which include KL regularization as a special case since the reference policy is assumed to have sufficient coverage. Hence it is likely that the upper bound for CB is already known (given that CB is a special case of MDP). For RLHF, its difference from CB mainly comes from the additional reward learning step, so I expect there could be more explanation on why $O(1/\epsilon)$ samples are sufficient for reward learning from preference data. However, the current version of the paper seems to lack such comparisons/remarks, which in my opinion are necessary for understanding the mechanism of RLHF (just as strong convexity of KL divergence for CB). For baselines, previous literature suggests that reward learning takes $O(1/\epsilon^2)$ samples. 

[1] Lan, Guanghui. "Policy mirror descent for reinforcement learning: Linear convergence, new sampling complexity, and generalized problem classes." Mathematical programming 198, no. 1 (2023): 1059-1106.

[2] Zhu, Banghua, Michael Jordan, and Jiantao Jiao. "Principled reinforcement learning with human feedback from pairwise or k-wise comparisons." In International Conference on Machine Learning, pp. 43037-43067. PMLR, 2023.

---

confidence: 2
contribution: 2
initial_rating: 3
presentation: 2
questions: - In definition 2.7, line 219 – 221, and definition 2.8, line 231 – 233, what does ``x sampled from d_0’’ mean in the sup? 

- Is there any benefit of using more than 2-stages of sampling?

- Can the authors provide more intuition why the number of samples needed are different in the first and the second stages of the algorithm, as presented in Theorem 3.3?
soundness: 2
strengths: The paper provides sharper analysis of the sampling complexity of KL-regularized contextual bandit problems, and provides novel results on the dependency of the sampling complexity on the data coverage.
summary: The paper provides a novel lower bound of Omega(1/epsilon) for the sampling complexity of finding epsilon-suboptimal solutions in KL-regularized contextual bandit problems. 

The paper then models the online KL-regularized RLHF problem as the KL-regularized contextual bandit problem and proposes a two-stage sampling algorithm. Using the strong convexity of the KL-regularization, the paper shows that the algorithm has sampling complexity O(1/epsilon), with an additive term that depends on the coverage coefficient of the reference policy.
weaknesses: The paper might benefit from some polishing. For instance, the definitions of some key terms are not rigorous (see questions below). In addition, the proof may lack rigor. For instance, line 320 – 321 uses Taylor expansion, and thus, if my understanding is correct, the equality (and the following inequality) does not hold.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Is the scaling with $D^2$ really necessary in Theorem 3.3 and Theorem 4.4 or can this be reduced to $C_{GL}$?
2. For $\eta$ very large (corresponding to no regularization), one would hope to recover the standard complexity bounds for contextual bandits, but this is not the case in any of the upper bounds (all of which will continue to increase as $\eta$ increases). I suspect a more refined analysis may allow one to obtain the minimum of the current complexity and the standard contextual bandit complexity. Could the authors comment on this?
soundness: 3
strengths: To my knowledge, this is the first work to obtain a tight bound on the sample complexity of RLHF with KL regularization. This is a commonly used setting in practice and as such it is important that we understand the sample complexity. It is interesting that one can obtain a $1/\epsilon$ rate as compared to the $1/\epsilon^2$ that might be expected.
summary: This paper studies the RLHF problem for contextual bandits, and aims to obtain a tight sample complexity on the problem. They consider both the reward observation setting and preference observation setting, and provide upper and lower bounds for each. Interestingly, they are able to show that including the KL constraint in the problem allows them to obtain a sample complexity that scales as $1/\epsilon$ instead of the more familiar $1/\epsilon^2$.
weaknesses: 1. The main coverage condition (Definition 2.6) is extremely strong. This will scale at least with the number of contexts: if we take $\theta = \theta’$, and choose $b(x) = 0$ for $x \neq x_0$ and $b(x_0) = B$ where $x_0$ is the minimum probability context under $d_0$, then the expression given in Definition 2.6 will scale as $1/d_0(x_0) \ge M$ for $M$ the number of contexts. Thus, the main sample complexity results of the paper (Theorem 3.3 and Theorem 4.4) really scale with the size of the context space. In general it is not acceptable to obtain a sample complexity scaling with the size of the context space, as this is typically extremely large, so these results are only meaningful asymptotically as $\epsilon \rightarrow 0$.
2. Furthermore, the result is also not tight in the regime where $\eta$ is very large.
3. The statement on Line 294 that $D^2 \le C_{GL}$ is not correct as a result of this ($C_{GL}$ in general will not scale with context size).
4. The statement in Theorem 3.1 and Theorem 4.3 that the coverage condition is $O(1)$ is also then incorrect—it should be $\log N_{\mathcal{R}}(\epsilon)$ I believe.
5. The problem setting could be clarified somewhat. In particular, it should be made more explicit that when a policy is $\epsilon$-optimal, this is with respect to $Q(\pi)$, the reward + KL objective, rather than just the reward. The latter is typically more standard for RL, so it should be made clear that the objective considered here is different.
6. The writing could be improved. There are various unclear or poorly worded statements (the following list is not exhaustive—please go through the paper carefully and resolve other such issues):
	* Line 59: “DPO suffers from a drop of chosen probability”. I am not sure what this means.
	* Line 62: “the learned model is easy to be hacked and become biased”. Grammatically incorrect, revise wording.
	* Line 64: The sentence starting with “Hence, the KL-regularization..” Is the first time KL regularization is mentioned. It seems like it needs to be introduced earlier for this sentence to read well.
	* Line 283: “identifically” is not a word.

7. There are also several issues with informal technical statements being made that are not necessarily correct:
	* Line 76: RLHF has been demonstrated to outperform offline methods because “it has further interactions with human or preference oracle”. It is not clear that this the reason (or what exactly is even meant by this sentence).
	* Line 273: It is much more standard in modern offline RL to obtain bounds under single policy concentrability (ie only one policy is covered). Thus, the claim that global coverage is standard in offline RL is too strong. 
	* Definition 2.6: What is the $\pi$ referred to here? It is not clear. 
	* Line 291: The supremum is over $x \sim d_0$. What does this mean? That the sup is taken over all $x$ in the support of $d_0$? This should be clarified (the same notation is used elsewhere as well).
	* Remark 3.2: The final sentence in this remark is not justified by Theorem 3.1. Simply showing a lower bound that is smaller than the lower bound for the standard contextual bandit does not imply that the true sample complexity is lower as the lower bound may just be loose—an upper bound is required to show this (which at this point in the paper has not been stated). Therefore, I would suggest removing this sentence.

Meta-review: This work obtains a tight bound on the sample complexity of RLHF with KL regularization where they show that the sample complexity scales with 1/epsilon instead of 1/epsilon^2. Unfortunately a couple of weaknesses were pointed out by the reviewers including the strength of the coverage condition, as well as several inconsistencies in the technical results. We encourage the authors to revise their promising manuscript to fix these issues.

Decision: Reject