MemOPD: On-Policy Distillation through Memory State Alignment for Long-Horizon Agents
August 11, 2026 ยท View on GitHub
MemOPD combines PPO with full-vocabulary OPD for long-horizon agents. The student generates each trajectory, AlignPack reconstructs the exact call-local model state, and a frozen teacher scores the student's sampled action tokens under the same state. Context, retrieved observations, padding, and retained context copies are excluded from both PPO and distillation losses.
Quick Start
Environment
The configuration uses Python 3.9, PyTorch 2.4.0, vLLM 0.6.3, and four GPUs.
conda create -n memopd python=3.9
conda activate memopd
pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121
pip install vllm==0.6.3
pip install -r requirements.txt
cd Mem1/train
pip install -e .
pip install flash-attn --no-build-isolation
The local retriever additionally requires faiss-gpu, datasets, transformers, fastapi, and uvicorn.
Retriever
Q2/Q8/Q16 and Wiki-RAG use the same dense Wikipedia retriever. Download the Wikipedia corpus and E5 index:
python setup/download.py --save_path /path/to/downloaded/documents
cat /path/to/downloaded/documents/part_aa \
/path/to/downloaded/documents/part_ab \
> /path/to/downloaded/documents/e5_Flat.index
gzip -dk /path/to/downloaded/documents/wiki-18.jsonl.gz
Set file_path in Mem1/train/retrieval_launch.sh to /path/to/downloaded/documents, then start the server:
cd Mem1/train
bash retrieval_launch.sh
Cold-Start SFT
cd Mem1/train
CUDA_VISIBLE_DEVICES=0,1,2,3 \
BASE_MODEL_PATH=/path/to/Qwen2.5-3B \
SFT_TRAIN_FILE=/path/to/your_sft_train.parquet \
SFT_VAL_FILE=/path/to/your_sft_val.parquet \
SFT_OUTPUT_DIR=/path/to/cold_start \
bash train_sft.sh
MemOPD Training
Prepare an SFT student, a teacher, and the RL data, then launch training:
cd Mem1/train
CUDA_VISIBLE_DEVICES=0,1,2,3 \
DATA_DIR=/path/to/disjoint_q2_rl_data \
ACTOR_MODEL_PATH=/path/to/sft_student \
TEACHER_MODEL_PATH=/path/to/teacher \
RETRIEVER_URL=http://127.0.0.1:8013/retrieve \
bash train_memopd.sh
Evaluation
Evaluation follows the standalone vLLM rollout used for the paper:
cd Mem1/inference
CUDA_VISIBLE_DEVICES=0 \
MODEL_PATH=/path/to/model \
bash start_vllm.sh
In another shell, generate and score trajectories:
MEMOPD_MAX_TURNS=6 \
MEMOPD_VLLM_URL=http://127.0.0.1:8014 \
MEMOPD_SEARCH_URL=http://127.0.0.1:8013/retrieve \
MODEL_PATH=/path/to/model \
python generate_rollout.py \
--model /path/to/model \
--data_file /path/to/test.parquet \
--output_file trajectories.jsonl
python eval.py --trajectory_jsonl_path trajectories.jsonl