Search papers, labs, and topics across Lattice.
This paper introduces single-rollout proximal policy optimization (SR-PPO) for token-level credit assignment in reinforcement learning applied to language models, addressing the high computational cost and challenges in credit assignment associated with traditional methods. By utilizing a calibrated token-level credit critic that predicts Pass@k success probabilities from a single rollout, the authors demonstrate that SR-PPO provides a more selective learning signal, effectively prioritizing challenging prefixes. The results indicate that SR-PPO achieves stable learning dynamics and significant improvements in Pass@128 success rates on mathematical reasoning benchmarks, showcasing its potential for enhancing RL efficiency in language models.
SR-PPO achieves significant gains in reasoning tasks by effectively assigning credit to individual tokens from a single rollout, transforming how we approach reinforcement learning in language models.
Estimating token-level advantages in reinforcement learning (RL) for language models remains challenging because scaling up episodic experience collection is expensive. The difficulty intensifies for baseline advantage estimation methods, where repeated sampling causes trajectories to diverge into substantially different reasoning prefixes. In this context, RL algorithms such as GRPO prove limited: an outcome reward is too sparse to be attributed to specific actions like intermediate steps, and comparisons across sampled traces are non-trivial because they are heterogeneous. To mitigate both the computational cost of repeated sampling and the difficulty of credit assignment, we study single-rollout proximal policy optimization (SR-PPO) featuring token-level credit assignment in RL for language models. Instead of estimating advantages by normalizing episodic returns within the candidate group, we train a calibrated token-level credit critic using Monte Carlo outcomes from one rollout per prompt. Specifically, we use the critic to predict the Pass@k success probability at the prompt prefix, which is derived from a Pass@1 attempt. This choice yields a more selective learning signal than Pass@1: it discounts easily solved prefixes while prioritizing hard ones whose success probability remains marginal. We show that as $k$ increases, Pass@k converges to a reachability indicator, reflecting whether a prefix can lead to at least one successful continuation. In an explicit state graph, the limit ($k \rightarrow \infty$) can be computed in $O(|V|+|E|)$ time, offering a promising surrogate for direct credit assignment without the need to sample contrastive traces. As an initial validation, SR-PPO exhibits stable learning dynamics, along with consistent gains in Pass@128 success rates on mathematical reasoning benchmarks such as HMMT26 and AIME24.