Title: Sharp Analysis for KL-Regularized Contextual Bandits and RLHF
Venue: NeurIPS 2025
Accepted: yes
Keywords: Contextual bandits

regularization
Abstract: Reverse-Kullback-Leibler (KL) regularization has emerged to be a predominant technique 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 of KL-regularization has been empirically demonstrated in various practical scenarios, current theoretical analyses of KL-regularized RLHF still yield the same $\mathcal{O}(1 / \epsilon^2)$ sample complexity as ones without KL-regularization. To understand the fundamental distinction between 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 $\mathcal{O}(1 / \epsilon)$ sample complexity when $\epsilon$ is sufficiently small. We also prove matching lower bounds for both settings. More specifically, we study how the coverage of the reference policy affects the sample complexity of KL-regularized online contextual bandits and RLHF. We show that with sufficient coverage from the reference policy, a simple two-stage mixed sampling algorithm can achieve an $\mathcal{O}(1 / \epsilon)$ sample complexity with only an additive dependence on the coverage coefficient, thus proving the benefits of online data even without explicit exploration. Our results provide a comprehensive understanding of the roles of KL-regularization and data coverage in online decision making, shedding light on the design of more efficient algorithms.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studied contextual bandits and RLHF problem with a KL-regularized objective function, and the novelty is mainly on the theoretical side of sharp analysis for sample complexity. The paper designed a two-stage mixed sampling algorithm to achieve a good upper bound, which is the lower bound for sample complexity for contextual bandits and the analysis can be extended to RLHF to get good results.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\nThe paper provided a solid theoretical contribution, and they also provided experimental results to verify their findings.\n\nWeaknesses:\nThe writing and notation are confusing. For example, in line 217, what is $R(\\theta_*,x,\\pi^*)$? Please see the limitations part for more details.\n\nQUESTIONS:\n1. The assumption of boundness for reward is widely used in related literature, but I didn't see the necessity of the non-negative assumption for reward. Could you clarify it further?\n2.  The authors discuss offline and online RLHF settings. But the writing is not clear, like which section or algorithm is for online and which part is for offline. For example, algorithm 1 seems like a hybrid setting. Could you discuss further?\n3. In line 216, $\\mathbb{E}_{x \\sim d_0} \\mathbb{E}_{a \\sim \\pi_o(\\cdot \\mid x)}\\left[\\left(R(\\widehat{\\theta}, x, a)-R\\left(\\theta_*, x, a\\right)\\right)^2\\right]=O(1 / n)$ is typically called generalization error. I didn't get why the authors called them in-sample error. Could you give further explanation?\n\nI really like the solid theoretical contribution of the paper. But I hope the authors have clearer writing to make the paper readable for more readers. Thus, depending on the rebuttal, I will consider changing my score.\n\nLIMITATIONS:\n1. The author mentioned the reverse KL regularization in the abstract. It will be better if the authors discuss forward KL v.s.  reverse KL regularization somewhere.\n2. I'd like to suggest that the authors give some proper citation for Definition 2.4.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides sharper sample complexity bounds for KL-regularized contextual bandits and RLHF. The authors propose a two-stage mixed sampling algorithm: in the first stage, the reward model and policy are updated using samples from a reference policy; in the second stage, the resulting intermediate policy is used to collect additional samples to estimate the optimal policy. With this approach, the algorithm achieves an O(1/e) sample complexity for small e, improving upon the standard O(1/e^2) complexity.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The paper provides a thorough theoretical analysis, including tight upper and lower bounds, and covers both global and local coverage assumptions.\n- The proposed two-stage mixed sampling algorithm is theoretically sound.\n- The algorithm is easily implementable, requiring only one additional policy update and online sampling step.\n\nWeakness:\n- The provided sample complexity is not adequately discussed or compared with prior work. For example, the number of samples required in the initial stage may explode in practice due to its dependence on the coverage coefficient, potentially limiting the algorithm’s applicability only to asymptotical analysis (with extremely small epsilon).\n- Empirical validation is limited, as the paper only compares the proposed method with an offline algorithm, where it is somewhat obvious that the online algorithm would perform better.\n- The potential limitations of the proposed algorithm are not thoroughly discussed.\n\nQUESTIONS:\n1. How does the coverage coefficient $D^2$ scale with the size of the state-action space? Could the authors provide any additional comparisons (either analytical or empirical) with existing online methods in the literature regarding sample complexity?\n2. Regarding Definition 3.5 (data coverage assumption for RLHF), the description is unclear. Could you provide additional explanation or intuition for the function $b(x)$? Is it dependent on $(\\theta, \\theta’)$?\n3. While higher KL regularization coefficients lead to sharper sample complexity bounds in suboptimality gap, it is unclear whether this can be called the \"power\" of regularization. Since increasing regularization changes the optimal policy (closer to the reference policy), isn't this just a trade-off?\n4. Minor issues:\n- Line 258: In $$\\pi: \\mathcal{X} \\to \\mathcal{A}$$, should the policy be deterministic or can it be stochastic?\n- Line 275: whether -> Whether\n- Definition 3.1 (Preference Oracle): Is the output of the mapping P {0, 1} or [0, 1]? It is unclear whether P is the sampled output or the probability.\n\nLIMITATIONS:\nThe paper does not sufficiently discuss its limitations, which should be more explicitly addressed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the RLHF problem in the contextual bandit setting, with the goal of establishing tight sample complexity bounds. The authors analyze both the RL and RLHF settings, deriving matching upper and lower bounds for each. Notably, they demonstrate that incorporating a KL constraint into the problem yields a sample complexity that scales as $1/\\epsilon$, rather than $1/\\epsilon^2$ as in prior works.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\nThe work proposes a tight sample complexity for RL/RLHF with KL-regularization. The paper is well written and all the theorems seem correct.\n\n**Weaknesses**\n\n1. Definition 2.5 is too strong. Essentially, by collect N data from $\\pi_0$, we can immediately find $\\theta$ such that $\\mathbb{E}_{x, a\\sim \\pi_0}[(R(\\theta, x, a) - R(\\theta^\\star, x, a))^2] \\le c  \\frac{\\log |\\mathcal{R}|/\\delta }{N} $ with probability $1-\\delta$. By taking $N$ as some log order term, the right hand side can be bounded by constant. In this regard, Definition 2.5 actually claim that by collecting log number of data, we can find $\\theta$ such that $R(\\theta, x, a) - R(\\theta^\\star, x, a)$ can be well bounded for all $x, a$. This is impossible in practice.\n2. The results in the paper only holds when $\\eta$ is small, as a small $\\eta$ ensures that $\\pi(a|x)/\\pi^\\star(a|x)$ can be well bounded given $R(\\theta, x, a) - R(\\theta^\\star, x, a)$ is well bounded. This makes the proposed tight bound practically meaningless, considering that $\\eta$ is usually a very small value – for example, it is typically set to 1e-3 or 1e-4 in PPO or GRPO.\n3. The analysis under RL from preference feedback is actually trivial. Given that the reward is bounded by B, it is straightforward that using exp(B) times more preference feedback is equivalent to directly obtaining reward feedback.\n\nQUESTIONS:\n1. See weakness.\n2. Could you give a more detailed comparison between the things in Appendix B and [1]?\n\n\n\n[1] Zhao, H., Ye, C., Xiong, W., Gu, Q., & Zhang, T. (2025). Logarithmic Regret for Online KL-Regularized Reinforcement Learning. arXiv preprint arXiv:2502.07460.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proves that for KL-regularized contextual bandits and RLHF, the correct sample complexity is of order $\\tilde{\\Theta}(1/\\epsilon)$, not $\\tilde{\\Theta}(1/\\epsilon^2)$ ignoring other factors. The lower bounds are proven via standard info-theoretic arguments, and the upper bounds are shown via proposing a two-stage approach. The first stage collects samples from the reference policy, and the second stage collects samples from the learned policy, where the learning happened from first stage's collected samples.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- Well-written\n\n**Weaknesses**\n- Lower bound w.r.t. policy coverage is lacking\n- References to the Appendix are lacking in the main text, e.g., Appendix D (experiments).\n- Moreover, discussions regarding the dependency of the coverage coefficient on $\\epsilon$ are lacking. Is it always negligible to claim that it is ignorable w.r.t. $1 / \\epsilon$? or can it be $\\frac{1}{\\epsilon} \\log\\frac{1}{\\epsilon}$, and so the current bound is tight up to logarithmic factors?\n\nQUESTIONS:\n1. In Rmk 3.9, the authors state that $e^B$ is probably unavoidable. Can the authors comment on whether the techniques from [1] can be adopted to remove such dependency?\n2. What parts of this paper can be easily extended to general preference learning scenarios?\n3. Related to weakness, can the authors comment on the (educated) conjectured \"correct\" dependency on the policy coverage?\n4. What is the intuition behind the fast rate of $1/\\epsilon$, in \"layman's terms\"? In learning theory, the intuition that I'm familiar with is regarding the curvature of the loss and mixability properties of the loss function. Do we have something similar here? If yes, then the authors should include such intuition in the main text.\n\n\n[1] https://arxiv.org/abs/2502.00666\n\n[2[ https://jmlr.org/papers/v16/vanerven15a.html\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}