🤖 machine learning

Rollout Pass-Rate Control: Steering Binary-Reward RL Toward Its Most Informative Regime

This paper proposes Prefix Sampling, a method that reconstructs trajectory prefixes to steer binary-reward agentic RL toward an optimal 50% pass rate, thereby maximizing reward entropy and contrastive signals to achieve significant wall-clock speedups and improved performance on benchmarks like SWE-bench and AIME.

Original authors: Tianshu Zhu, Wenyu Zhang, Xiaoying Zuo, Lun Tian, Haotian Zhao, Yucheng Zeng, Jingnan Gu, Daxiang Dong, Jianmin Wu, Dawei Yin, Dou Shen

Published 2026-05-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tianshu Zhu, Wenyu Zhang, Xiaoying Zuo, Lun Tian, Haotian Zhao, Yucheng Zeng, Jingnan Gu, Daxiang Dong, Jianmin Wu, Dawei Yin, Dou Shen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Wasting Energy on "Too Easy" or "Too Hard" Tasks

Imagine you are a coach training a team of student athletes to solve complex puzzles. You give them a batch of 8 puzzles at a time.

The "Too Easy" Batch: 7 or 8 students solve the puzzle instantly. They are bored, and you learn nothing new because they all got it right.
The "Too Hard" Batch: 0 or 1 student solves it. Everyone else is stuck. You learn nothing because there's no successful example to study.
The "Just Right" Batch: 4 students solve it, and 4 fail. This is the sweet spot. You have a perfect mix of success and failure to learn from. The students who failed can see exactly what the successful ones did differently.

In the world of AI, this is called Reinforcement Learning (RL). The AI generates many "rollouts" (attempts) to solve a problem. The paper argues that current AI training wastes a massive amount of computer power on the "Too Easy" and "Too Hard" batches because they don't provide useful learning signals.

The Solution: "Prefix Sampling" (The "Head Start" and "Handicap" Trick)

The researchers propose a clever method called Prefix Sampling to fix this. Instead of just throwing away the bad batches or asking the AI to try again from scratch (which is slow and expensive), they use a "time-travel" trick.

Think of the AI's attempt as a long story being written.

The "Too Hard" Scenario: The AI tries to solve a hard problem and mostly fails.
- The Trick: The system finds the one successful attempt the AI did make in that batch. It takes the first half of that successful story (the "prefix") and forces the AI to start its next attempt from that exact point.
- The Analogy: It's like giving a struggling student a head start. "You got stuck here, but look, you actually solved the first part correctly. Start your next attempt from here." This makes the hard problem easier, pushing the success rate up toward 50%.
The "Too Easy" Scenario: The AI solves the problem too easily.
- The Trick: The system finds the one failed attempt in that batch. It takes the first half of that failure and forces the AI to start from there.
- The Analogy: It's like giving a genius student a handicap. "You solved this too fast. Let's pretend you made a mistake at the beginning. Start your next attempt from this error." This makes the easy problem harder, pushing the success rate down toward 50%.

Why 50% is the Magic Number

The paper does some math to prove that 50% success is the most informative point.

Entropy (Confusion): If you know an AI will always win or always lose, there is no surprise. If it wins half the time, there is maximum "surprise" or information to learn from.
Contrast: To learn, you need to compare a "win" against a "loss." If everyone wins, you have no losses to compare. If everyone loses, you have no wins. You need a 50/50 split to get the best contrast.

How It Works in the Real World (The "Stateful" Part)

This is the hardest part to explain, but it's crucial. In simple math problems, an AI just writes text. But in software engineering (like fixing code), the AI is an "agent" that opens files, runs commands, and changes the computer's state.

Usually, if you want to replay a previous attempt, you have to reset the whole computer environment, which is slow.

The Innovation: The researchers built a system that can "replay" the AI's actions step-by-step to reconstruct the exact state of the computer (the code, the files, the history) without starting over.
The Masking: When the AI continues from this replayed state, the system tells the AI: "You didn't write the first part (we replayed that for you). You only get credit for the new part you write." This ensures the AI learns from its own new decisions, not the old ones.

The Results: Faster and Smarter

The researchers tested this on two types of tasks:

