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

2026-08-07

AI Summary

This work addresses the performance degradation in long-horizon agents caused by context accumulation and the inefficacy of sparse rewards in guiding memory updates. To resolve the misalignment between teacher scoring states induced by memory compression, the authors propose an online policy distillation mechanism with memory alignment, which reconstructs the original model invocation states to restore token positions and causal visibility. The approach integrates context packing, full-vocabulary supervision signals, and joint optimization with PPO. Compared to conventional history concatenation methods, it achieves a 7.0% absolute improvement in F1 score; relative to standard PPO, it yields up to a 416.2% performance gain while accelerating actor computation by 1.63× during training.

Abstract

Long-horizon agents accumulate growing contexts during interaction, impairing performance and stability. Compact memory mitigates this problem by compressing and rewriting the history retained between model invocations. Learning what to retain typically relies on proximal policy optimization (PPO) with final task rewards, but sparse rewards provide little guidance for individual memory updates. This limitation motivates on-policy distillation (OPD), which supplies dense teacher supervision on student rollouts. For such supervision to be valid, the teacher must evaluate each sampled action under the same state in which it was generated. However, the context rewriting performed during memory compression can break this alignment. When sampled responses are retained and re-encoded for later invocations, flattening the interaction into a persistent history may cause the teacher to score the action under a state that the student never visited during rollout. The action therefore remains on-policy by provenance, but not necessarily by state. We therefore propose Memory-Aligned On-Policy Distillation (MemOPD). MemOPD records the inputs and sampled outputs of each model invocation, restores its original token positions and causal visibility, and packs the reconstructed invocations for efficient teacher scoring. The teacher provides full-vocabulary supervision at the sampled action positions, while PPO preserves the final task objective. Experiments verify state alignment across several context updates and show that it improves F1 by 7.0% over persistent-history teacher scoring in a matched control. Overall, MemOPD-3B improves F1 over PPO by up to 416.2%, while packing yields up to a 1.63x speedup in actor computation during training. The code for this work is publicly available at: https://github.com/TPssp/MemOPD.

Problem

• long-horizon agents
• memory compression
• on-policy distillation
• state alignment
• sparse rewards

Innovation

• on-policy distillation
• memory state alignment
• long-horizon agents
• context compression
• teacher-student supervision