Exploratory Preference Optimization: Harnessing Implicit Q*-Approximation for Sample-Efficient RLHF

Introduction

Reinforcement learning from human feedback (RLHF) is a central tool to align language models to human values and elicit useful behavior (Christiano et al., 2017; Bai et al., 2022; Ouyang et al., 2022). Using human-labeled preference data, RLHF achieves enhanced capabilities using a modest amount of data compared to unsupervised pre-training (on the order of tens of millions versus trillions of tokens) by treating the language model as a “policy” and optimizing it with reinforcement learning techniques.

Even though RLHF is typically only applied with preference data from humans or other language models, one might hope that it has potential to produce super-human capabilities because recognizing novel behavior and insights is typically easier than generating novel behavior. Indeed, it is often much easier to verify correctness of a given proof or program than it is to produce one from scratch. By repeatedly generating new proposals and labeling them with human feedback, a language model could gradually push beyond the boundary of human capabilities. Unfortunately, even with the great disparity in difficulty between generation and verification, a major barrier to achieving enhanced capabilities via RLHF is the volume of human feedback, i.e., sample complexity, required by existing methods. Thus, a promising research direction is to develop sample-efficient methods for RLHF.

A natural way to address the sample efficiency problem for RLHF is to augment algorithms with online exploration. Online exploration exploits interactive access to human or AI feedback by deliberately encouraging the model to produce diverse, novel responses. RLHF algorithms that exploit online feedback have received limited investigation, and in spite of encouraging initial results, existing approaches either do not update the language model (Dwaracherla et al., 2024), or engage in purely passive exploration (Guo et al., 2024; Gao et al., 2024), with no mechanism to encourage novelty or diversity. Passive exploration is intuitively insufficient, as we are unlikely to generate novel and correct proofs by chance; we make this precise in Proposition 2.1. Thus, the full potential of online exploration as a new paradigm for language model training has yet to be realized.

The central challenge in equipping language models with deliberate exploration is to efficiently navigate the vast, combinatorially large space of token sequences to find responses for which feedback will be maximally informative. The contemporary theory of reinforcement learning offers—at a conceptual level—solutions to this problem, providing algorithm design principles for exploration that can optimally take advantage of problem structure and achieve sample efficiency to the best extent one can hope for (Jiang et al., 2017; Agarwal et al., 2019; Foster and Rakhlin, 2023). However, the most powerful approaches in this space are computationally intractable in the general reinforcement learning setting (Jiang et al., 2017; Jin et al., 2021; Foster et al., 2021), and prior attempts to adapt them to RLHF either make unrealistic modeling assumptions (i.e., do not allow for general function approximation) (Xu et al., 2020; Novoseller et al., 2020; Pacchiano et al., 2021; Wu and Sun, 2023; Zhan et al., 2023b; Du et al., 2024; Das et al., 2024), or are computationally inefficient and not feasible to faithfully implement (Chen et al., 2022; Wang et al., 2023; Ye et al., 2024). Can we, perhaps by specializing to language modeling, develop practical, provable, and empirically efficient online exploration methods for RLHF?

We propose a new algorithm for online exploration in RLHF, Exploratory Preference Optimization (XPO), which is simple and practical—a one-line change to (online) Direct Preference Optimization (DPO; Rafailov et al. (2023); Guo et al. (2024))—yet enjoys the strongest known provable guarantees and promising empirical performance. XPO augments the DPO objective with a novel and principled exploration bonus, empowering the algorithm to explore outside the support of the initial model. We show that XPO is provably sample-efficient, and converges to a near-optimal language model policy under natural exploration conditions (Jin et al., 2021; Xie et al., 2023; Zhong et al., 2022). Critically, and in contrast to prior work, our theory holds irrespective of whether the initial model is sufficiently exploratory on its own. To summarize:

XPO offers the first practical and provably sample-efficient online exploration algorithm for RLHF with general function approximation.

Our design and analysis of XPO uses previously disparate techniques from language modeling and theoretical reinforcement learning, combining them in a serendipitous fashion through the perspective of KL-regularized Markov decision processes (Neu et al., 2017).

First, generalizing Rafailov et al. (2024), we observe that DPO can be viewed as implicitly performing Bellman error minimization (Xie and Jiang, 2020) to approximate the optimal value function $Q^{\star}$ in a KL-regularized MDP. We use this to provide a novel KL-regularized regret decomposition.

Then, we show that global optimism (Jiang et al., 2017; Jin et al., 2021; Xie et al., 2023), a powerful RL exploration technique that has classically been viewed as computationally intractable (Dann et al., 2018; Kane et al., 2022; Golowich et al., 2024), can be implemented in any KL-regularized MDP with deterministic transitions (generalizing language modeling) by adding a surprisingly simple exploration bonus to the DPO objective. This yields the XPO objective.

We expect our analysis techniques and perspective to be useful more broadly. In particular, the guarantees for XPO hold not just for language models, but for any reinforcement learning problem with a stochastic starting state and (potentially unknown) deterministic transition dynamics (“Deterministic Contextual MDP”).

In Section 3.4, we perform a proof-of-concept experiment to validate our theory, and find that XPO can match the performance of DPO variants (Xu et al., 2023; Tran et al., 2024; Dong et al., 2024) based on passive or heuristic exploration using significantly less preference data. These initial findings suggest that augmenting language models with online exploration may indeed lead to benefits over passive exploration.

Two concurrent and independent works posted to arXiv just before this preprint, Cen et al. (2024); Zhang et al. (2024), propose algorithms that equip DPO with exploration bonuses similar to XPO. On the theoretical side, both works are restricted to the contextual bandit formulation of RLHF, and do not consider the general reinforcement learning framework in this work or make the connection to $Q^{\star}$-approximation and KL-regularized MDPs. Compared to our results, which give provable sample complexity guarantees with general function approximation, Zhang et al. (2024) do not provide sample complexity guarantees, while Cen et al. (2024) provide guarantees only for linear contextual bandits. In addition, and importantly, the sample complexity guarantees in Cen et al. (2024) have exponential dependence on the KL regularization parameter, which our results avoid. Empirically, both works find benefits from exploration.

2 Paper Organization

Section 2 presents background on RLHF, online feedback, and the necessity of exploration. Section 3 presents our algorithm and main theoretical guarantees, including motivation behind the algorithm design and a proof sketch. Section 3.4 presents experimental results, and we conclude with discussion in Section 4. Proofs and additional results are deferred to the appendix.

Background

This section contains necessary background to present our main results. We begin by recalling the standard formulation of reinforcement learning from human feedback from offline data (Section 2.1), then introduce the online feedback model and highlight the need for systematic exploration (Section 2.2).

We study RLHF in a general reinforcement learning formulation which subsumes the token-level MDP formulation considered in prior work (Rafailov et al., 2024), but is somewhat broader.

In the context of language modeling, the main object of interest is the token-level MDP (Rafailov et al., 2024). Here, $s_{1}\sim\rho$ represents a prompt, each action $a_{h}$ represents a token (with $\mathcal{A}$ representing the vocabulary), and the state $s_{h}=(s_{1},a_{1},\ldots,a_{h-1})$ is the prompt and sequence of tokens so far. The language model is represented by a policy $\pi$, which maps the current context $s_{h}=(s_{1},a_{1},\ldots,a_{h-1})$ to a distribution over the next token $a_{h}$. The trajectory $\tau=(s_{1},a_{1}),\ldots,(s_{H},a_{H})$ produced by this process can be interpreted as the language model’s response to the prompt $s_{1}$; we will occasionally use the terms “trajectory” and “response” synonymously in this context.

Our main results apply to any Deterministic Contextual MDP (DCMDP) for which the initial state is stochastic, but the subsequent transition dynamics are deterministic and potentially unknown. This formulation encompasses but strictly generalizes the token-level MDP.

Based on the preference dataset $\mathcal{D}_{\mathsf{pref}}$, the goal is to learn a policy $\widehat{\pi}$ with high reward. Following prior theoretical works on RLHF, we consider a KL-regularized reward objective (Xiong et al., 2023; Ye et al., 2024), defined for a regularization parameter $\beta>0$, via

We aim to compute a policy $\widehat{\pi}$ such that

for some small $\varepsilon>0$. Such a guarantee means that $\widehat{\pi}$ near-optimally maximizes reward, yet stays relatively close to $\pi_{\mathsf{ref}}$ (as a function of $\beta$). The choice of $\beta>0$, which is important for safety and reliability, is typically viewed as a domain specific hyperparameter (Tang et al., 2024). Our main focus in this paper is the small-$\beta$ regime, which allows $\widehat{\pi}$ to meaningfully deviate from $\pi_{\mathsf{ref}}$ and generate potentially novel responses. Notably, by taking $\beta$ sufficiently small, it is possible to translate suboptimality bounds for the regularized reward into bounds for the unregularized reward (e.g., Zhu et al., 2023; Zhan et al., 2023a).

We refer to this setting as offline RLHF because the algorithm relies only on the offline dataset $\mathcal{D}_{\mathsf{pref}}$ for training, and does not perform any active data collection.