Software Engineering (SWE-bench): Fixing real-world code bugs.
Math (AIME 2025): Solving hard math olympiad problems.

The Gains:

Speed: The AI reached the same high performance level in half the time (up to 2x faster) on the larger models.
Efficiency: It wasted less computer power on useless "all-win" or "all-loss" attempts.
Performance: The final AI was actually better at solving problems than the standard training method, reaching higher scores.

Summary

The paper introduces a "traffic controller" for AI training. Instead of letting AI run wild and wasting time on tasks that are too easy or too hard, it actively steers the training batches toward a 50% success rate. It does this by giving struggling AI a "head start" from a previous success and giving over-confident AI a "handicap" from a previous failure. This keeps the learning process in the "Goldilocks zone," making training significantly faster and more effective.

Technical Summary: Rollout Pass-Rate Control via Prefix Sampling

1. Problem Statement

In binary-reward Reinforcement Learning with Verifiable Rewards (RLVR), such as software engineering (SWE-bench) and mathematical reasoning, online policy-gradient methods like Group Relative Policy Optimization (GRPO) rely on sampling groups of rollout trajectories to generate contrastive learning signals. However, a significant inefficiency arises when these sampled groups exhibit skewed pass rates (e.g., all successes or all failures, or heavily imbalanced ratios like 1/8 or 7/8).

In such skewed regimes, binary rewards provide weak contrastive signals. Specifically:

All-fail (0/N) or all-pass (N/N) groups are typically discarded by group filtering heuristics (e.g., in GRPO or DAPO) because they offer no relative advantage signal.
Skewed groups (e.g., 1/8 or 7/8) that survive filtering still provide suboptimal learning signals compared to balanced groups.
Compute Waste: Generating these uninformative trajectories consumes substantial computational resources and wall-clock time without yielding proportional training benefits.

Existing methods attempt to mitigate this through implicit filtering, dynamic sampling, or injecting external guidance (hints, oracle prefixes, or teacher traces). However, these approaches often fail to explicitly control the pass rate of retained groups toward the most informative regime, or they alter the resource budget and protocol by introducing external data.

2. Methodology: Prefix Sampling (PS)

The authors propose Rollout Pass-Rate Control (Rollout-PRC) as an explicit objective, identifying a 50% pass rate as the optimal operating point for binary-reward RL. They then introduce Prefix Sampling (PS), a bidirectional framework that steers skewed rollout groups toward this 50% regime without requiring external data or changing the rollout budget.

2.1 Theoretical Basis: Why 50%?

The paper demonstrates that a 50% pass rate ( $p=0.5$ ) maximizes four complementary metrics of learning signal strength:

Reward Entropy: Maximizes the information content of binary outcomes.
Survival Probability: Maximizes the probability that a group survives group filtering (i.e., is not all-fail or all-pass).
RLOO Advantage Energy: Maximizes the variance of Leave-One-Out (RLOO) advantage estimates under GRPO.
Contrastive Pairs: Maximizes the number of success–failure pairs available for contrastive learning within a group.

2.2 The Prefix Sampling Framework

Prefix Sampling operates by replaying trajectory prefixes to adjust the difficulty of the continuation for skewed groups:

For "Hard" Groups (Low Pass Rate, e.g., 1/8 or 2/8): The system selects a successful trajectory from the same batch (or a saved prefix) and replays its prefix. This acts as a "head start," providing the current policy with a correct partial solution to continue from, thereby increasing the likelihood of a successful completion.
For "Easy" Groups (High Pass Rate, e.g., 6/8 or 7/8): The system selects a failing trajectory and replays its prefix. This acts as a "handicap," forcing the policy to navigate a difficult path it previously failed, thereby increasing the likelihood of failure and creating a contrastive signal.

2.3 Implementation in Stateful Environments

A key innovation of PS is its ability to function in stateful agent environments (e.g., coding agents with tool use and environment interaction):

