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

arXiv:2608.07068 · cs.AI · Submitted 2026-08-07 · Read on arXiv

Listen

Transcript

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.

Zhiyuan Liu, Tinghong Ye, Chenghao Liu, Yizhuo Li, Songfang Huang

Peking University · Zhejiang University · Harbin Institute of Technology

cs.AI

Submitted: 2026-08-07

Updated: 2026-08-10

Code: https://github.com/TPssp/MemOPD

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 77/100

The gist: The paper addresses the challenge of training long-horizon language agents that accumulate growing contexts during interaction, which "impairs performance and stability." Compact memory systems

Key concepts

On-policy distillation: A training technique where a teacher model provides supervision on actions sampled by the student's own policy, rather than on a fixed expert script. This ensures the teacher scores decisions the student actually made, but only if the state used for scoring matches the state during rollout.
Memory state alignment: The requirement that the exact tokenized input, token positions, causal visibility, and prediction positions used during training match those from the original rollout. Compact memory rewriting can change these, so identical text may represent different states. MemOPD reconstructs these states to keep training valid.
Rollout context equivalence (RCE): A verification check that compares the logits from a packed training sequence to those from running each model invocation independently. If they match within numerical tolerance, the packed computation is equivalent to the original rollout, ensuring no hidden corruption.

Terminology

Distillation — 24× in this paper · explained in 2 episodes
Alignment — 23× in this paper · explained in 1 episode
Actor — 14× in this paper · explained in 3 episodes
On-Policy Distillation — 13× in this paper · explained in 3 episodes
RAG — 7× in this paper · explained in 3 episodes
Benchmark — 5× in this paper · explained in 2 episodes
Provenance — 5× in this paper · explained in 4 episodes
Empirical Methods in Natural Language Processing — 4× in this paper
Knowledge Distillation — 4× in this paper · explained in 11 episodes
Proximal Policy Optimization (PPO) — 4× in this paper · explained in 5 episodes
Regularization — 4× in this paper · explained in 3 episodes
Reinforcement Learning — 3× in this paper · explained in 4 episodes
Verification — 3× in this paper · explained in 2 episodes
Ablation — 2× in this paper · explained in 1 episode
Alignment Verification — 2× in this paper · explained in 1 episode
International Conference on Learning Representations — 2× in this paper
Transformer — 2× in this paper · explained in 2 episodes
AI Agents — 1× in this paper · explained in 2 episodes
Ablation Studies — 1× in this paper · explained in 1 episode
Benchmarking — 1× in this paper · explained in 3 episodes
Fidelity — 1× in this paper · explained in 2 episodes
Introduces Rollout Context Equivalence — 1× in this paper
KL Regularization — 1× in this paper · explained in 2 episodes
Large Language Models (LLMs) — 1× in this paper · explained in 6 episodes
Regret — 1× in this paper · explained in 3 episodes

Summary

The paper addresses the challenge of training long-horizon language agents that accumulate growing contexts during interaction, which impairs performance and stability. Compact memory systems mitigate this problem by compressing and rewriting the history retained between model invocations. However, 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.

The central problem identified is a state alignment issue: 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. During interaction, the same content can have different roles, positions, and visibility — e.g., a response token that was sampled becomes later context, so flattening changes the autoregressive prefix, causal visibility, and prediction locations. This causes the training objective to optimize a different conditional distribution from the one that produced the behavior."

The paper quantifies this problem in a 3B model audit: persistent-history reconstruction produced a p99 log probability error of 1.774, changed the top prediction at 651 sampled action positions, and falsely triggered PPO clipping for 13.29% of the actions. A key insight is that student-generated actions alone are insufficient for valid on-policy distillation, and the authors introduce memory state alignment as a necessary condition for applying OPD to long-horizon agents.

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. Specifically, the framework:

Reconstructs exact call states: During rollout, MemOPD records the exact tokenized input zt = [q, xt] and the exact sampled token IDs of each response at, rather than re-decoding. A memory state is formally defined as σt = (zt, ρt, Vt, γt), where ρt gives token positions, Vt specifies causal visibility, and γt(k) identifies the position whose output distribution predicts action token at,k. The paper emphasizes that identical decoded text can represent different memory states.
Packs reconstructed invocations: The compiler places recorded invocations into one physical sequence y = [q, x0, a0, x1, a1,..., xT−1, aT−1], treating physical adjacency and causal visibility separately. The stable prefix q is shared (stored once) because its token IDs and positions are identical across invocations, while tokens from different calls have blocked visibility between them. Positions restart from the indices used by the corresponding independent invocation instead of increasing across the packed sequence. When a response is retained for a later invocation, it is represented twice: once as the exact sampled action and once as context with the positions and visibility of that later invocation.
Defines the sampled action domain: Retained response copies and observations condition future actions but were not sampled at those physical positions. The mask mact j = 1[yj lies in the sampled occurrence of a recorded at] selects the true policy decisions. Although a later context occurrence may repeat the same token IDs, it occupies different physical positions and receives mact = 0.
Introduces Rollout Context Equivalence (RCE): RCE requires the packed representation to produce the same action logits as independently executing each model invocation: ‖lpack g(t,k) − lind t,k‖∞ ≤ ϵnum for every sampled action token. RCE therefore tests whether packed training reproduces the conditional computation used during rollout rather than merely forming a valid tensor.

MemOPD combines local teacher guidance with task-level policy optimization:

Teacher guidance: Uses full-vocabulary reverse KL divergence, LOPD = (1/Σ j mT j) Σ j mT j Σ v pθ,j(v) log[pθ,j(v)/pT,j(v)], where the frozen teacher and student receive the same tokens, positions, and visibility. The teacher supplies dense local preferences, but it does not determine whether a sequence of memory updates and retrieval actions succeeds.
PPO task objective: PPO retains the final task reward, defined over the sampled action domain, with a reference KL penalty KLref j = [lold,j − lref,j]+ and reward rej = rj − βKLref j. GAE advances across ordered sampled actions and skips positions used only for context. The actor objective is Lactor = LPPO − αH + λLOPD.

The paper stresses that the reference policy is separate from both the teacher and the behavior snapshot — the reference is frozen at SFT initialization while πold is the rollout-generating policy.

Benchmarks: Multi-objective QA from HotpotQA and Natural Questions (Q2, Q8, Q16 with 2, 8, 16 questions each), trained on Q2 and tested on transfer to longer horizons, plus single-objective Wiki-RAG.
Models: Qwen2.5 3B student; gpt-oss-120b generates training trajectories (20,036 turn-level examples); a frozen 7B teacher with λ = 0.02; five independent seeds.

Multi-objective retrieval: "MemOPD-3B achieves the highest EM and F1 at all horizons. Averaged over five seeds, MemOPD improves F1 over PPO by 14.3%, 283.2%, and 416.2% on Q2, Q8, and Q16, respectively. On Q16, it reduces peak context by 33.0% and inference time by 14.1%. The growing advantage shows that aligned teacher guidance is more useful at longer horizons."
Wiki-RAG transfer: MemOPD improves EM and F1 over PPO by 6.1% and 7.4%, respectively... It also has the lowest dependency and inference time in the table. Relative to PPO, MemOPD reduces dependency by 31.6% and inference time by 28.5%.
State alignment verification: In an audit of 64 trajectories, 199 invocations, and 12,083 sampled tokens, persistent history changes the top prediction at 651 action positions and falsely clips 13.29% of the actions even before any policy update. Reconstructed packing instead matches independent invocations at the numerical floor. After memory updates, top-1 agreement falls to 81.20%, the p99 sampled log probability difference reaches 19.854, and the mean KL divergence reaches 0.793, changing the teacher prediction at 444 action positions.
Matched teacher-state control: "Persistent-history teacher scoring improves F1 and EM over PPO by 5.6% and 4.4%, respectively, showing that dense teacher guidance remains useful despite the state mismatch. Reconstructing the teacher state further improves F1 and EM by 7.0% and 10.0%, respectively."
RCE across context updates: The compiler preserves RCE under full response retention, suffix retention, summary replacement, sliding windows, and retrieval refresh — all achieving p99 log-probability errors around 3.4–3.8 × 10−5 with 100% top-1 agreement — showing memory state alignment depends on the actual computation presented to the model rather than the particular form of memory rewriting.
Packing efficiency: Packing achieves up to a 1.63× actor speedup while preserving RCE. This speedup comes from shared computation without changing the optimized states (1.17–1.63× forward, 1.08–1.63× forward+backward across nine configurations).

