MemOPD: On-Policy Distillation through Memory State Alignment for Long-Horizon Agents
- URL: http://arxiv.org/abs/2608.07068v1
- Date: Fri, 07 Aug 2026 10:18:31 GMT
- Title: MemOPD: On-Policy Distillation through Memory State Alignment for Long-Horizon Agents
- Abstract summary: Long-horizon agents accumulate growing contexts during interaction, impairing performance and stability.<n>We propose Memory-Aligned On-Policy Distillation (MemOPD) to mitigate this problem.<n>MemOPD records the inputs and sampled outputs of each model invocation, restores its original token positions, and packs the reconstructed invocations for efficient teacher scoring.
- Score: 48.05798760668219
- License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
- 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.
Related papers
- Rethinking On-Policy Distillation of Large Language Models II: One Training Example [31.843793429278808]
On-policy distillation combines student-generated rollouts with dense token-level supervision from a teacher.<n>One-shot OPD keeps improving for hundreds of steps and recovers most of full-data OPD's gain across task domains and model families.<n>We explain this result through the states visited during training and the rate at which the student aligns with the teacher.
arXiv Detail & Related papers (2026-09-03T17:54:38Z) - OPTD: On-Policy Transition Distillation with Consistency-Guided Adaptive Compression for Few-Step Diffusion Language Models [39.15861870749305]
We propose OPTD, On-Policy Transition Distillation with consistency-guided adaptive compression.<n>It samples partial states from the few-step student's own trajectories, uses a frozen, question-only teacher to identify outcome-aligned future candidates, and orders them by current-state confidence.<n>It consistently improves the quality--efficiency trade-off and attains the strongest overall quality-constrained AUP among the evaluated few-step baselines.
arXiv Detail & Related papers (2026-08-03T23:09:43Z) - TFP: Temporally Conditioned Memory-Fusion Policies for Visuomotor Learning [11.149498165285491]
Vision--Language--Action (VLA) policies such as $_0.5$ and OpenVLA perform well on many tasks, but they are often reactive.<n>We argue that such tasks require not only memory, but dynamics-aware belief updates.
arXiv Detail & Related papers (2026-07-09T09:24:30Z) - Multi-Turn On-Policy Distillation with Prefix Replay [73.10000453999088]
We propose Replayed-Prefix On-Policy Distillation (ReOPD) for agentic tasks.<n>ReOPD reuses pre-collected teacher trajectories as replayed prefixes.<n>It preserves or improves OPD-level accuracy, uses zero tool calls during student training, and is at least 4$times$ faster per training step than OPD.
arXiv Detail & Related papers (2026-07-06T07:56:53Z) - Procedural Memory Distillation: Online Reflection for Self-Improving Language Models [51.938248694868584]
Reinforcement learning with verifiable rewards (RLVR) evaluates each rollout against a verifier and updates the policy from that episode-level signal.<n>Across episodes and epochs, the model repeatedly encounters related problems under a changing policy, producing cross-episode signals that episode-local updates cannot capture.<n>We propose Procedural Memory Distillation (PMD), which converts these crossepisode signals into reusable procedural memory and distills it into the policy's weights during training.
arXiv Detail & Related papers (2026-07-01T21:20:57Z) - Prune-OPD: Efficient and Reliable On-Policy Distillation for Long-Horizon Reasoning [66.52232008796294]
Prune-OPD dynamically aligns training budgets with supervision quality.<n>It reduces training time by 37.6%--68.0% while preserving, and often improving, performance on challenging benchmarks.
arXiv Detail & Related papers (2026-05-08T14:38:53Z) - TCOD: Exploring Temporal Curriculum in On-Policy Distillation for Multi-turn Autonomous Agents [55.27396165691312]
We identify a key limitation of vanilla OPD in multi-turn agent settings, which we term Trajectory-Level KL Instability.<n>We propose TCOD, a framework that controls the trajectory depth exposed to the student and progressively expands it with a curriculum schedule.<n> Experimental results across four student-teacher pairs show that TCOD mitigates KL escalation and enhances KL stability throughout training, improving agent performance by up to 18 points over vanilla OPD.
arXiv Detail & Related papers (2026-04-27T03:38:27Z) - Offline Reinforcement Learning with Implicit Q-Learning [85.62618088890787]
Current offline reinforcement learning methods need to query the value of unseen actions during training to improve the policy.
We propose an offline RL method that never needs to evaluate actions outside of the dataset.
This method enables the learned policy to improve substantially over the best behavior in the data through generalization.
arXiv Detail & Related papers (2021-10-12T17:05:05Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.