Skip to content

kiminh/PCTD

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

3 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

PCTD: Preference-Guided Counterfactual Task Decomposition for Agent Tool Retrieval


Overview

PCTD is a task decomposition training and inference framework for intelligent assistants (e.g. phone voice assistants). Built on an open-source LLM (this project uses Qwen3-8B/4B), it uses a two-stage pipeline — SFT (Supervised Fine-Tuning) → RL (GPDO) — to teach the model to rewrite and decompose complex user instructions into single-turn, single-task atomic tasks separated by <br>.

Example — Given a multi-turn, multi-task instruction, PCTD rewrites it into single-turn, single-task atomic tasks (separated by <br>), each of which retrieves a precise tool:

Input  (multi-turn · multi-task):
  用户: 帮我查查火车票
  助手: 请问出发地和目的地分别是哪里?
  用户: 从南京到杭州的,还有帮我打开蓝牙

Output (atomic tasks · `<br>`-separated):
  查从南京到杭州的火车票<br>打开蓝牙

Tools recalled per sub-task:
  查从南京到杭州的火车票  →  train_service
  打开蓝牙                →  bluetooth

Features

  • Two-stage training: SFT warm-start → RL (GRPO/GDPO) refinement on the merged SFT checkpoint.
  • Dual reward signals for RL: a preference reward model (quality) and a counterfactual tool-recall reward (utility) — directly tying decomposition quality to downstream tool retrieval.
  • Counterfactual tool-recall reward: compares NDCG of tool recall from decomposed sub-tasks vs. the original query; decomposition that improves recall gets a positive reward.
  • LLM + state-machine corpus generator (mtd_tool_data): auto-produces multi-turn, multi-task dialogue data across 11 execution states.
  • End-to-end pipeline: Data generation → SFT → Preference RM → RL → VLLM inference → GSB / Tool-Recall Evaluation.

MTDTool Benchmark construction

The data generation follows a four-stage process. (1) Initialization: Tools and scenarios are sampled to establish initial constraints, and two LLM generators produce the initial user query q0 and assistant output o0. (2) State Evolution & Context Decision: At each turn t ≥ 1, the system decides whether to inherit the previous context or switch to a new topic, and updates the pending intent pt,j and execution state st,j for each active tool. (3) Turn Generation: Conditioned on the dialogue history Ht, current states, and active tools, the framework generates the user query qt, derives the atomic task sequence At, and produces the state-aligned assistant output ot. (4) Per-Turn Annotation & Storage: The framework maps the atomic task sequence At to the used tool set T used t , assigns scenario labels and attribution information, and stores the structured record. This pipeline preserves the coherence of multi-turn intent evolution while providing fine-grained annotations for task decomposition and tool retrieval.

PCTD Pipeline

MTDTool JSON
   │
   ├─► Step 1·SFT Data ─► Step 2·SFT Train ─► Step 3·Merge + chat_template
   │                                                      │
   └─► Step 4·RL Data ────────────────────────────────────┤
                                                          ▼
                          Step 6·RL (GPDO) ◄── Step 5·Tool Emb (opt)
                              │               ◄── Step 5b·Pref RM
                              │   reward = 0.2·reward_model + 0.8·tool_recall
                              ▼
                    Step 7·vLLM Inference ─► Step 8·Evaluation (GSB + Tool Recall)

Step 0 (mtd_tool_data) generates the MTDTool corpus upstream. See each step below for details.


Step 1 · SFT Data Generation

Converts MTDTool raw JSON (containing messages dialogue history + gt-plan task decomposition annotations) into LLaMA-Factory's Alpaca format, and injects a system prompt.

Input: MTDTool raw JSON

{
  "messages": [
    {"role": "user", "content": "帮我查查火车票"},
    {"role": "assistant", "content": "请问出发地和目的地分别是哪里?"},
    {"role": "user", "content": "从南京到杭州的,还有帮我打开蓝牙"}
  ],
  "gt-plan": "查从南京到杭州的火车票<br>打开蓝牙",
  "gt-tools": "train_service<br>bluetooth"
}

Running:

python sft-model/src/gen_sft_model_dataset.py --out_name train \
 --infile  data/mtd_tool/your_train.json

Output files:

sft-model/data/origin_data/{task}/{version}/
├── train.json                # Copied raw data
├── train_alpaca.json         # Alpaca format training set

Step 2 · SFT Training (LoRA)

LoRA fine-tuning of Qwen3-8B using the LLaMA-Factory framework.

Running:

# Modify path variables in sft-model/sh/train_lora_sft.sh (PROJECT_DIR, MODEL_PATH, etc.)
bash sft-model/sh/train_lora_sft.sh