Initial approaches to offline RLHF (Christiano et al., 2017; Ouyang et al., 2022) proceed by first estimating a reward function $\widehat{r}$ from $\mathcal{D}_{\mathsf{pref}}$ using the Bradley-Terry model, then optimizing an estimated version of the KL-regularized objective in Eq. 2 using policy optimization methods like PPO, i.e.,

The starting point for our work is an alternative approach introduced by Rafailov et al. (2023), Direct Preference Optimization (DPO). DPO is motivated by a closed-form solution for the policy that optimizes the KL-regularized objective in Eq. 2, and condenses the two-step process above into a single policy optimization objective, removing the need for reward function estimation. Concretely, DPO solvesWe adopt the convention that the value of the DPO objective is $+\infty$ if $\pi$ does not satisfy $\pi\ll{}\pi_{\mathsf{ref}}$.

for a user-specified policy class $\Pi$, where $\sigma(x)\vcentcolon={}\frac{\exp(x)}{1+\exp(x)}$ is the sigmoid function.

2 Online Feedback and Exploration in RLHF

DPO and other offline RLHF methods have achieved great success in language model alignment, but are fundamentally limited to behaviors that are well-supported by the initial model $\pi_{\mathsf{ref}}$ and preference data $\mathcal{D}_{\mathsf{pref}}$. RLHF with online feedback offers a promising approach to move beyond this limitation by collecting feedback from responses sampled from the model during training (Guo et al., 2024).

3 The Necessity of Deliberate Exploration

Existing approaches to online RLHF adapt offline techniques by applying them iteratively. As an example, Online DPO (Guo et al., 2024) proceeds as follows:The closely related Iterative DPO approach (Xu et al., 2023; Tran et al., 2024) proceeds in the same fashion, but samples a large batch of preference pairs from each policy $\pi^{{\scriptscriptstyle(t)}}$ instead of a single pair, and performs fewer updates.

Compute $\pi^{{\scriptscriptstyle(t)}}$ by solving the DPO objective in Eq. 4 with the current preference dataset $\mathcal{D}_{\mathsf{pref}}^{{\scriptscriptstyle(t)}}$.

Sample $\tau^{{\scriptscriptstyle(t)}},\widetilde{\tau}^{{\scriptscriptstyle(t)}}\sim{}\pi^{{\scriptscriptstyle(t)}}\mid{}s_{1}^{{\scriptscriptstyle(t)}}$, then label as $(\tau_{+}^{{\scriptscriptstyle(t)}},\tau_{-}^{{\scriptscriptstyle(t)}})$ and update $\mathcal{D}_{\mathsf{pref}}^{{\scriptscriptstyle(t+1)}}\leftarrow\mathcal{D}_{\mathsf{pref}}^{{\scriptscriptstyle(t)}}\cup\{(\tau_{+}^{{\scriptscriptstyle(t)}},\tau_{-}^{{\scriptscriptstyle(t)}})\}$.

We refer to such an approach as passive exploration, as the responses are sampled directly from the policy $\pi^{{\scriptscriptstyle(t)}}$ without an explicit mechanism to encourage diversity. The following proposition shows that passive exploration is insufficient to discover novel behavior: Unless the initial policy $\pi_{\mathsf{ref}}$ has good coverage, Online DPO can fail to learn a near-optimal policy.

Fix $\beta\in(0,\tfrac{1}{8}\log(2))$, and consider the bandit setting ($H=1$, $\mathcal{S}=\varnothing$, and $\left\lvert\mathcal{A}\right\rvert=2$). There exists a reference policy $\pi_{\mathsf{ref}}$ such that for all $T\leq{}\tfrac{1}{2}\exp(\frac{1}{8\beta})$, with constant probability, all of the policies $\pi^{{\scriptscriptstyle(1)}},\ldots,\pi^{{\scriptscriptstyle(T+1)}}$ produced by Online DPO satisfy

That is, the sample complexity required by Online DPO is exponential in $\frac{1}{\beta}$, which is unacceptable in the small-$\beta$ regime; inspecting the proof, it is straightforward to see that the same conclusion holds for Iterative DPO and purely offline DPO. The idea behind Proposition 2.1 is simple: If $\pi_{\mathsf{ref}}$ places small probability mass on the optimal action, Online DPO may fail to ever explore this action until the number of iterations is exponentially large. This reflects the intuition that in the small-$\beta$ regime, more deliberate exploration is required to discover behaviors or capabilities not already covered by $\pi_{\mathsf{ref}}$.

Various empirical works have suggested that offline DPO can under-perform relative to vanilla RLHF with PPO due to a lack of on-policy sampling (Xiong et al., 2023; Guo et al., 2024; Dong et al., 2024; Tang et al., 2024). Proposition 2.1 highlights a conceptually distinct phenomenon, where both of the aforementioned algorithms (as well as online variants of DPO) fail due to poor coverage from $\pi_{\mathsf{ref}}$, in spite of on-policy sampling.

Exploratory Preference Optimization

We now present our main algorithm XPO, which addresses the limitations of existing alignment methods by augmenting DPO with active exploration. We first describe the algorithm and motivation (Section 3.1), then present theoretical guarantees (Section 3.2), and sketch the analysis (Section 3.3).

XPO (Exploratory Preference Optimization) is displayed in Algorithm 1. The algorithm takes as input a user-specified policy class $\Pi$ and proceeds in almost the same fashion as Online DPO. For each step $t\in[T]$, given the current policy $\pi^{{\scriptscriptstyle(t)}}$ and an initial state $s_{1}^{{\scriptscriptstyle(t)}}$, the algorithm begins by sampling a pair of trajectories $\tau^{{\scriptscriptstyle(t)}}\sim\pi^{{\scriptscriptstyle(t)}}\mid{}s_{1}^{{\scriptscriptstyle(t)}}$ and $\widetilde{\tau}^{{\scriptscriptstyle(t)}}\sim\pi_{\mathsf{ref}}\mid{}s_{1}^{{\scriptscriptstyle(t)}}$, which are labeled as $(\tau_{+}^{{\scriptscriptstyle(t)}},\tau_{-}^{{\scriptscriptstyle(t)}})$ based on the preference feedback and used to update the preference dataset via $\mathcal{D}_{\mathsf{pref}}^{{\scriptscriptstyle(t+1)}}\leftarrow\mathcal{D}_{\mathsf{pref}}^{{\scriptscriptstyle(t)}}\cup\{(\tau_{+}^{{\scriptscriptstyle(t)}},\tau_{-}^{{\scriptscriptstyle(t)}})\}$. The most important step is 7, which updates the policy to $\pi^{{\scriptscriptstyle(t+1)}}$ via the following optimistic variant of the DPO objective:

Here, $\alpha\geq{}0$ is an optimism parameter; for $\alpha=0$, the algorithm nearly equivalent to Online DPO, except that we sample $\tau^{{\scriptscriptstyle(t)}}\sim\pi^{{\scriptscriptstyle(t)}}\mid{}s_{1}^{{\scriptscriptstyle(t)}}$ and $\widetilde{\tau}^{{\scriptscriptstyle(t)}}\sim\pi_{\mathsf{ref}}\mid{}s_{1}^{{\scriptscriptstyle(t)}}$ instead of sampling $(\tau^{{\scriptscriptstyle(t)}},\widetilde{\tau}^{{\scriptscriptstyle(t)}})\sim\pi^{{\scriptscriptstyle(t)}}\mid{}s_{1}^{{\scriptscriptstyle(t)}}$ at each iteration. As we will see now, for $\alpha>0$, the term

in Eq. 6 encourages the policy to behave optimistically, and produce diverse responses $\tau$.

Optimism in the face of uncertainty is a widely used technique in reinforcement learning theory (Agarwal et al., 2019; Lattimore and Szepesvári, 2020; Foster and Rakhlin, 2023). In its most standard form, the optimism principle is usually stated as follows: One should explore by choosing their actions according to the most optimistic view of the world, given all of the data that has already been observed. The idea is that if we choose a decision according to this principle, one of two good things can happen: (i) the optimistic view is correct, and we receive large reward; or (ii) the optimistic view is incorrect, but we receive useful information that will help to better estimate the state of the world in subsequent iterations.

In other words, $\pi^{\star}_{\beta}$ implements an accurate internal reward model. From this viewpoint:

The standard DPO term in Eq. 6 encourages the policy $\pi$ to build an accurate internal model for rewards under the Bradley-Terry model; this can be viewed as a form of implicit $Q^{\star}$-approximation, since we are implicitly minimizing the Bellman errors in Eq. 8.

In light of Equation 9 it is natural to approximate $V_{\beta}^{\pi}(s_{1})$, the regularized value function for $\pi$, by $r(\tau)-\beta\log\frac{\pi(\tau)}{\pi_{\mathsf{ref}}(\tau)}$. Using this approximation, the first term in Equation 6 biases the policy toward a large value function such that $V^{\star}_{\beta}\lesssim V^{\pi}_{\beta}$, implementing implicit (global) optimism in the face of uncertainty (up to an inconsequential difference in on-policy rewards). The fact that this suffices to drive exploration is quite subtle, and leverages non-trivial properties of the KL-regularized MDP, including the fact that Eq. 8 holds on a per-trajectory basis.

