Dataset Reset Policy Optimization for RLHF

1 Introduction

• In standard RLHF, an offline preference dataset trains a reward model, then online RL optimizes it without reusing that dataset.The introduction asks whether reusing offline data could improve learning efficiency.Source: §1 · pp. 1–2
• DR-PO assumes the optimizer can reset to any state, an assumption the authors say is naturally satisfied when fine-tuning generative models.For text generation, a reset corresponds to supplying a partial sentence and the initial prompt to the policy.Source: §1 · p. 2
• DR-PO integrates offline preference data into online RLHF by resetting policy optimization to states from the dataset.The paper reports strong performance guarantees and computational tractability compared with prior theoretical RLHF work.Source: §1 · p. 2
• When PPO is the policy optimizer, DR-PO outperforms PPO and DPO on TL;DR summarization and Anthropic HH without additional computation or memory overhead to PPO.The comparison is reported across two standard RLHF benchmarks.Source: §1 · p. 2

2 Related Work

• Several prior provable online RLHF methods require tabular MDPs or restricted function classes, while DR-PO targets general function approximation with offline preference data.The cited approaches also include posterior sampling, complicated exploration bonuses, confidence sets, or non-convex optimization.Source: §2 · p. 2
• Prior theoretical offline RLHF methods cited here rely on confidence sets and constrained maximin optimization, which the authors characterize as computationally inefficient.One cited setting learns reward from behavior rather than preference feedback, so it differs from this paper's setting.Source: §2 · p. 3
• Unlike work on expert demonstrations in tabular or linear MDPs, DR-PO considers preference-based offline data that need not come from a high-quality expert.Its dataset-reset exploration avoids optimism-based strategies, which the authors describe as computationally tractable with general function approximation.Source: §2 · p. 3
• Dataset resets are presented as complementary to empirical online RLHF algorithms and potentially applicable to methods including P3O and APA.The authors also note potential relevance to hybrid methods that add online generations to DPO.Source: §2 · p. 3
• The paper combines dataset resets with KL regularization to limit deviation from the offline data and reduce over-optimization of the learned reward model.The authors connect reward-model over-optimization to reward hacking and say the combined ideas yield strong performance in theory and practice.Source: §2 · p. 3

3 Preliminaries

• The setting is an episodic, time-inhomogeneous MDP with horizon H, state and action spaces, bounded rewards, and a common initial state.The cumulative reward along a trajectory is bounded by r_max.Source: §3 · p. 3
• A policy induces state-action and trajectory visitation distributions, while value functions represent expected cumulative reward from a state or state-action pair.The paper seeks an epsilon-optimal policy relative to a high-quality target, which need not be globally optimal.Source: §3 · pp. 3–4
• Human preferences are observed on offline trajectory pairs sampled from a given supervised-fine-tuned policy, while the true reward is unobservable.The preference probability is modeled through an increasing link function of the difference in trajectory rewards.Source: §3 · p. 4
• The Bradley-Terry-Luce model uses a sigmoid link, and the reward model is learned from preference data by maximum likelihood estimation.The link's nonlinearity is used to characterize the difficulty of learning the reward model.Source: §3 · p. 4
• An unlabeled trajectory dataset sampled from the supervised-fine-tuned policy is also available and may be larger than the human-labeled dataset.The paper specifies that this dataset is distinct from the preference-labeled trajectory-pair data.Source: §3 · p. 4
• The paper assumes the agent can reset to a chosen state and continue a rollout, an ability it identifies in applications such as language-model fine-tuning.In text generation, resetting means feeding a partial sentence to the transformer policy to generate subsequent tokens.Source: §3 · p. 4

4 Dataset Reset Policy Optimization

