MemoryAgent Arena

A unified, extensible, and easy-to-use evaluation framework for memory agents.

This project is designed to evaluate memory-augmented agents under a common protocol: given a stream of memory chunks, an agent incrementally builds memory, answers questions at specified positions, and is scored with benchmark-specific metrics in a shared output format.

The framework supports both simple baselines such as concat/RAG and tool-using memory agents such as MemAgent, MemAlpha, MemT, UMA, and SMA. SMA training is maintained in a separate reproduction project; this repository stays focused on benchmark execution, scoring, and result analysis.

Why This Framework?

Memory-agent results are often hard to compare because different projects use different data formats, prompts, generation scripts, judge prompts, output files, and metrics. This repository aims to make those comparisons cleaner.

Core goals:

• Unified evaluation protocol: all agents consume the same EvalData format and emit the same responses_*.jsonl / evaluated_*.jsonl files.
• Extensible benchmark registry: benchmarks declare their loader, evaluator, metric profile, and agent data source in one place.
• Extensible agent registry: agents declare import path, adapter behavior, and runtime defaults in one place.
• Reusable runners: generation, evaluation, resume, sharding, and JSONL repair are handled by shared runner/runtime modules.
• Easy baseline execution: maa generate runs baselines with explicit task lists or --tasks all and automatic service management.
• Clean runtime boundary: training code is kept out of the arena, while lightweight UMA/SMA evaluation loops remain available through external/memory_agent.

Repository Layout

agents/                 Memory-agent implementations and agent registry/adapters
benchmarks/             BenchmarkSpec and central benchmark registry
data/                   Dataset builders and processed benchmark data
evaluation/             Evaluator implementations and evaluator registry
scripts/experimental/   Experiment-specific helper scripts
scripts/patches/        Environment compatibility patch helpers
scripts/run/            Optional shell wrappers around the `maa` CLI
scripts/tools/          Standalone utility scripts
external/               External agent runtimes and third-party benchmark code
docs/                   Analysis notes and implementation docs
assets/                 Figures

evaluate_async.py       CLI entry point for generation/evaluation
evaluation_runner.py    Shared generation/evaluation runners
benchmark_runtime.py    JSONL output, resume, and metric policy helpers
generate_stats.py       Result aggregation and plotting

Architecture

The framework is organized around three registries/adapters:

• benchmarks/registry.py: defines benchmark behavior through BenchmarkSpec.
• agents/registry.py: defines available agents through AgentSpec, including runtime defaults such as env vars, model, and concurrency.
• agents/adapters.py: adapts different agent APIs into one runner-facing protocol.
• evaluation/registry.py: maps benchmark specs to evaluator implementations.

At runtime, the flow is:

BenchmarkSpec -> load EvalData
AgentSpec + AgentAdapter -> generate responses
JsonlOutputStore -> write/resume responses_*.jsonl
Evaluator -> score responses
MetricPolicy -> summarize evaluated_*.jsonl

This means adding a benchmark or agent should not require editing the core runner.

Installation

Prerequisites

• Python 3.10+
• uv
• CUDA-capable GPU for local vLLM/SGLang inference

Main Environment

uv sync
source .venv/bin/activate

CUDA/SGLang compatibility patches used by this project:

python scripts/patches/patch_triton.py
python scripts/patches/patch_sglang.py

If your environment needs a prebuilt FlashAttention wheel, install it after uv sync.

Environment Variables

Create a .env file when using local or remote OpenAI-compatible APIs. Common variables:

LOCAL_API_BASE=http://127.0.0.1:8000/v1
OPENAI_API_KEY=EMPTY
OPENROUTER_API_KEY=...

For UMA/SMA agent loops:

PROMPT_TEMPLATE_PATH=external/memory_agent/memory_agent_runtime/sma/prompts/prompt_template.yaml
UMA_PROMPT_TEMPLATE_PATH=external/memory_agent/memory_agent_runtime/uma/prompts/prompt_template.yaml
VERL_TOKENIZER_PATH=/path/to/tokenizer

Quick Start

Start a local OpenAI-compatible inference server, for example vLLM:

source .venv/bin/activate
vllm serve Qwen/Qwen3-4B-Instruct-2507 \
  --host 0.0.0.0 \
  --port 8000 \
  --max-model-len 32768

Run one agent on one task:

python evaluate_async.py \
  --task locomo \
  --agent concat \
  --output-dir results/qwen3-4b/locomo \
  --generate-only

Score an existing responses file:

python evaluate_async.py \
  --task locomo \
  --input-file results/qwen3-4b/locomo/responses_concat.jsonl \
  --output-dir results/qwen3-4b/locomo

Use DeepSeek as the judge. DeepSeek scores are written into the same evaluated_*.jsonl files as metric.deepseek_llm_score:

python evaluate_async.py \
  --task locomo \
  --input-file results/qwen3-4b/locomo/responses_concat.jsonl \
  --output-dir results/qwen3-4b/locomo \
  --deepseek-eval

Command Line

MemoryAgent Arena provides a unified CLI, similar in spirit to lm_eval.

Without reinstalling the environment, use the module form:

python -m memory_agent_arena.cli tasks
python -m memory_agent_arena.cli tasks --detail
python -m memory_agent_arena.cli agents

After uv sync installs the project scripts, use either command name:

memory-agent-arena tasks
ma_eval tasks
maa agents
maa wizard

Use tasks to inspect every registered benchmark. Use tasks --detail to show processed dataset cache status and dataset statistics, and tasks build <task|all> to build missing processed data:

maa tasks --detail
maa tasks build ledgerqa[ss2]
maa tasks build 'locomo[refined]'
maa tasks build all --fail-fast

Task aliases use family[config] syntax. For example, ledgerqa[ss2] expands to the canonical task ledgerqa-ss2, and locomo[refined] expands to locomor. Result directories and evaluate_async.py --task values still use canonical task ids.

Teacher-style LoCoMo prompts are not benchmark variants. They are private SMA experiment settings; public benchmark runs should use locomo or locomo[refined]. If you need the teacher-style SMA memory prompt for an ablation, pass --teacher with the SMA agent only.

Use maa wizard for an interactive command builder. It supports generate, score, and run, prints the exact CLI command(s), and asks before execution. For a nicer selector UI, install the optional dependency with:

uv sync --extra wizard

Generate responses:

maa generate \
  --tasks locomo,locomor \
  --agent sma \
  --agent-id SMA_EXAMPLE \
  --model /path/to/model-or-checkpoint \
  --results-dir results/qwen3-4b \
  --concurrency 4

Score existing responses:

maa score --results-dir results/qwen3-4b
maa score \
  --tasks locomo,locomor \
  --agent-id SMA_EXAMPLE \
  --results-dir results/qwen3-4b \
  --judge local

maa score scans responses_*.jsonl files and evaluates only candidates whose evaluated_*.jsonl file is missing or incomplete. Use --tasks and --agent-id as filters; --agent-id refers to the response file suffix, not an agent registry key. With --judge deepseek, scoring fills metric.deepseek_llm_score in the same evaluated_*.jsonl files.

Run generation, scoring, and optional aggregation:

maa run \
  --tasks locomo \
  --agent sma \
  --agent-id SMA_EXAMPLE \
  --model /path/to/model-or-checkpoint \
  --results-dir results/qwen3-4b \
  --stats

Run multiple agents sequentially:

maa run \
  --tasks locomo,ledgerqa-ss2 \
  --agents concat,rag,sma \
  --agent-id concat,rag,SMA_EXAMPLE \
  --model Qwen/Qwen3-4B-Instruct-2507,Qwen/Qwen3-4B-Instruct-2507,/path/to/model \
  --results-dir results/qwen3-4b \
  --stats

When --agents is used, agents are executed in order and each agent runs all requested tasks before the next agent starts. If --agent-id or --model is provided, it must have the same number of values as --agents; otherwise agent ids default to agent names and models default to each agent profile. A failed agent is recorded and later agents continue to run; use --fail-fast to stop at the first failed agent. With multi-agent run --stats, aggregation runs once after all agents finish.

Summarize results:

maa stats --results-dir results/qwen3-4b --save-txt
maa stats ledgerqa --results-dir results/qwen3-4b

By default, stats uses .agentignore and .benchignore found under the results directory or its parents. maa stats ledgerqa switches to .agentignore_ledgerqa and .benchignore_ledgerqa, while keeping the same glob matching rules such as DS_*.

Manage local services:

maa generate --tasks locomo --agent sma --model /path/to/model
maa score --tasks locomo --agent-id SMA_EXAMPLE
maa run --tasks locomo --agent sma --model /path/to/model