Visibility corruption is the larger source of error compared to position corruption, but incorrect positions alone still change 260 top predictions and falsely clip 4.97% of the actions.
The exact action domain selects 12,083 sampled action tokens, whereas trajectory and response masks would add 66,830 and 74,357 spurious positions, respectively — demonstrating that state reconstruction is distinct from identifying which tokens are decisions.

The paper concludes: "Compact memory can separate sampled actions from the states later used to train them. MemOPD resolves this mismatch by reconstructing every invocation, distinguishing sampled actions from later context copies, and combining state-aligned teacher guidance with task-level PPO. RCE verifies the packed computation, while the matched control shows that teacher guidance already improves F1 over PPO by 5.6% and state alignment adds a further 7.0%. Across five seeds, MemOPD improves F1 over PPO by up to 416.2%, while packing yields up to a 1.63× speedup in actor computation during training."

Improvements for AI systems

An AI system improved with MemOPD can:

Train long-horizon agents without state mismatch from memory compression. During rollout, the system records exact tokenized inputs and sampled token IDs per invocation; during training, it restores each invocation’s original token positions, causal visibility, and sampled-action locations. This prevents the teacher or PPO from scoring an action under a context the student never actually visited.
Use a valid action mask that separates sampled decisions from later context copies. Only truly sampled action tokens receive policy-gradient and distillation losses; repeated copies of the same content in later observations are treated purely as context, avoiding spurious gradients on positions that were never policy outputs.
Provide dense, state-aligned local supervision. A frozen teacher computes full-vocabulary reverse KL against the student’s own distribution on the exact same state (same tokens, positions, and causal mask), giving per-action learning signal for memory-update and retrieval decisions instead of relying only on sparse final-task rewards.
Combine dense teacher guidance with task-level PPO. The system retains a PPO objective over the final reward with a reference KL penalty computed against a frozen SFT policy that is separate from both the teacher and the behavior snapshot, preserving exploration stability while following local preferences.
Guarantee that packed training is computationally equivalent to independent invocations. Before training, the system can check Rollout Context Equivalence (RCE): packed logits for every sampled action must match independently executed invocation logits to within numerical error. This catches any corruption from batching, position resets, or visibility blocking before gradients are computed.
Train on compact memory systems more efficiently. The compiler packs all recorded invocations into one physical sequence by sharing the stable prefix and blocking cross-invocation visibility, achieving 1.17–1.63× forward speedups and 1.08–1.63× forward+backward speedups while preserving exact rollout equivalence.
Scale to longer horizons without performance collapse. After training on short tasks (2-question QA), the system transfers to much longer horizons (8 and 16 questions): F1 improves over PPO by 283.2% on Q8 and 416.2% on Q16 in the paper’s multi-objective retrieval benchmark, with 33.0% lower peak context and 14.1% lower inference time.
Handle diverse memory-rewriting mechanisms robustly. Whether memory is updated by full response retention, suffix retention, summary replacement, sliding windows, or retrieval refresh, the system maintains RCE and numerically identical student/teacher states, so the training objective stays aligned regardless of the compression strategy.
Avoid false PPO clipping and incorrect top-prediction changes. In contrast to persistent-history flattening—which in the paper’s audit caused p99 log-probability error of 1.774, changed top predictions at 651 sampled action positions, and falsely clipped 13.29% of actions—the reconstructed-packing approach matches independent invocation outputs at the numerical floor.
Achieve superior task performance on downstream long-horizon agent benchmarks. The improved system demonstrates higher EM/F1 on multi-objective retrieval and Wiki-RAG transfer with fewer dependencies and lower inference time, while retaining the capability to optimize memory-update and retrieval policies rather than merely answering isolated questions.