Output: LoRA adapter checkpoint, saved at sft-model/data/output_models/{task}/{version}/checkpoint-{step}/


Step 3 · SFT Model Merge & chat_template Replacement

RL training starts from the SFT model, so the LoRA adapter produced during the SFT phase must first be merged into the base model. Additionally, since SFT training data uses the Alpaca format template (qwen_alpaca), which doesn't match the base model's (Qwen3) default chat_template, the tokenizer's chat_template must be replaced after merging to ensure apply_chat_template behavior during RL training and inference is consistent with SFT.

3.1 Merge LoRA adapter

python sft-model/src/merge.py \
    --pretrained_model_path /path/to/Qwen3-8B \
    --lora_model_path sft-model/data/output_models/{task}/{version}/checkpoint-{step}

The merged model is output to {lora_model_path}/merged/ by default.

3.2 Replace chat_template

After merging, the model uses the original Qwen3 tokenizer's chat_template. This needs to be replaced with the Alpaca template consistent with SFT training:

from transformers import AutoTokenizer

merged_path = "sft-model/data/output_models/{task}/{version}/checkpoint-{step}/merged"
tokenizer = AutoTokenizer.from_pretrained(merged_path)

# Replace with Alpaca format chat_template
alpaca_chat_template = (
    "{% set roles = {'user': '用户','human': '用户','assistant': '系统答复'} %}{% if messages|length == 1 and messages[0]['role'] == 'system' %}<|im_start|>{{ messages[0]['content'] }}<|im_end|>{% else %}<|im_start|>你是一个多轮改写与任务拆解专家,根据用户指令对单轮、多轮、多任务进行改写和拆解,目标是分解成单轮单任务。\n\n【任务定义】\n- 你需要结合用户的上下文对话历史来改写成单轮对话\n- 你需要把多任务拆解为单任务\n\n【改写任务定义】\n- 无需改写场景: 单轮或者上下文无关任务, 示例: (Q1: 打开蓝牙) (A1: 好的已打开) (Q2: 打开微信) -> 输出: 打开微信\n- 需要改写场景: 支持单轮和多轮对话,包含指代消解,关键信息补充等场景。\n -- 指代消解:解析代词(如""、"这个")或省略信息,结合历史对话补全当前任务的完整信息;示例: (Q1: 打开蓝牙)(A1: 好的已打开) (Q2: 关闭它) -> 输出: 关闭蓝牙\n -- 关键信息补充:根据对话上下文补充缺失的必要信息; 示例:(Q1: 你能帮我查一下明天的天气吗?) (A1: 请问你想查询哪个城市的天气?) (Q2: 北京) ->输出: 查询北京明天的天气\n\n【多任务拆解定义】\n- 一句话多任务: 一句话中包含多个任务指令; 示例1: (Q: 对比杭州、上海、南京三个城市的天气) -> 输出:查询杭州的天气<br>查询上海的天气<br>查询南京的天气<br>对比三个城市的天气; 示例2: (Q: 打开wifi、蓝牙和深色模式) -> 输出:打开wifi<br>打开蓝牙<br>打开深色模式\n- 多句话多任务: 多句话中包含多个任务指令; 示例1: (Q:打开蓝牙,关闭微信) -> 输出:打开蓝牙<br>关闭微信; 示例2: (Q:查一下我爸爸的生日,再查一下我妈妈的) -> 输出:查一下我爸爸的生日<br>再查一下我妈妈的生日\n\n【输出范式】\n任务1<br>任务2<br>任务3<br>任务4;任务之间用`<br>`分隔。\n\n{% for message in messages[:] %}{{ roles[message['role']] }}: {{ message['content'] }}\n{% endfor %}任务分解:<|im_end|>{% endif %}"
)
tokenizer.chat_template = alpaca_chat_template
tokenizer.save_pretrained(merged_path)

Why this step is critical:

  • The RL training framework verl internally uses tokenizer.apply_chat_template() to convert the messages list into a prompt string, which must be consistent with the system_prompt + dialogue concatenation format used during SFT data generation
  • RL inference (rl-model/src/infer_vllm.py) also calls apply_chat_template to construct the prompt
  • Without replacing the chat_template, the prompt format would not match the SFT phase, leading to severely degraded inference performance

Output: The merged complete model is located at {lora_model_path}/merged/, ready to be used as the base model for RL training.


Step 4 · RL Data Generation

Converts MTDTool raw JSON into the Parquet format required by the verl/GPDO framework.

Running:

python rl-model/src/gen_rl_model_data.py  --out_name train \
 --infile  data/mtd_tool/your_train.json

