AWA-RL: Abstention-Aware Reinforcement Learning for Search Agents

Paper (Arxiv） • Model Weights (Hugging Face)

This repository contains the inference and evaluation code for To Answer or to Abstain: Mitigating Search-Agent Hallucinations via Abstention-Aware Reinforcement Learning.

AWA-RL trains search agents to balance capability and reliability: the agent should answer when the retrieved evidence supports an answer, and abstain with not enough information when the evidence is insufficient.

Results

The results above are from the AWA-RL paper. They show that AWA-RL can reliably control a model's refusal rate while balancing answer recall (reported as accuracy) and precision. In particular, adjusting the courage factor provides a stable control knob for the refusal rate: higher or lower encouragement shifts the model's willingness to answer, allowing the operating point to be selected for the desired balance between coverage and reliability.

What Is Included

• Minimal search-agent inference loop.
• Evaluation scripts for MuSiQue, HotpotQA, and 2WikiMultiHopQA validation subsets.
• OpenAI-compatible LM judge evaluation for Accuracy, Precision, Reliability-Aware F1 (RA-F1), and refusal rate.
• Three 1k validation datasets in data/.

Environment

Install the runtime dependencies:

pip install -r requirements.txt

The inference script uses vLLM and assumes enough GPUs are available for the released checkpoint.

External Services

This repository assumes that two services are already running.

FlashRAG Search Server

Run a FlashRAG server compatible with the Search-R1 retrieval setup. The evaluation code expects an HTTP endpoint such as:

http://127.0.0.1:1563

The search server should return top-k Wikipedia passages for the agent's <tool_call> queries. For retrieval setup details, follow the Search-R1 FlashRAG server instructions.

LM Judge Server

For RA-F1, deploy an OpenAI-compatible judge model endpoint. For example:

MODEL_PATH=/path/to/gpt-oss-120b \
TP_SIZE=4 \
PORT=6002 \
bash scripts/serve_lmjudge_gpt_oss_120b.sh

The evaluation script will call this endpoint with a correctness-judging prompt and expects a JSON response containing "judgement": "correct" or "judgement": "incorrect".

Datasets

The following validation subsets are included:

data/musique-val-1k.parquet
data/hotpotqa-val-1k.parquet
data/2wiki-val-1k.parquet

Each file follows the same Qampari-style structure used by the original experiments, with question, reward_model.ground_truth, and extra_info.id fields.

Run Evaluation

Set the model checkpoint and service endpoints:

export MODEL_PATH=/path/to/awa-rl-checkpoint
export SEARCH_URL=http://127.0.0.1:1563
export LMJUDGE_URL=http://127.0.0.1:6002/v1
export LMJUDGE_MODEL=/path/to/gpt-oss-120b

Run MuSiQue validation:

bash scripts/run_musique_val1k.sh

Or call the Python entry point directly:

PYTHONPATH=src python -m search_eval.eval \
  --model-path "$MODEL_PATH" \
  --dataset musique-val-1k \
  --search-url "$SEARCH_URL" \
  --lmjudge-url "$LMJUDGE_URL" \
  --lmjudge-model "$LMJUDGE_MODEL" \
  --output-dir outputs/musique-val-1k \
  --tp-size 4 \
  --pp-size 1

Evaluate the other included datasets by changing --dataset:

PYTHONPATH=src python -m search_eval.eval --model-path "$MODEL_PATH" --dataset hotpotqa-val-1k --search-url "$SEARCH_URL" --lmjudge-url "$LMJUDGE_URL" --lmjudge-model "$LMJUDGE_MODEL" --output-dir outputs/hotpotqa-val-1k --tp-size 4

PYTHONPATH=src python -m search_eval.eval --model-path "$MODEL_PATH" --dataset 2wiki-val-1k --search-url "$SEARCH_URL" --lmjudge-url "$LMJUDGE_URL" --lmjudge-model "$LMJUDGE_MODEL" --output-dir outputs/2wiki-val-1k --tp-size 4

For a quick smoke test:

MAX_QUESTIONS=5 SKIP_EVAL=1 bash scripts/run_musique_val1k.sh

Metrics

Let N_c, N_r, and N_w denote the number of correct, refused, and wrong responses, with N_total = N_c + N_r + N_w.

Accuracy  = N_c / N_total
Precision = N_c / (N_c + N_w)
RA-F1     = 2 * Precision * Accuracy / (Precision + Accuracy)
Refusal   = N_r / N_total

The evaluation output is written to:

<output-dir>/<dataset>_predictions.jsonl
<output-dir>/<dataset>_metrics.json
<output-dir>/<dataset>_lmjudge.json

Citation

@misc{zhang2026answerabstainmitigatingsearchagent,
      title={To Answer or to Abstain: Mitigating Search-Agent Hallucinations via Abstention-Aware Reinforcement Learning}, 
      author={Fengji Zhang and Tianyu Fan and Yuxiang Zheng and Xinyao Niu and Chengen Huang and Jacky Keung and Bei Chen},
      year={2026},
      eprint={2607.10738},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2607.10738}, 
}