As remarked above, another difference between XPO and online/iterative DPO is that instead of sampling the preference pairs via $(\tau^{{\scriptscriptstyle(t)}},\widetilde{\tau}^{{\scriptscriptstyle(t)}})\sim\pi^{{\scriptscriptstyle(t)}}$, we sample $\tau^{{\scriptscriptstyle(t)}}\sim\pi^{{\scriptscriptstyle(t)}}\mid{}s_{1}^{{\scriptscriptstyle(t)}}$ and $\widetilde{\tau}^{{\scriptscriptstyle(t)}}\sim\pi_{\mathsf{ref}}\mid{}s_{1}^{{\scriptscriptstyle(t)}}$. This small change is important: it is possible to show that in general, sampling $(\tau^{{\scriptscriptstyle(t)}},\widetilde{\tau}^{{\scriptscriptstyle(t)}})\sim\pi^{{\scriptscriptstyle(t)}}$ can lead to degenerate behavior in which the algorithm fails to adequately explore in the small-$\beta$ regime, even when $\pi_{\mathsf{ref}}$ itself has good coverage.

While we use $\widetilde{\tau}^{{\scriptscriptstyle(t)}}\sim\pi_{\mathsf{ref}}\mid{}s_{1}^{{\scriptscriptstyle(t)}}$ in Algorithm 1, XPO is significantly more general, and leads to provable guarantees for any fixed sampling policy $\widetilde{\tau}^{{\scriptscriptstyle(t)}}\sim\widetilde{\pi}\mid{}s_{1}^{{\scriptscriptstyle(t)}}$, as well as certain data-dependent sampling schemes (e.g., sampling $\widetilde{\tau}^{{\scriptscriptstyle(t)}}\sim{\sf unif}(\pi^{{\scriptscriptstyle(1)}},\ldots,\pi^{{\scriptscriptstyle(t)}})\mid{}s_{1}^{{\scriptscriptstyle(t)}}$); different choices may have different tradeoffs and benefits in practice. A general version of XPO which leaves the sampling distribution for $\widetilde{\tau}^{{\scriptscriptstyle(t)}}$ as a free parameter is given in Section C.1 (Algorithm 2).

XPO is highly practical, and can easily be incorporated into existing language modeling and RLHF pipelines as a drop-in replacement for Online DPO (a one-line change to existing code). The theoretical guarantees for the algorithm continue to hold under standard modifications such as (i) incorporating additional preference data from $\pi_{\mathsf{ref}}$ or another reference policy; and (ii) performing a smaller number of iterations, but collecting a larger batch of preference data from $\pi^{{\scriptscriptstyle(t)}}$ (as in Iterative DPO).

2 Theoretical Guarantees

To provide sample complexity guarantees for XPO, we make some standard statistical assumptions. The first assumption asserts that the policy class $\Pi$ is powerful enough to represent the optimal KL-regularized policy.

The policy class $\Pi$ satisfies $\pi^{\star}_{\beta}\in\Pi$.

Policy realizability is a minimal assumption for sample-efficient reinforcement learning (Agarwal et al., 2019; Lattimore and Szepesvári, 2020; Foster and Rakhlin, 2023); through Eq. 9, it is equivalent to a form of reward/value realizability. For language modeling, $\Pi$ will typically correspond to a class of language models with fixed architecture but variable weights. Next, we make a regularity assumption on the policies in $\Pi$ (Rosset et al., 2024).

For all $\pi\in\Pi$ and trajectories $\tau=(s_{1},a_{1}),\ldots,(s_{H},a_{H})$,

Note that $V_{\mathsf{max}}$ is measurable and controllable in practice; our guarantees scale polynomially with this parameter. For log-linear policies where $\pi(a\mid{}s)\propto\exp\left(\nicefrac{{f(s,a)}}{{\beta}}\right)$, we expect $V_{\mathsf{max}}\lesssim{}R_{\mathsf{max}}$.

The trajectory-level coverability coefficient is given by

Assumption 3.2 implies a trivial bound of C_{\mathsf{cov}}(\Pi)\lesssim{}\exp\big{(}\frac{V_{\mathsf{max}}}{\beta}\big{)}. Indeed, $C_{\mathsf{cov}}(\Pi)$ measures coverage with respect to the best possible distribution $\mu$, while the bound implied by Assumption 3.2 takes $\mu=\pi_{\mathsf{ref}}$, so we expect $C_{\mathsf{cov}}(\Pi)\ll\exp(V_{\mathsf{max}}/\beta)$ when $\pi_{\mathsf{ref}}$ does not provide adequate coverage on its own (e.g., the example in Proposition 2.1). This is precisely the setting where we expect deliberate exploration to be helpful. We also note that there is a trivial bound $C_{\mathsf{cov}}(\Pi)\leq|\mathcal{A}|^{H}$, but because coverability depends on the structure of the (restricted) class $\Pi$, the value can be significantly smaller in general (e.g., if policies $\pi\in\Pi$ are highly correlated or stochastic).

The main sample complexity guarantee for XPO is as follows.

Let us discuss some key features of this result.

Theorem 3.1 shows that XPO converges to a near-optimal policy with sample complexity polynomial in the coverability coefficient $C_{\mathsf{cov}}(\Pi)$; in particular, to learn an $\varepsilon$-optimal policy $T=\widetilde{O}\left(\frac{C_{\mathsf{cov}}(\Pi)\log\lvert\Pi\rvert}{\varepsilon^{2}}\right)$ episodes are required.We state the result for finite classes ($\log\lvert\Pi\rvert<\infty$) to simplify presentation, following the standard in RL theory (Agarwal et al., 2019; Foster and Rakhlin, 2023); the result readily extends to infinite classes through standard arguments. By scaling with $C_{\mathsf{cov}}(\Pi)$, Theorem 3.1 can be viewed as a strict improvement over offline RLHF (Zhu et al., 2023; Zhan et al., 2023a), as well as prior works on online RLHF that rely on passive exploration (Xiong et al., 2023; Gao et al., 2024; Chang et al., 2024). In particular, these works scale with coverage parameters for $\pi_{\mathsf{ref}}$, the simplest of which take the form $C_{\mathsf{conc}}(\Pi)\vcentcolon={}\sup_{\tau\in(\mathcal{S}\times\mathcal{A})^{H}}\sup_{\pi\in\Pi}\frac{\pi(\tau)}{\pi_{\mathsf{ref}}(\tau)}$. Under Assumption 3.2, we have that $C_{\mathsf{conc}}(\Pi)=\exp(V_{\mathsf{max}}/\beta)$ which, as discussed above, upper bounds $C_{\mathsf{cov}}(\Pi)$ but can be much larger when $\pi_{\mathsf{ref}}$ has poor coverage. The dependence on $C_{\mathsf{cov}}(\Pi)$ in Theorem 3.1 reflects the fact that XPO can explore responses not covered by $\pi_{\mathsf{ref}}$. Many works consider more general notions of coverage that account for reward function structure, in the same vein as SEC, as well as single-policy variants; both can be problematic for similar reasons.

In Appendix C, we give a generalization of Theorem 3.1 (Theorem 3.1′) which scales with a more comprehensive exploration parameter, the Sequential Extrapolation Coefficient (SEC), matching (for DCMDPs) the most general results in prior work on exploration in RLHF, but with a significantly simpler algorithm (Chen et al., 2022; Wang et al., 2023; Ye et al., 2024). The SEC also leads to polynomial sample complexity for tabular and linear MDPs, a common setting considered in prior work (Xu et al., 2020; Novoseller et al., 2020; Pacchiano et al., 2021; Wu and Sun, 2023; Zhan et al., 2023b; Das et al., 2024). See Appendix A for a detailed comparison. We emphasize that Theorem 3.1 applies to any DCMDP (including but not limited to the token-level MDP), even if the dynamics are unknown; as such, the result meaningfully extends beyond the contextual bandit formulation of RLHF found in many prior works (Zhu et al., 2023; Xiong et al., 2023; Das et al., 2024; Ye et al., 2024).

By avoiding explicit dependence on $\exp(\frac{1}{\beta})$, XPO provably improves upon Online DPO when $\beta$ is small; per Proposition 2.1, the latter must pay $\exp(\frac{1}{\beta})$ even when $C_{\mathsf{cov}}(\Pi)\leq{}2$. This improvement stems from the fact that KL-regularization does not automatically lead to exploration or grant meaningful control of coverability in the small-$\beta$ regime.

Most prior approaches to RL with general function approximation that incorporate global forms of optimism similar to Eq. 7 (Jiang et al., 2017; Sun et al., 2019; Du et al., 2021; Jin et al., 2021; Xie et al., 2023; Liu et al., 2024) are known to be computationally intractable to implement in general (Dann et al., 2018), and involve solving non-convex, non-differentiable constrained optimization problems. Thus, it is natural to ask why our result is not too good to be true. The answer is that even though the objective in Eq. 6 is simple, it is still non-convex in general, even if one employs log-linear policies of the form

Separately, we mention in passing that we believe it should be possible to derive tighter sample complexity bounds for large $\beta>0$, in the vein of Tiapkin et al. (2023a).