• After data collection, DR-PO sends the online batch to a policy-optimization oracle, which may use PG, NPG, PPO, or actor-critic updates.The PPO implementation fits a critic for advantage estimation and updates the policy using clipping.Source: §4 · pp. 4–5
• For each online episode, DR-PO samples a state from an offline trajectory and resets the current policy there instead of starting at the initial state.The resulting rollouts form the online batch used for policy optimization.Source: §4 · p. 5
• Each rollout records learned-model rewards and a log policy ratio relative to the supervised-fine-tuned policy, which can estimate a KL penalty.The penalty is used in theory and practice to limit policy deviation so the learned reward remains a good approximation under the learned policy's trajectory distribution.Source: §4 · p. 5
• The experiments use PPO as the policy optimizer, whereas the theory uses NPG; the authors describe PPO as a scalable heuristic approximation to NPG.They report that hyperparameters that work well for PPO also work for DR-PO.Source: §4 · p. 5

5 Theoretical Analysis

• The meta-algorithm separates reset-based online data collection from policy optimization, which can use an oracle such as PG, NPG, actor-critic, or PPO-style updates.The analysis instantiates that oracle with NPG; when direct reset is unavailable, a πSFT roll-in and policy roll-out can implement it, though direct reset reduces computation.Source: §5 · pp. 5–6
• The NPG oracle fits a Q-function estimate from online data, then updates the policy through KL-based Mirror Descent.The KL penalty keeps the updated policy from deviating too much from πSFT and keeps its support within πSFT’s support.Source: §5 · pp. 5–6

5.1 Theoretical Sample Complexity

• The analysis assumes realizable, bounded reward and Q-function classes; handling model misspecification is left for future work.The reward and Q-function classes must contain the corresponding true reward and policy Q-functions.Source: §5.1 · p. 6
• The coverage assumption requires πSFT to cover the comparator policy both trajectory-wise and state-action-wise.Trajectory-wise coverage is needed because human feedback is trajectory-wise; otherwise reward learning may not recognize high-performing policy traces at test time.Source: §5.1 · pp. 6–7
• The bound trades estimation error against NPG optimization error: increasing T and decreasing λ reduce optimization error but can amplify estimation error through larger CSFT.The analysis attributes the bound’s three terms to reward estimation, Q-function estimation, and NPG optimization, respectively.Source: §5.1 · p. 7
• DR-PO with NPG can find an ε-optimal policy with polynomial sample complexity in labeled trajectory pairs and unlabeled trajectories.The guarantee assumes reward and Q-function realizability and single-policy concentrability; finite classes can be extended using covering numbers.Source: §5.1 · pp. 6–7

6.1 How well can DR-PO optimize the RLHF objective?

• DR-PO achieves higher reward-model scores than all baselines and the highest GPT-4 win rate against human references, with slightly higher KL than PPO.Its reward–KL frontier also gives higher scores than PPO for most KL values.Source: §6.1 · pp. 8–9, Table 1 · p. 9, Figure 1 · p. 9

6.2 Analysis of Dataset Reset Proportion

• Even a small reset proportion yields higher-scoring generations with lower KL than PPO, and any reset proportion improves win rate; 100% resets perform best overall.β is the share of online generations starting from dataset resets; β = 0 recovers PPO.Source: §6.2 · p. 9, Table 2 · p. 10, Figure 2 · p. 10

6.3 DR-PO Transfer Performance

• On zero-shot CNN/DailyMail transfer, DR-PO is the only tested policy to outperform the supervised-finetuned model on every reported metric.The evaluation uses Pythia 2.8B; GPT-4 evaluates the win rate.Source: §6.3 · pp. 9–10, Table 3 · p. 10

6.4 DR-PO Scaling Performance on Anthropic HH

• On Anthropic HH, DR-PO showed scaling improvements similar to PPO across 125M, 1B, and 6.9B model scales, yet performed strictly better and produced more-preferred generations than all baselines.Figure 3 reports GPT-4 win rates as means and standard deviations across three seeds.Source: §6.4 · p. 11, Figure 3 · p. 11

7 Conclusion

• DR-PO exploits generative models’ ability to reset from offline data to enhance RLHF, with effectiveness demonstrated in theory and practice.The conclusion says dataset resets can be implemented in other online data-collection components, beyond the PPO-style optimizer tested in the experiments.Source: §7 · p. 11

Algorithm 3 DR-PO with NPG update

• Algorithm 3 partitions the unlabeled dataset across iterations and samples a state from each trajectory as a reset point for policy evaluation.It then collects a trajectory from that state; the supplied algorithm excerpt does not include all subsequent update details.Source: Algorithm 3 DR-PO with NPG update · p. 15

