Reuse your FLOPs: Scaling RL on Hard Problems by Conditioning on Very Off-Policy Prefixes

TL;DR

This work introduces PrefixRL, where it is proved that the PrefixRL objective is not only consistent with the standard RL objective but also more sample efficient, and discovers back-generalization: training only on prefixed problems generalizes to out-of-distribution unprefixed performance.

Abstract

Typical reinforcement learning (RL) methods for LLM reasoning waste compute on hard problems, where correct on-policy traces are rare, policy gradients vanish, and learning stalls. To bootstrap more efficient RL, we consider reusing old sampling FLOPs (from prior inference or RL training) in the form of off-policy traces. Standard off-policy methods supervise against off-policy data, causing instabilities during RL optimization. We introduce PrefixRL, where we condition on the prefix of successful off-policy traces and run on-policy RL to complete them, side-stepping off-policy instabilities. PrefixRL boosts the learning signal on hard problems by modulating the difficulty of the problem through the off-policy prefix length. We prove that the PrefixRL objective is not only consistent with the standard RL objective but also more sample efficient. Empirically, we discover back-generalization: training only on prefixed problems generalizes to out-of-distribution unprefixed performance, with learned strategies often differing from those in the prefix. In our experiments, we source the off-policy traces by rejection sampling with the base model, creating a self-improvement loop. On hard reasoning problems, PrefixRL reaches the same training reward 2x faster than the strongest baseline (SFT on off-policy data then RL), even after accounting for the compute spent on the initial rejection sampling, and increases the final reward by 3x. The gains transfer to held-out benchmarks, and PrefixRL is still effective when off-policy traces are derived from a different model family, validating its flexibility in practical settings.

Figures (16)

• Figure 1: PrefixRL: On-Policy RL Conditioned on Off-Policy Prefixes. We leverage previously spent compute ( ) on hard problems in the form of correct off-policy traces rejection sampled from the base LLM we start RL from. Off-policy traces could also come from other model families or previous RL runs. We append prefixes of a single correct off-policy trace per problem to the original problem, creating prefixed problems ( ). Then, we simply run on-policy RL on prefixed and no-prefix (original) problems ( ). PrefixRL places the RL policy in higher-rewarding states, which boosts the learning signal. Performance transfers from the prefixed to no-prefix problems via a phenomenon we call back-generalization.
• Figure 2: PrefixRL affords a self-improvement pipeline that recycles RL flops on hard problems. We instantiate PrefixRL for self-improvement by collecting a dataset of off-policy traces through large-scale rejection sampling on the base LLM (distilled Llama3.1-8B). In FLOPs-matched training, PrefixRL outperforms the strongest baseline (SFT on rejection-sampled data + RL): 2$\times$ higher compute efficiency (including rejection-sampling cost) and $>$45$\%$ (over 3$\times$ relative) higher final training accuracy on no-prefix training problems (left), with gains transferring to standardized evals such as AIME ’25 (right).
• Figure 3: Supervising the policy on $\mathcal{D}_\mathrm{off}$ can cause diversity collapse or training instabilities during RL: (a, b) Warm-starting the RL run by running SFT on $\mathcal{D}_\mathrm{off}$ before (mid-training) reduces token entropy (a) and hurts exploration during RL (worse Pass@64 performance in (b)). (c, d) Directly using $\mathcal{D}_\mathrm{off}$ during online RL by updating the current RL policy with importance-weighted off-policy traces (in addition to on-policy traces) leads to training instabilities. We see the gradient norm (clipped at 1.0) blow up during training (d) and this leads to optimization collapse (c).
• Figure 4: Prefixing on off-policy trace improves probability of future success. When we condition on prefixed problems, we increase the accuracy by placing the policy at key strategy-revealing states (Erdös Lemma in example ). For five problems (P1, $\ldots$, P5) we plot accuracy when conditioning on prefixes of varying lengths, as a proportion of the full off-policy prefix length ( ). The accuracy is near zero until a key state is visited, after which it jumps sharply.
• Figure 5: Back-generalization (train-test mismatch): Using Llama3.1-8b-instruct, we run RL only on prefixed problems with prefix length lies in the shaded interval. We evaluate different training step checkpoints across the full range of prefix lengths, including no-prefix problems. Training on longer prefixes improves performance on shorter prefixes and can eventually improve no-prefix, indicating back-generalization (a,b). When training uses only very long prefixes (severe mismatch with no-prefix), back-generalization to no-prefix problems takes more training steps (e.g., 800 iterations) (c).

Theorems & Definitions (10)

• Theorem 3.2: Consistency of the PrefixRL objective
• Theorem 3.3: Suboptimality gap of PrefixRL
• Proposition 3.4: Worst-case separation with standard RL
• proof
• proof
• Lemma B.4: Performance difference lemma; kakade2002approximately
• proof
• Lemma B.5: Three-point identity for KL
• proof
• Lemma B.6: Lemma 15 in song2022hybrid