本文提出’Trajectory Policy Gradient Theorem’，从理论上证明在LLM在线强化学习中仅用响应级别奖励即可无偏估计token级奖励的策略梯度，并基于此设计了TRePO算法，简化PPO设计并具备token级建模能力。
Reinforcement Learning, Large Language Model, Instruction Tuning, Reasoning, Efficiency
Shenghua He, Tian Xia, Xuan Zhou, Hui Wei
Amazon, PAII Inc., UC Merced
Generated by grok-3
Background Problem
大语言模型（LLM）的强化学习（RL）在后训练阶段（如RLHF）中面临一个常见挑战，即’零奖励假设’（Zero-Reward Assumption），其中非终止动作（即中间token生成）没有即时任务特定奖励，只有最终token获得整个响应的奖励。这种假设源于实践中难以获取精确的token级奖励，而响应级别奖励更容易获得。现有方法分为两类：一类（如GRPO, ReMax）简化critic网络，认为中间奖励为零对critic影响较小；另一类试图通过响应级别奖励模型近似token级奖励，但缺乏理论支持。本文试图解决这一问题，通过提出’Trajectory Policy Gradient Theorem’，从理论上证明仅用响应级别奖励即可无偏估计token级奖励的策略梯度，为LLM对齐提供更高效、实用的RL方法。
Method
本文提出了’Trajectory Policy Gradient Theorem’，其核心思想是：在LLM应用场景中，基于REINFORCE和Actor-Critic类算法，即使只有响应级别奖励（RM(W)），也可以无偏估计真实未知token级奖励的策略梯度。具体推导如下：
- 通过Lemma 1，重新定义了LLM中状态（W0,t-1）和动作（wt），并推导出策略梯度公式，表明可以直接从策略πθ采样轨迹，而无需依赖经典RL中的稳态分布。
- Theorem 1进一步证明，通过在每个时间步t采样额外轨迹W(t+1)，计算响应级别奖励的期望值EW(t+1)[RM(W(t+1))]，即可无偏估计策略梯度，而无需明确计算token级奖励r(wt)或折扣因子γ。
- Theorem 2引入了近优基线（baseline）来降低梯度方差，基线为EW(t)[RM(W(t))]，通过采样估计。 基于此理论，作者提出了Token-Reinforced Policy Optimization (TRePO)算法：
- 在选定时间步D上，通过调整推理温度（从0开始逐渐增大）采样多个完整轨迹W(t)，估计期望奖励。
- 对非选定时间步，通过线性插值估计优势函数（advantage function）。
- 使用类似PPO的clipped surrogate loss进行优化，但无需critic网络，内存效率与GRPO相当。 批判性思考：理论推导看似严谨，但对采样数量和质量的依赖可能导致实际应用中估计偏差，尤其是在轨迹长度较长或奖励分布复杂的情况下。此外，TRePO的采样策略（调整温度、选择高概率区域）缺乏具体优化指导，可能在不同任务中表现不稳定。
Experiment
本文未提供具体的实验数据或结果来验证提出的理论和TRePO算法的有效性，仅在理论分析和文献回顾中引用了现有工作的实验结果（如Xu等人对PPO和DPO的比较）。作者提到未来工作将包括在数学推理、多轮医疗对话和多智能体RL等任务中评估TRePO与PPO、GRPO等方法的性能对比。 批判性思考：缺乏实验支持是本文的一大缺陷。理论上，TRePO通过采样估计期望奖励可能比PPO更准确地捕捉token级信息，但其计算成本（多次采样）和工程复杂性（温度调整、轨迹选择）未被实际验证。文献中引用的PPO优于DPO的结果（如在CodeContests数据集上PPO的通过率达22.4%，而DPO仅为0.0%）虽支持作者对PPO理论优势的推测，但无法直接推及TRePO的实际效果。此外，作者未讨论采样数量M和时间步选择|D|对性能的影响，实验设计的合理性和全面性存疑。
Further Thoughts
本文提出的’Trajectory Policy Gradient Theorem’在理论上为LLM的RL训练提供了一个有趣的视角，即通过响应级别奖励间接建模token级信息，这可能启发未来在奖励模型设计上的创新。例如，可以探索如何结合多模态数据（如文本和图像）构建更强大的响应级别奖励模型，而无需关注token级分解。此外，TRePO在多轮交互任务（如医疗对话）中的潜在应用值得关注，因为其对中间步骤奖励的估计能力可能比DPO等方法更适合处理长序列依赖。然而，实际应用中，采样带来的计算成本可能限制其在大规模LLM训练中的普及，未来可以考虑与高效推理技术（如vLLM提到的PagedAttention）结合，优化TRePO的实现效率。另一个值得思考的方向是，是否可以通过自适应采样策略（如基于奖励方差动态调整采样数量）来平衡精度和成本，这可能需要结合在线学习或元学习方法来实现。