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

Introduction to the show: ident: Paper Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "MemOPD: On-Policy Distillation through Memory State Alignment for Long-Horizon Agents".

Jane: The paper was written by Zhiyuan Liu, Tinghong Ye, Chenghao Liu, Yizhuo Li and Songfang Huang from Peking University and Zhejiang University and Harbin Institute of Technology.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Paper summary: Tom: Good to be back, folks. Jane and I have a new paper on the table, and we brought friends.

Jane: Lu's here, Meng's here, and Lalam is dialing in from the big-picture desk.

Lu: Happy to be here. This one grabbed me fast.

Meng: Me too. It's on training agents that have to work over many steps without their context ballooning.

Tom: Right — long-horizon agents. They search, they retrieve, they answer, and every step adds tokens to the prompt.

Jane: That grows without bound, so some agents compress what they remember between calls. That's compact memory.

Lu: And the paper shows a subtle trap: when you train such an agent, you can get the history wrong.

Tom: Exactly. The agent samples an action under one context, but after memory compression, the training data may show that same action under a different context.

Jane: The tokens might be identical. The positions shift, the visibility changes, and suddenly the teacher is scoring a state the student never visited.

Meng: That's the core claim — an action can be on-policy by origin but not by state.

Tom: Their fix is MemOPD. It records each model invocation, reconstructs the exact state used for each sampled action, and packs those states efficiently for training.

Jane: They also build a verification check called RCE, rollout context equivalence, to prove the packed computation matches running each call independently.

Lu: And the numbers are wild. On the long Q16 benchmark, F1 jumps 416 percent over plain PPO training.

Meng: On Wiki-RAG, the gains are smaller but still real — about seven percent F1, with faster inference and lower dependency.

Lalam: Why it matters is bigger than one benchmark. Distillation and reinforcement learning both assume the teacher sees what the student saw. This paper shows how to honor that assumption when memory rewriting breaks it.

Tom: So it's a training-correctness paper with a practical engine behind it.

Jane: And it makes you question every flattened transcript you've ever trained on.

Lu: I love that. A genuine "wait, is my batch even valid?" moment.

Tom: We'll walk through the paper page by page, starting where it starts — with the problem statement.

Jane: And spoiler, the problem is sneakier than it looks. Let's go.

Page 1: Tom: So page one sets up the pain: agents that accumulate context get slower and less stable.

Jane: Every invocation adds reasoning, observations, tool output. Transformers pay attention cost, and models drown in their own history.

Lu: The paper points at compact memory as the fix — learn what to keep, rewrite the retained context between calls.

Meng: And it cites MEM1 as an example, a system that learns a compact internal state for both memory and reasoning.

Jane: But here's the training problem. These agents are usually tuned with PPO, and the reward only arrives at the very end of the task.

Tom: Sparse reward. A memory update in the middle gets no direct feedback on whether it helped.

Lu: So the paper turns to on-policy distillation. A teacher model can provide dense supervision on the student's own sampled actions.

Meng: That's the standard trick — teacher scores what the student actually generated, not a fixed expert script.

Jane: But there's a catch that the paper names on page one. For teacher supervision to be valid, the teacher has to score each action under the same state where the student produced it.

Tom: And compact memory breaks that. A response gets generated, then parts of it get retained and re-encoded into a later context.

Lu: Flattening the whole interaction into one transcript changes token positions, causal visibility, prediction locations.

Meng: So you get a training state that the student policy never actually visited during rollout.

Jane: The action stays on-policy by origin. But not by state. That's the crux.

Tom: The paper's answer, briefly: record the exact inputs and sampled outputs of every call, restore their original positions and visibility, and pack those reconstructed calls for scoring.

Lu: Then PPO keeps the final task objective while the teacher adds dense full-vocabulary guidance at the sampled action positions.

Meng: And the preview numbers are already on page one — F1 up 416.2 percent over PPO on the longest horizon, plus a 1.63x speedup from packing.

Jane: The abstract also gives a matched control number: 7 percent F1 improvement just from aligning teacher states versus persistent-history scoring.

Tom: I like how they frame the contribution — state alignment as a necessary condition for on-policy distillation, not an optional nicety.

Lalam: It's a correctness argument, really. The whole field of agent training quietly assumes the recorded context is the lived context. Page one says: check that assumption.

Meng: And they put the code online, so people can actually go check.

Jane: But the sharpest evidence comes on page two, where they show how badly flattening corrupts real rollouts.