Replay-through-Execution: Instead of simply concatenating text, the system reconstructs the agent's state (code, conversation history, tool outputs, environment state) by feeding the saved prefix responses through the standard action parser and environment executor.
Prefix Masking: To ensure credit assignment remains on-policy, the loss is masked for all tokens in the replayed prefix. Optimization is restricted only to the continuation tokens generated by the current policy. This prevents the model from receiving credit for actions it did not generate while still benefiting from the reconstructed state.
Adaptive Control: An adaptive controller maintains an exponential moving average (EMA) of the rerollout pass rate for each "bucket" of skew (1/8, 2/8, 6/8, 7/8). It dynamically adjusts the prefix ratio (the length of the prefix to replay) to steer the pass rate of each bucket toward the 0.5 target.

3. Key Contributions

Formulation of Rollout-PRC: The paper identifies and formalizes pass-rate control as a principled efficiency objective for binary-reward RL, proving that 50% pass rate maximizes reward entropy, advantage energy, and contrastive structure.
Prefix Sampling Algorithm: A bidirectional framework that uses self-generated prefixes (no external oracles) to steer skewed groups. It uniquely supports stateful agent environments via replay-through-execution and prefix masking.
Empirical Validation: Demonstrated across multiple scales (Qwen3-4B to 32B) and domains (SWE-bench and AIME 2025), showing that PS achieves baseline-level performance in fewer steps and delivers significant end-to-end speedups.

4. Experimental Results

The authors evaluated PS on SWE-bench Verified (agentic RL) and AIME 2025 (mathematical reasoning) using Qwen3 models (4B, 8B, 14B, 32B).

4.1 Benchmark Performance

SWE-bench Verified:
- Qwen3-14B: Reached the baseline's peak performance (0.273) at step 170 (vs. 300 for baseline) and achieved a new peak of 0.295 (+4.7 percentage points).
- Qwen3-32B: Reached the baseline peak at step 290 (vs. 410) and achieved a peak of 0.422 (+5.4 percentage points).
AIME 2025:
- Qwen3-4B: Reached baseline peak at step 140 (vs. 230) with a +7.3 pp gain.
- Qwen3-8B: Reached baseline peak at step 190 (vs. 260) with a +10.8 pp gain.

4.2 Efficiency and Speedups

End-to-End Speedup: PS delivered 2.01× speedup on Qwen3-14B and 1.55× on Qwen3-32B.
Training Efficiency: The number of training steps to reach the baseline's ordinary-task score was reduced by 1.51× to 1.92× across models.
Valid Groups: The number of valid (non-degenerate) rollout groups per batch increased significantly (e.g., from 25.9 to 36.7 for 14B), indicating more efficient use of the rollout budget.
Wall-Clock Time: In stateful SWE runs, prefix replay reduced per-step wall-clock time by reusing reconstructed states rather than regenerating prefixes, contributing to the overall speedup.

4.3 Ablation Studies

Ablations on the 4B math run confirmed the decomposition of gains:

Replay: The primary driver of convergence speed (1.61× speedup with fixed replay).
Adaptive Control: Essential for keeping rerollout pass rates near 0.5; fixed ratios led to suboptimal pass rates (e.g., 0.36 or 0.65) and slower convergence.
Bidirectional Coverage: Controlling both hard and easy buckets is necessary; controlling only hard buckets yielded significantly lower gains.

5. Significance and Claims

The paper claims that Prefix Sampling offers a principled, resource-efficient method to improve binary-reward RL training without relying on external supervision or increasing compute budgets.

Mechanism: The gains are attributed to steering skewed, partially successful groups toward the high-signal 50% regime, thereby increasing reward entropy and contrastive pairs.
Stateful Capability: Unlike prior methods that rely on text-only prefix concatenation, PS successfully reconstructs stateful agent environments (code, tools, history) through execution replay, making it applicable to complex agentic tasks.
Modesty: The authors explicitly state that their claims are scoped to binary-reward RLVR with grouped rollouts. They do not claim that 50% is optimal for all curriculum designs or that the method eliminates the need for diverse data; rather, it optimizes the utilization of the existing rollout budget. The method is presented as a "drop-in" controller compatible with existing GRPO++ pipelines.

In summary, the paper demonstrates that by actively managing the pass rate of rollout groups to target 50%, one can significantly accelerate the training of agentic RL systems while preserving or improving final performance, all while operating within the same computational constraints as standard baselines.