Towards a Belief-Based World Model for LLM Agents

Consider an autonomous system that plans and executes a series of actions to achieve some goal. It’s hypothesized that creating models of the world (or world models) can improve the learning of the system’s decision-making function, or policy. Instead of needing costly real-world interactions, a policy can be trained on diverse, imagined future trajectories from a world model (Hafner et al., 2020; Ha and Schmidhuber, 2018). World modeling training objectives can also be used alongside reinforcement learning (RL) to improve the policy itself (Yu et al., 2026; Schwarzer et al., 2021; Copet et al., 2025). Recent world modeling efforts have primarily gone towards improving the policy at training time (Bruce et al., 2024; Agarwal et al., 2025; Ren et al., 2025). Once trained, the policy does not use the world model within a control loop (i.e., model-free control). While this is a promising direction, it is unclear how well such policies generalize to novel situations at test time. To help with this, the policy may also have access to a world model at test-time (i.e., model-based control), where foresight from the world model enables test-time reasoning (e.g., planning) in difficult or novel situations. This direction has received comparatively less attention and is the focus of this work.

Recent model-based control methods (Maes et al., 2026; Zhou et al., 2025; Hao et al., 2023; Janner et al., 2022; Schrittwieser et al., 2020) allow the policy to interact with a world model through a simulation interface: given a representation of the current state and a candidate action (or action sequence), what future trajectory follows? We raise a broader question: is simulation the only interface that a world model should expose to a policy? While simulation may be sufficient for training model-free policies, we argue that it is inadequate for model-based control. Different types of actions require different information from a world model. Consider pragmatic actions, which primarily advance the agent toward its goal (e.g., picking up a mug in order to drink coffee). Simulation is well-suited to these actions because it provides foresight by evaluating the consequences of acting. In contrast, epistemic actions are information-gathering actions that primarily reduce uncertainty about the current state of the environment (Kaelbling et al., 1998; Kirsh and Maglio, 1994). In this setting, the policy benefits from a belief—what is known and what is uncertain—about the current state to reason about whether epistemic actions are needed. Simulation can only expose uncertainty indirectly through variation across predicted futures, which may conflate uncertainty about the present state with stochasticity in the environment’s state transition dynamics. Moreover, accurately recovering current-state uncertainty from simulation may require many samples and may still miss low-probability outcomes. Thus, we argue that the current-state belief should be exposed directly rather than reconstructed indirectly from future simulations (exemplified in Fig. 1).

Recently, large language model (LLM) agents have emerged as general-purpose policies in open-ended environments, from software engineering to robotics (Yang et al., 2024; Drouin et al., 2024; Huang et al., 2022; Ahn et al., 2022). Given their general-purpose abilities, LLMs may be able to implicitly world model while planning actions. However, prior works have found that LLMs struggle with maintaining and updating beliefs in complex, long-horizon partially observable tasks (Rahmani, 2026; Jha et al., 2026; Zou et al., 2026). Recent world models for LLMs do not adequately alleviate this burden; they predominantly expose action-conditioned foresight while leaving belief estimation to the agent. The agent must therefore implicitly maintain uncertainty about the current state while simultaneously reasoning about action selection, which is likely a suboptimal division of responsibilities. This motivates a modular setting: the trained LLM performs general reasoning and action selection, while a separate world model performs current and future state belief estimation. The remaining question is how these independent components should communicate.

Separating belief estimation from action selection has precedence in decision-making under partial observability (Kaelbling et al., 1998) and has subsequently been explored in deep reinforcement learning (RL), where dedicated belief estimation can substantially improve decision-making (Hafner et al., 2020; Igl et al., 2018; Singh et al., 2021; Wang et al., 2023; Chen et al., 2022). However, these approaches typically train the policy around a particular belief representation, and the control loop—the manner in which the policy and world model interact—is determined by the algorithm designer. This largely avoids the need to design an explicit interface between world model and policy. The LLMs considered in this work present a different setting. The LLM may be a pretrained, frozen, or even proprietary general-purpose reasoner that cannot readily be retrained around an arbitrary belief representation. At the same time, the agent’s general reasoning abilities make a fixed control loop unnecessary: rather than specifying how the agent must use the world model, we can expose world model capabilities through an interpretable interface and allow the LLM itself to decide when they are useful. We therefore propose exposing belief in natural language, enabling an LLM to initiate queries about current or future states without being trained around a particular belief representation.