Tom: Let's look at that. The audit numbers are brutal.

Page 2: Tom: Page two shows the damage with real numbers.

Jane: They audit a 3B model's rollouts and compare three ways of reconstructing the same actions.

Lu: Persistent history — just flattening everything into one transcript — produces a p99 log probability error of 1.774.

Meng: That's huge. It means the model's own predicted probabilities get distorted at the tail.

Tom: And the top prediction changed at 651 sampled action positions.

Jane: Plus 13.29 percent of actions would falsely trigger PPO clipping. The ratio looks out of bounds even though the policy never changed.

Lu: That's the scariest part. PPO's clipping mechanism assumes the ratio is wrong because the policy moved. Here it's wrong because the state is wrong.

Meng: The figure on page two explains the mechanism visually. A sampled response becomes retained memory, then a later call sees that memory in a different position.

Jane: Same content, different role. What was a prediction location becomes a conditioning token.

Tom: The paper calls it the "why flattening misaligns" story — and it's the heart of the motivation.

Lu: Then they list contributions. State mismatch formulation, the MemOPD reconstruction framework, and RCE as a verification criterion.

Meng: And the related work section is genuinely useful. It situates this between memory management systems and distillation.

Jane: Right — MemGPT, A-MEM, prompt compression, those manage what the agent remembers.

Tom: But none of them check whether the training interface can still score a sampled decision after rewriting.

Lu: And there's a nice nod to DAgger. DAgger says: query the expert on states the learner actually visits.

Meng: Exactly. But the paper pushes further — even if you collected supervision on learner states, a stored trajectory might not reproduce those states later.

Jane: So the stored state itself has to be re-verified. That's the gap.

Tom: Also they mention generalized knowledge distillation, where a teacher scores student-generated sequences. Standard OPD assumes the autoregressive prefix stays available.

Lu: Context rewriting breaks that assumption. A response reappears later as context with different visibility.

Meng: So the teacher scores the wrong conditional distribution. The action domain is right, but the conditioning is wrong.

Jane: That sets up page three, where they stop hand-waving and write down the formalism.

Tom: Time to define what a memory state actually is.

Page 3: Tom: Page three gets formal. They define the rollout loop with equations.

Jane: The behavior policy samples a response given the task prefix and the mutable context. The environment returns an observation and reward. Then a context update rule builds the next input.

Lu: That update rule is where the magic — or the corruption — happens.

Meng: They instantiate it with MEM1's compact memory protocol. Only retained content plus the newest observation survives into the next call.

Jane: So a three-call interaction looks like: task prefix, then a call with retained memory and observation, never an accumulated transcript.

Tom: That distinction is central. A sampled response and its later context copy are two different computational events.

Lu: And they remind us the pipeline starts with SFT to teach the format, then PPO with a frozen behavior snapshot, a reference policy, and a critic.

Meng: Before any objective scores an action, MemOPD reconstructs the invocation that produced it.

Jane: Then comes the definition I like — a memory state is not just the text.

Tom: Right. It's a tuple: the tokenized input, the token positions, the causal visibility, and the prediction position that maps to each action token.

Lu: So identical decoded text can correspond to two different memory states if the positions or visibility differ.

Meng: That's the punchline. Two strings that look the same to a human are different states to a transformer.

Jane: And the rollout state for a particular action token also includes the earlier tokens of the same response.

Tom: State alignment requires the reconstructed state to equal the rollout state for every sampled action token.

Lu: If that equality fails, the behavior likelihoods, the PPO ratios, and the teacher targets all describe a decision that was never made.

Meng: Provenance alone doesn't save you. The action came from the student, sure, but under a different context.

Jane: So they've turned a vague intuition into a precise requirement.

Lalam: And that precision is what makes the rest of the paper testable. You can't argue with an equality condition.

Tom: The next page shows how to actually satisfy it — reconstruction, packing, and the RCE test.

Lu: That's where the engineering gets clever.

Jane: Let's dig in.

Page 4: Tom: Page four is the engineering heart.

Jane: They record the exact token IDs at rollout time. No decoding and re-tokenizing, because that could change the sampled action.

Lu: Then a compiler arranges all invocations into one physical sequence.

Meng: The task prefix is stored once, with its original positions, and made visible to every invocation block.

Jane: Other tokens see only their own call's context. Attention across calls is blocked.

Tom: Positions restart per call instead of marching forward across the packed sequence. That's crucial.