Our results are limited to MDPs with deterministic dynamics and stochastic start state (DCMDPs). We believe that without further modifications, the DPO objective is not suitable for stochastic dynamics, as Eq. 9 no longer holds on a per-trajectory basis.

A related point concerns trajectory coverability. In the standard (as opposed to preference-based) RL setting, it is possible to achieve guarantees that scale with state-action coverability (Xie et al., 2023), defined via:

3 Proof Sketch for \crtcrefthm:main

Our starting point for the proof of Theorem 3.1 is the following regret decomposition, which is proven as a consequence of the implicit $Q^{\star}$-approximation result in Eq. 9.

For any pair of policies $\pi$ and $\nu$, it holds that

This result decomposes the error of any policy into two pairs of terms: The first pair in Eq. 13 measures the extent to which the policy’s internal reward model overestimates the optimal value, and directly informs the notion of optimism in XPO, while the second pair in Eq. 14 measures the reward model’s predictive accuracy. Critically, as a consequence of the fact that Eq. 9 holds uniformly for all trajectories, the regret decomposition measures error under (i) the policy $\pi$ itself (on-policy error), and (ii) an arbitrary reference policy $\nu$, which we will instantiate as the historical data distribution.

Let $\boldsymbol{\mu}^{{\scriptscriptstyle(t)}}\vcentcolon={}\frac{1}{t-1}\sum_{i<t}\pi^{{\scriptscriptstyle(i)}}\otimes\pi_{\mathsf{ref}}$ denote the policy that, given $s_{1}$, samples $\tau\sim{}\pi^{{\scriptscriptstyle(i)}}$ for $i\sim{\sf unif}([t-1])$ and samples $\widetilde{\tau}\sim\pi_{\mathsf{ref}}$, with the convention that $\boldsymbol{\mu}^{{\scriptscriptstyle(1)}}$ is arbitrary. Observe that $\min_{t\in[T+1]}J_{\beta}(\pi^{\star}_{\beta})-J_{\beta}(\pi^{{\scriptscriptstyle(t)}})\leq{}\frac{1}{T}\sum_{t=1}^{T}J_{\beta}(\pi^{\star}_{\beta})-J_{\beta}(\pi^{{\scriptscriptstyle(t)}})$. For each step $t$, applying Lemma 3.1 with $\pi=\pi^{{\scriptscriptstyle(t)}}$ and $\nu=\pi_{\mathsf{ref}}$ gives

The reward estimation error term in Eq. 15 samples $\tau\sim{}\pi^{{\scriptscriptstyle(t)}}\mid{}s_{1}$ and $\widetilde{\tau}\sim{}\pi_{\mathsf{ref}}\sim{}s_{1}$ (on-policy). To relate this to the purely off-policy objective in 7 of XPO, we use a potential argument based on coverability (Xie et al., 2023) which, for any $\alpha>0$, allows us to bound the above expression by

The XPO objective in 7 minimizes an empirical analogue of this quantity (up to a standard translation between log-loss and square loss under the Bradley-Terry model), so a concentration argument (Lemma D.5) allows us to conclude that the iterates of XPO satisfy $\Psi_{\texttt{XPO}}^{{\scriptscriptstyle(t)}}(\pi^{{\scriptscriptstyle(t)}})\lesssim{}\alpha^{-1}\frac{\log\lvert\Pi\rvert}{t}+\sqrt{\frac{\log\lvert\Pi\rvert}{t}}$ with high probability. Plugging this bound into Eq. 16 yields

4 Empirical Validation

To close this section, we provide a preliminary empirical evaluation of XPO in real-world RLHF experiments. To implement XPO, we use the iterative DPO (Xu et al., 2023; Tran et al., 2024; Dong et al., 2024) pipeline from Dong et al. (2024) with 3 total iterations (that is, we set $T=3$, but draw a large batch of pairs from $\pi^{{\scriptscriptstyle(t)}}$), and augment the DPO objective with the optimism term in XPO. We use the same base model (which we refer to as Llama-3-8B-Flow-SFT),https://huggingface.co/RLHFlow/LLaMA3-SFT. prompt sets for each iteration,https://huggingface.co/datasets/RLHFlow/iterative-prompt-v1-iter1-20K, https://huggingface.co/datasets/RLHFlow/iterative-prompt-v1-iter2-20K, https://huggingface.co/datasets/RLHFlow/iterative-prompt-v1-iter3-20K. and preference model (to generate preference feedback),https://huggingface.co/RLHFlow/pair-preference-model-LLaMA3-8B. as Dong et al. (2024), which makes our results generally comparable to theirs. Over all three iterations, we fix $\pi_{\mathsf{ref}}$ to be the base model, Llama-3-8B-Flow-SFT.

In Table 1, we compare XPO with the following baselines: 1) iterative DPO with the same setup (i.e., XPO with $\alpha=0$), 2) Llama-3-8B-Flow-Final, the final model from Dong et al. (2024), and 3) the industry-level instruction-tuned model Llama-3-8B-it,https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct on various academic and chat benchmarks (Zhong et al., 2023; Nie et al., 2020; Hendrycks et al., 2021; Rein et al., 2023; Cobbe et al., 2021; Dubois et al., 2024; Li et al., 2024; Clark et al., 2018; Lin et al., 2022; Zellers et al., 2019; Sakaguchi et al., 2021). We compute all baseline numbers ourselves with the same configuration for a fair comparison. A key distinction between our experimental setup and that of Dong et al. (2024) is that we construct preference pairs from only two responses, whereas Dong et al. (2024) use best/worst-over-8-responses for preference pair construction as a heuristic exploration strategy. In other words, the final models we obtain (XPO-iter3, and a baseline, DPO-iter3 in Table 1) use only $\nicefrac{{1}}{{4}}$ the number of generated responses compared the final model (Llama-3-8B-Flow-Final)https://huggingface.co/RLHFlow/LLaMA3-iterative-DPO-final from Dong et al. (2024).

We find that the model obtained by XPO improves over the non-exploratory baseline (DPO-iter) on the chat benchmarks (which offer roughly $\sim$90% agreement and/or Spearman correlation to Chatbot Arena (Chiang et al., 2024)), and attains performance comparable to the industry-level (Llama-3-8B-it) or 4$\times$data-usage (Llama-3-8B-Flow-Final) models. At the same time, XPO also improves over the non-exploratory baseline on most of the academic benchmarks, again achieving comparable performance with the industry-level and 4$\times$data-usage models, and does not introduce significant performance regression in any benchmark. In contrast, we observe that the iterative DPO baseline (without exploration) causes obvious regression in the math (GSM8K) benchmark. However, we caution that conducting separate training runs with different random seeds can yield results with relatively high variance (e.g., the difference in win rates can be up to $3\%\sim 4\%$) for both chat benchmarks; due to resource limitations, we defer a more comprehensive evaluation to future work. See Appendix F for additional results.

Discussion

Our work provides the first practical and provably sample-efficient online exploration algorithm for RLHF with general function approximation, a step toward fully realizing the potential of online exploration for aligning language models. Our results also show that viewing DPO as a form of implicit $Q^{\star}$-approximation can directly inform new algorithmic interventions (e.g., implicit optimism), and offer an example of fruitful interplay between language modeling and theoretical reinforcement learning. Building on this viewpoint, an exciting direction for future work is to import the broader set of tools from the literature on reinforcement learning theory (e.g., more powerful exploration principles (Foster et al., 2021)) and harness them for language modeling and alignment; in this context, we expect our analysis techniques based on the KL-regularized MDP to find broader use.

From a reinforcement learning perspective, interesting technical directions for future work include (i) providing instance-dependent sample complexity bounds for XPO; and (ii) supporting RL settings beyond deterministic contextual MDPs. On the practical side, immediate followup directions include extending XPO to support general preference models (Munos et al., 2023; Swamy et al., 2024) or more general feedback modalities (Ethayarajh et al., 2024).

References

Appendix A Related Work

Theoretical analysis of algorithms for RLHF is becoming an active area of research. Much of this research focuses on purely offline RLHF (Zhu et al., 2023; Zhan et al., 2023a), which is complementary to our work. Many works also consider a so-called hybrid RLHF setting, where the algorithm has access to online feedback, but requires the initial policy $\pi_{\mathsf{ref}}$ to have good coverage (e.g., bounded concentrability or related quantities) (Xiong et al., 2023; Gao et al., 2024; Chang et al., 2024).To our knowledge, all prior works in this space require uniform notions of concentrability as opposed to single-policy concentrability. Gao et al. (2024) state guarantees in terms of single-policy concentrability under the assumption that certain regression errors can be bounded, but this cannot be achieved in general without further coverage or exploration-like conditions. These hybrid algorithms do not engage in systematic exploration (i.e., they explore passively), and hence cannot provide meaningful guarantees if $\pi_{\mathsf{ref}}$ does not adequately cover the optimal policy (e.g., for the setting in Proposition 2.1).

For online RLHF, the most relevant related work can be summarized as follows:

Most prior work (Xu et al., 2020; Novoseller et al., 2020; Pacchiano et al., 2021; Wu and Sun, 2023; Zhan et al., 2023b; Du et al., 2024; Das et al., 2024) gives algorithms and sample complexity guarantees for the special case of tabular or linear MDPs; these algorithms use exploration bonuses that are tailored to linear models, and are not suitable for the general function approximation setting we consider (e.g., for LLMs). Nonetheless, we obtain polynomial sample complexity guarantees for tabular and linear MDPs (Examples D.1 and D.2), though our results are restricted to deterministic dynamics (we believe that moving beyond the DPO objective is likely required to handle stochastic dynamics).

More relevant to our work is Ye et al. (2024), who give algorithms and sample complexity guarantees for online RLHF with general function approximation for the special case of contextual bandits ($H=1$). For contextual bandits, their sample complexity guarantees scale with a complexity measure, the eluder coefficient, which is equivalent to the Sequential Extrapolation Coefficient in our most general result, Theorem 3.1′. However, their exploration algorithm requires solving a rather complicated optimization problem, and it is unclear whether it is possible to implement it faithfully for language models (in particular, their experiments use an alternative, heuristic approach to exploration which is only loosely inspired by the theory).

Lastly, Chen et al. (2022); Wang et al. (2023) give guarantees for RLHF with general function approximation based on eluder dimension-like complexity measures which are incomparable to, but in some cases more general than Theorem 3.1′. However, these works require model-based function approximation (as opposed to the model-free setup we consider), and do not lead to efficient or practical algorithms when specialized to language modeling.

A difference worth highlighting between our work and some (but not all) of the works above (Zhu et al., 2023; Xiong et al., 2023; Das et al., 2024; Ye et al., 2024) is that we model RLHF as a general reinforcement learning problem as opposed to a contextual bandit problem. The problem of autoregressive sequence prediction can equivalently be formulated as RL in the token-level MDP, or as a contextual bandit problem (RL with horizon $H=1$) in which the “action space” consists of all possible token sequences. However, because our work supports general deterministic contextual MDPs (DCMDPs) with unknown dynamics and not just the token-level MDP, it is strictly more general than the contextual bandit formulation.

Recent work of Rafailov et al. (2024) shows that DPO, when applied to the token-level MDP can be viewed as estimating the KL-regularized value function $Q^{\star}_{\beta}$; their work does not consider sample complexity or online exploration. Our results extend their observation to any deterministic contextual MDP and—more importantly—show that it is possible to harness this perspective to provide provable end-to-end sample complexity guarantees.

Online exploration in RLHF has received limited exploration so far, with notable examples including Online DPO (Guo et al., 2024) and Iterative DPO (Xu et al., 2023; Tran et al., 2024; Pang et al., 2024; Mitra et al., 2024; Dong et al., 2024). As discussed in Section 2, these methods engage in purely passive exploration, meaning that sample from the current model $\pi^{{\scriptscriptstyle(t)}}$ without an explicit mechanism to encourage diverse, exploratory responses.

Dwaracherla et al. (2024) perform a dedicated empirical evaluation of active exploration for language models. However, this work does not actually train the language model, and thus cannot be viewed as a form of RLHF; instead the authors train a reward model iteratively, and use this in tandem with various active sampling schemes to accept or reject responses proposed by $\pi_{\mathsf{ref}}$. Nevertheless, the positive results achieved by Dwaracherla et al. (2024) in this limited setting are suggestive of the potential power of online exploration in RLHF. Similarly, Ye et al. (2024) perform a limited evaluation of empirical exploration schemes inspired by theoretical RL, but only report results for reward modeling benchmarks, not language modeling.

Most closely related, Xiong et al. (2023); Dong et al. (2024) perform an extensive empirical evaluation of Iterative DPO variants, and find that Iterative DPO with passive exploration can already have significant benefits over offline DPO. These works also incorporate a “best/worst-over-$n$” trick for preference pair construction, which can be viewed as a heuristic to promote exploration, but does not have provable guarantees. See Section 3.4 for further discussion.

Outside the context of language models, an active line of research provides structural complexity measures and algorithms that enable sample-efficient exploration in reinforcement learning in general settings (Russo and Van Roy, 2013; Jiang et al., 2017; Sun et al., 2019; Wang et al., 2020; Du et al., 2021; Jin et al., 2021; Foster et al., 2021; Xie et al., 2023; Foster et al., 2023; Liu et al., 2024). The techniques from this line of research that support general function approximation, while sample-efficient, are computationally intractable to implement in general (Dann et al., 2018), involving non-convex and non-differentiable constrained optimization problems. We use the unique structure of the KL-regularized MDP formulation and deterministic contextual MDP (DCMDP) to derive the exploration objective in XPO which—while still non-convex—is differentiable and directly amenable to a practical implementation with language models.

First introduced in Ziebart et al. (2008); Ziebart (2010), a number of recent works provide sample complexity guarantees for reinforcement learning in KL-regularized or entropy-regularized MDPs (Kozuno et al., 2022; Tiapkin et al., 2023b, a), mainly focusing on the special case of tabular (finite-state/action) MDPs. To the best of our knowledge, the optimistic objective in XPO is novel in this context.

Appendix B Technical Tools

Let $(X_{t})_{t\leq{T}}$ be a sequence of real-valued random variables adapted to a filtration $(\mathscr{F}_{t})_{t\leq{}T}$. If $\left\lvert X_{t}\right\rvert\leq{}R$ almost surely, then with probability at least $1-\delta$,

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

Appendix C Proof of \crtcrefthm:main

This section is organized as follows. First, in Section C.2, we present a more general version of XPO, which makes use of an arbitrary, user-specified sampling policy for the second response $\widetilde{\tau}$. Then, in Section C.2, we state a more general version of Theorem 3.1 (Theorem 3.1′), and show how it implies Theorem 3.1. Examples are then given in Section D.3.

In the remainder of the section, we prove Theorem 3.1′. We first prove a number of intermediate results:

In Section D.4, we state preliminaries regarding the KL-regularized MDP, and use them to prove the implicit $Q^{\star}$-approximation lemma (Lemma D.3).

In Section D.5, we prove the central regret decomposition lemma (Lemma 3.1).

In Section D.6, we prove a key concentration result used within Theorem 3.1′.

Finally, in Section D.7, we prove Theorem 3.1′, with proofs for supporting lemmas deferred to Section D.8.

Algorithm 2 presents a general version of XPO. The algorithm is identical to Algorithm 1, except that it makes use of an arbitrary, user-specified user-specified sampling policy for the second response $\widetilde{\tau}$.

In more detail, the algorithm takes as input a sampling strategy $\boldsymbol{\pi}_{\mathsf{samp}}$ which, at step $t$, computes a sampling policy $\widetilde{\pi}^{{\scriptscriptstyle(t)}}$ via $\widetilde{\pi}^{{\scriptscriptstyle(t)}}\leftarrow\boldsymbol{\pi}_{\mathsf{samp}}(\pi^{{\scriptscriptstyle(1)}},\ldots,\pi^{{\scriptscriptstyle(T)}})$. The algorithm then samples the response pair $(\tau^{{\scriptscriptstyle(t)}},\widetilde{\tau}^{{\scriptscriptstyle(t)}})$ via $\tau^{{\scriptscriptstyle(t)}}\sim\pi^{{\scriptscriptstyle(t)}}\mid{}s_{1}^{{\scriptscriptstyle(t)}}$ and $\widetilde{\tau}^{{\scriptscriptstyle(t)}}\sim\widetilde{\pi}^{{\scriptscriptstyle(t)}}\mid{}s_{1}^{{\scriptscriptstyle(t)}}$. Algorithm 1 is a special case of this scheme in which $\widetilde{\pi}^{{\scriptscriptstyle(t)}}=\pi_{\mathsf{ref}}$ for all $t$.

A secondary difference from Algorithm 1 is that Algorithm 2 assumes access to a dataset $\mathcal{D}_{\mathsf{opt}}^{{\scriptscriptstyle(t)}}$ consisting of $t$ responses sampled from $\widetilde{\pi}^{{\scriptscriptstyle(t)}}$, which are used to compute the optimistic term in 9. In Algorithm 1, because $\widetilde{\pi}=\pi_{\mathsf{ref}}$ is static, we can simply re-use the responses $\widetilde{\tau}^{{\scriptscriptstyle(1)}},\ldots,\widetilde{\tau}^{{\scriptscriptstyle(t)}}$ for this task, setting $\mathcal{D}_{\mathsf{opt}}^{{\scriptscriptstyle(t)}}=\left\{\widetilde{\tau}^{{\scriptscriptstyle(1)}},\ldots,\widetilde{\tau}^{{\scriptscriptstyle(t)}}\right\}$. However, for general time-varying sampling scheme, it may be necessary to draw a fresh dataset of responses from $\widetilde{\pi}^{{\scriptscriptstyle(t)}}$ to compute $\mathcal{D}_{\mathsf{opt}}^{{\scriptscriptstyle(t)}}$.

As a practical example, Algorithm 3—displayed below—instantiates the general scheme in Algorithm 2 by setting $\widetilde{\pi}^{{\scriptscriptstyle(t)}}={\sf unif}(\pi^{{\scriptscriptstyle(1)}},\ldots,\pi^{{\scriptscriptstyle(t)}})$ to sample from the historical data distribution at step $t$. For this scheme, it suffices to set $\mathcal{D}_{\mathsf{opt}}^{{\scriptscriptstyle(t)}}=\left\{\tau^{{\scriptscriptstyle(1)}},\ldots,\tau^{{\scriptscriptstyle(t)}}\right\}$, re-using the responses sampled from $\pi^{{\scriptscriptstyle(1)}},\ldots,\pi^{{\scriptscriptstyle(t)}}$.