generate, score, and run automatically check the ports declared by the agent profile. If the expected model is already served, it is reused. If a port is occupied by a different model, the old service is stopped and the expected one is launched. Use --no-auto-serve to skip service management when you want to manage vLLM or embedding servers manually. Remote endpoints such as OPENAI_API_BASE=http://10.x.x.x:8000/v1 are detected and skipped with a message. Each CLI invocation creates one service log session under runtime/services/logs/ by default. Override the root with --runtime-dir or MAA_RUNTIME_DIR. Within the same session, repeated service restarts append to the same role/port log with a separator; separate CLI invocations use separate timestamped directories. Agent workspaces and other intermediate memory artifacts default to runtime/agents/<agent_id>/<task>/ when an agent needs a writable workspace.

Services are not stopped automatically after standalone generate or score, so one agent can be reused across multiple benchmarks. run stops local agent vLLM ports before local judge scoring to avoid GPU conflicts. By convention, 8000 and 8001 are agent model ports, 8002 is the local judge port, and 8080 is the embedding service port. Extra vLLM flags can be supplied with MAA_VLLM_EXTRA_ARGS for all vLLM services, or MAA_VLLM_JUDGE_EXTRA_ARGS for the local judge only. For example, MAA_VLLM_JUDGE_EXTRA_ARGS="-tp 4 --gpu-memory-utilization 0.5" runs the judge with a smaller tensor-parallel group and lower memory target. By default, agent models use MAA_AGENT_CUDA_VISIBLE_DEVICES=4,5,6,7. A single agent model uses that whole list. Two-agent-model setups split that list in half: 8000 gets the first half and 8001 gets the second half. Tensor parallel size is inferred from the visible GPU count. Local judge also defaults to GPUs 4,5,6,7 with -tp 4. Default --gpu-memory-utilization is 0.7. Default agent and QA model context window is --max-model-len 65536; local judge uses 16384. vLLM services also default to VLLM_USE_FLASHINFER_SAMPLER=0, unless you export a different value before running the CLI.

For agents with baseline-specific settings, maa generate and maa run apply the runtime defaults declared in agents/registry.py. These defaults set missing environment variables such as EMBEDDING_SERVICE_ENDPOINT, OPENAI_API_BASE, JUDGE_API_BASE, MEMT_*, UMA_*, or GAM_*, but they do not overwrite variables you have already exported.

Inspect or disable profiles:

maa agents
maa generate --agent memt --tasks locomo --print-env --dry-run
maa generate --agent memt --tasks locomo --no-profile --model EdwinYue/Mem-T-4B

Baselines

Run baselines directly through maa generate. This keeps task expansion, service startup, model defaults, and resume behavior in one CLI path:

maa generate --agent concat --tasks locomo
maa generate --agent memt --tasks "locomo hotpotqa"
maa generate --agent uma --tasks all

--tasks all expands to every registered benchmark. For a smaller run, pass an explicit comma- or space-separated task expression.

More examples:

maa generate --agent memagent --tasks all
maa generate --agent memagent_woq --tasks locomo
maa generate --agent rag --tasks ledgerqa-ss2
maa run --agent sma --tasks locomo --stats

Supported Benchmarks

The central benchmark registry is in benchmarks/registry.py.

Currently supported tasks include:

banking77 clinic convomem hotpotqa knowmebench
locomo locomor
ledgerqa-* longmemeval longmemeval_v2_small
memalpha msc nlu perltqa pubmed_rct squad
trec_coarse trec_fine

Dynamic task names are also supported:

• ledgerqa-<suffix>, for example ledgerqa-ss10
• hotpotqa_<num_docs>, for example hotpotqa_10

Supported Agents

The central agent registry is in agents/registry.py.

Common agent keys:

concat memagent memalpha memt rag gam amem mem1 rlm uma sma

Notes:

• uma points to the older UMA agent implementation.
• sma points to SMAAgent.

Output Format

Generation writes JSONL files:

results/<run>/<task>/responses_<agent_id>.jsonl

Each row contains:

• qid
• query
• expected_answer
• response
• generation_time
• optional intermediate_paths
• optional tool_call_stats

Evaluation writes:

results/<run>/<task>/evaluated_<agent_id>.jsonl

Metrics include task-dependent fields such as:

• official_score
• llm_score
• f1_score
• exact_match
• sub_em
• rouge_score

Use generate_stats.py to aggregate results:

python generate_stats.py --results-dir results/qwen3-4b

Extending the Framework

Add a New Benchmark

1. Implement or register a loader that returns list[EvalData].
2. Add a BenchmarkSpec in benchmarks/registry.py.
3. If needed, add an evaluator in evaluation/evaluators.py and map it in evaluation/registry.py.

Minimal shape:

BenchmarkSpec(
    name="my_benchmark",
    loader=load_my_benchmark,
    evaluator_key="default",
    metric_profile="qa",
    data_source="my_benchmark",
)