Lu: And if a response was retained as memory, it appears twice — once as the sampled action, once as later context with the later call's positions and visibility.

Meng: Two occurrences, two roles, two different computational identities.

Jane: If a sequence limit would drop a conditioning token, the compiler rejects the example outright.

Tom: No silent changes to the training state. That's a strong design choice.

Lu: Then they define RCE — rollout context equivalence. For every sampled action token, the packed logits must match the independent-call logits within numerical tolerance.

Meng: That's a testable guarantee. You can literally run both and compare.

Jane: They test it against independent invocation logits, not just "looks like a valid tensor." That distinction matters.

Tom: Then the action domain. Correct states don't tell you which tokens are decisions.

Lu: The sampled action mask marks only the original sampled occurrence of each response token.

Meng: A later context copy with the same token IDs gets a zero. It conditions future actions; it isn't one.

Jane: They also define the teacher mask, which can select a subset of the sampled action domain.

Tom: And they're careful to say the action domain can't fix broken logits, and correct logits can't fix a wrong decision mask.

Lu: Both have to be right. State reconstruction preserves conditioning; the mask preserves which tokens count.

Meng: The figure pulls it all together — recording, separating, packing, then a shared pipeline for PPO and teacher guidance.

Jane: And the speedup comes from sharing that stable prefix across calls without changing the optimized states.

Tom: So alignment isn't just correct. It's efficient.

Lalam: That's the detail I value most. The field won't adopt a correctness fix that costs double. Packing makes it a win on both axes.

Lu: The next pages ask whether it pays off in actual benchmarks.

Jane: And the answer is a big yes.

Page 5: Tom: Page five layers the objectives on top of the aligned states.

Jane: Teacher guidance is a full-vocabulary reverse KL at the sampled action positions.

Lu: That exposes the teacher's whole distribution, not just the sampled token.

Meng: And PPO stays in charge of the final task reward. The teacher can't override whether the task actually succeeds.

Jane: There's also a reference policy — frozen at the SFT initialization — that penalizes drift with a token-level KL term.

Tom: The reward becomes task reward minus that reference penalty.

Lu: Clever detail: GAE advances across the ordered sampled actions, skipping context-copy positions.

Meng: So value learning respects the decision structure, not the storage layout.

Jane: The combined actor objective is PPO minus entropy plus a weighted teacher term.

Tom: Then the experiments. They build a multi-objective QA benchmark from HotpotQA and Natural Questions.

Lu: Each query packs several questions. Q2, Q8, Q16 — 2, 8, 16 questions — so retrieval gets longer and memory updates pile up.

Meng: Training happens on Q2. Q8 and Q16 test transfer to longer horizons. That's a tough test.

Jane: The student is Qwen2.5 3B. Trajectory data comes from a large teacher model, 20,036 turn-level examples after filtering.

Tom: PPO and MemOPD share the same initialization, the same data order, the same masks, the same evaluation protocol.

Lu: That's the right way to run a controlled comparison. Five seeds each.

Meng: And Table 1 is impressive. MemOPD beats PPO by 14.3 percent F1 on Q2, 283.2 percent on Q8.

Jane: Then 416.2 percent on Q16. The longer the horizon, the bigger the gain.

Tom: It also cuts peak context by 33 percent and inference time by 14 percent on Q16.

Lu: So the model isn't just scoring better — it's remembering more efficiently.

Meng: That growing advantage is the paper's best argument. Alignment matters more when memory rewriting happens more often.

Lalam: And the trend is exactly what you'd predict from the theory. More rewrites, more misalignment, more room for the fix to shine.

Jane: They're not done. Page six checks whether the gains transfer to a single-objective setting.

Tom: Plus the audit that proves the mechanism.

Page 6: Tom: Page six starts with Wiki-RAG, a single-objective retrieval benchmark.

Jane: MemOPD beats PPO by 6.1 percent EM and 7.4 percent F1.

Lu: And it wins on efficiency — 31.6 percent lower dependency, 28.5 percent faster inference.

Meng: Peak context creeps up 2.2 percent, which is a fair trade for much better answers.

Tom: But the real meat is the alignment audit.

Jane: They take 64 real trajectories, 199 model invocations, 12,083 sampled action tokens.

Lu: And they compare persistent-history scoring against independent calls and reconstructed packing.

Meng: Persistent history changes the top prediction at 651 positions and falsely clips 13.29 percent of actions.

Jane: Reconstructed packing matches the independent calls down at the numerical floor — 3.43e-5 in p99 log probability error.

