Search Agent Distillation with On-Policy Relabeling

Real-model search-agent distillation with Qwen2.5-0.5B-Instruct + LoRA.

This repo studies a simple question:

What happens when a student is distilled from clean offline trajectories, but deployment requires multi-step search under noisy context and distribution shift?

The short answer is the core result of this project:

• offline-only distillation can collapse
• on-policy relabeling recovers the states the student actually visits
• constrained action selection makes the improvement stable enough to measure

Method

The pipeline is:

1. generate synthetic multi-hop search worlds
2. collect oracle offline traces
3. train a Qwen + LoRA student policy
4. roll out the current student and relabel visited states with the oracle
5. retrain on offline + on-policy data

The student predicts constrained Choice: k actions instead of fully free-form tool strings. That keeps the setup agent-like while avoiding evaluation noise from formatting drift and open-world hallucinations.

Main Result

Representative results:

Setting	Offline final success	OPD final success
single-action low-data	0.00	1.00
ReAct + candidates	0.22	0.69
choice scoring	0.91	0.97

The strongest evidence comes from the low-data regime:

• offline only briefly improves and then collapses
• offline + OPD reaches perfect final success

That is the main claim of the repo: on-policy relabeling is most valuable when rollout states drift away from the clean teacher distribution.

Setup

Install dependencies:

pip install -r requirements.txt

This project was developed with a local copy of:

/root/autodl-tmp/models/Qwen2.5-0.5B-Instruct

If direct Hugging Face access is blocked, a mirror such as hf-mirror.com is sufficient for downloading the model ahead of time.

Run

Smoke test:

python train.py \
  --train-worlds 64 \
  --eval-worlds 16 \
  --rollout-worlds 32 \
  --offline-warmup-epochs 1 \
  --opd-epochs 1 \
  --batch-size 2 \
  --grad-accum-steps 4 \
  --output-dir runs/smoke

Representative low-data run:

python train.py \
  --model-path /root/autodl-tmp/models/Qwen2.5-0.5B-Instruct \
  --train-worlds 8 \
  --eval-worlds 32 \
  --rollout-worlds 8 \
  --offline-warmup-epochs 1 \
  --opd-epochs 2 \
  --batch-size 2 \
  --grad-accum-steps 4 \
  --output-dir runs/low_data

Each run writes:

• summary.json
• metrics_baseline.json
• metrics_opd.json
• training_curves.png
• baseline_adapter/
• opd_adapter/

Repository Layout

qwen_search_distill/
├── train.py
├── requirements.txt
├── README.md
└── figures/
    ├── make_project_figure.py
    ├── method_overview.png
    ├── project_overview.png
    └── project_overview.svg

Scope

This repo is:

• a compact real-model demonstration of search-agent distillation
• a minimal OPD-style relabeling experiment
• a controlled testbed for distribution shift in agent trajectories

This repo is not:

• a production search stack
• a full web-RAG system
• a benchmark-scale agent training framework

Citation

@misc{search_agent_opd_demo,
  title={Search Agent Distillation with On-Policy Relabeling},
  author={JayBuendia},
  year={2026},
  note={GitHub repository}
}