Task Memory Engine (TME): Enhancing State Awareness for Multi-Step LLM Agent Tasks

📄 arXiv: 2504.08525v4 📥 PDF

作者: Ye Ye

分类: cs.AI, cs.CL

发布日期: 2025-04-11 (更新: 2025-08-22)

备注: 14 pages, 5 figures. Preprint prepared for future submission. Includes implementation and token-efficiency analysis. Code at https://github.com/biubiutomato/TME-Agent

🔗 代码/项目: GITHUB

💡 一句话要点

提出任务记忆引擎TME，增强LLM Agent多步任务中的状态感知能力

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: LLM Agent 多步任务 任务记忆 状态感知 提示工程 知识图谱 上下文学习

📋 核心要点

  1. 现有LLM Agent在多步任务中缺乏结构化的任务状态理解，导致性能不稳定、幻觉和长期连贯性差。
  2. 提出任务记忆引擎TME，通过分层任务记忆树TMT跟踪任务执行，存储任务步骤的输入、输出和状态。
  3. 实验表明，TME能显著提高多步Agent任务的完成准确性和可解释性，且实现开销小。

📝 摘要（中文）

大型语言模型（LLMs）越来越多地被用作多步任务的自主Agent。然而，现有的大多数框架未能保持对任务状态的结构化理解，通常依赖于线性提示连接或浅层记忆缓冲区。这导致性能脆弱、频繁出现幻觉以及较差的长期连贯性。本文提出了任务记忆引擎（TME），这是一个轻量级且结构化的记忆模块，它使用分层任务记忆树（TMT）跟踪任务执行。树中的每个节点对应于一个任务步骤，存储相关的输入、输出、状态和子任务关系。我们引入了一种提示合成方法，该方法基于活动节点路径动态生成LLM提示，从而显著提高执行一致性和上下文基础。通过多步Agent任务的案例研究和对比实验，我们证明了TME以最小的实现开销实现了更好的任务完成准确性和更可解释的行为。核心TME组件的参考实现可在https://github.com/biubiutomato/TME-Agent获得，包括基本示例和结构化内存集成。虽然当前的实现使用基于树的结构，但TME被设计为具有图感知能力，支持可重用的子步骤、收敛的任务路径和共享依赖项。这为未来的基于DAG的记忆架构奠定了基础。

🔬 方法详解

问题定义：现有LLM Agent在执行多步任务时，通常采用线性拼接prompt或浅层记忆缓冲，缺乏对任务状态的结构化理解。这导致Agent在复杂任务中容易出现幻觉，难以保持长期连贯性，最终影响任务完成的准确率。现有方法难以有效追踪任务的执行过程和状态变化，无法充分利用上下文信息指导后续步骤的执行。

核心思路：TME的核心思路是构建一个结构化的任务记忆，以树状结构（TMT）来组织任务的执行过程。每个节点代表一个任务步骤，记录该步骤的输入、输出、状态以及与其他子任务的关系。通过维护这个任务记忆树，TME能够清晰地追踪任务的执行路径，并基于当前节点路径动态生成LLM的prompt，从而增强LLM对任务状态的感知能力。

技术框架：TME主要包含两个核心组件：任务记忆树（TMT）和提示合成模块。TMT是一个分层树状结构，用于存储任务的执行过程信息。提示合成模块则根据当前TMT中的活动节点路径，动态生成LLM的prompt。具体流程如下：1. Agent执行一个任务步骤；2. TME将该步骤的输入、输出和状态信息存储到TMT的相应节点中；3. 当Agent需要执行下一个步骤时，提示合成模块会根据当前TMT的活动节点路径，提取相关的上下文信息；4. 将提取的上下文信息与当前步骤的指令组合成prompt，输入给LLM；5. LLM根据prompt生成下一步的动作或输出。

关键创新：TME最关键的创新在于其结构化的任务记忆表示和动态提示合成方法。与传统的线性prompt拼接或浅层记忆缓冲相比，TME能够更有效地组织和利用任务的上下文信息，从而提高LLM Agent的执行一致性和上下文基础。此外，TME的树状结构也为未来的图结构扩展奠定了基础，可以支持更复杂的任务依赖关系和子任务复用。

关键设计：TMT的节点设计是关键。每个节点需要包含足够的信息来描述任务步骤的状态，包括输入、输出、状态、子任务关系等。提示合成模块的设计也至关重要，需要能够根据TMT的活动节点路径，有效地提取相关的上下文信息，并将其与当前步骤的指令组合成prompt。此外，TME的设计考虑了可扩展性，未来的图结构扩展可以支持更复杂的任务依赖关系和子任务复用。目前实现采用树结构，但设计上支持DAG（有向无环图）结构，为后续扩展提供了基础。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文通过案例研究和对比实验，验证了TME的有效性。实验结果表明，TME能够显著提高多步Agent任务的完成准确性，并降低幻觉的发生。与基线方法相比，TME在任务完成准确率方面取得了显著提升，同时保持了较低的实现开销。论文还提供了TME核心组件的参考实现，方便其他研究者和开发者使用。

🎯 应用场景

TME具有广泛的应用前景，可应用于各种需要多步推理和规划的LLM Agent任务，例如智能客服、自动化流程、游戏AI等。通过增强LLM Agent的状态感知能力，TME可以提高任务完成的准确性和效率，降低幻觉的发生，并提升Agent的长期连贯性。未来，TME有望成为构建更可靠、更智能的LLM Agent的重要组成部分。

📄 摘要（原文）

Large Language Models (LLMs) are increasingly used as autonomous agents for multi-step tasks. However, most existing frameworks fail to maintain a structured understanding of the task state, often relying on linear prompt concatenation or shallow memory buffers. This leads to brittle performance, frequent hallucinations, and poor long-range coherence. In this work, we propose the Task Memory Engine (TME), a lightweight and structured memory module that tracks task execution using a hierarchical Task Memory Tree (TMT). Each node in the tree corresponds to a task step, storing relevant input, output, status, and sub-task relationships. We introduce a prompt synthesis method that dynamically generates LLM prompts based on the active node path, significantly improving execution consistency and contextual grounding. Through case studies and comparative experiments on multi-step agent tasks, we demonstrate that TME leads to better task completion accuracy and more interpretable behavior with minimal implementation overhead. A reference implementation of the core TME components is available at https://github.com/biubiutomato/TME-Agent, including basic examples and structured memory integration. While the current implementation uses a tree-based structure, TME is designed to be graph-aware, supporting reusable substeps, converging task paths, and shared dependencies. This lays the groundwork for future DAG-based memory architectures.