Learning When to Plan: Efficiently Allocating Test-Time Compute for LLM Agents


TL;DR

本文提出了一种两阶段训练方法(监督微调+强化学习),使大型语言模型(LLM)智能体能够学习在序贯决策任务中动态地决定何时进行规划,从而以更高效的计算成本实现更优的性能和更强的可控性。

关键定义

本文提出了一个用于动态规划的 conceptual framework (概念框架),其核心定义如下:

  1. 动态规划智能体 (Dynamic Planning Agent):这是一种能够根据当前状态的需要,灵活决定是否分配测试时计算资源 (test-time compute) 进行规划的智能体。它摒弃了“总是规划”或“从不规划”的固定策略,寻求一种更高效的决策模式。

  2. 概念性策略分解:本文将单一LLM的行为概念上分解为三个协同工作的策略,它们通过统一的输出格式实现:
    • 决策策略 (Decision Policy, $\phi_{\theta}$): 判断当前是否需要生成新计划。这通过模型是否选择输出 \(<plan>\) 标志来隐式实现。
    • 规划策略 (Planning Policy, $\psi_{\theta}$): 当决定要规划时,生成新的自然语言计划 $p_t$。
    • 行动策略 (Acting Policy, $\pi_{\theta}$): 根据当前上下文和计划(新生成的或已有的)来生成具体行动 $a_t$。
  3. 规划优势 (Planning Advantage, $A_{plan}$):指生成一个新计划所带来的预期未来回报的提升量,与继续使用旧计划相比。其定义为:

    \[A_{plan}(c_{t})=\mathbb{E}_{p_{t}\sim\psi_{\theta}(\cdot\mid c_{t},d_{t}=1)}[V^{\pi_{\theta}}(c_{t},p_{t})-V^{\pi_{\theta}}(c_{t},p_{t-1})]\]

    其中 $V^{\pi_{\theta}}(c_t, p_t)$ 是在给定上下文 $c_t$ 和计划 $p_t$ 下的预期未来回报。

  4. 规划成本 (Cost of Planning, $C_{plan}$):规划行为带来的总成本,包括:

    \[C_{plan}=C_{tokens}+C_{latency}+C_{noise}\]
    • 计算成本 ($C_{tokens}$): 生成计划所耗费的 token 数量。
    • 延迟成本 ($C_{latency}$): 规划所花费的真实时间,在本文的实验环境中可忽略。
    • 不稳定性成本 ($C_{noise}$): 过度或不一致的重新规划可能导致行为不稳定(如目标摇摆、无效回溯),从而降低任务成功率。这是一个概念性成本,其负面影响会通过降低任务回报而间接被优化。
  5. 计划漂移 (Plan Drift):指随着智能体与环境交互,现有计划的有效性随时间推移而衰减的现象。计划的抽象层次、模型的准确性和环境的动态性都会影响漂移速度。这个概念解释了为何智能体需要周期性地重新评估是否需要规划。

相关工作

当前,在序贯决策领域,相关工作主要分为两类。一类是经典的规划方法,如蒙特卡洛树搜索 (Monte Carlo Tree Search, MCTS) 和模型预测控制 (Model Predictive Control, MPC),它们强调在行动前进行显式展望。另一类是大型语言模型(LLM)的推理方法,如 ReAct,它通过提示词引导模型在每一步行动前都进行思考和规划。

然而,现有方法存在一个关键瓶颈:“总是规划”的策略在计算上是昂贵的,并且在长时序任务中可能因引入行为不稳定性而降低性能;而“从不规划”则限制了智能体解决复杂问题的能力。目前缺乏一种有效的方法,能让LLM智能体学会在何时分配宝贵的测试时计算资源(即规划)是真正有益的。

本文旨在解决这个具体问题:如何训练LLM智能体,使其能够在序贯决策任务中学会动态、高效地分配测试时计算资源,即学会“何时应该规划”。

本文方法

本文的核心方法是提出一个概念框架,并基于此设计了一个两阶段训练流程,以教会LLM智能体动态规划。

Dynamic Planning Agent Architecture 图2:动态规划智能体架构。智能体是一个单一、整体的LLM,其概念性策略通过统一的输出格式实现。规划决策($\phi_{\theta}$)由模型选择是否以 \(<plan>\) 标志开始其生成来隐式做出。这个单一输出字符串随后被解析以提取动作($a_{t}$)和(如果存在的话)新计划($p_{t}$),从而执行行动($\pi_{\theta}$)和规划($\psi_{\theta}$)策略。

概念框架与优化目标

本文将智能体的决策过程建模为一个成本效益分析:仅当规划带来的预期收益(规划优势 $A_{plan}$)超过其总成本(规划成本 $C_{plan}$)时,才触发规划行为。

为了让智能体隐式地学会这种权衡,本文设计了一个强化学习目标。智能体的参数 $\theta$ 通过最大化以下目标函数进行优化:

\[\theta^{\*}=\arg\max_{\theta}\mathbb{E}_{\tau\sim\theta}\left[\sum_{t=0}^{H}\gamma^{t}\left(R_{task}(s_{t},a_{t})-d_{t}\cdot C_{tokens,t}\right)\right]\]

其中,$R_{task}$ 是环境提供的任务奖励,$d_{t}$ 是一个二元变量(如果第 $t$ 步进行了规划则为1,否则为0),$C_{tokens,t}$ 是规划产生的 token 数量的成本。这个目标函数直接鼓励智能体在任务回报的提升(隐式地代表了 $A_{plan}$)和规划的计算成本之间做出权衡。而由频繁规划导致的行为不稳定等负面效应(即 $C_{noise}$)会自然地导致任务回报降低,从而被间接惩罚。

两阶段训练流程

为了实现上述目标,本文提出一个包含监督微调(SFT)和强化学习(RL)的两阶段训练流程:

第一阶段:监督微-调预训练 (SFT Priming)

第二阶段:强化学习微调 (RL Fine-Tuning)

创新点

实验结论

本文通过在 \(POGS\) 和 \(Crafter\) 环境中的一系列实验,验证了动态规划的有效性。

Zero-shot, SFT, and RL results 图1:(a-b) Zero-shot结果显示,在Crafter和POGS中存在最佳的“金发姑娘”规划频率。(c-d) SFT结果表明,规划智能体在与基础模型的KL散度较低时表现更佳。(e-f) RL结果显示,经过SFT预训练的规划智能体比非规划基线更具样本效率,并能更稳定地达成复杂目标。

Zero-Shot 评估结论:

SFT 预训练结论:

RL 微调结论:

Best-of-N comparison on Crafter 图5:“最好中的最好”(Best-of-N,N=100)在Crafter上的比较图。结果显示,人类协作(N=20)取得了最强的进展,其次是SFT+RL动态规划和SFT+RL无规划,而基础RL基线则表现落后。

最终结论: 实验证明,训练LLM智能体动态地决定何时规划,相比于固定的规划策略,是一种更有效、更高效的方法。本文提出的两阶段SFT+RL训练范式是实现这一目标的有效路径,它不仅提升了智能体的自主性能和效率,还增强了其与人类协作的可控性,为构建更强大、更安全的智能体系统开辟了新的方向。