In short, we argue that the simulation interface of current world models for LLM agents should be augmented to explicitly expose beliefs over the current state. To achieve this, we introduce Belief-Based World Models (BB-WMs), which should maintain a belief over the current state, update it as new observations arrive, and propagate it forward during action-conditioned simulation. This affords an LLM agent two complementary interfaces: natural language queries ask about the world model’s current belief when reasoning about epistemic actions, while simulation evaluates the consequences of pragmatic actions. Before designing methods for learning scalable Belief-Based World Models, this paper focuses on a more fundamental question: if a world model represents a belief over the state, does exposing this belief to the LLM agent improve decision-making? To isolate this question, we intentionally design and study hand-crafted, benchmark-specific BB-WMs, abstracting away the separate challenge of learning accurate beliefs from experience.

World Models for LLMs: There have been efforts to equip LLM agents with separate world models that provide foresight into the consequences of candidate actions. WMA (Chae et al., 2025) finds that out-of-the-box LLMs cannot accurately simulate the outcome of actions in web interaction tasks, so they learn a separate world model to provide action-conditioned foresight. WALL-E (Zhou et al., 2025) replaces generative world models with a neurosymbolic world model, which checks the validity of an LLM-proposed action before it acts; if the action is predicted to be invalid, the LLM is asked to replan. DreamPhase (Hamidi et al., 2026) learns a latent world model that generates hypothetical future observations from predicted future latent states, scores the generated futures with a learned value function, and distills feedback into natural language to condition a frozen LLM agent. These methods predominantly expose an action-conditioned simulation interface to the policy.

Interestingly, recent work by Qian et al. (Qian et al., 2026) shows that providing an LLM agent with a ground-truth simulator alone does not translate into improved decision-making at test-time; agents often fail to invoke simulation when useful, misuse its outputs, or even degrade when simulation is enforced. Their analysis suggests that a key challenge is in when an agent decides to query a world model and how it integrates the resulting information into its reasoning. These findings motivate our efforts to extend world models with belief modeling capabilities, allowing the LLM to benefit from a different type of information.

Belief Representation in LLM Agents: Some works have sought to improve LLM agents by maintaining estimates of the current environment state rather than relying solely on an agent’s ability to reason over the interaction history. Approaches such as Statler, QuBE, StateAct, and ABBEL (Yoneda et al., 2024; Kim et al., 2024; Rozanov and Rei, 2024; Lidayan et al., 2025) maintain compact textual or structured representations of task-relevant state that are updated as new observations become available. These representations emphasize a single estimate or summary of the current state rather than explicitly preserving uncertainty over many possible states. Some recent work has begun to model this uncertainty directly. BeliefMem (Liao et al., 2026) retains multiple candidate conclusions together with probabilities, allowing competing interpretations of past evidence to coexist and updating them as additional evidence is observed. Agent-BRACE (Singh et al., 2026) more directly adopts the POMDP notion of belief for LLM agents, representing the current state uncertainty and conditioning the policy on this structured belief. These works demonstrate growing interest in explicitly representing uncertainty about the current state and making this information available to the agent during decision making. However, these approaches are not world models. They focus on estimating and maintaining the environment’s current state but do not provide an action-conditioned simulation interface for reasoning about the future. BB-WMs seek to bridge simulation with belief estimation: the policy can directly access the world model’s belief when reasoning about uncertainty in the current state, while continuing to obtain foresight through simulation.

Preliminaries: An agent, at any timestep $t$, receives an observation $o_{t}$, takes action $a_{t}$ in the environment, and receives the subsequent observation $o_{t+1}$. The agent selects actions in order to achieve some goal $g$. A world model maintains some internal representation $x_{t}$ of the environment’s true latent state $s_{t}$ based on the agent’s interaction history. Most world models used for model-based control expose a similar interface to the policy. Specifically, the policy queries the world model with $K$ candidate actions $\{a_{t}^{(i)}\}_{i=1}^{K}$ (in practice, the policy may query with a sequence of actions). The world model predicts (or samples) the representation of the future state after taking the action: $\hat{x}^{(i)}_{t+1}=Predict(x_{t},a_{t}^{(i)})$. The predicted internal state need not itself be exposed to the policy. We denote by $R(\hat{x}^{(i)}_{t+1})$ the readout of the predicted future, which may take many different forms, including a latent representation, natural-language description, symbolic state, or generated observation. The policy uses these readouts to determine which action to execute in the real environment. Prior methods may differ in terms of their internal representation $x_{t}$ and the response given $\mathcal{R}(\hat{x}_{t+1})$ to the policy, but this simulation interface underlies many world model families, which we show in Fig. 2.

