PCTD: Preference-Guided Counterfactual Task Decomposition for Agent Tool Retrieval
For developers of AI agents that rely on tool retrieval from ambiguous instructions, PCTD improves generalization to unseen tools by mitigating reward hacking.
PCTD addresses reward hacking in reinforcement learning-based task decomposition for agent tool retrieval by introducing a counterfactual reward that quantifies marginal causal gain, cutting off spurious correlations. It achieves state-of-the-art performance in retrieval, decomposition quality, and out-of-domain generalization.
Task decomposition aims to transform ambiguous instructions into executable atomic subtasks, thereby guiding high-precision tool retrieval. However, our analysis reveals that directly adopting tool retrieval metrics, i.e., Recall or NDCG, as rewards for task decomposition can easily induce reward hacking in reinforcement learning-based methods. Specifically, models tend to maximize retrieval matching through strategies such as repetitive decomposition. This spurious correlation between the shallow features of decomposition results and retrieval metric impairs generalization in Out-of-Domain (OOD) scenarios involving unseen tools. To address this issue, we propose PCTD, a Preference-guided Counterfactual Task Decomposition framework. PCTD quantifies the marginal causal gain of decomposition on retrieval ranking through a counterfactual reward, thereby cutting off spurious correlations at their source. Meanwhile, it introduces a preference reward to impose fine-grained structural supervision on logical coherence and atomicity, encouraging the model to generate high-quality decompositions. In addition, we construct MTDTool, the task decomposition benchmark specifically designed for mobile multi-turn interactions. Extensive experiments demonstrate that PCTD alleviates repetitive decomposition and surpasses SOTA methods in retrieval, decomposition quality, and OOD generalization.