C.2 General Version of \crtcrefthm:main

Our most general sample complexity guarantee for XPO (Algorithm 1 and Algorithm 2), Theorem 3.1′, is stated in terms of the following preference-based analogue of the Sequential Extrapolation Coefficient (SEC) from Xie et al. (2023) (also known as an eluder coefficient or decoupling coefficient (Zhong et al., 2022; Ye et al., 2024)). Recall that for a trajectory $\tau=(s_{1},a_{1}),\ldots,(s_{H},a_{H})$, we define

For a pair of policies $\pi$ and $\widetilde{\pi}$, we define $\pi\otimes\widetilde{\pi}$ as the joint policy that, given $s_{1}$, samples $\tau\sim{}\pi\mid{}s_{1}$ and $\widetilde{\tau}\sim\widetilde{\pi}\mid{}s_{1}$. We write $(\tau,\widetilde{\tau})\sim{}\pi\otimes\widetilde{\pi}\mid{}s_{1}$ as shorthand for this process.

For a policy class $\Pi$, sampling strategy $\boldsymbol{\pi}_{\mathsf{samp}}$, and entropy regularization parameter $\beta>0$, we define the Sequential Extrapolation Coefficient via

where $\widetilde{\pi}^{{\scriptscriptstyle(t)}}=\boldsymbol{\pi}_{\mathsf{samp}}(\pi^{{\scriptscriptstyle(1)}},\ldots,\pi^{{\scriptscriptstyle(t)}})$, and where we define $\boldsymbol{\mu}^{{\scriptscriptstyle(t)}}\vcentcolon={}\frac{1}{t-1}\sum_{i<t}\pi^{{\scriptscriptstyle(i)}}\otimes\widetilde{\pi}^{{\scriptscriptstyle(i)}}$, with the convention that $\mu^{{\scriptscriptstyle(1)}}$ is arbitrary.

Note that for Algorithm 1, which sets $\widetilde{\pi}^{{\scriptscriptstyle(t)}}=\pi_{\mathsf{ref}}$ for all $t$, we can simplify the definition above to

where $\boldsymbol{\mu}^{{\scriptscriptstyle(t)}}\vcentcolon={}\frac{1}{t-1}\sum_{i<t}\pi^{{\scriptscriptstyle(i)}}\otimes\pi_{\mathsf{ref}}$.

Our general sample complexity guarantee is as follows.

As a special case, if we set $\alpha=c\cdot{}\frac{\beta}{(V_{\mathsf{max}}+R_{\mathsf{max}})e^{2R_{\mathsf{max}}}}\cdot\sqrt{\frac{\log(\lvert\Pi\rvert T\delta^{-1})\log(T)}{T\cdot{}\mathsf{SEC_{RLHF}}(\Pi,T,\beta;\pi_{\mathsf{ref}})}}$ for an absolute constant $c>0$, then Algorithm 1 ensures that with probability at least $1-\delta$,

The following result shows that the SEC is always bounded by the coverability coefficient in Definition 3.1.

Theorem 3.1 follows immediately by combining Theorem 3.1′ with Lemma D.1.

D.3 Additional Examples for \crtcrefthm:main_general

for a given value function class $\mathcal{F}\subseteq(\mathcal{S}\times\mathcal{A}\to R_{\mathsf{max}})$. Note that for such a class, we can take $V_{\mathsf{max}}\leq R_{\mathsf{max}}$, and that $Q^{\star}_{\beta}\in\mathcal{F}$ implies that $\pi^{\star}_{\beta}\in\Pi_{\mathcal{F}}$.

The following lemma bounds the SEC for log-linear policy classes in terms of a preference-based analogue of the value function SEC in Xie et al. (2023).

For any value function class $\mathcal{F}\subseteq(\mathcal{S}\times\mathcal{A}\to R_{\mathsf{max}})$, we have that $\mathsf{SEC_{RLHF}}(\Pi,T,\beta;\boldsymbol{\pi}_{\mathsf{samp}})\leq\mathsf{SEC_{RLHF}}(\mathcal{F},T;\boldsymbol{\pi}_{\mathsf{samp}})$, where

where $\pi^{{\scriptscriptstyle(t)}}\vcentcolon={}\pi_{f^{{\scriptscriptstyle(t)}}}$, $\widetilde{\pi}^{{\scriptscriptstyle(t)}}=\boldsymbol{\pi}_{\mathsf{samp}}(\pi^{{\scriptscriptstyle(1)}},\ldots,\pi^{{\scriptscriptstyle(t)}})$, and $\boldsymbol{\mu}^{{\scriptscriptstyle(t)}}\vcentcolon={}\frac{1}{t-1}\sum_{i<t}\pi^{{\scriptscriptstyle(i)}}\otimes\widetilde{\pi}^{{\scriptscriptstyle(i)}}$ (with the convention that $\mu^{{\scriptscriptstyle(1)}}$ is arbitrary), and where $\mathcal{T}_{\beta}$ is the KL-regularized Bellman operator defined in Section D.4.

Proof of Lemma D.2. This is an immediate corollary of Lemma D.4. ∎

We first apply this bound to give a polynomial bound on the SEC in tabular DCMDPs where $\mathcal{S}$ and $\mathcal{A}$ are finite.

Suppose that $\boldsymbol{\pi}_{\mathsf{samp}}$ sets $\pi^{{\scriptscriptstyle(t)}}=\widetilde{\pi}$ for all $t$ for some fixed policy $\widetilde{\pi}$. When $\mathcal{F}=\left\{f:\mathcal{S}\times\mathcal{A}\to R_{\mathsf{max}}\right\}$ consists of all functions over tabular state and action spaces with $\lvert\mathcal{S}\rvert,\lvert\mathcal{A}\rvert<\infty$, we have $\mathsf{SEC_{RLHF}}(\mathcal{F},T;\boldsymbol{\pi}_{\mathsf{samp}})\leq\widetilde{O}(H\lvert\mathcal{S}\rvert\lvert\mathcal{A}\rvert)$ and $\log\lvert\Pi_{\mathcal{F}}\rvert\lesssim\widetilde{O}(\lvert\mathcal{S}\rvert\lvert\mathcal{A}\rvert)$. It follows that XPO (Algorithm 1) achieves

Example D.1 is a corollary of the following more general result.

In a Linear MDP (Jin et al., 2020), we have

for $B=O(\sqrt{d})$ and $R=O(R_{\mathsf{max}})$, then $\pi^{\star}_{\beta}\in\Pi_{\mathcal{F}}$, satisfying Assumption 3.1. For this setting, when $\boldsymbol{\pi}_{\mathsf{samp}}$ sets $\pi^{{\scriptscriptstyle(t)}}=\widetilde{\pi}$ for all $t$ for some fixed policy $\widetilde{\pi}$, we have $\mathsf{SEC_{RLHF}}(\mathcal{F},T;\boldsymbol{\pi}_{\mathsf{samp}})\leq\widetilde{O}(d)$ and $\log\lvert\Pi_{\mathcal{F}}\rvert\lesssim\widetilde{O}(d)$. It follows that XPO (Algorithm 1) achieves

In this section, we give some basic background on value functions and dynamic programming for the KL-regularized MDP (Ziebart et al., 2008; Ziebart, 2010), then use these properties to prove Lemmas D.3 and D.4, which show that the optimal KL-regularized policy implicitly performs models rewards and performs $Q^{\star}$-approximation.

and that the policy that obtains the maximum above is

From here, beginning with $Q^{\star}_{\beta}(s_{H},a_{H})\vcentcolon={}r(s_{H},a_{H})$, $\pi^{\star}_{\beta}(a_{H}\mid{}s_{H})=\pi_{Q^{\star}_{\beta}}(a_{H}\mid{}s_{H})$, and $V^{\star}_{\beta}(s_{H})=V_{Q^{\star}_{\beta}}(s_{H})$ for $s_{H}\in\mathcal{S}_{H}$, for each $s_{h}\in\mathcal{S}_{h}$, we can inductively define for each $h\in[H]$:

The following lemma, generalizing Rafailov et al. (2024), shows that the optimal KL-regularized policy $\pi^{\star}_{\beta}$ can be viewed as implicitly modeling rewards.

For any DCMDP, it holds that for all admissibleWe use “admissible” to a refer to a trajectory generated by executing an arbitrary policy $\pi:\mathcal{S}\to\Delta(\mathcal{A})$ in the MDP. trajectories $\tau=(s_{1},a_{1}),\ldots,(s_{H},a_{H})$,

where $V^{\star}_{\beta}$ is the KL-regularized value function defined in Eq. 26.

Proof of Lemma D.3. Let $\tau=(s_{1},a_{1}),\ldots,(s_{H},a_{H})$, and recall that for any DCMDP, all state transitions except for $s_{1}\sim\rho$ are deterministic. Then we have