Belief-Based World Models: In a BB-WM (see bottom right of Fig. 2), the internal representation $x_{t}=(b_{t},\xi_{t})$ includes an explicit belief $b_{t}$ over the environment’s current latent state $s_{t}$, with $\xi_{t}$ denoting any additional information. The BB-WM specifies an interface that allows the belief to be directly exposed to the policy. Specifically, given all past observations and actions $h_{t}=(o_{1:t},a_{1:t-1})$, the belief is $b_{t}\equiv P(s_{t}|h_{t})$. Before executing $a_{t}$, the prior representation of the next timestep is $x^{-}_{t+1}=(b^{-}_{t+1},\xi^{-}_{t+1})=Predict(x_{t},a_{t})$. After executing $a_{t}$ and observing the resulting $o_{t+1}$, the posterior representation is $x_{t+1}=(b_{t+1},\xi_{t+1})=Update(x^{-}_{t+1},o_{t+1})$. Then, the policy has access to two interfaces.

Simulation interface: For a candidate action $a_{t}^{(i)}$, the world model predicts the resulting next state $\hat{x}^{(i)}_{t+1}$. The policy receives a policy-facing readout: $y_{sim}^{(i)}=R_{sim}(\hat{x}^{(i)}_{t+1})$.

Belief-query interface: Unlike current world modeling methods, with a BB-WM the policy can directly query the current posterior belief without supplying a hypothetical action. Specifically, the policy can ask some query $q$ over belief $b_{t}$ through the interface: $y_{bel}=R_{bel}(b_{t},q)$.

Before taking an action, the policy may issue multiple queries across both interfaces, and the responses can be used as conditioning for the policy to decide the next action.

Recall our research question: “if a world model represents a belief over the state, does exposing this belief to the LLM agent improve decision-making?" To isolate the effect of exposing the belief to the agent, we deliberately abstract away the separate challenge of learning an accurate BB-WM by studying simple text-based environments and leveraging prior knowledge about the game environment. This allows us to hand-specify the task-relevant state space, the prior belief prediction $Predict(.)$, and the posterior belief update $Update(.)$, instead of needing to learn them from environment interactions. Because the hand-designed state space is semantic and interpretable, we can define a natural-language belief readout $R_{bel}(b_{t},q)$, allowing a pretrained LLM policy to query the BB-WM without finetuning. Our goal is not to propose hand-crafted world models as a scalable solution, but to establish whether explicit belief access is useful in the first place. We describe our benchmark-specific instantiations next.

Our goal is to study whether exposing a belief over the current state to an LLM agent improves its decision-making. We therefore perform a controlled study in which the LLM agent is frozen and variants differ only in the information exposed by the world model. Our objective is not to maximize benchmark performance, but to isolate the relative benefit of belief access. For this reason, we also do not compare to concurrent methods for belief tracking (referenced in Sec. 2).

Methods Evaluated: On each benchmark, we evaluate three LLMs (Llama-3.1-8B-Instruct (Grattafiori et al., 2024), Qwen3-14B (Yang et al., 2025), and Sonnet 4.6) on two agentic frameworks (ReAct (Yao et al., 2023) and ReflAct (Kim et al., 2025)). For each agent, we evaluate several different world models. In Belief, we allow the agent to only ask queries about the current belief state (no simulation). In WALL-E, we pair the agent with WALL-E’s action-conditioned next-state predictor (no belief state, just a valid/invalid action prediction). In BB-WM, we compose the belief modeling with WALL-E, allowing the agent to both do simulation and ask questions on the current belief state.

Evaluation Budget: In ALFWorld/BabyAI, the agent can take up to 30/64 environment actions to achieve the task, following. ScienceWorld varies the step budget per task family. In both benchmarks, we terminate execution early if the agent seems to be derailed; specifically if it issues 10 consecutive invalid actions or 10 consecutive world model queries.

Metrics: An effective agent should complete the most number of tasks in the fewest number of steps. We use reward-within-budget to measure both. The average reward obtain so far is reported as a function of the fraction of steps taken out of the total action budget. ALFWorld’s and BabyAI’s return is binary (0/1) depending on task success and ScienceWorld’s reward is a score from [0,1], allowing partial rewards in intermediate steps for reaching certain milestones.

This work shows that belief access improves LLM agent decision-making and often compounds with simulation-based world modeling, motivating Belief-Based World Models (BB-WMs), where current-state belief and future-state prediction provide distinct and complementary information to an agent. This work has two important limitations. First, our BB-WMs are hand-crafted and benchmark-specific. Thus, our results establish that accurate beliefs are useful when exposed to an LLM policy but do not address how such beliefs should be learned in realistic, open-ended settings. Second, our simulation interface is based on WALL-E-style action-validity prediction. This is a relatively limited form of simulation and does not propagate the current belief forward into a belief over hypothetical future states. Both limitations are important directions for future work, and our results indicate that addressing them could yield meaningful benefits in agent decision-making under partial observability.