Policy-compatible joint decoding improves reasoning on hard problems
When All Rollouts Fail: Policy-Compatible Joint Decoding for Hard-Problem Reasoning
Machine LearningArtificial Intelligence
Summary
When language models try to solve very hard problems, they often fail every attempt, so they can't learn what the right answer looks like. The authors show that using expert examples that are too different from the model’s usual way of thinking doesn’t help because the model just imitates shallowly. They propose a new method, PICO, which combines expert guidance with the model’s current preferences step-by-step, making learning more effective. Their tests show this helps language models get better at reasoning tasks.
What this means in practice
- •For ai developers: Enhance reasoning capabilities of language models on difficult problems using PICO to guide training.
- •For chatbot designers: Improve chatbot reasoning accuracy by integrating policy-compatible expert guidance during model training.
Authors
Zehao Wang, Hanzhu Chen, Runquan Gui, Zhe Wang, Lin Yang
Abstract
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.