Processing pipeline:

  1. Read raw JSON, extract messages (truncated to the last user), gt-plan, gt-tools, recall_query
  2. Build a DataFrame with columns: query, user_messages, history, tool, gt
  3. Convert to HuggingFace Dataset, map to verl required format:
    • prompt: user_messages
    • reward_model.ground_truth: gt-plan
    • extra_info: containing split, question, tool, history, recall_query
  4. Save as .parquet file

Output:

rl-model/data/origin_data/{task}/{version}/
├── train.json
├── train.parquet
├── val.json
└── val.parquet

Step 5 · Tool Embedding Precomputation (Optional)

Offline precomputation of tool description text embedding vectors to accelerate tool recall service startup and avoid re-encoding on every launch.

Running:

# Single GPU
python common/tool_recall_service/generate_embeddings.py \
    --emb_model_path /path/to/Qwen3-Embedding-4B \
    --tool_api_path /path/to/tool_api.xlsx \
    --output_dir /path/to/tool_emb_cache \
    --gpu_id 0 --batch_size 8

# Multi-GPU parallel
python common/tool_recall_service/generate_embeddings.py \
    --emb_model_path /path/to/Qwen3-Embedding-4B \
    --tool_api_path /path/to/tool_api.xlsx \
    --output_dir /path/to/tool_emb_cache \
    --num_gpus 4 --batch_size 32

Output:

tool_emb_cache/
├── embeddings.npy    # (n_docs, hidden_dim) float32 array
└── metadatas.json    # [{tool_en_name, tool_description}, ...]

Step 5b · Preference Reward Model Training (pref_reward_model)

pref_reward_model trains a Preference Reward Model that scores the quality of "task decomposition rewrite results" for ranking, providing the reward_model reward signal for Step 6 RL training (i.e. the underlying model of the reward_model service deployed on port 8010). During RL training, the decomposition results generated by the model are scored by this reward model as preference reward feedback to the policy.

Data Pipeline

Preference annotation data (Excel) → gen_reward_data.py → sorted TSV → training.

Input: data/preference_reward_train.xlsx, containing columns such as dialog_id, round_id, rank (rewrite quality ranking, 1 is best), query, assistant, response_text (task decomposition rewrite result).

gen_reward_data.py groups by dialog_id + round_id, sorts candidates within each group by rank from high to low, filters out over-length samples (max_token_length=1280), requires at least 2 candidates per group, and outputs one JSON list per row:

["【当前轮对话】\n用户:... \n【任务分解改写结果】\n改写A(最优)", "改写B", "改写C(最差)"]

Text is concatenated by utils.build_rm_input(): system_prompt + 【历史轮对话】 + 【当前轮对话】用户:query + 【任务分解改写结果】response. The system_prompt defines 5-dimensional evaluation criteria: completeness / accuracy / coreference resolution / expression norms / context consistency, plus task decomposition specific rules.

Model and Loss

Item Description
Base Qwen2.5-1.5B-Instruct (AutoModelForSequenceClassification, num_labels=1)
Fine-tuning LoRA (rank=64, alpha=128, dropout=0.1), target layers q/k/v/o/gate/up/down_proj
Output Scalar reward score reward
Loss compute_rank_list_loss: pairwise margin loss, computes logsigmoid(r_high − r_low) summed and negated over all candidate pairs, ensuring higher-quality rewrites score higher than lower-quality ones

Training

# Modify MODEL_PATH / PROJECT_DIR / distributed parameters (nproc_per_node, etc.) in sh/train_reward_model.sh
bash pref_reward_model/sh/train_reward_model.sh

torchrun launches DDP distributed training. A custom collate_fn handles the case where "one sample contains multiple texts to rank" (padded to the maximum text count within the batch). Checkpoints are evaluated and saved every save_steps, retaining model_best by ranking accuracy.

Key parameters:

Parameter Value
Base model Qwen2.5-1.5B-Instruct
Fine-tuning method LoRA (rank=64, alpha=128)
Learning rate 1e-5, cosine, warmup 100 steps
Epochs 10
Batch size 1 per device, gradient_accumulation=4
Sequence length 1280
Precision bf16
GPU 8 cards (DDP)

Evaluation metrics: ranking perfect accuracy (acc), Top-1 accuracy (highest score for the best candidate). Training curves are output as static PNGs by iTrainingLogger.

Output: LoRA adapter + tokenizer, saved at data/output_models/{task}/{version}/model_best/.

Inference

infer.py loads model_best (LoRA adapter + base), scores multiple rewrites per row of the test set TSV, and outputs an Excel file containing query, history, rewrite, reward, reward_order, ranking_perfect_correct, top1_correct.

After the trained reward model is deployed as the reward_model service (port 8010), it can be used as the reward_model_score signal in Step 6 RL training.


Step 6 · RL Training (GPDO)