Add a New Agent

1. Implement an agent under agents/.
2. Register it in agents/registry.py.
3. If the agent follows the standard interface, use the default adapter.
4. If it needs special construction or return parsing, add an adapter in agents/adapters.py.

Standard agent interface:

class MyAgent(BaseAgent):
    async def add_memory_async(self, chunk: str) -> None:
        ...

    async def QA_batch_async(self, query_list: list[str]) -> list[str]:
        ...

Registry example:

"myagent": AgentSpec(
    name="myagent",
    module_name="agents.my_agent",
    class_name="MyAgent",
    default_model="Qwen/Qwen3-4B-Instruct-2507",
    default_concurrency=4,
    env_defaults={
        "OPENAI_API_KEY": "EMPTY",
        "OPENAI_API_BASE": "http://127.0.0.1:8000/v1",
    },
)

If the agent needs special behavior:

"myagent": AgentSpec(
    name="myagent",
    module_name="agents.my_agent",
    class_name="MyAgent",
    adapter_key="my_adapter",
)

SMA Training

SMA training has been split out of this arena repository. Use this project for benchmarking memory agents and scoring generations; use the separate sma-training reproduction project for verl-based training scripts, dataset conversion, checkpoint merging, and paper-specific launchers.

Data Preparation

Processed datasets are built through the unified task registry:

maa tasks build <task>
maa tasks build all
maa tasks build locomo[base] hotpotqa[200]

Raw benchmark datasets are not committed to this repository. Public datasets are downloaded by their builders when possible, and the downloaded files are cached under data/raw/.

Benchmark task	Data source behavior
locomo[base]	Auto-downloads data/raw/locomo10.json from the official LoCoMo GitHub URL. Override with LOCOMO_RAW_URL.
locomo[refined]	Uses data/raw/locomor/locomo_refined.json. Set LOCOMOR_RAW_URL to download it from your own URL, or place the file manually.
knowmebench	Auto-downloads the Hugging Face dataset realty2333/knowMe-Bench into data/raw/KnowMeBench/ after your HF account has access to the gated dataset. Set HF_TOKEN if needed; override with KNOWMEBENCH_HF_REPO.
hotpotqa[200] and other HotpotQA configs	Auto-downloads the corresponding file, for example data/raw/hotpotqa/eval_200.json.
convomem	Auto-downloads Salesforce/ConvoMem into data/raw/ConvoMem/; this is a larger multi-file download. Override the local source with CONVOMEM_RAW_DIR.
longmemeval	Auto-downloads the small LongMemEval JSON file into data/raw/.
longmemeval_v2_small	Requires LongMemEval-V2 files under data/raw/longmemeval-v2/, or set LONGMEMEVAL_V2_DATA_DIR.
msc	Auto-downloads MemGPT/MSC-Self-Instruct through Hugging Face datasets.
MemoryAgentBench tasks (nlu, banking77, clinic, trec_*)	Require a MemoryAgentBench parquet file. Default: data/raw/memoryagentbench/test.parquet; override with MEMORYAGENTBENCH_PARQUET_PATH.
MemAlpha tasks (perltqa, pubmed_rct, squad)	Require a MemAlpha parquet file. Default: data/raw/memalpha/test.parquet; override with MEMALPHA_PARQUET_PATH.

For Hugging Face downloads, set HF_ENDPOINT or HF_HUB_ENDPOINT if you need a mirror. After raw data is present or downloadable, rerun maa tasks build <task> to create the processed cache under data/.cache/.

Agent runtimes that are needed for integrated baselines live under external/. For example, deltamem defaults to external/delta-Mem, while gam loads the vendored external/gam/src package. Delta-Mem model adapters are still external artifacts; pass --adapter-dir when running that agent.

Development Notes

Recommended local checks:

python -m py_compile evaluate_async.py evaluation_runner.py benchmark_runtime.py
python -m py_compile agents/registry.py agents/adapters.py
python -m py_compile benchmarks/registry.py evaluation/registry.py evaluation/evaluators.py
bash -n scripts/run/run.sh scripts/run/run_generation.sh scripts/run/run_score.sh
bash -n scripts/experimental/run_memt_remote_shards.sh

For a quick SMA smoke test:

python tests/test_sma_agent.py

License

MemoryAgent Arena is distributed under the license in LICENSE. Third-party source snapshots under external/ retain their own licenses and attribution; see external/THIRD_PARTY.md and the license files in each vendored package.