where the second equality uses that $(\mathcal{T}_{\beta}f)(s_{h},a_{h})=r(s_{h},a_{h})+V_{f}(s_{h+1})$ for any admissible trajectory in a deterministic MDP, and the third equality uses the explicit form for $\pi^{\star}_{\beta}$ in terms of $V^{\star}_{\beta}$ and $Q^{\star}_{\beta}$ given in Equation 25. Rearranging yields the result. ∎

We can also prove the following, more general version of Lemma D.4.

Proof of Lemma D.4. Let $\tau=(s_{1},a_{1}),\ldots,(s_{H},a_{H})$. Then we have

where the first equality uses the definition of $V_{f}$, the second equality uses that $(\mathcal{T}_{\beta}f)(s_{h},a_{h})=r(s_{h},a_{h})+V_{f}(s_{h+1})$ for any admissible trajectory in a deterministic MDP, and the third equality uses that $\pi_{f}(a\mid{}s)=\pi_{\mathsf{ref}}(a\mid{}s)e^{\frac{f(s,a)-V_{f}(s)}{\beta}}$. Rearranging yields the result. ∎

D.5 Regret Decomposition

In this section we prove the central regret decomposition for XPO, restated below.

Proof of Lemma 3.1. It follows immediately from the definition of the KL-regularized reward that

However, since $\beta\log\frac{\pi^{\star}_{\beta}(\tau)}{\pi_{\mathsf{ref}}(\tau)}-r(\tau)=V^{\star}_{\beta}(s_{1})$ for all admissible trajectories by Lemma D.3, we have that

for all policies $\nu$, as the initial state $s_{1}$ does not depend on the policy under consideration. The result now follows by rearranging

D.6 Concentration Lemmas

Recall that we define $\boldsymbol{\mu}^{{\scriptscriptstyle(t)}}=\frac{1}{t-1}\sum_{i<t}\pi^{{\scriptscriptstyle(i)}}\otimes\widetilde{\pi}^{{\scriptscriptstyle(i)}}$. For a given policy $\pi$, define

The following lemma is our central concentration guarantee for Algorithm 1.

Suppose that Assumptions 3.2 and 3.1 hold. Then Algorithm 1 guarantees that with probability at least $1-\delta$, for all steps $t\in[T]$,

for $\kappa\vcentcolon={}(8(R_{\mathsf{max}}+V_{\mathsf{max}})e^{2R_{\mathsf{max}}})^{-2}$.

Proof of Lemma D.5. Let $t\in\left\{2,\ldots,T+1\right\}$ be fixed.

and $\widehat{B}^{{\scriptscriptstyle(t)}}(\pi)=\alpha\sum_{\tau\in\mathcal{D}_{\mathsf{opt}}^{{\scriptscriptstyle(t-1)}}}\log\pi(\tau)$. Then we can equivalently write

For a given policy $\pi$, recall that we define

Then, in light of Lemma D.3, under the Bradley-Terry model (Eq. 1), we have that for all $t$,

For any fixed $t\geq{}1$, with probability at least $1-\delta$, all $\pi\in\Pi$ satisfy

Rearranging Lemma D.6, with probability at least $1-\delta$, all $\pi\in\Pi$ satisfy

Hence, as long as $\pi^{\star}_{\beta}\in\Pi$ (Assumption 3.1), the definition of $\pi^{{\scriptscriptstyle(t)}}$ in Algorithm 2 implies that

We next appeal to another basic concentration result.

For any fixed $t\geq{}1$, with probability at least $1-\delta$, all $\pi\in\Pi$ satisfy

Combining Lemma D.7 with Eq. 32, we conclude that with probability at least $1-2\delta$,

To conclude, we further simplify the expression via

where the last inequality uses that for $x,y\geq{}0$, $(x-y)^{2}\leq{}4(x+y)(\sqrt{x}-\sqrt{y})^{2}$.

Finally, using Lemma D.3, we have $f_{\pi^{\star}_{\beta}}\in\left[-R_{\mathsf{max}},R_{\mathsf{max}}\right]$ almost surely, while $f_{\pi^{{\scriptscriptstyle(t)}}}\in\left[-V_{\mathsf{max}},V_{\mathsf{max}}\right]$ by Assumption 3.2. We appeal to the following lemma.

If $x\in\left[-X,X\right]$ and $y\in[-Y,Y]$ for $X\geq{}0$, $Y\geq{}1$, then

This proves the result after taking a union bound over all steps $t$.

Next, using Eq. 30 and a somewhat standard argument from van de Geer (2000); Zhang (2006), we calculate that

Since $D^{2}_{\mathsf{H}}\left(\cdot,\cdot\right)\leq{}2$ and $-\log(1-x)\geq{}x$ for $x\leq{}1$, we conclude that

Proof of Lemma D.7. Let $\tau^{{\scriptscriptstyle(1)}},\ldots,\tau^{{\scriptscriptstyle(t-1)}}$ denote the trajectories in $\mathcal{D}_{\mathsf{opt}}^{{\scriptscriptstyle(t-1)}}$. Let $\widehat{b}^{{\scriptscriptstyle(i)}}(\pi)=\alpha\log\pi(\tau^{{\scriptscriptstyle(i)}})$, and let

which implies that $\left\lvert Z^{{\scriptscriptstyle(i)}}(\pi)\right\rvert\leq{}2\frac{\alpha}{\beta}V_{\mathsf{max}}$. From here, the result follows immediately by applying Lemma B.1 with the sequence $(Z_{i}(\pi))$ and taking a union bound over $\pi\in\Pi$. ∎

Proof of Lemma D.8. We consider three cases. First, if $x\in\left[-2Y,2Y\right]$, then

for some $z\in\left[-2Y,2Y\right]$. In this regime, we have $\sigma^{\prime}(z)\geq{}\sigma^{\prime}(2Y)=e^{2Y}/(1+e^{2Y})^{2}\geq{}(4e^{2Y})^{-1}$. Next, if $x\geq{}2Y>0$, we can directly bound

where the last line holds whenever $Y\geq{}1$. We conclude in this case that

Finally, we consider the case where $x\leq{}-2Y\leq{}0$. In this case, we can similarly lower bound

as long as $Y\geq{}1$. From here, proceeding in the same fashion as the second case yields the result. ∎

D.7 Proof of \crtcrefthm:main_general

Proof of Theorem 3.1′. Before diving into the proof, we re-state two central technical lemmas. The first lemma, generalizing Rafailov et al. (2024), shows that the optimal KL-regularized policy $\pi^{\star}_{\beta}$ can be viewed as implicitly modeling rewards.

This lemma allows us to view the DPO objective as a form of implicit $Q^{\star}$-approximation. Building on this lemma, we prove the following regret decomposition.

This result shows that the (regularized) regret of any policy $\pi$ can be decomposed into two terms. The term in Eq. 14 measures the extent to which $\pi$ (implicitly) models the reward; by Lemma D.3, this term is zero when $\pi=\pi^{\star}_{\beta}$. Meanwhile, the term in Eq. 13 measures the extent to which the policy $\pi$ over-estimates the internal reward; we will control this term using optimism. Importantly, the regret decomposition in Lemma 3.1 holds for an arbitrary roll-in policy $\nu$. This will facilitate minimizing the terms in the regret decomposition in a data-driven fashion. Before proceeding, we remark that Lemma D.3 and Lemma 3.1 together imply that

For each step $t$, we apply Lemma 3.1 with $\pi=\pi^{{\scriptscriptstyle(t)}}$ and $\nu=\widetilde{\pi}^{{\scriptscriptstyle(t-1)}}$, which gives

Next, recall that we define $\boldsymbol{\mu}^{{\scriptscriptstyle(t)}}=\frac{1}{t-1}\sum_{i<t}\pi^{{\scriptscriptstyle(t)}}\otimes\widetilde{\pi}^{{\scriptscriptstyle(t)}}$ Consider a fixed step $t\geq{}2$, and define

Then, using the AM-GM inequality, for any $\eta>0$ we can bound

Note that by definition, we have that $\sum_{t=1}^{T}\mathcal{I}^{{\scriptscriptstyle(t)}}\leq\mathsf{SEC_{RLHF}}(\Pi,T,\beta;\boldsymbol{\pi}_{\mathsf{samp}})$. Hence, by plugging Eq. 36 into Eq. 35 and summing, we conclude that

above. Let $f_{\pi}(\tau,\widetilde{\tau})\vcentcolon={}\beta\log\frac{\pi(\tau)}{\pi_{\mathsf{ref}}(\tau)}-\beta\log\frac{\pi(\widetilde{\tau})}{\pi_{\mathsf{ref}}(\widetilde{\tau})}$. By Lemma D.3, we have that for any pair of admissible trajectories $(\tau,\widetilde{\tau})$ that share the initial state $s_{1}$, $f_{\pi^{\star}_{\beta}}(\tau,\widetilde{\tau})=r(\tau)-r(\widetilde{\tau})$, so we can rewrite Eq. 38 as

We now recall the central concentration lemma for XPO (Lemma D.5).

It follows that if we set $\eta=\frac{\beta\kappa}{\alpha{}T}\leq{}\frac{\beta\kappa}{\alpha(t-1)}$, then with probability at least $1-\delta$, for all $t\in[T]$,

