🤖 AI Summary
This paper addresses the pervasive “zero-reward assumption” challenge in LLM reinforcement learning—namely, the difficulty of obtaining token-level immediate rewards, leaving only sparse, response-level rewards. We introduce the Trajectory Policy Gradient Theorem, which rigorously proves for the first time that response-level rewards can yield unbiased estimates of the true token-level policy gradient. Methodologically, we establish a mathematical equivalence between response-level rewards and token-level gradients, revealing that mainstream algorithms—including PPO and GRPO—are inherently compatible with this setting. Leveraging this insight, we propose TRePO, a lightweight, memory-efficient algorithm. Key contributions include: (1) providing a unified theoretical foundation for response-level RL; (2) substantially reducing reward engineering overhead in LLM alignment; and (3) achieving competitive performance with simplified implementation and improved training efficiency.
📝 Abstract
We study a common challenge in reinforcement learning for large language models (LLMs): the Zero-Reward Assumption, where non-terminal actions (i.e., intermediate token generations) receive zero task-specific immediate reward, while only the final token receives a reward for the entire response. This assumption arises frequently in practice, as precise token-level rewards are often difficult or infeasible to obtain in LLM applications. In this work, we provide a unifying theoretical perspective. We introduce the Trajectory Policy Gradient Theorem, which shows that the policy gradient based on true, unknown token-level rewards can be unbiasedly estimated using only a response-level reward model, regardless of whether the Zero-Reward Assumption holds or not, for algorithms in the REINFORCE and Actor-Critic families. This result reveals that widely used methods such as PPO, GRPO, ReMax, and RLOO inherently possess the capacity to model token-level reward signals, offering a theoretical justification for response-level reward approaches. Our findings pave the way for more practical, efficient LLM fine-tuning, allowing developers to treat training algorithms as black boxes and focus on improving the response-level reward model with auxiliary sub-models. We also offer a detailed analysis of popular RL and non-RL methods, comparing their theoretical foundations and practical advantages across common LLM tasks. Finally, we propose a new algorithm: Token-Reinforced Policy Optimization (TRePO), a theoretically grounded method that is simpler than PPO, matches GRPO in memory efficiency, and holds promise for broad applicability.