Memory-R2: Fair Credit Assignment for Long-Horizon Memory-Augmented LLM Agents

1 Introduction

• Memory changes inherited states across rollouts, making GRPO comparisons unfair and trajectory-level credit signals noisy or biased.Downstream failures are difficult to attribute to current operations, earlier corrupted memory, or later updates that overwrite useful information.Source: §1 · p. 2
• LoGo-GRPO combines trajectory-level global rewards with local rerollouts from identical intermediate memory states for fairer comparisons and cleaner memory-operation supervision.It also introduces session-wise attribution signals while preserving end-to-end long-horizon optimization.Source: §1 · p. 2
• A shared LLM backbone with role-specific prompts supports cooperative fact extraction and memory management, jointly targeting memory formation and evolution.The extractor identifies salient information; the manager decides whether to insert, update, or delete entries.Source: §1 · p. 2
• Chunk-wise alternating extractor and manager decisions make memory construction a multi-step process, while training horizons progress from 8 to 16 to 32 sessions.The curriculum is intended to stabilize long-horizon optimization by building reliable short-horizon behavior first.Source: §1 · p. 3
• The framework reports strong gains over prior memory-agent baselines using only two training conversations, with generalization across benchmarks, model scales, and answer agents.The stated contributions also include fairer session-level credit assignment through global-local optimization.Source: §1 · p. 3

2.1 Memory Agent Architectures

• Memory-agent architectures include graph- or structure-based substrates and system-inspired designs, while many rely on heuristic or prompt-based storage and editing policies.Memory-R2 retains a modular extractor–manager architecture but optimizes the memory lifecycle with reinforcement learning.Source: §2.1 · p. 3

2.2 Reinforcement Learning for Memory Agents

• RL-based memory methods show promise but largely use outcome-level rewards, do not explicitly address credit assignment across diverging memory states, and leave joint formation, evolution, and retrieval underexplored.Memory-R2 addresses these gaps with multi-step extractor–manager training, shared-parameter co-learning, and a global-local GRPO objective.Source: §2.2 · p. 3

3.1 Problem Formulation: Multi-step Memory Bank Construction

• Memory construction proceeds chunk by chunk: a fact extractor proposes salient content, then a memory manager selects operations such as INSERT, UPDATE, or DELETE.The two cooperative roles use role-specific prompts on a shared LLM backbone, and the shared policy is optimized using downstream task performance.Source: §3.1 · pp. 3–4

3.2 Length-Normalized Step-level RL with Shared Extractor–Manager Policy

• Each extractor or manager call is treated as one generation step, giving steps comparable weight regardless of output length.The step-level ratio and advantage are used in the LoGo-GRPO objective.Source: §3.2 · p. 4

3.3 LoGo-GRPO for Multi-session Credit Assignment

• LoGo-GRPO combines full-trajectory global learning with local rerollouts from shared intermediate memory states for lower-bias session-level credit assignment.It samples a subset of sessions for local rerolls, balancing long-horizon supervision with fairer comparisons at manageable cost.Source: §3.3 · pp. 4–5
• QA quality is measured with token-level F1, while a compression penalty discourages memory tokens from exceeding a fixed fraction of cumulative session tokens.Global rollouts evaluate terminal memory and attribute reward to sessions according to where required evidence occurs.Source: §3.3 · p. 5

3.4 Curriculum Learning for Long-Horizon Credit Assignment

• Training horizon increases in stages from 8 to 16 to 32 sessions, with each stage's best validation checkpoint initializing the next.The stages progress from basic memory operations to stronger inter-session dependencies and full long-horizon optimization.Source: §3.4 · p. 6

4.1 Experiment Setup

• Experiments train on LoCoMo with a 2:1:7 train/validation/test split and evaluate OOD transfer on LongMemEval, MSC-Self-Instruct, and MemBench using F1, BLEU-1, judge scores, and M-Fail.The extractor and manager share a Qwen2.5-7B-Instruct backbone; the answer agent is held fixed during training for stable reward signals.Source: §4.1 · p. 6

4.2 Main Results