Tom: So the corruption isn't a rounding artifact. It's structural.

Lu: The paper also shows the teacher divergence grows over time. Before the first memory update, exact and flattened states agree perfectly.

Meng: After later invocations, top-1 agreement between teacher states drops to 81.2 percent.

Jane: And the p99 log probability difference explodes to 19.854.

Tom: Mean KL of 0.793. The teacher flipped its prediction at 444 action positions.

Lu: Those are positions where the student gets guidance for a distribution it never actually faced.

Meng: Table 3 sums it up: persistent history fails, batching at the numerical floor passes, reconstructed packing passes.

Jane: Then Table 4 asks whether the method is tied to MEM1's specific memory scheme.

Tom: They test five controlled context updates — full response retention, suffix retention, summary replacement, sliding windows, retrieval refresh.

Lu: Every single one passes RCE at the numerical floor.

Meng: That's a strong generality claim. The compiler only cares about realized token contexts, not symbolic memory roles.

Jane: So the alignment interface is portable across memory designs.

Tom: Which sets up page seven: the ablations that isolate where the gains actually come from.

Lu: And those results are beautifully clean.

Page 7: Tom: Page seven runs the controlled experiments that pin down the mechanism.

Jane: First, a matched Q2 control. PPO without teacher, persistent-history teacher, and MemOPD. Everything else identical.

Lu: Persistent teacher still helps — 5.6 percent F1 and 4.4 percent EM over PPO.

Meng: So dense guidance is useful even when the state is wrong. That's an important result on its own.

Jane: Then state alignment adds another 7 percent F1 and 10 percent EM on top.

Tom: So the full package beats PPO by 13 percent F1 and 14.8 percent EM in that matched control.

Lu: That decomposition is lovely. Teacher helps, alignment helps more.

Meng: Next they corrupt states deliberately. They break visibility only, positions only, or both.

Jane: Visibility is the bigger error source. But wrong positions alone still flip 260 top predictions and falsely clip 4.97 percent of actions.

Tom: So both matter. You can't skip either.

Lu: They also separate state reconstruction from action-domain correctness.

Meng: The exact action domain selects 12,083 sampled tokens. Trajectory masks would add 66,830 spurious positions, response masks 74,357.

Jane: Wrong masks would count the same decision twice — once as the action, once as context.

Tom: Then the efficiency ablation. Packing speeds up the actor by up to 1.63x while preserving RCE.

Lu: And the discussion section draws the philosophical line: provenance versus state validity.

Meng: A rollout guarantees the student produced the action. It doesn't guarantee the recorded context is the lived context.

Jane: Teacher guidance and PPO serve different roles — local preference versus global task success.

Tom: And the persistent-teacher result shows the risk: you can get real gains and still be leaving a chunk of performance on the table.

Lu: Because part of your supervision is quietly optimizing the wrong distribution.

Meng: The transportable interface point lands too — any memory rewrite scheme can plug into this reconstruction compiler.

Lalam: Which means the contribution outlives this specific benchmark. It's a protocol for honest agent training, not a one-off trick.

Jane: So the method survives contact with different memory designs.

Tom: That's the whole arc. Now let's wrap it up.

Conclusion: Tom: So the paper leaves us with a clean lesson.

Jane: Sampled actions and valid states are separate things. Both must be checked during training.

Lu: The persistent-history shortcut looks fine on the surface, but the audits show it corrupts logits, teacher targets, even PPO clipping.

Meng: And the fix isn't exotic. Record what actually happened, restore it, verify it.

Jane: RCE gives the field a concrete way to audit training representations instead of trusting their shape.

Tom: The numbers back it up: 416 percent F1 gain on the longest horizon, 7 percent gain purely from alignment, 1.63x speedup.

Lalam: I think this nudges the whole agent-training ecosystem toward state-level honesty. If your batch doesn't reproduce the rollout, your objective is fiction.

Lu: And the benchmarks reward it. Longer tasks, bigger wins. That's exactly where memory rewriting gets aggressive.

Meng: The code is public, so we can all stress-test it on our own memory schemes.

Jane: For us, the takeaway is simple: when you train an agent, ask what state each token was actually sampled under.

Tom: If you can't answer that, your teacher might be grading a test the student never took.

Lu: Great image. We're stealing that.

Jane: We'll miss this paper, but the habit of asking the question will stick.

Tom: Time to move on — the next paper is already waiting.

Jane: Until then, keep your states aligned.