PBSD: Privileged Bayesian Self-Distillation for Long-Horizon Credit Assignment
2026-06-08 • Machine Learning
Machine LearningComputation and Language
AI summaryⓘ
The authors address the problem of figuring out which specific steps in a long sequence of actions help achieve a final goal when only the final result is rewarded, a challenge especially tricky in multi-turn tasks. They introduce PBSD, a method that uses Bayesian reasoning to break down the overall outcome into signals for each step, helping identify which steps contributed positively or negatively. This approach turns sparse final feedback into detailed guidance for learning, improving the agent's performance and ability to generalize to new situations. Their experiments show that PBSD helps agents learn better policies and transfer knowledge more effectively.
Reinforcement LearningCredit AssignmentBayesian InferenceSelf-DistillationMulti-turn AgentsSparse RewardsPolicy OptimizationTrajectoryLikelihood RatioGeneralization
Authors
Yang Tian, Rui Wang, Xumeng Wen, Junjie Li, Shizhao Sun, Lei Song, Jiang Bian, Bo Zhao
Abstract
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.