• Memory-R2 achieves the best overall F1 and BLEU-1 among training-free and trained LoCoMo baselines; versus Memory-R1, F1 rises from 43.14 to 50.60 and B1 from 36.44 to 44.01.The backbone-controlled comparison uses Qwen2.5-7B-Instruct; Memory-R2 also reaches judge score 80.99.Source: §4.2 · p. 7, Table 1 · p. 7
• Zero-shot OOD gains span all reported metrics, and F1 rises from 10.3 to 46.8 for Qwen2.5-3B, with memory-module training driving the dominant gains across answer agents.On LongMemEval-oracle, F1 improves from 27.88 to 50.60; the reported gains are achieved when trained on only two LoCoMo conversations.Source: §4.2 · pp. 7–8, Figure 2 · p. 7

4.3 Ablation Studies

• On LoCoMo, replacing LoGo-GRPO with GRPO lowers F1 from 49.67 to 46.62 and B1 from 43.77 to 40.97; removing curriculum lowers F1 to 24.12 and raises M-Fail to 46.5%.Direct 32-session training falls from validation F1 0.47 to 0.27 and raises M-Fail above 70%, while curriculum training stabilizes around F1 0.50 with M-Fail under 7%.Source: §4.3 · p. 8, Table 2 · p. 8
• Length-normalized step weighting outperforms token-level loss, which yields F1 43.53 and B1 38.10, and is reported to prevent output-length bias under the shared policy.The paper attributes overall gains to fair credit assignment, curriculum learning, multi-step construction, and shared extractor–manager co-learning.Source: §4.3 · pp. 8–9
• Architecture ablations support explicit extractor–manager roles and parameter sharing: merging roles yields 39.14 F1, disjoint parameters 44.31 F1, and training only the extractor 28.30 F1.Alternative interaction depths score 40.39, 41.37, and 37.61 F1 for N=4/8/10; training only the manager yields 45.34 F1.Source: §4.3 · pp. 8–9

4.4 More Analysis: Latency and Compression

• Memory-R2 improves F1 while reducing per-conversation inference latency for Qwen2.5-3B and Qwen2.5-7B; shorter generations drive the 3B gain, while shorter and more stable generations contribute at 7B.The reported latency changes accompany reduced decoding work and more stable memory construction at 7B.Source: §4.4 · p. 9
• Across Qwen2.5-3B and Qwen2.5-7B, λcomp = 0.3 achieves the best F1 and BLEU-1, while smaller penalties may retain noisy memories and stronger compression may remove useful evidence.The authors use λcomp = 0.3 as the default setting.Source: §4.4 · p. 9, Figure 4 · p. 9

5 Conclusion

• Memory-R2 targets fair credit assignment under diverging memory states in multi-session reinforcement learning.Its LoGo-GRPO method combines trajectory-level optimization with local rerollouts from shared intermediate memory states.Source: §5 · p. 9
• The framework jointly trains memory formation and evolution through a shared extractor–manager policy, uses chunked multi-step construction, and applies a session-horizon curriculum.These components accompany the fair-credit-assignment method.Source: §5 · p. 9
• Experiments report performance over prior memory-agent baselines and generalization across OOD benchmarks, model scales, and answer agents; the authors call credit assignment a key ingredient for robust long-horizon agents.They position the work as a foundation for future memory-centric RL research.Source: §5 · pp. 9–10

• Memory-bank construction alternates extraction and management over chunked sessions.Source: Figure 5 · p. 13
• Global rollouts optimize end-to-end performance using rewards from the final memory.Source: Figure 5 · p. 13
• Local rerollouts from shared memory states provide low-bias credit assignment, and both signals enter one GRPO-style objective.Curriculum learning enables stable long-horizon training.Source: Figure 5 · p. 13

• For each question, the answer agent receives each speaker’s top-30 memory entries by embedding similarity, subject to a 0.3 threshold.The prompt asks for step-by-step reasoning over timestamped memories and a final answer inside an <answer> tag.Source: §A.1 · p. 13
• Training examples come from the full memory-construction pipeline on two LoCoMo training conversations using GPT-4o at T=0, with samples retained only when token-level F1 is at least 0.25.Retained samples are randomly split 90%/10% into training and validation sets.Source: §A.1 · p. 13
• The answer-level reward uses standard SQuAD-style token F1 after lowercasing, punctuation and article removal, and whitespace tokenization.The reward is aligned with the answer-level evaluation metric.Source: §A.1 · p. 13
• GRPO training uses eight rollouts per prompt, runs for five epochs, and applies a token-level KL penalty with coefficient 0.001 in the actor loss.Evaluation and checkpointing occur every five optimizer steps.Source: §A.1 · p. 13

A.2 RL Training for Fact Extraction and Memory Management

• The fact extractor and memory manager share one Qwen2.5-7B-Instruct backbone, trained with a LoCoMo curriculum expanding trajectories from 8 to 16 to 32 sessions.Training uses two conversations with 328 QA pairs; GRPO samples 16 global trajectories and uses a 0.5 local sampling fraction with 4 resampled turns.Source: §A.2 · p. 14

B.1 Prompt Template for Fact Extraction (Memory Formation)

• Fact extraction emits atomic, self-contained third-person facts naming the subject, preserving temporal details, and tagging each fact with its originating dia_id.The prompt directs the extractor to produce one JSON object per durable fact for downstream memory management.Source: §B.1 · pp. 14–16

B.2 Prompt Template for the Memory Manager (Memory Evolution)

• The memory manager selects insert, update, delete, or no operation while enforcing atomic entries, fact-preserving updates, and verification of noisy retrieved candidates.Updates must not silently discard prior factual claims, and embedding-similarity candidates may be unrelated to the incoming fact.Source: §B.2 · p. 17
• The manager prompt requires each memory item to represent one fact and prohibits updates from removing existing factual information.Figure 7 describes the prompt’s JSON edit output and its constraints against fact loss, duplicate entries, and noisy retrieval-driven overwrites.Source: §B.2 · p. 19, Figure 7 · p. 22
• The manager checks for contradiction, semantic duplication, same-story refinement or progression, and otherwise inserts, in that order.Updates retain prior factual information; a later development of the same entity’s story can update an existing entry.Source: §B.2 · pp. 19, 21

B.3 Prompt Template for Answer Agent (Memory Usage)

• The answer agent resolves relative time expressions using the supporting memory’s timestamp and separates memories into speaker-labeled blocks to avoid speaker confusion.It returns an <answer>...</answer> span, which is extracted and scored against the gold answer with SQuAD-style token-level F1.Source: §B.3 · p. 22

C.1 LLM-as-a-Judge

• LLM-as-a-Judge assigns binary Correct/Wrong labels to assess semantic equivalence, with leniency for formatting differences and more verbose answers.The evaluator is gpt-4o-mini, chosen as a cost-quality trade-off for scoring the full evaluation set.Source: §C.1 · pp. 23–24, Figure 9 · p. 24

C.2 Memory-Failure Rate (M-Fail)

• M-Fail is the fraction of required gold evidence absent from the memory bank; zero means all required evidence is present.The metric isolates memory-construction errors from downstream retrieval or answer-generation effects.Source: §C.2 · p. 24

D Limitations and Future Work

• The evidence covers text-only multi-session dialogue and downstream QA; effectiveness in image-, video-, or embodied settings remains unexplored.The authors say practical effectiveness in richer environments still requires investigation.Source: D Limitations and Future Work · p. 24
• More effective memory construction may support personalized assistants, long-term educational support, and tools that maintain user context across interactions.The authors describe these as potential positive impacts, not established deployment outcomes.Source: D Limitations and Future Work · p. 24
• Persistent memory may retain sensitive information or allow stored errors to affect later interactions, raising privacy and security concerns.The authors identify privacy-preserving design, secure storage and access control, and safer memory updates as future directions.Source: D Limitations and Future Work · p. 25

E Additional Ablations and Generalization Results

• LoGo-GRPO consistently outperforms standard GRPO across question types and curriculum stages, supporting the benefit of local rerollouts.The comparison is reported across all evaluated categories and stages.Source: E Additional Ablations and… · p. 25
• Memory-R2 improves F1 while reducing per-conversation latency, with gains linked to more concise generation and shorter output-length tails.The diagnostics also report reduced total completion tokens, especially for Qwen2.5-3B.Source: E Additional Ablations and… · p. 25, Figure 12 · p. 26
• Mean and standard deviation over three runs confirm that Memory-R2’s reported improvements are robust across random seeds.Table 3 reports results averaged across three independent runs with different random seeds.Source: E Additional Ablations and… · p. 25
• Curriculum training stabilizes validation F1 near 0.50, while direct 32-session training falls from 0.47 to 0.27; M-Fail stays below 7% versus over 70%.Curriculum memory-bank size grows 194→416→483 items at stage ends, while the direct run inflates and truncates erratically; per-session F1 is 0.61/0.58/0.45 versus 0.23 for the direct run.Source: E Additional Ablations and… · p. 25, Figure 11 · p. 26