[ICLR 2026] Look Back to Reason Forward: Revisitable Memory for Long-Context LLM Agents

This repo contains the official implementation of ICLR 2026 paper ReMemR1: Look Back to Reason Forward: Revisitable Memory for Long-Context LLM Agents.

News

• [Apr 2026] Added .claude/ AI-assisted development support with skill files for algorithm, reward, training, evaluation, and task extension 🤖.
• [Apr 2026] Our Project Page is now live 🌐.
• [Feb 2026] Our Checkpoint has been released on Huggingface🤗 🎉.
• [Jan 2026] Our paper get accepted by ICLR 2026 🎉🎉🎉
• [Jan 2026] Nvidia propose GDPO, which shares the same design logic as our multi-level reward aggregation
• [Sep 2025] Our paper is released on Huggingface and Alphaxiv. Please upvote our paper if you like this work :)

Overview

• Conceptual Example: wheat gleaning in the field

• Q1: How we address the constraints in linear doc scan?
• A1: We introduce Callback Mechanism to allow non-linear memory re-visit.

• Q2: How we precisely reward callback/update behaviors?
• A2: We introduce Multi-Level Rewarding, aggregated at advantage level.

Installation

conda create -n rememr1 python=3.11
conda activate rememr1
pip install httpx==0.23.1 aiohttp -U ray[serve,default] vllm

pip install nltk pyyaml beautifulsoup4 html2text wonderwords tenacity fire
pip install vllm==0.9 --index-url https://download.pytorch.org/whl/cu126
pip install "sglang==0.4.6"
pip install hydra-core accelerate tensordict torchdata wandb "tensordict<=0.6.2"

Data Processing

Trianing Data: This research use the same training data as MemAgent. The data files are publicly available, and can be downloaded from huggingface. After the download is finished, put hotpotqa_train_32k.parquet and hotpotqa_dev.parquet under data/train/.

Data for Evaluation: The data for evaluation is sourced from HotpotQA and 2WikiMultiHopQA. To process the data for long-context QA, simply run:

bash scripts/0_run_data_process.sh

Training

Prepare for Multi-Node Training

You may skip this step if you only want to start a single-node trianing.

To start multi-node training, you should first ray servers on the head and worker nodes.

# head node
ray start --head --dashboard-host=0.0.0.0

# worker node
ray start --address=<head_node_address>

More references can be found in ray's documentation here.

Start the Training

After all the nodes are ready (or you prefer single-node training), the training can be launched via:

bash scripts/1_run_train_ReMemR1_3B.sh
bash scripts/1_run_train_ReMemR1_7B.sh

You might adjust the N_NODE variable to match your number of devices.

Evaluation

Once the training is converged, use scripts/merge_ckpt.sh to merge the checkpoints before evaluation. For example,

bash scripts/merge_ckpt.sh "results/memory_agent/ReMemR1_3B/global_step_200/actor"

This will automatically put the merged checkpoint into results/memory_agent/ReMemR1_3B/global_step_200/actor/hf_ckpt.

Once the checkpoint merging is done, run the below script for evaluation:

bash scripts/2_run_eval_ReMemR1.sh

Licensing

This project is licensed under the MIT License. It includes components from MemAgent, licensed under the Apache License 2.0. Thanks for their awesome work!

AI-Assisted Development

This project includes a .claude/ directory with structured guidance for AI coding agents (Claude Code, Codex, etc.):

• .claude/CLAUDE.md — Project overview, architecture, key abstractions, and coding conventions
• .claude/skills/modify-algorithm.md — Guide for modifying the MemoryAgent and callback mechanism
• .claude/skills/modify-reward.md — Guide for modifying the multi-level reward system
• .claude/skills/run-training.md — Guide for configuring and launching training runs
• .claude/skills/run-evaluation.md — Guide for running evaluation and interpreting results
• .claude/skills/add-new-task.md — Guide for adding new tasks and datasets

These files help AI agents understand the codebase structure and make targeted modifications without extensive context exploration.

Citation

@article{rememr1,
  author       = {Yaorui Shi and
                  Yuxin Chen and
                  Siyuan Wang and
                  Sihang Li and
                  Hengxing Cai and
                  Qi Gu and
                  Xiang Wang and
                  An Zhang},
  title        = {Look Back to Reason Forward: Revisitable Memory for Long-Context {LLM}
                  Agents},
  journal      = {CoRR},
  volume       = {abs/2509.23040},
  year         = {2025},
  eprinttype    = {arXiv},
  eprint       = {2509.23040},
}