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 可以在脑中模拟成千上万种未来,以选择最佳计划——而无需承担现实环境中试错的成本或风险。
研究人员尝试了多种方法来构建此类模型:
专用神经网络 — 速度快且可训练,但在理解复杂文本指令或泛化到新情境方面能力有限。它们无法轻易利用 LLM 所包含的广泛、非结构化知识。
多模态模型 (例如文本到视频) — 预测结果逼真,但对于实时规划而言速度太慢。单次预测可能需要几秒,这对于需要快速反应的 Agent 来说不可用。
将 LLM 作为直接模拟器 — 用当前状态提示 LLM,并要求它以文本形式给出下一个状态。这往往产生不可靠的数值,查询速度慢,而且在规划循环中反复调用的成本高昂。
作者们看到了一个更好的媒介:** 代码**。代码精确、执行速度极快,并且人类可解释。但让 LLM 编写复杂且正确的模拟代码并非易事。
代码世界模型 (CWM) 简介
关键思想是将世界建模任务转化为程序合成。
给定:
- 一个 RL 环境的自然语言描述 (就像一本规则手册) ,
- 一小组示例转移 (轨迹) ,
目标是生成一个能完美模拟该环境的 Python 类。
这个类必须实现一个 step(state, action)
方法:
生成的代码会与提供的轨迹进行验证。如果预测结果与真实结果匹配,我们就得到了一个成功的 CWM;否则,错误信息会被反馈进入代码编写流程进行改进。
图 1: CWM 的迭代工作流程: 生成候选代码,与已知轨迹进行验证,并不断精炼,直到代码完全匹配环境。
为了严格评估这一想法,作者创建了代码世界模型基准 (CWMB)——包含 18 个多样化的 RL 环境 (从 CartPole 到复杂的 MuJoCo 物理任务) ,每个环境都有清晰的文本描述和少量轨迹数据集。
但如何让 LLM 编写、调试和精炼如此复杂的代码?一次性提示显然不够。
GIF-MCTS 登场。
GIF-MCTS: 更智能的代码编写方式
构建正确代码是一个编写、测试和修复的迭代过程。GIF-MCTS 使用蒙特卡洛树搜索 (MCTS) 来统筹这一过程。
MCTS 因 AlphaGo 而闻名,它在可能的决策路径上搜索,平衡对新思路的探索与对已知优法的利用,以寻找最优策略。
对于 GIF-MCTS:
- 节点 = 程序的某个版本 (部分或完整)
- 动作 = 用于修改代码的特定 LLM 提示
- 节点价值 = 准确率 — 正确预测的测试轨迹比例
当算法不断尝试生成新代码片段、修复错误、优化逻辑以提高准确率时,搜索树会逐步扩展。
图 2: 在 GIF-MCTS 中,每个节点是一个按准确率评分的程序。动作 (生成、改进、修复) 在反馈引导下扩展搜索树,直到找到高性能的 CWM。
GIF-MCTS 使用三种专门的编辑动作:
1. 生成新代码行
主要的探索动作。LLM 获取现有代码片段,并被要求扩展它,有时会从零开始。从同一节点重复生成可探索多种补全方式。
2. 改进预测
针对能运行但结果错误的代码。LLM 会得到:
- 完整程序
- 出错输入
- 错误输出
- 正确期望输出
它会解释错误,并基于改进的逻辑重写代码——就像调试单元测试。
3. 修复错误
有些输出会导致崩溃 (语法/运行时错误) 。修复提示会将错误信息和代码反馈给 LLM 进行修复,而不是直接弃用,从而挽救潜力逻辑。
在 MCTS 中平衡这三种动作,GIF-MCTS 能系统地搜索代码空间,逐步逼近准确的 CWM。
实验与结果
作者测试了 GIF-MCTS 在通用代码生成和专门的 CWM 创建任务上的表现。
APPS: 编程竞赛挑战
为验证 GIF-MCTS 的代码生成能力,他们使用了 APPS 基准——一套高难度编程问题——并聚焦于最困难的“竞赛”子集。
表 1: 在 APPS 竞赛问题上,GIF-MCTS 使用相同 LLM 主干模型取得了当前最佳性能。
GIF-MCTS (Llama 3–70B) 解决了 28.3% 的问题,超过了思维链提示和类似 WorldCoder 算法,成为高难度代码合成任务的领先方法。
CWMB: 构建世界模型
主要测试: 为 CWMB 基准中的 18 个环境生成 CWM。
评估的两项指标:
准确率 \(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) \]归一化回报 \(\mathcal{R}\):
\[ \mathcal{R}(\text{CWM}) = \frac{R(\pi_{\text{CWM}}) - R(\pi_{\text{rand}})}{R(\pi_{\text{true}}) - R(\pi_{\text{rand}})} \]
衡量使用 CWM 的规划器与随机策略及完美预言机规划器的表现差距。
表 2: 在离散和连续环境中,GIF-MCTS 生成的 CWM 准确率更高,规划回报也优于 WorldCoder。
GIF-MCTS 全面优于 WorldCoder,在离散与连续动作空间均带来更好规划效果。GPT-4 Turbo 在复杂连续任务中尤为出色,包括成功建模具有挑战性的 Humanoid 环境。
RTFM: 语言与逻辑结合
RTFM 网格世界要求阅读一份描述怪物弱点和目标的“手册”,再据此采取行动,是对规则转化为可执行逻辑能力的严苛考验。
表 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秒——速度提升百万倍。
表 8: CWM 通过用即时代码执行替代重复 LLM 调用,实现了巨大的速度提升。
搜索的力量: GIF-MCTS 的引导式探索比单纯精炼效果更佳——结合生成、改进、修复让方法更稳健。
准确率与回报差距: 高准确率并不总意味着高回报。遗漏稀有但关键的状态 (如成功终止状态) 会严重影响规划器表现。未来工作应重视建模这些边缘案例。
局限性: CWM 目前假设环境为确定性且完全可观测,并依赖清晰文本描述。扩展到随机、部分可观测或图像定义的任务是重要下一步。
结论
生成代码世界模型提出了一个令人振奋的设想: 将 LLM 从缓慢的模拟器转型为程序合成器,从而打造快速、精确、可解释且样本高效的世界模型。
GIF-MCTS 证明,结构化的搜索驱动方法可引导 LLM 解决超越简单提示的复杂推理与编码任务。随着 LLM 编码能力提升,类似 GIF-MCTS 的方法可为曾被视为不可解的任务生成先进世界模型。
这是让 AI Agent 真正理解、预测并规划复杂世界的重要一步——由它们亲手编写的代码驱动。