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

2026-07-17

AI Summary

This work addresses the limitations of existing reinforcement learning–based task decomposition approaches, which often suffer from reward hacking—such as repetitive subtask generation—due to the direct use of retrieval metrics as rewards, leading to poor out-of-domain generalization in tool-augmented settings. To mitigate this, the authors propose a preference-guided counterfactual task decomposition framework that introduces counterfactual causal reasoning into task decomposition for the first time. By employing counterfactual rewards to quantify the causal contribution of each decomposition step to retrieval performance, the method severs spurious correlations between superficial features and retrieval metrics. Additionally, a structured preference reward provides fine-grained supervision over the logical coherence and atomicity of decomposed steps. Experiments on the newly introduced mobile multi-turn interaction benchmark, MTDTool, demonstrate that the proposed approach substantially alleviates repetitive decomposition and consistently outperforms state-of-the-art methods in retrieval effectiveness, decomposition quality, and out-of-domain generalization.

Abstract

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.

Problem

• task decomposition
• reward hacking
• tool retrieval
• out-of-domain generalization
• spurious correlation

Innovation

• Counterfactual Reward
• Preference-Guided Learning
• Task Decomposition
• Tool Retrieval
• Out-of-Domain Generalization