Credit Assignment with Resets in Language Model Reasoning
For researchers and practitioners training language models on multi-step reasoning tasks, this work provides a simple yet effective method to improve credit assignment without external supervision, leading to better performance.
The paper addresses the problem of uniform credit assignment in reinforcement learning for language model reasoning, which ignores which steps contributed to success or failure. The proposed Self-Reset Policy Optimization (SRPO) method, which self-localizes erroneous steps and resets there, consistently outperforms standard GRPO and Random-Reset Policy Optimization (RRPO) across models and reasoning benchmarks.
Contemporary reinforcement learning with verifiable reward methods post-train language models on multi-step reasoning by assigning a single outcome reward uniformly across all tokens in a trajectory. Such uniform assignment ignores which steps contributed to success or failure. Improving credit assignment can address this limitation by enabling targeted refinement of faulty reasoning steps, rather than updating entire trajectories uniformly. Resets are one such simple mechanism, enabling more precise credit assignment by returning to an intermediate state and resampling counterfactual continuations, so that outcome differences can be attributed to decisions made at that point. We propose two such methods: Random-Reset Policy Optimization (RRPO), where reset states are drawn randomly from reasoning steps, and Self-Reset Policy Optimization (SRPO), where the model self-localizes the erroneous step in an incorrect trajectory and resets there. We analyze these methods within the Conservative Policy Iteration (CPI) framework. Extending CPI with a credit-assignment oracle that targets improvable states yields provable improvements over random resets. Across models and reasoning benchmarks, SRPO consistently outperforms standard GRPO and RRPO by sampling multiple suffix continuations at a self-localized reset and learning from their rewards, using only the model itself with no external supervision.