Source-linked AI summary
When All Rollouts Fail: Policy-Compatible Joint Decoding for Hard-Problem Reasoning
Zehao Wang, Hanzhu Chen, Runquan Gui, Zhe Wang, Lin Yang
TL;DR
When all on-policy rollouts fail, RLVR and group-relative methods may lack an informative learning signal, while expert trajectories can be poorly matched to the student policy. PICO generates verified, policy-compatible trajectories through joint decoding and learns from them alongside on-policy rollouts. It reports consistent gains over standard RLVR and expert-guided baselines across reasoning and agentic tasks.
Problem
When every rollout fails on a difficult problem, RLVR receives no positive reward signal, while expert trajectories may lie in low-probability regions under the student policy.
Source: §1 · p. 1Method
PICO uses teacher–student joint decoding to generate verified trajectories compatible with the student, then optimizes them with on-policy rollouts using a mixed-behavior objective.
Source: §1 · p. 2Results
PICO consistently outperforms standard RLVR and representative expert-guided baselines across tasks and model scales; with Qwen3-4B-Base, it improves average score over GRPO by 4.1 percentage points.
Source: §4.2 · p. 7, §6 · p. 9Takeaways & Limitations
The results highlight policy compatibility as important to learning from expert-guided trajectories across mathematical, general reasoning, and agentic tasks.
Source: §6 · p. 9
Abstract
from arXiv · showhide
Reinforcement Learning with Verifiable Rewards (RLVR) derives its learning signal from binary rewards via on-policy exploration and becomes a powerful paradigm for improving the reasoning of large language models. Yet it faces a learning cliff on problems beyond model's current ability: when every rollout fails, the samples reveal nothing about a correct solution, and group-relative estimators (e.g., GRPO, RLOO) collapse the update to zero advantage. Expert trajectories offer a remedy, but they are incompatible with the current policy and often lie in the student's low-probability regions, so the surviving updates demand large parameter shifts and decay into surface-level imitation rather than transferable reasoning. We argue that a useful learning signal on hard problems must be both verifiably correct and sufficiently compatible with the current policy. We propose Policy-Compatible Reinforcement Learning (PICO), which performs distribution-constrained joint decoding: expert data guides decoding at the token level while preserving the student's relative token probabilities, and a dynamic top-k mechanism confines the sampling distribution under a finite divergence budget, yielding trajectories that are verified and policy-compatible by construction. These trajectories are optimized jointly with on-policy rollouts under a mixed-behavior objective that mitigates their distribution mismatch via importance sampling. Experiments across reasoning diverse datasets and agentic benchmarks demonstrate that PICO consistently outperforms standard RLVR and prior off-policy baselines.
1 INTRODUCTION
RLVR can provide no informative update on difficult problems when every rollout fails, while expert trajectories can be misaligned with the student policy. PICO uses jointly guided, policy-compatible trajectories to recover learning signals, and experiments report gains across reasoning and agentic tasks.
1 INTRODUCTION
- When every sampled rollout fails, binary rewards provide no positive signal and group-relative advantages can collapse to zero.Source: §1 · p. 1
- Expert trajectories can provide useful reasoning signals, but low-probability demonstrations may cause policy drift, surface imitation, or mismatch-related instability.Source: §1 · pp. 1–2
- PICO uses teacher–student joint decoding to generate verified trajectories compatible with the student, then optimizes them alongside on-policy rollouts.The method combines token-level guidance with failure-triggered mixed-behavior optimization.Source: §1 · p. 2
- PICO consistently outperforms standard RLVR and expert-guided baselines across mathematics, general reasoning, and agentic tasks.The paper also reports generalizability to agentic tasks and validates key components through ablations.Source: §1 · p. 2
2 PRELIMINARIES
RLVR optimizes expected binary verifier reward on policy-generated responses. GRPO derives relative advantages from rewards within rollout groups and updates the policy with a clipped, KL-regularized objective.
2 PRELIMINARIES
- RLVR scores a response with binary reward 1 when its extracted answer matches the ground truth, and 0 otherwise.Source: §2 · p. 2
- Figure 1 depicts PICO replacing one failed student rollout with a verified joint trajectory generated from intersecting teacher and student top-k candidates.Source: Figure 1 · p. 3
- GRPO standardizes each trajectory’s reward within its rollout group to compute relative advantages, without a separately trained value model.Source: §2 · p. 3
- The policy-gradient objective uses clipping to limit large updates and a KL penalty to constrain divergence from a reference policy.The KL penalty can be downweighted or omitted when larger policy shifts are desired.Source: §2 · p. 3
3 METHOD
PICO combines teacher and student token preferences under a constraint on distribution shift, then uses verified joint trajectories to reconstruct all-failure rollout groups. Importance sampling accounts for the joint trajectory’s behavior-policy mismatch during optimization.
3 METHOD
- PICO combines teacher and student token preferences to generate rollouts, using dynamic top-k selection and importance sampling in its two method components.Source: §3 · pp. 3–4
3.1 DISTRIBUTION-CONSTRAINED JOINT DECODING
- The consensus support set intersects the teacher’s and student’s top-k tokens, excluding candidates implausible under either distribution.The teacher contributes through candidate selection while the student determines compatibility among retained tokens.Source: §3.1 · p. 4
- The joint-decoding objective balances expected teacher confidence against reverse KL divergence from the student policy, with λ controlling their trade-off.A smaller λ better preserves the student’s relative token preferences; a larger λ gives more influence to the teacher.Source: §3.1 · p. 4
3.2 DYNAMIC TOP-k JOINT DECODING
- The closed-form policy restricts support to jointly selected tokens and reweights student probabilities by teacher confidence; the final policy sets λ = 0.This preserves student relative probabilities within the support while leaving teacher guidance to candidate selection.Source: §3.2 · pp. 4–5
- Dynamic top-k selection chooses the smallest teacher cutoff whose consensus set retains at least Zmin student probability mass, falling back to the student policy if no cutoff works.Source: §3.2 · p. 5
3.3 MIXED-BEHAVIOR POLICY OPTIMIZATION
- When every trajectory in a rollout group is incorrect, PICO samples a joint trajectory and replaces one failed trajectory if verification succeeds, restoring reward variation.Source: §3.3 · p. 5
- The reconstructed group mixes on-policy and joint trajectories, and an off-policy importance ratio accounts for the joint trajectory’s token-level behavior mismatch.Source: §3.3 · p. 5
4 EXPERIMENTS
PICO is evaluated on mathematical, general-reasoning, and agentic tasks against on-policy and expert-guided baselines. It leads the reported comparisons, while ablations highlight the roles of policy compatibility, dynamic top-k selection, and importance sampling.
4.1 EXPERIMENTAL SETUP
- Experiments use Qwen3-4B-Base, Qwen3-8B-Base, and Qwen3-4B-Instruct; agentic tasks use Qwen3-4B-Instruct as both student and teacher.Training uses DAPO-Math-17K for mathematical and general reasoning and ALFWorld and ScienceWorld for agentic tasks; evaluations include AIME 2025, AIME 2026, HMMT 2026, MATH-500, GPQA-Diamond, and HLE.Source: §4.1 · pp. 6–7
4.2 MAIN RESULTS
- PICO improves average scores over GRPO by 4.1 points with Qwen3-4B-Base and 3.6 points with Qwen3-8B-Base, and beats the strongest competitor by 2.2 and 1.8 points, respectively.Gains on GPQA-Diamond and HLE show generalization beyond the mathematics-focused training distribution.Source: §4.2 · p. 7, Table 1 · p. 7
- Across both model scales, PICO maintains higher training rewards and fewer solve-none questions than GRPO and LUFFY, which tend to plateau.PICO also produces longer responses, alongside higher verified rewards and fewer solve-none questions.Source: §4.2 · p. 7
4.3 AGENTIC TASKS
- PICO achieves a 74.3% average success rate across ALFWorld and ScienceWorld, outperforming all baselines.It improves over GRPO by 6.8 percentage points and over LUFFY and Prefix-RFT by 4.9 and 5.2 points on average, respectively.Source: §4.3 · pp. 7–8, Table 2 · p. 8
4.4 ABLATION STUDY
- Dynamic top-k decoding achieves the highest ablation average score of 37.2%, while removing the importance-sampling ratio reduces it to 35.3%.Direct teacher sampling yields 33.8%; among fixed settings, k = 10 performs best.Source: §4.4 · p. 8, Table 3 · p. 8
5 RELATED WORK
Prior work incorporates expert trajectories through direct optimization or rollout guidance, but these approaches face distribution mismatch or may rely on expert-induced states. PICO instead uses expert data for token-level guidance while constraining generated trajectories to remain compatible with the student policy.
5.1 ON- AND OFF-POLICY LEARNING
- Directly optimizing on expert trajectories can create substantial distribution mismatch, excessive policy drift, and unstable optimization when those trajectories are unlikely under the student policy.LUFFY mixes expert demonstrations with on-policy rollouts, ReLIFT interleaves RL and supervised fine-tuning, and SRFT and Chord unify supervised and reinforcement learning.Source: §5.1 · p. 9
5.2 EXPERT-GUIDED EXPLORATION AND DECODING
- Prefix-based guidance can steer exploration but may leave students unable to reach expert-induced states from the original prompt; PICO instead guides generation at the token level.OC-GRPO applies importance sampling to account for mismatch between guided and unguided rollouts.Source: §5.2 · p. 9
6 CONCLUSION
PICO uses expert data to generate verified trajectories while controlling deviation from the student policy, then selectively adds them to all-failure groups. Experiments show gains across reasoning and agentic tasks, and the appendix states a uniqueness result for its decoding optimization.
6 CONCLUSION
- PICO outperforms standard RLVR and representative expert-guided baselines across mathematical reasoning, general reasoning, and agentic tasks, without teacher guidance at inference.Ablations validate the framework’s components; the paper identifies policy compatibility as a key principle for incorporating expert knowledge into RLVR.Source: §6 · p. 9
A.1.1 OPTIMIZATION PROBLEM FORMULATION AND CLOSED-FORM SOLUTION
- The joint-decoding optimization has a unique global minimizer because its KL term is strictly convex and its teacher-confidence term is linear.The passage identifies this stationary point with Equation 7.Source: §A.1.1 · p. 14
A.1.2 PROOF OF THE KL DECOMPOSITION
The proof shows that hard joint decoding minimizes distribution shift for a fixed support by using the student distribution conditioned on that support. It also specifies the limits of the action-level correction and the failure-triggered training procedure.
A.1.2 PROOF OF THE KL DECOMPOSITION
- For fixed support F_t, the minimum student-policy divergence is −log Z_t, attained by the hard policy; soft teacher reweighting strictly increases shift unless q_t is constant on F_t.The method therefore uses the minimum-shift endpoint, letting the teacher determine support rather than reweighting retained tokens.Source: §A.1.2 · p. 14
- When every guided step retains mass Z_t ≥ Z_min and infeasible steps fall back to the student, the per-token constraint yields a trajectory-level KL bound over horizon T_max.Variable-length responses use an absorbing end-of-sequence state and equal policies after termination.Source: §A.1.2 · pp. 14–15
A.1.3 IMPORTANCE-SAMPLING CORRECTION
- The action-level importance-sampling correction is exact conditional on the prefix and retained support, but cannot recover contributions from tokens outside F_t.At data collection, the retained-mass constraint controls both KL shift and residual support truncation, while the correction factor is the token-independent Z_t.Source: §A.1.3 · p. 15
- The mixed-behavior objective is a surrogate, not an unbiased estimator of on-policy RLVR: hard support may remove student tokens, and PPO clipping introduces bias.The action-level identity does not correct prefixes sampled under joint decoding into prefixes sampled under the student policy.Source: §A.1.3 · p. 15
A.2 ALGORITHMIC PROCEDURE
- The teacher stays frozen, π_old is fixed during rollout collection, and privileged guidance is provided only to the teacher.The student evaluates the original query and the shared generated response prefix.Source: §A.2 · p. 15
B.1 PASS@k RESULTS
PICO improves Pass@16 over GRPO at both model scales. Its average results are best among baselines for the 4B model and match the best baseline for the 8B model.
B.1 PASS@k RESULTS
- 8.8 and 10.8 percentage points: PICO raises average Pass@16 over GRPO on Qwen3-4B-Base and Qwen3-8B-Base, respectively.These gains indicate a higher likelihood of finding at least one correct solution under repeated sampling.Source: §B.1 · p. 15
- 43.3%: PICO achieves the best average Pass@16 on Qwen3-4B-Base, exceeding the strongest baseline by 6.6 points.Source: §B.1 · p. 15
- 45.5%: PICO reaches this average Pass@16 on Qwen3-8B-Base, matching the best baseline.Source: §B.1 · p. 15
B.2 TRAINING CURVES
The training curves track how joint decoding guides rollout generation while preserving proximity to the student policy. Reported trends show retained student probability mass increasing and the associated KL divergence decreasing during training.
Behavior-Policy Divergence
- Retained student probability mass rises from approximately 0.91 to 0.94 as token-level KL divergence falls from roughly 0.10 to 0.06.Figure 5 presents raw measurements and a smoothed trend for Qwen3-4B-Base.Source: §B.2 · pp. 16–17, Figure 5 · p. 18
Dynamic Top-k Joint Decoding
- Dynamic top-k decoding intersects student and teacher top-k token sets, increasing teacher-side k until retained student mass reaches Zmin before sampling from the restricted student distribution.Figure 3 illustrates this candidate-set construction at each decoding step.Source: Figure 3 · p. 17
Question Rescue
- Figure 4 tracks unsolved-question counts under student self-sampling versus joint decoding, with the shaded gap marking rescued questions across mathematical tasks and ALFWorld.The caption identifies Qwen3-4B-Base, Qwen3-8B-Base, and Qwen3-4B-Instruct as the evaluated models.Source: Figure 4 · p. 18
C EXPERIMENTAL DETAILS
The experiments use three Qwen3 student configurations across mathematical reasoning and agentic tasks, with Qwen3-4B-Instruct also serving as the teacher. Mathematical and agentic training use different datasets and have reported training times.
Models and Training Setup
- Qwen3-4B-Base and Qwen3-8B-Base are the mathematical-reasoning students, while Qwen3-4B-Instruct is used for agentic experiments and serves as teacher across settings.The teacher conditions on expert solutions to provide token-level guidance.Source: C EXPERIMENTAL DETAILS · p. 18
Training Data and Hardware
- Mathematical reasoning methods share the DAPO-Math-17K corpus and hyperparameter configuration, while agentic training uses ALFWorld and ScienceWorld separately.All experiments use four NVIDIA H20 GPUs.Source: C EXPERIMENTAL DETAILS · p. 18
Training Time
- Training takes approximately 1.5–2 days for either mathematical-reasoning model and approximately 1 day per agentic task for Qwen3-4B-Instruct.Source: C EXPERIMENTAL DETAILS · p. 18
C.1 TRAINING DETAILS
PICO adds teacher-guided decoding during training, increasing time per step while improving rollout utilization. Its reported utilization-per-time ratio remains comparable to GRPO.
Training Efficiency
- PICO increases average time per training step from 28.3 to 43.6 minutes, but improves rollout utilization from 57.2% to 91.9%, allowing more sampled groups to provide effective learning signals.Source: §C.1 · p. 18
- Utilization per minute is 2.11 percentage points for PICO versus 2.01 for GRPO, despite PICO's additional computation.Source: §C.1 · p. 18
Training Configuration
- Table 5 lists implementation and hyperparameter configurations for the mathematical-reasoning experiments.Source: Table 5 · p. 19
C.2 EVALUATION DETAILS
Evaluation covers mathematical, cross-domain reasoning, and interactive agentic benchmarks, using task-specific success measures and sampling protocols. The prompts distinguish student-visible task context from teacher-only expert guidance and specify separate formats for mathematical and ALFWorld tasks.
C.2 EVALUATION DETAILS
- Evaluation covers in-domain AIME 2025, AIME 2026, and HMMT 2026, broader MATH-500, out-of-domain GPQA-Diamond and text-only HLE, and interactive ALFWorld and ScienceWorld.Reasoning benchmarks use 16 responses per problem; agentic evaluation uses up to 20 interaction steps per episode and reports mean episode success over eight episodes per task instance.Source: §C.2 · pp. 19–20
D PROMPTS
- The student receives task and interaction context, while the teacher additionally receives expert solutions or trajectories to produce token-level predictions; expert data is not exposed to the student.Student probabilities determine relative sampling probabilities among retained candidates.Source: D PROMPTS · p. 21
D.1 MATHEMATICAL REASONING PROMPT
- The mathematical prompt asks the student to reason step by step and place its final answer within a boxed format; the teacher prompt supplies private guidance for a complete solution.The teacher prompt requests equations, intermediate conclusions, and justifications.Source: §D.1 · p. 21
D.2 ALFWORLD PROMPTS
- The ALFWorld prompt directs the teacher to align the known-correct expert trajectory with current observations and history, then select an admissible next action.The agent must issue exactly one command through env_step, following the admissible-command list.Source: §D.2 · p. 23