paper_title: 由蒙特卡洛树搜索引导的大语言模型生成代码世界模型 blog_title: 教 LLM 编写世界代码: 通往更智能 AI Agent 的新路径

想象一下教一个机器人下棋。你可以给它看数百万盘棋局,希望它能像许多深度学习模型一样学会其中的模式。或者,你也可以给它一本规则手册。有了规则,机器人就不仅仅是模仿过去的棋局——它能够对任何可能的棋盘局面进行推理,预测结果,并战略性地规划自己的走法。

这本规则手册,就是 AI 研究者所说的世界模型——一个关于世界如何运作的内部模拟。

对于一个 AI Agent 来说,拥有一个准确的世界模型就像拥有了超能力。它能让 Agent 提前规划,理解行动的后果,并以极高的效率适应新任务。但构建这些模型充满挑战。传统方法需要海量数据,而虽然大语言模型 (LLM) 拥有广博的知识和推理能力,但直接将它们用作模拟器通常缓慢、昂贵且不可靠

如果我们能两全其美呢?

这就是一篇引人入胜的新论文背后的核心思想: 由蒙特卡洛树搜索引导的大语言模型生成代码世界模型。研究者们不再要求 LLM 成为世界模型,而是让它编写一个程序来充当世界模型。这个程序——一个代码世界模型 (CWM)——精确、可解释,且运行速度快如闪电。

为了构建 CWM,他们引入了一种强大的代码生成方法,名为 **GIF-MCTS **(使用蒙特卡洛树搜索进行生成、改进和修复) ,通过编写、调试和精炼代码的迭代循环来引导 LLM,直到它能精确地模拟环境。

在本文中,我们将剖析代码世界模型的工作原理,探索 GIF-MCTS 如何精心打造它们,并讨论为什么这种方法可能重塑样本高效、智能 Agent 的创建方式。


世界模型的问题所在