GPDO (Group Policy Distance Optimization) training using the verl framework. Training starts from the SFT checkpoint and optimizes task decomposition quality through dual reward signals.

Prerequisites:

  • Step 2 SFT training completed with an available checkpoint
  • Step 3 SFT model merge and chat_template replacement completed
  • Step 4 RL data generation completed
  • Step 5b preference reward model training completed and deployed as reward_model service (port 8010)

Running:

# Modify path variables in rl-model/sh/train_gdpo_lora.sh
# The script automatically starts the tool_recall service (port 8020) and reward_model service (port 8010)
bash rl-model/sh/train_gdpo_lora.sh

Key parameters:

Parameter Value Description
Base model Merged SFT checkpoint model Continues training on top of the Phase 1 SFT model
Algorithm GPDO adv_estimator=gdpo
Fine-tuning method LoRA (rank=64, alpha=128) Inherits SFT LoRA configuration
Learning rate 1e-5, cosine scheduler warmup 10%
Epochs 10
Sampling count 8 rollout.n=8, 8 samples per prompt
KL coefficient 0.001 Constrains policy from deviating too far from the reference model
Batch size train=128, val=1024
GPU 8 cards

Reward signals (weights controlled by environment variables REWARD_WEIGHT and TOOL_WEIGHT):

Signal Default Weight Computation
reward_model_score 0.2 Calls the reward_model service, scores the generated result paired with GT, outputs after sigmoid
tool_recall_score 0.8 Counterfactual tool recall reward: compares the NDCG of tool recall for decomposed sub-tasks vs. the NDCG of direct recall from the original query

Counterfactual tool recall reward (counterfactual_reward.py):

  1. Call the tool recall service with the original query, compute NDCG@k, full coverage
  2. Call the tool recall service with the model's decomposition result, compute NDCG@k, full coverage
  3. Counterfactual reward = cf_plus_weight × (NDCG_decomposed − NDCG_original) + full_plus_weight × (coverage_decomposed − coverage_original)
  4. Reward range [-1, 1], positive values indicate the decomposition improved tool recall quality

Output: LoRA adapter checkpoint, saved at rl-model/checkpoints/{task}/{version}/global_step_{N}/actor/lora_adapter/


Step 7 · Inference

Step 7a · SFT Model Inference

Batch inference on the SFT model using vLLM (no need to merge LoRA, uses the base model directly):

bash sft-model/sh/infer_lora_vllm.sh

Pipeline:

  1. Automatically detect idle GPUs
  2. split_file.py splits the test set by GPU count
  3. Each GPU runs infer_vllm.py (temperature=0, greedy decoding)
  4. combine_file.py merges results from multiple GPUs

Step 7b · RL Model Inference

Before RL model inference, the LoRA adapter must be merged into the base model:

bash rl-model/sh/infer_grpo_lora_vllm.sh

Pipeline:

  1. Automatically detect idle GPUs
  2. merge.py merges the LoRA adapter into the base model, outputting the merged complete model
  3. split_file.py splits the test set by GPU count
  4. Each GPU runs infer_vllm.py (temperature=0)
  5. combine_file.py merges results from multiple GPUs

Step 8 · Evaluation

Step 8a · GSB Evaluation (LLM-as-Judge)

Uses an LLM as a judge to compare the decomposition quality of model predictions against ground truth:

python evaluation/gsb/runners/evaluate.py \
    --config_dir ./evaluation/gsb/config \
    --input_file /path/to/predictions.json \
    --output_dir ./evaluation/gsb/outputs \
    --mode absolute \
    --pred_field "predict1,predict2,...."

Two evaluation modes:

  • pairwise (pairwise comparison): outputs Win / Lose / Tie rates
  • absolute (absolute scoring): outputs a 5-point composite score

This evaluation uses the absolute mode.

Configuration files:

File Purpose
evaluation/gsb/config/evaluation.yaml Evaluation parameters (mode, concurrency, etc.)
evaluation/gsb/config/judge_models.yaml Judge LLM configuration
evaluation/gsb/config/prompts.yaml Evaluation prompt templates

Step 8b · Tool Recall Evaluation

Evaluates the tool recall quality of the model's decomposed sub-tasks:

# SFT model tool recall evaluation
bash sft-model/sh/evl_sft_tool_recall.sh

# RL model tool recall evaluation
bash rl-model/sh/evl_rl_too_recall.sh

Evaluation pipeline:

  1. Merge inference results (containing pred_plan) with raw data (containing gt-tools)
  2. Call the tool recall service (port 8020) to recall top-k tools using pred_plan as the query
  3. Compute metrics: NDCG@K, Completion@K, Recall@K

About

PCTD: Preference-Guided Counterfactual Task Decomposition for Agent Tool Retrieval

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages