生成代码世界模型：结合大型语言模型和蒙特卡罗树搜索

引言

在人工智能领域，能够对世界进行建模是智能体实现目标的关键能力。当面对新的环境时，智能体必须迅速理解其机制，以便制定计划并实现目标。在这种情况下，自然语言条件可以帮助智能体将当前观察与过去的知识联系起来，从而更好地理解世界。因此，通过自然语言向智能体传递关于新任务的信息，具有特别的前景。

近年来，大型语言模型（LLMs）在自然语言处理领域取得了革命性的进展，并为世界建模提供了巨大的机会。然而，将LLMs与世界模型结合的最佳方式尚不明确。在这篇文章中，我们介绍了一种新的方法，即通过编写Python代码来生成RL（强化学习）世界模型。

代码世界模型（Code World Models）

什么是代码世界模型？

代码世界模型（CWM）是由LLMs生成的一种世界模型，这些模型以Python代码的形式存在，专门用于基于模型的强化学习（RL）。与直接使用LLMs进行规划相比，调用代码具有精准、可靠、可解释和高效的优点。然而，编写合适的代码世界模型需要理解复杂的指令、生成具有非平凡逻辑的精确代码，以及通过单元测试和环境轨迹反馈自我调试长程序。

为了解决这些挑战，我们提出了GIF-MCTS（Generate, Improve and Fix with Monte Carlo Tree Search），这是一种新的代码生成策略，用于优化LLMs的代码生成能力。

代码世界模型的框架

在模型基础的强化学习问题设置中，我们将环境表示为具有状态空间、动作空间、转移函数和标量奖励函数的马尔可夫决策过程。世界模型的任务是准确表示转移函数和奖励函数。

我们假设：

环境是确定的且完全可观察的。
提供的环境自然语言描述足够详细，可以推断出观察空间以及转移和奖励函数的逻辑。

代码世界模型基准（CWMB）

为了全面测试不同环境下的世界模型生成，我们定义了一个新的基准，包括18个难度不同的RL环境。每个环境都有相应的自然语言描述和精心设计的轨迹数据集。这个基准旨在通过提供的数据学习准确合成代码世界模型，并在不同复杂度的环境中评估不同的代码生成方法。

GIF-MCTS：生成、改进和修复

我们提出的GIF-MCTS方法使用蒙特卡罗树搜索（MCTS）来生成和优化代码世界模型。GIF-MCTS的主要流程包括选择、扩展、评估和价值回溯。

动作类型

GIF-MCTS方法包括三种主要动作类型：

生成新行：利用LLM的随机采样能力，在树的不同分支中生成不同的代码片段。
改进预测：LLM会被提示改进现有代码，并提供错误预测的输入示例及其预期输出。
修复错误：当代码包含语法或运行时错误时，LLM会尝试修复这些错误，并提供修复逻辑。

实验

基准

我们在三个基准上评估了GIF-MCTS的性能：新的CWMB、流行且具有挑战性的编码基准APPS，以及语言条件网格世界RTFM。GIF-MCTS在所有三个基准上都超过了现有方法，包括APPS的“Competition”部分，CWMB的更难连续环境以及RTFM。

实验结果

使用GIF-MCTS生成的代码世界模型能够成功用于规划，生成的基于模型的RL智能体在样本效率和推理速度上显著提高。我们展示了在多个环境中使用合成的CWM进行成功的规划，结果表明CWM在推理速度上比直接查询LLM快四到六个数量级，同时在CWM准确时，其性能可以与访问真实世界模型的oracle规划器相匹敌。

讨论

限制

尽管代码世界模型（CWM）框架是一个令人兴奋的基于模型规划的方向，但我们仍依赖于环境的确定性和完全可观察性的假设。未来的工作可以探索如何将该方法扩展到包含随机性和部分可观察环境中。

此外，提供能够合理转换为Python函数的环境描述（例如记录关键变量的手册）也是一个潜在问题。在没有可用描述的情况下，预处理技术（例如图像到文本模型）可能是解决这一问题的一种方法。

结论

我们提出了一种新的框架，即代码世界模型（CWM），通过结合大型语言模型（LLM）和蒙特卡罗树搜索（MCTS），为基于模型的强化学习（RL）智能体构建世界模型。我们的实验表明，GIF-MCTS方法在多个基准上都表现优异，能够生成高效且可解释的代码世界模型，并显著提高基于模型的RL智能体的样本效率和推理速度。

未来的研究可以继续完善代码生成方法，优化底层LLM模型，并探索如何将CWM应用于更复杂的环境中。我们相信，CWM框架将推动开发出更快、更可解释且样本效率更高的基于模型的RL智能体。

致谢

本文的研究得到了芬兰研究理事会（旗舰计划：芬兰人工智能中心FCAI，资助编号352986、358246）和欧盟（H2020资助101016775和NextGenerationEU）的支持。我们感谢CSC为本项目提供的LUMI超级计算机的访问权限，该计算机由EuroHPC联合企业所有，由CSC（芬兰）和LUMI联盟主办。

参考文献

Ha, D. , & Schmidhuber, J. (2018). World models.✅
Hendrycks, D. , Basart, S., Kadavath, S., et al. (2021). Measuring coding challenge competence with apps.✅
Kocsis, L. , & Szepesvári, C. (2006). Bandit based monte-carlo planning.✅
Lin, J. , Du, Y., Watkins, O., et al. (2023). Learning to model the world with language.✅
Tang, H. , Key, D., & Ellis, K. (2024). Worldcoder, a model-based llm agent: Building world models by writing code and interacting with the environment.✅
Zhong, V. , Rocktäschel, T., & Grefenstette, E. (2020). RTFM: Generalising to new environment dynamics via reading.✅

这篇文章不仅介绍了代码世界模型和GIF-MCTS方法的理论基础，还通过实验验证了其在多个基准上的优越性能。希望未来的研究能够进一步提升这一领域的发展，为人工智能和强化学习带来更多突破。