AMemGym: Interactive Memory Benchmarking for Assistants in Long-horizon Conversations

This repo contains the code and data for the paper: AMemGym: Interactive Memory Benchmarking for Assistants in Long-horizon Conversations.

---

Overview

AMemGym is the first interactive, on-policy evaluation framework for conversational memory in LLM-based assistants. Unlike traditional static benchmarks that rely on pre-generated conversations, AMemGym enables realistic evaluation by allowing assistants to generate their own responses and learn from environmental feedback—bridging the gap between evaluation and real-world deployment.

Key Features

• Realistic Evaluation: Assistants actively participate in conversations with simulated users adapting to their responses
• Fine-Grained Diagnostics: Pinpoints failures in Write, Read, and Utilization operations
• Optimization Feedback: Enables autonomous agent self-evolution through environmental feedback
• Fully Automated: Scalable generation of diverse, high-quality scenarios spanning 128K-512K+ context lengths

---

Installation

# Clone the repository
git clone https://github.com/your-org/amemgym.git
cd amemgym

# Create and activate a virtual environment (optional but recommended)
uv venv
source .venv/bin/activate 

# Install with uv (recommended)
uv sync

# Install the package with pip
uv pip install -e .

Set up LLM API access by creating a .env file or exporting environment variables:

export OPENAI_API_KEY=your_api_key
export OPENAI_BASE_URL=https://api.openai.com/v1

---

Quick Start

Running On-Policy Evaluation

Prepare Environment Data

First, prepare the environment data. You can use the provided v1.base dataset (aligned with the paper) or create your own.

# Generate environment data using the provided configuration that you can customize
uv run python -m amemgym.env.gen \
    --data_dir data/v1.base \
    --config_path configs/env/v1.base.json \
    --persona_path data/personas/nemotron.parquet

Note that data/personas/nemotron.parquet contains user personas from Nemotron-Personas used in the environment data generation.

Running Main Evaluation

# Run main evaluation with a specific agent configuration
uv run python -m amemgym.eval.overall \
    --agent_config configs/agent/awi.json \
    --env_data data/v1.base/data.json \
    --output_dir eval-output/v1.base/overall

Available Agent Configurations:

You can customize your own agent configurations or even create new agent types by following the interface in amemgym.assistants.base.BaseAgent.

Running Upper-Bound and Random Baselines

Use the following commands to run upper-bound and random baseline evaluations for normalized memory scores.

# Run upper-bound evaluation
uv run python -m amemgym.eval.upperbound \
    --agent_config <a specific agent config, e.g., configs/agent/awi.json> \
    --env_data data/v1.base/data.json \
    --output_dir eval-output/v1.base/upperbound

# Run random baseline evaluation
uv run python -m amemgym.eval.random \
    --env_data data/v1.base/data.json \
    --output_file eval-output/v1.base/random_metrics.json

(Optional) Running Fine-Grained Diagnostics

# Run fine-grained diagnostics for a specific agent configuration
uv run python -m amemgym.eval.diagnosis \
    --agent_config <a specific agent config, e.g., configs/agent/awi.json> \
    --env_data data/v1.base/data.json \
    --output_dir eval-output/v1.base  # consistent with overall evaluation output dir

Running Evolution Experiments

For self-evolution experiments (Table 3 in paper):

uv run python -m amemgym.eval.evolution \
    --agent_config configs/agent/awi-evolve/complete.json

Available Evolution Configurations:

---

Citation

If you find AMemGym useful for your research, please cite our paper:

@inproceedings{
    jiayang2026amemgym,
    title={{AM}emGym: Interactive Memory Benchmarking for Assistants in Long-Horizon Conversations},
    author={Cheng Jiayang and Dongyu Ru and Lin Qiu and Yiyang Li and Xuezhi Cao and Yangqiu Song and Xunliang Cai},
    booktitle={The Fourteenth International Conference on Learning Representations},
    year={2026},
    url={https://openreview.net/forum?id=sfrVLzsmlf}
}