基于模型的强化学习 (RL) 中,Agent 学习预测其环境接下来会发生什么。给定当前状态 \(s\) 和一个选定的动作 \(a\),世界模型会输出:

  • 下一个状态 \(s'\)
  • 奖励 \(r\)
  • 任务是否结束 (done 标志)

借助这种预测能力,Agent 可以在脑中模拟成千上万种未来,以选择最佳计划——而无需承担现实环境中试错的成本或风险。

研究人员尝试了多种方法来构建此类模型:

  1. 专用神经网络 — 速度快且可训练,但在理解复杂文本指令或泛化到新情境方面能力有限。它们无法轻易利用 LLM 所包含的广泛、非结构化知识。

  2. 多模态模型 (例如文本到视频) — 预测结果逼真,但对于实时规划而言速度太慢。单次预测可能需要几秒,这对于需要快速反应的 Agent 来说不可用。

  3. 将 LLM 作为直接模拟器 — 用当前状态提示 LLM,并要求它以文本形式给出下一个状态。这往往产生不可靠的数值,查询速度慢,而且在规划循环中反复调用的成本高昂。

作者们看到了一个更好的媒介:** 代码**。代码精确、执行速度极快,并且人类可解释。但让 LLM 编写复杂且正确的模拟代码并非易事。


代码世界模型 (CWM) 简介

关键思想是将世界建模任务转化为程序合成

给定:

  • 一个 RL 环境的自然语言描述 (就像一本规则手册) ,
  • 一小组示例转移 (轨迹) ,

目标是生成一个能完美模拟该环境的 Python 类

这个类必须实现一个 step(state, action) 方法:

\[ (\hat{s}', \hat{r}, \hat{d}) = \text{code\_environment.step}(s, a) \]

生成的代码会与提供的轨迹进行验证。如果预测结果与真实结果匹配,我们就得到了一个成功的 CWM;否则,错误信息会被反馈进入代码编写流程进行改进。

代码世界模型框架概览。一个 LLM、环境描述和轨迹被输入到 GIF-MCTS 代码生成过程中。这会产生一个候选的代码世界模型,并与轨迹进行验证。如果失败,反馈将被送回生成器以精炼代码。如果成功,它就可以用于规划。

图 1: CWM 的迭代工作流程: 生成候选代码,与已知轨迹进行验证,并不断精炼,直到代码完全匹配环境。

为了严格评估这一想法,作者创建了代码世界模型基准 (CWMB)——包含 18 个多样化的 RL 环境 (从 CartPole 到复杂的 MuJoCo 物理任务) ,每个环境都有清晰的文本描述和少量轨迹数据集。

但如何让 LLM 编写、调试和精炼如此复杂的代码?一次性提示显然不够。

GIF-MCTS 登场。


GIF-MCTS: 更智能的代码编写方式

构建正确代码是一个编写、测试和修复的迭代过程。GIF-MCTS 使用蒙特卡洛树搜索 (MCTS) 来统筹这一过程。

MCTS 因 AlphaGo 而闻名,它在可能的决策路径上搜索,平衡对新思路的探索与对已知优法的利用,以寻找最优策略。

对于 GIF-MCTS:

  • 节点 = 程序的某个版本 (部分或完整)
  • 动作 = 用于修改代码的特定 LLM 提示
  • 节点价值 = 准确率 — 正确预测的测试轨迹比例

当算法不断尝试生成新代码片段、修复错误、优化逻辑以提高准确率时,搜索树会逐步扩展。

GIF-MCTS 树示例。算法探索不同的代码生成路径,并利用反馈 (预测准确率)

图 2: 在 GIF-MCTS 中,每个节点是一个按准确率评分的程序。动作 (生成改进修复) 在反馈引导下扩展搜索树,直到找到高性能的 CWM。

GIF-MCTS 使用三种专门的编辑动作:

1. 生成新代码行

主要的探索动作。LLM 获取现有代码片段,并被要求扩展它,有时会从零开始。从同一节点重复生成可探索多种补全方式。

2. 改进预测

针对能运行但结果错误的代码。LLM 会得到:

  • 完整程序
  • 出错输入
  • 错误输出
  • 正确期望输出

它会解释错误,并基于改进的逻辑重写代码——就像调试单元测试。

3. 修复错误

有些输出会导致崩溃 (语法/运行时错误) 。修复提示会将错误信息和代码反馈给 LLM 进行修复,而不是直接弃用,从而挽救潜力逻辑。

在 MCTS 中平衡这三种动作,GIF-MCTS 能系统地搜索代码空间,逐步逼近准确的 CWM。


实验与结果

作者测试了 GIF-MCTS 在通用代码生成和专门的 CWM 创建任务上的表现。


APPS: 编程竞赛挑战

为验证 GIF-MCTS 的代码生成能力,他们使用了 APPS 基准——一套高难度编程问题——并聚焦于最困难的“竞赛”子集。

表格显示不同方法在 APPS 基准上的表现。GIF-MCTS 取得最高 28.3% 严格准确率。

表 1: 在 APPS 竞赛问题上,GIF-MCTS 使用相同 LLM 主干模型取得了当前最佳性能。

GIF-MCTS (Llama 3–70B) 解决了 28.3% 的问题,超过了思维链提示和类似 WorldCoder 算法,成为高难度代码合成任务的领先方法。


CWMB: 构建世界模型

主要测试: 为 CWMB 基准中的 18 个环境生成 CWM。

评估的两项指标:

  1. 准确率 \(A\):

    \[ A = \frac{1}{N} \sum_{i=1}^{N} \left( \frac{1}{3} \mathbf{1}[s'_i, \hat{s}'_i] + \frac{1}{3} \mathbf{1}[r_i, \hat{r}_i] + \frac{1}{3} \mathbf{1}[d_i, \hat{d}_i] \right) \]
  2. 归一化回报 \(\mathcal{R}\):

    \[ \mathcal{R}(\text{CWM}) = \frac{R(\pi_{\text{CWM}}) - R(\pi_{\text{rand}})}{R(\pi_{\text{true}}) - R(\pi_{\text{rand}})} \]


    衡量使用 CWM 的规划器与随机策略及完美预言机规划器的表现差距。

代码世界模型基准 (CWMB) 的结果表

表 2: 在离散和连续环境中,GIF-MCTS 生成的 CWM 准确率更高,规划回报也优于 WorldCoder。

GIF-MCTS 全面优于 WorldCoder,在离散与连续动作空间均带来更好规划效果。GPT-4 Turbo 在复杂连续任务中尤为出色,包括成功建模具有挑战性的 Humanoid 环境。


RTFM: 语言与逻辑结合

RTFM 网格世界要求阅读一份描述怪物弱点和目标的“手册”,再据此采取行动,是对规则转化为可执行逻辑能力的严苛考验。

RTFM 环境结果表,GIF-MCTS 表现显著优于 WorldCoder

表 3: 在 RTFM 中,当给予更多精炼预算时,GIF-MCTS 引导 GPT-4 生成了完美模型。

在此任务中,GIF-MCTS 的准确率显著超过 WorldCoder。通过 50 次 LLM 调用,GPT-4 生成了 100% 精确的 CWM,使规划器获得归一化回报 1.0——与预言机表现持平。


关键要点与未来方向

这项研究不仅是新算法——更改变了我们对 LLM 在 RL 推理角色的认知

  • 样本与推理效率: CWM 所需环境交互极少 (测试为 10 条短轨迹) ,且运行速度比逐步调用 LLM 快几个数量级
    Humanoid-v4 中,CWM 步仅耗时 0.0001秒,而 GPT-4 需要 146.7秒——速度提升百万倍。

推理时间比较表,CWM 比直接查询 GPT-4 快 4 至 6 个数量级

表 8: CWM 通过用即时代码执行替代重复 LLM 调用,实现了巨大的速度提升。

  • 搜索的力量: GIF-MCTS 的引导式探索比单纯精炼效果更佳——结合生成改进修复让方法更稳健。

  • 准确率与回报差距: 高准确率并不总意味着高回报。遗漏稀有但关键的状态 (如成功终止状态) 会严重影响规划器表现。未来工作应重视建模这些边缘案例。

局限性: CWM 目前假设环境为确定性且完全可观测,并依赖清晰文本描述。扩展到随机、部分可观测或图像定义的任务是重要下一步。


结论

生成代码世界模型提出了一个令人振奋的设想: 将 LLM 从缓慢的模拟器转型为程序合成器,从而打造快速、精确、可解释且样本高效的世界模型。

GIF-MCTS 证明,结构化的搜索驱动方法可引导 LLM 解决超越简单提示的复杂推理与编码任务。随着 LLM 编码能力提升,类似 GIF-MCTS 的方法可为曾被视为不可解的任务生成先进世界模型。

这是让 AI Agent 真正理解、预测并规划复杂世界的重要一步——由它们亲手编写的代码驱动。