WorldLLM: Improving LLMs' world modeling using curiosity-driven theory-making

TL;DR

WorldLLM enhances large language models' ability to predict and understand structured environments by combining Bayesian inference, active exploration, and hypothesis refinement, leading to more accurate and interpretable world models.

Contribution

This paper introduces WorldLLM, a novel framework that integrates Bayesian inference, reinforcement learning, and LLMs to improve structured world modeling and hypothesis refinement.

Findings

Improves predictive accuracy in environment modeling.

Generates human-interpretable theories of environment dynamics.

Effectively explores environments using curiosity-driven policies.

Abstract

Large Language Models (LLMs) possess general world knowledge but often struggle to generate precise predictions in structured, domain-specific contexts such as simulations. These limitations arise from their inability to ground their broad, unstructured understanding in specific environments. To address this, we present WorldLLM, a framework that enhances LLM-based world modeling by combining Bayesian inference and autonomous active exploration with reinforcement learning. WorldLLM leverages the in-context learning abilities of LLMs to guide an LLM-based world model's predictions using natural language hypotheses given in its prompt. These hypotheses are iteratively refined through a Bayesian inference framework that leverages a second LLM as the proposal distribution given collected evidence. This evidence is collected using a curiosity-driven reinforcement learning policy that explores the environment to find transitions with a low log-likelihood under our LLM-based predictive model using the current hypotheses. By alternating between refining hypotheses and collecting new evidence, our framework autonomously drives continual improvement of the predictions. Our experiments demonstrate the effectiveness of WorldLLM in a textual game environment that requires agents to manipulate and combine objects. The framework not only enhances predictive accuracy, but also generates human-interpretable theories of environment dynamics.