B Proof of Theorem 5.7

• Theorem B.6 states that Algorithm 1 with an NPG update returns a policy satisfying a high-probability guarantee under its assumptions.The supplied passage does not include the guarantee’s bound.Source: B Proof of Theorem 5.7 · pp. 15–17
• The proof proceeds by bounding Q-function estimation error, analyzing NPG performance, and accounting for reward uncertainty in the final suboptimality gap.The proof description identifies these as its three steps.Source: B Proof of Theorem 5.7 · p. 17

B.1 Q function Estimation Error

• The subsection identifies Q-function estimation error as a quantity that must be controlled in the proof.The excerpt provides the setup and lemma statement, but not the derivation or numerical form of the error bound.Source: §B.1 · pp. 17–18
• Under Assumption 5.3, Lemma B.7 states that its conclusion holds with probability at least 1 − δ1 for any δ1 ∈ (0, 1], for all t ∈ [T] and π ∈ {πt, π⋆}.Source: §B.1 · p. 17

B.2 NPG Analysis

• NPG analysis aims to find a near-optimal policy with respect to the estimated reward, accounting for the evaluation error defined in Lemma B.7.The excerpt also introduces a bound on the KL divergence between the output policy and πSFT.Source: §B.2 · pp. 18, 20

B.3 Unregularized Suboptimality Gap w.r.t. r⋆

• In proving Theorem 5.7, the authors bound term (1) using the MLE guarantee (Lemma C.2) with probability at least 1 −δ2, term (2) via Lemma B.8 conditioned on E1, and term (3) via Lemma B.11.Source: §B.3 · pp. 20–21

C.1 Least Sqaures Guarantee

• The least-squares guarantee is stated for bounded real-valued functions, independent samples, and targets whose conditional expectation equals the target function.The supplied statement is attributed to Song et al. (2022).Source: §C.1 · p. 21

C.2 Maximum Likelihood Estimation Guarantee

• The maximum-likelihood estimation subsection states a high-probability guarantee under Assumption 5.2.The supplied passages do not include the guarantee’s formula.Source: §C.2 · p. 21

C.3 Performance Difference

• The performance-difference lemma is stated for any two policies and reward function.The appendix restates and proves this lemma.Source: §C.3 · p. 22

D.1.1 Task Details

• Training data for both datasets came from OpenAI’s summarize-from-feedback repository, with dataset-specific details listed in Table 4.Source: §D.1.1 · p. 22

D.2 Dataset Reset Implementation Details

• A LogitsProcessor uses β to select generations for dataset resets and teacher-forces reference scores at selected roll-in positions.β is the proportion of generations in a batch selected for resets.Source: §D.2 · pp. 22–23

D.2.1 Computation

• Reset generation has the same computational requirements as generation from the initial state distribution; experiments used four of 16 A6000 GPUs per run.PPO and DR-PO used the same per-device batch size.Source: §D.2.1 · p. 23

D.3 Details on GPT4 Winrate

• The GPT-4 prompt asks which summary best captures a post’s important points based on accuracy, coverage, and coherence, and requests a one-sentence comparison plus an A/B choice.Source: §D.3 · p. 24

D.3.1 Win Rate Example

• In the example, GPT-4 preferred DR-PO’s more detailed summary because it better represented the post’s tone and content, including that the girlfriend knew about the feelings.The explanation says the alternative was too abbreviated and omitted that detail.Source: §D.3.1 · p. 25

Prompt 1

• The appendix illustrates prompts about pursuing veterinary school from finance and addressing a charter-school principal’s treatment of teachers.Source: Prompt 1 · pp. 25–26

D.5 Hyperparameters

• For TL;DR and CNN/DailyMail, PPO and DR-PO share hyperparameters except for mixing proportion; all processes share decoding, LoRA, and tokenizer settings.Source: Table 5 · p. 27

• DPO without LoRA had preference accuracy on CNN/DailyMail comparable to the reward model, so it was used for comparisons against PPO and DR-PO.Source: §D.6 · p. 28