论文

MemOPD:面向长程智能体的记忆状态对齐在线策略蒸馏

MemOPD: On-Policy Distillation through Memory State Alignment for Long-Horizon Agents

模型优化模型训练上下文与知识强化学习上下文工程记忆蒸馏Agentic RLAgent记忆

摘要

长程智能体在交互过程中累积不断增长的上下文,损害性能与稳定性。紧凑记忆通过压缩与改写模型调用之间保留的历史来缓解这一问题。学习应保留什么通常依赖带最终任务奖励的近端策略优化(PPO),但稀疏奖励对单次记忆更新几乎没有指导作用。这一局限促成了在线策略蒸馏(OPD),它在学生的采样轨迹上提供密集的教师监督。要使这种监督有效,教师必须在生成每个采样动作时所处的相同状态下对其进行评估。然而,记忆压缩期间执行的上下文改写会破坏这种对齐:当采样的响应被保留并重新编码以供后续调用时,把交互扁平化为持久历史可能导致教师在学生采样时从未到过的状态下为该动作打分。因此,该动作按来源仍是在线的(on-policy),但按状态则未必。为此我们提出记忆对齐在线策略蒸馏(MemOPD)。MemOPD记录每次模型调用的输入与采样输出,恢复其原始token位置与因果可见性,并将重建的调用打包以实现高效的教师打分。教师在采样动作位置提供全词表监督,同时PPO保留最终任务目标。实验验证了多种上下文更新下的状态对齐,并在匹配对照中显示其相比持久历史教师打分将F1提升7.0%。总体而言,MemOPD-3B相比PPO最多将F1提升416.2%,而打包使训练中的actor计算最多加速1.63倍。代码已开源于 https://github.com/TPssp/MemOPD。

MemOPD:面向长程智能体的记忆状态对齐在线策略蒸馏:论文配图
图2:MemOPD概览。此处,一次调用指一次模型调用。(1) MemOPD重构每次调用,将采样动作与其后在上下文中的出现区分开,并将对齐后的计算打包到共享任务前缀之下。(2) RCE在数值容差内比较打包与独立的完整词表logits。(3) 对齐后的批次支持在采样动作域上进行PPO,以及在教师监督掩码范围内进行完整词表的教师引导。