Plugging this bound back into Eq. 37, we have that

where the last line uses that $\kappa\leq V_{\mathsf{max}}^{-2}$. It follows that by choosing

Finally, we note that $(V_{\mathsf{max}}+\kappa^{-1/2})=O((V_{\mathsf{max}}+R_{\mathsf{max}})e^{2R_{\mathsf{max}}})$.

D.8 Proofs for SEC Bounds

be the distribution that achieves the value of the coverability coefficient in Definition 3.1. Let us abbreviate $C_{\mathsf{cov}}\equiv C_{\mathsf{cov}}(\Pi)$. For a trajectory $\tau$, let

Letting $\mathcal{T}\vcentcolon={}(\mathcal{S}\times\mathcal{A})^{H}$, we can further bound this by

so that $\text{(I)}\leq{}32C_{\mathsf{cov}}$.

where the last inequality is by Cauchy-Schwarz. We conclude that

To proceed, we restrict our attention to the case where $\widetilde{\pi}^{{\scriptscriptstyle(t)}}=\widetilde{\pi}$ for all $t$ for some fixed $\widetilde{\pi}$. We observe that in this case, for all $t$,

since $\tau$ and $\widetilde{\tau}$ are conditionally independent given $s_{1}$, and since $d^{\pi,\pi^{\prime}}(\tau,\widetilde{\tau})=0$ if $\tau,\widetilde{\tau}$ do not share the same $s_{1}$. It follows that

Finally, by Lemma 4 of Xie et al. (2023), we have that for all $\tau\in\mathcal{T}$, $\sum_{t=1}^{T}\frac{d^{\pi^{{\scriptscriptstyle(t)}}}(\tau)}{\sum_{i<t}d^{\pi^{{\scriptscriptstyle(i)}}}(\tau)+C_{\mathsf{cov}}{}\nu(\tau)}\leq{}O(\log(T))$, which yields $\text{(II)}\leq O(C_{\mathsf{cov}}\log(T))$. This proves the result.

Proof for Example D.2. We claim for any pair of trajectories $\tau,\widetilde{\tau}$ and function $f\in\mathcal{F}$, we can write

With this definition, we observe that in the case where $\widetilde{\pi}^{{\scriptscriptstyle(t)}}=\widetilde{\pi}$ for all $t$, we can write the value of $\mathsf{SEC_{RLHF}}$ for a sequence of policies $\pi^{{\scriptscriptstyle(1)}},\ldots,\pi^{{\scriptscriptstyle(T)}}$ as

Appendix E Additional Proofs

This section contains proofs for supporting results found throughout Section 2 and Section 3.

Proof of Proposition 2.1. Consider the bandit setting where $H=1$, $\mathcal{S}=\varnothing$, and $\mathcal{A}=\{\mathfrak{a},\mathfrak{b}\}$. Let $\beta>0$ be given. We consider the reward function $r$ given by $r(\mathfrak{a})=1$ and $r(\mathfrak{b})=\frac{1}{2}$. We choose the reference model to set $\pi_{\mathsf{ref}}(\mathfrak{a})=\varepsilon$ and $\pi_{\mathsf{ref}}(\mathfrak{b})=1-\varepsilon$ for a parameter $\varepsilon\vcentcolon={}\exp(-\frac{c}{\beta})$, where $c>0$ is an absolute constant whose value will be chosen at the end of the proof. We choose $\Pi=\{\pi_{\mathsf{ref}},\pi^{\star}_{\beta}\}$, which we note satisfies Assumption 3.1 and Assumption 3.2 with $V_{\mathsf{max}}=O(1)$.

Specialized to the bandit setting, Online DPO takes the following simplified form:

Sample pair of actions $a^{{\scriptscriptstyle(t)}},\widetilde{a}^{{\scriptscriptstyle(t)}}\sim\pi^{{\scriptscriptstyle(t)}}$.

Label the actions as $(a_{+}^{{\scriptscriptstyle(t)}},a_{-}^{{\scriptscriptstyle(t)}})$ according the Bradley-Terry model:

and update $\mathcal{D}_{\mathsf{pref}}^{{\scriptscriptstyle(t+1)}}\leftarrow{}\mathcal{D}_{\mathsf{pref}}^{{\scriptscriptstyle(t)}}\cup\{(a_{+}^{{\scriptscriptstyle(t)}},a_{-}^{{\scriptscriptstyle(t)}})\}$.

Compute $\pi^{{\scriptscriptstyle(t+1)}}$ via

Our construction uses the fact that depending on the preference dataset $\mathcal{D}_{\mathsf{pref}}^{{\scriptscriptstyle(t)}}$, the minimizer in Eq. 45 may not be uniquely defined. Let $\mathcal{E}^{{\scriptscriptstyle(t)}}$ denote the event that at iteration $t$, $a^{{\scriptscriptstyle(t)}}=\widetilde{a}^{{\scriptscriptstyle(t)}}=\mathfrak{b}$. We appeal to a technical lemma.

Suppose we initialize with $\pi^{{\scriptscriptstyle(1)}}=\pi_{\mathsf{ref}}$. As long as $c\leq\frac{1}{8}$, $\varepsilon\leq{}1/2$, the following properties hold:

Whenever $\mathcal{E}^{{\scriptscriptstyle(1)}},\ldots,\mathcal{E}^{{\scriptscriptstyle(t)}}$ hold, we can choose the policy $\pi^{{\scriptscriptstyle(t+1)}}$ to satisfy $\pi^{{\scriptscriptstyle(t+1)}}=\pi_{\mathsf{ref}}$, which has

By Lemma E.1 and the union bound, we have that

as long as $T\leq{}\frac{1}{2\varepsilon}$. It follows that whenever this occurs, $\max_{\pi}J_{\beta}(\pi)-J_{\beta}(\pi^{{\scriptscriptstyle(t)}})\geq{}\frac{1}{8}$ for all $t\in[T+1]$.

Note that since online DPO selects $\pi^{{\scriptscriptstyle(t)}}=\pi_{\mathsf{ref}}$ for all $t$ in our counterexample above, this also immediately implies a lower bound for offline DPO (interpreting $\pi^{{\scriptscriptstyle(T+1)}}$ as the policy returned by offline DPO).

Proof of Lemma E.1. We prove this claim inductively. Let $t\in[T]$ be fixed, and suppose the claim holds for $1,\ldots,t-1$. If we assume $\mathcal{E}^{{\scriptscriptstyle(1)}},\ldots,\mathcal{E}^{{\scriptscriptstyle(t-1)}}$ hold, then we have $\pi^{{\scriptscriptstyle(t)}}=\pi_{\mathsf{ref}}$ inductively. In this case,

Now, for the second part of the claim, suppose that $\mathcal{E}^{{\scriptscriptstyle(1)}},\ldots,\mathcal{E}^{{\scriptscriptstyle(t+1)}}$ hold. Then for all $t^{\prime}\in[t+1]$, $a_{+}^{{\scriptscriptstyle(t^{\prime})}}=a_{-}^{{\scriptscriptstyle(t^{\prime})}}=\mathfrak{b}$, which implies that

for all $\pi\in\Pi$ such that $\pi\ll\pi_{\mathsf{ref}}$. It follows that $\pi^{{\scriptscriptstyle(t+1)}}=\pi_{\mathsf{ref}}$ is a valid minimizer for Eq. 45.

Finally, we compute that as long as $\varepsilon\leq{}1/2$ and $c\leq\frac{1}{8}$

Appendix F Experiments: Additional Results and Details

Table 2 displays the performance of XPO and the comparator models described in Section 3.4 on additional reasoning tasks. We observe that the XPO model outperforms Llama-3-8B-it, and still comparable to Llama-3-8B-Flow-Final, which uses 4x more data than XPO. The average over all academic benchmarks is 59.94 for Llama-3-8B-Flow-Final vs. 59.61 for XPO-iter3, in addition to the performance gain from XPO-iter3 in the chat benchmarks (Table 1).

The experiments were conducted on 8 x Nvidia H100 GPUs. In our implementation, we mainly follow the general version of XPO (Algorithm 2), and we pick $\widetilde{\pi}^{\scriptscriptstyle(t)}=\pi^{\scriptscriptstyle(t)}$ and $\mathcal{D}_{\mathsf{opt}}^{{\scriptscriptstyle(t)}}=\mathcal{D}_{\mathsf{pref}}^{\scriptscriptstyle(t)}$. In each iteration, we fix the base model (Llama-3-8B-Flow-SFT) as $\pi_{\mathsf{ref}}$, set $\beta=0.1$, use a global batch size of $16$, and use a learning rate of $5\times 10^{-7}$ with cosine scheduling. The $\alpha$ parameter follows the schedule $\{1\times 10^{-5},5\times 10^{-6},0\}$ for the three iterations. We clip the $\log\frac{\pi(\tau)}{\pi_{\mathsf{ref}}(\tau)}$ term for both positive and negative trajectories to $$, but only for the exploration term, in order to enhance stability. This is motivated by Assumption 3.2. The number of training epochs for each iteration is 2, and the warmup ratio is 0.03.