ATOD: Annealed Turn-Aware On-Policy Distillation for Multi-Turn Agentic Tasks

🔥 Overview

We introduce ATOD (Annealed Turn-aware On-policy Distillation), a hybrid online distillation algorithm that combines OPD and GRPO with a smoothly annealed schedule and turn-level disagreement-uncertainty reweighting (T-DUR) for training small language-model agents on long-horizon, multi-turn tasks.

📖 Results

ATOD achieves substantial improvements over OPD, SOD, TCOD, SDAR, and GRPO baselines on ALFWorld, WebShop, and Search-QA across all student scales, and even surpasses the corresponding teacher models.

🛠️ Installation

Python environment

conda create -n atod python==3.12 -y
conda activate atod

pip3 install vllm==0.11.0

pip3 install flash-attn==2.7.4.post1 --no-build-isolation --no-cache-dir
pip install -e .

Log in to Weights & Biases if you use WandB logging (scripts pass trainer.logger=['console','wandb'] in many cases):

export WANDB_API_KEY=your_key_here

Install Supported Environments

1. ALFWorld

Install with pip:

pip3 install gymnasium==0.29.1
pip3 install stable-baselines3==2.6.0
pip3 install alfworld

Download PDDL & Game files and pre-trained MaskRCNN detector (will be stored in ~/.cache/alfworld/):

alfworld-download -f

2. WebShop

WebShop requires Python <=3.10, so begin by creating a new environment:

conda create -n verl-webshop python==3.10 -y
conda activate verl-webshop

Install WebShop:

cd ./agent_system/environments/env_package/webshop/webshop
./setup.sh -d all

After WebShop is installed, return to the root directory and install the verl package:

cd repo_root/
pip3 install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124
pip3 install flash-attn==2.7.4.post1 --no-build-isolation
pip3 install -e .
pip3 install vllm==0.8.2
# spacy 3.7.2 requires typer<0.10.0,>=0.3.0, but you have typer 0.15.2 which is incompatible.
# weasel 0.3.4 requires typer<0.10.0,>=0.3.0, but you have typer 0.15.2 which is incompatible.

The warnings can be safely ignored.

3. Search

cd ./agent_system/environments/env_package/search/third_party
pip install -e .
pip install gym==0.26.2

Prepare dataset (data will be saved at ~/data/searchR1_processed_direct):

cd repo_root/
python examples/data_preprocess/preprocess_search_r1_dataset.py

Since faiss-gpu is not available via pip, we setup a separate conda environment for the local retrieval server. Running this server will use around 6GB of GPU memory per GPU, so make sure to account for this in your training run configuration. Build Retriever environments:

# Create and activate the retriever environment with Python 3.10
conda create -n retriever python=3.10 -y
conda activate retriever

# Install PyTorch (with GPU support) and related libraries
conda install numpy==1.26.4 # needed to stop incompatible version of numpy from being installed via pip
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu124

# Install other Python packages
pip install transformers datasets pyserini huggingface_hub

# Install the GPU version of faiss
conda install faiss-gpu==1.8.0 -c pytorch -c nvidia -y

# Install the API service framework
pip install uvicorn fastapi

Download the index:

conda activate retriever

local_dir=~/data/searchR1
python examples/search/searchr1_download.py --local_dir $local_dir
cat $local_dir/part_* > $local_dir/e5_Flat.index
gzip -d $local_dir/wiki-18.jsonl.gz

Start the local flat e5 retrieval server:

conda activate retriever

# redirect the output to a file to avoid cluttering the terminal
# we have observed outputting to the terminal causing spikes in server response times
bash examples/search/retriever/retrieval_launch.sh > retrieval_server.log 

Training

Step 1. Train the teacher (GRPO). Scripts live under examples/grpo_teacher_trainer/ ({alfworld, search, webshop} × {4B, 30B-A3B}) and assume the repo root as working directory. Edit the model path in the script to point to your local checkpoint.

# Teacher training — {alfworld, search, webshop} × {4b, 30ba3b}
bash examples/grpo_teacher_trainer/run_alfworld_grpo_qwen3_4b.sh
bash examples/grpo_teacher_trainer/run_search_grpo_qwen3_30ba3b.sh
bash examples/grpo_teacher_trainer/run_webshop_grpo_qwen3_4b.sh

Step 2. Train the student with ATOD. Scripts live under examples/atod_trainer/. Each script sets the student model, the GRPO-trained teacher (point teacher_model_path to the teacher trained in Step 1), the annealing schedule (κ, ρ, T), and the T-DUR parameters.

# ATOD (ours) — {alfworld, search, webshop} × {0.6b, 1.7b, 4b}
bash examples/atod_trainer/run_alfworld_atod_qwen3_1.7b.sh
bash examples/atod_trainer/run_search_atod_qwen3_1.7b.sh
bash examples/atod_trainer/run_webshop_atod_qwen3_1.7b.sh

⭐️ Citation

If you find this project useful, welcome to cite us.

@misc{atod2026,
      title={ATOD: Annealed Turn-Aware On-Policy Distillation for Multi-Turn Agentic Tasks}, 
      year={2026},
      eprint={2606.27814},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2606.27814}, 
}

🤝 Acknowledgement

This project builds on SDAR, verl-agent, veRL, ALFWorld, SkillRL, and Search-R1. We thank the authors of those projects.