Source-linked AI summary
PBSD: Privileged Bayesian Self-Distillation for Long-Horizon Credit Assignment
Yang Tian, Rui Wang, Xumeng Wen, Junjie Li, Shizhao Sun, Lei Song, Jiang Bian, Bo Zhao
TL;DR
Long-horizon search agents receive sparse outcome rewards that poorly distinguish helpful from misleading intermediate actions. PBSD converts Bayesian answer evidence into turn-level credit signals, improving performance, generalization, and transfer from short-context training to long-context inference.
Problem
Long-horizon search agents receive coarse trajectory-level rewards that poorly distinguish helpful from misleading intermediate reasoning and tool interactions.
Method
PBSD uses Bayes’ rule to convert an intractable posterior-to-prior answer ratio into an estimable privileged answer-conditioned action-likelihood ratio for turn-level credit.
Results
PBSD consistently outperforms outcome-only RLVR and strong baselines, with improved in-domain performance, cross-dataset generalization, and long-context transfer.
Takeaways & Limitations
Bayesian evidence calibration provides an effective mechanism for improving credit assignment in long-horizon search-agent reinforcement learning.
Takeaways & Limitations
PBSD depends on verified final answers and sufficiently capable, well-calibrated answer-conditioned models, limiting applicability when answers are ambiguous or calibration is poor.
Abstract
from arXiv · showhide
Long-horizon agentic tasks pose a fundamental credit assignment challenge for outcome-base reinforcement learning: trajectory-level rewards verify final correctness but provide limited guidance on which intermediate reasoning steps or tool interactions contribute to the outcome. The difficulty is especially pronounced in multi-turn search agents, where successful trajectories may contain misleading actions and failed trajectories may contain valuable evidence-gathering steps. We propose PBSD (Privileged Bayesian Self-Distillation), a Bayes-calibrated self-distillation method for fine-grained credit assignment under sparse final rewards. PBSD measures trajectory quality through the posterior-to-prior probability ratio of the verified answer and applies Bayes' rule to convert this hard-to-estimate answer-side ratio into a tractable likelihood ratio between a standard student model and a privileged answer-conditioned teacher model. Autoregressive decomposition of this Bayesian evidence score yields turn-level signals that identify whether each intermediate turn supports or undermines the verified outcome. Consequently, PBSD provides a principled and elegant reweighting scheme that transforms sparse outcome supervision into Bayes-calibrated turn-level credit signals, while remaining fully compatible with standard policy optimization. Experiments demonstrate that PBSD consistently enhances performance across both in-domain and out-of-domain settings, and effectively transfers knowledge from short-context training to long-context inference, suggesting that its fine-grained credit assignment mechanism facilitates more effective policy learning and yields improved generalization.
1 Introduction
PBSD addresses the severe credit-assignment challenge in long-horizon RLVR search agents by converting sparse verified outcomes into Bayes-calibrated, turn-level evidence signals. It preserves outcome-based learning while differentially weighting intermediate actions and improves performance and generalization in 30B MoE experiments.
- Motivation: Long-horizon RLVR search agents face severe credit-assignment difficulties because successful trajectories can contain misleading actions and failed trajectories can contain useful evidence-gathering steps.The challenge arises when trajectory-level verification provides limited guidance about which intermediate interactions contributed to the final outcome.
- Method: PBSD evaluates each intermediate action by its posterior-to-prior evidence ratio for the verified final outcome.This formulates turn-level credit assignment as Bayesian evidence estimation.
- Method: PBSD uses Bayes’ rule to transform the intractable answer-side evidence ratio into an estimable privileged action-likelihood ratio.Direct estimation would require marginalizing over all possible future continuations from a given prefix.
- Method: PBSD calibrates turn-level advantages from the Bayesian evidence score without introducing an independent loss function.The terminal verifiable reward remains the global learning signal, while intermediate actions receive differential weights according to their evidence for the verified outcome.
- Results: PBSD demonstrates superior performance with substantially fewer training steps than outcome-only RLVR and strong cross-dataset generalization on a 30B MoE model.These results suggest that PBSD learns transferable credit signals for long-horizon agentic behavior.
2 Related Work
Related work identifies long-horizon agentic reinforcement learning as a difficult credit-assignment setting involving extended heterogeneous reasoning and tool-use sequences. Existing search-agent and on-policy distillation approaches motivate PBSD’s lightweight Bayesian turn-level signal, which avoids handcrafted rewards, external evaluators, and costly search-based estimation.
- Credit Assignment: Long-horizon agentic reinforcement learning requires credit assignment across extended sequences of reasoning steps, tool calls, observations, and decisions.These tasks differ from single-turn reasoning through heterogeneous, multi-step interactions.
- Bayesian Credit Assignment: PBSD derives a lightweight Bayesian calibration signal directly from sampled trajectories, enabling turn-level credit assignment without handcrafted process rewards, external evaluators, or costly search-based estimation.This contrasts with methods whose per-turn credit estimation may require numerous additional rollouts or value evaluations in long-context search-agent settings.
- On-Policy Distillation: OPD combines on-policy trajectory collection with token-level teacher supervision, but depends on teacher capability and compatibility with the student’s exploration distribution.Different teacher strategies or action preferences can weaken this compatibility.
- Search Agents: Search agents iteratively issue queries, inspect retrieved evidence, and synthesize answers; Search-R1 and DeepResearcher apply outcome-based reinforcement learning to acquire or scale such behavior.Search-R1 demonstrates multi-turn search training without large-scale supervised trajectories, while DeepResearcher scales reinforcement learning in real-world web environments.
3 Methodology
PBSD converts sparse trajectory-level outcome supervision into calibrated turn-level credit by estimating Bayesian evidence through privileged teacher–student likelihood ratios. It reweights inherited policy advantages while filtering unreliable evidence and preserving standard policy optimization.
- From trajectory-level rewards to turn-level credit: Standard GRPO assigns every token in a trajectory the same outcome-level advantage, making credit assignment coarse and noisy for heterogeneous long-horizon turns.The multi-turn gradient decomposition exposes this shared-advantage scheme and its weakness in search-agent settings.
- From trajectory-level evaluation to turn-level evidence: PBSD assesses trajectory support using the posterior-to-prior probability ratio of the verified answer, then applies Bayes’ rule to rewrite it as a likelihood ratio.Direct posterior and prior estimation is difficult because it requires marginalizing over unobservable latent paths.
- From trajectory-level evaluation to turn-level evidence: The method compares standard student likelihood p(τi | xi) with answer-conditioned teacher likelihood p(τi | xi, y⋆), using ground-truth answers only through teacher conditioning.The teacher–student trajectory likelihood ratio is decomposed autoregressively across assistant turns.
- From trajectory-level evaluation to turn-level evidence: Each turn receives si,t = log p(ai,t | xi, τi,<t, y⋆) −log p(ai,t | xi, τi,<t), measuring its incremental contribution to verified-answer evidence.This converts trajectory-level correctness into a fine-grained signal for intermediate-turn credit.
- Evidence-Calibrated Advantage Reweighting: PBSD uses Bayesian evidence to modulate inherited outcome advantages, down-weighting opposing turns in successful trajectories and strengthening their penalties in failed trajectories.A tanh-shaped weight with symmetric clipping limits outlier influence while preserving the original GRPO reinforcement or penalty direction.
- Stability and implementation: PBSD filters low-magnitude, noise-dominated evidence near zero and, for R3-routed MoE models, recomputes scores with non-replay likelihoods to avoid routing artifacts.The complete workflow converts outcome supervision into turn-level credit without introducing an additional optimization objective.
4 Experiment
PBSD achieves the strongest reported performance across evaluation settings, improving over GRPO on in-domain and BrowseComp subsets while generalizing from 64K-token training to 256K-token evaluation. Its gains are accompanied by broader sampled-solution coverage, faster and more stable optimization, and a shift toward more frequent, focused interactions with fewer generated tokens.
- Comparison with other methods: PBSD achieves the best results across all evaluation settings, scoring 40.87 on in-domain validation and 35.83 on BC(300).It improves over GRPO by 2.62 points on validation and 3.50 points on BC(300).
- Comparison with other methods: PBSD shows more favorable optimization than GRPO, with faster improvement, higher final accuracy, and more stable convergence across BC subsets.These gains are especially pronounced on the easy and medium subsets, supporting stronger robustness on challenging out-of-domain instances.
- Cross-Benchmark Generalization under Long-Context Evaluation: Despite 64K-token training, PBSD is evaluated at 256K tokens to test transfer beyond both the training distribution and context length.The long-context setting is particularly challenging on BrowseComp because solving problems can require longer exploration and more tool invocations.
- Capability Boundary Analysis: SFT+PBSD consistently outperforms SFT across all Pass@k values, indicating improvements to both likely responses and the density of valid sampled reasoning paths.Higher gains at larger sampling budgets suggest that PBSD expands the model’s effective capability boundary.
- Behavioral Dynamics of PBSD during Training: During training, PBSD increases interaction turns while substantially reducing generated tokens, shifting computation toward more frequent, focused, and evidence-oriented search behavior.The teacher–student likelihood gap also gradually decreases as the teacher increasingly recognizes student-generated actions as supportive of the verified answer, without direct KL optimization or privileged-policy distillation.
5 Conclusion
PBSD addresses credit assignment in RLVR-based long-horizon search agents by deriving turn-level credit from privileged answer-conditioned likelihoods. It uses Bayes’ rule to transform an intractable answer-side estimation problem into an estimable action-likelihood ratio.
- 5 Conclusion: PBSD is a Bayes-calibrated self-distillation method for credit assignment in RLVR-based long-horizon search agents.The method derives turn-level credit from privileged answer-conditioned likelihoods.
- 5 Conclusion: PBSD interprets trajectory quality as a posterior-to-prior evidence ratio.This interpretation provides the basis for its Bayesian credit-assignment procedure.
- 5 Conclusion: Bayes’ rule converts the intractable answer-side estimation problem into an estimable action-likelihood ratio.PBSD then uses this ratio for detached calibration.
6 Limitations and Future Work
PBSD depends on verified final answers to construct privileged answer-conditioned likelihoods, fitting RLVR settings with automatically checkable outcomes. In open-ended agentic tasks, ambiguous, incomplete, or multiple valid solutions make this conditioning signal harder to define.
- Limitation: PBSD requires verified final answers to construct privileged answer-conditioned likelihoods for Bayesian evidence estimation.This assumption is aligned with RLVR settings such as search QA and verifiable reasoning tasks.
- Limitation: Automatically checkable final outcomes make PBSD well aligned with search QA and verifiable reasoning tasks.
- Limitation: In open-ended agentic scenarios, ambiguous or incomplete ground-truth answers make PBSD’s privileged conditioning signal harder to define.
- Limitation: Multiple valid solutions in open-ended tasks further complicate defining the privileged conditioning signal.
A Search Agent System Prompt
The search agent system prompt defines the available MCP-style tools and states the high-level search objective.
- The system prompt specifies the available MCP-style tools and the high-level search objective.
Search Agent System Prompt
The search agent system prompt defines structured MCP tool-use formatting and instructs the agent to produce complete, accurate, well-reasoned answers using additional tools.
- Agent Objective: The agent’s objective is to provide complete, accurate, and well-reasoned answers using additional tools.The prompt also incorporates information from the previous tool use and the current date.
- Tool Use: The prompt provides a tool-use formatting section for structured MCP interactions.It specifies server and tool names, argument tags, and valid JSON formatting.
- Tool Use: Agents invoke server tools through the `use_mcp_tool` interface.The interface requires server_name, tool_name, and arguments fields enclosed in dedicated tags.
- Tool Use: The prompt supplies available tool functions through MCP server definitions in JSON Schema format.The definitions are referenced using `{mcp_server_definitions}`.
B Pseudocode
PBSD samples trajectories, computes verifier-based group-normalized advantages, and reweights token advantages using privileged answer-conditioned likelihood comparisons at the turn level. The resulting policy-gradient update is otherwise unchanged, and privileged answers are excluded at inference.
- Training procedure: The turn-level likelihood gap produces a multiplier that reweights every loss-token advantage within the corresponding assistant turn.Turns becoming more likely under privileged context are amplified or suppressed according to trajectory advantage, with clipped soft correction controlling the adjustment.
- Training procedure: The policy-gradient loss is updated using reweighted token advantages, while privileged information is used only for advantage reweighting and is unavailable at inference.At inference, the model receives only the original prompt without privileged answers.
- Training procedure: For each update, PBSD samples a group of trajectories from the policy and computes verifier rewards with normalized advantages.The algorithm takes policy πθ, prompts X, verifier V, and group size G as inputs.
- Training procedure: PBSD augments each prompt with its privileged answer, then scores generated tokens under the original and privileged contexts.The privileged context is formed separately for each trajectory, with token-position mapping used for aligned scoring.
C Case Study: Turn-Level Credit Assignment
The case studies show that PBSD assigns assistant-turn credit by comparing privileged teacher and student likelihoods, distinguishing useful reasoning from distracting or weakly supported steps. In both correct and incorrect rollouts, it separates locally valuable reasoning from misleading branches and erroneous decisions.
- Turn-level signal visualization: PBSD visualizes assistant-turn reasoning, teacher–student likelihood comparisons, and qualitative labels because rollout logs lack token-level likelihoods.Turns are labeled High-value, Low-value, or Filtered according to alignment, distraction, or weak likelihood gaps.
- Correct rollout: In the successful Amar Sangee rollout, PBSD marks clue resolution and candidate-title discovery as high-value while identifying redundant verification as low-value.Key examples include linking the forest-film clue to Satyajit Ray and surfacing Amar Sangee.
- Turn-level signal visualization: PBSD filters turns with weak likelihood gaps instead of assigning confident positive or negative credit.A Filtered turn is distinct from reasoning aligned with the teacher signal or a distracting or misleading branch.
- Incorrect rollout: In the incorrect rollout, PBSD distinguishes locally useful reasoning steps from misleading branches and the final erroneous decision despite the wrong answer.Figure 7 presents turn-level PBSD signals for this incorrect trajectory.