Source-linked AI summary
Towards a Belief-Based World Model for LLM Agents
Shubham Kumar, Harshit Kumar, Narendra Ahuja, Saurabh Jha
TL;DR
LLM agents struggle with belief maintenance in long-horizon partially observable tasks, while simulation alone incompletely communicates current-state uncertainty. The paper introduces BB-WMs, which maintain queryable beliefs alongside simulation, and finds that belief access improves decision-making and complements simulation-based interfaces.
Problem
Simulation-based world-model interfaces incompletely communicate uncertainty about the current state, which LLM agents may need for decision-making under partial observability.
Method
BB-WMs maintain an explicit belief over the current state, update it with observations, and expose it to LLM policies through natural-language queries alongside simulation.
Results
LLM agents benefit from access to an accurate belief over the current state, with belief queries and simulation queries providing complementary benefits.
Takeaways & Limitations
Belief-Based World Models motivate giving LLM agents direct natural-language access to what is known and uncertain while preserving simulation for future consequences.
Takeaways & Limitations
The study uses benchmark-specific BB-WM instantiations and deliberately hand-specifies task-relevant state spaces and belief transitions rather than learning scalable accurate BB-WMs.
Abstract
from arXiv · showhide
Large language models (LLMs) are being used as policies for autonomous decision-making and planning in many domains. Despite their strong reasoning capabilities, LLMs struggle with long-horizon tasks, especially under partial observability. World models are a promising way to enhance policy performance, both during training and inference. During inference, agents currently use world models to simulate the consequences of candidate actions before committing to an action, which can improve decision-making. However, we argue that simulation alone is an incomplete interface for decision-making under partial observability: simulation doesn't adequately capture uncertainty about the current state, which agents may need for accurate decision-making. We address this limitation with Belief-Based World Models (BB-WMs), which model and maintain a belief that LLMs can query to access information on what is known and uncertain about the current state. Before developing methods to learn accurate BB-WMs, we first ask a more fundamental question: does exposing a world model's belief directly to an LLM policy improve decision-making? Our results show that giving LLM agents access to world model beliefs improves task performance under partial observability, while remaining complementary to existing simulation-based world models. Code is released at https://github.com/skumar-ml/belief-world-models.
1 Introduction
World models help policies evaluate future consequences, but simulation alone does not directly expose uncertainty about the current state. The paper proposes exposing beliefs to LLM agents so they can reason about information-gathering actions alongside pragmatic actions.
- World models let policies anticipate the outcomes of candidate actions without costly real-world interactions.
- Simulation evaluates pragmatic actions, whereas epistemic actions require information about what is known and uncertain in the current state.
- Simulation can conflate current-state uncertainty with stochastic future transitions and become sample-inefficient over longer horizons.
- LLM agents are strong general-purpose policies but struggle to maintain and update beliefs in complex, long-horizon partially observable environments.
- Belief-Based World Models maintain and expose current-state beliefs through natural-language queries while retaining simulation for action consequences.
2 Related Works
Prior approaches either provide simulation-based foresight or maintain estimates of the current state. BB-WMs combine belief tracking with action-conditioned simulation through a directly queryable interface.
- Decoupled world models provide LLM agents with foresight into candidate-action consequences before execution.
- Ground-truth simulation does not reliably improve test-time decision-making because agents may fail to query it, misuse outputs, or degrade when simulation is enforced.
- Textual and structured state representations maintain estimates of task-relevant current state as new observations arrive.
- Unlike these state estimators, BB-WMs combine direct belief access with action-conditioned simulation for future reasoning.
3 Belief-Based World Models
BB-WMs extend simulation-based world models with an explicit belief over the current latent state. They update that belief from observations and expose both belief queries and simulated-future readouts to the policy.
- Existing world models simulate candidate actions and return policy-facing readouts of predicted future states.
- Prior world-model families return predicted rewards, latent states, observations, or natural-language descriptions after simulation.
- A BB-WM represents its internal state as an explicit belief over the environment’s current latent state plus additional information.
- The belief is updated by first predicting a prior representation after an action and then incorporating the next observation.
- BB-WMs let policies query the current posterior belief without supplying a hypothetical action.
- Policies may issue multiple belief and simulation queries before acting, using their responses to condition the next action.
4 Method
The method tests whether pretrained LLM policies benefit from querying explicit beliefs while using simulation to evaluate candidate actions in partially observed text environments. It uses hand-specified state representations, belief updates, natural-language queries, and rule-based or oracle action validation across ALFWorld and ScienceWorld.
- Method design: The study hand-specifies task-relevant state spaces, prior beliefs, posterior updates, and natural-language belief readouts so pretrained LLMs can query beliefs without finetuning.This isolates the value of exposing beliefs from the separate challenge of learning an accurate BB-WM.
- ALFWorld: ALFWorld models deterministic state components alongside categorical beliefs over object locations across receptacles.The belief starts uniform over possible receptacles and is updated by collapsing on a found object or renormalizing after a failed search.
- Method design: The ALFWorld simulation component is intentionally incomplete: WALL-E neither predicts the complete next state nor represents a belief.This design isolates the contribution of the belief-query interface.
- ALFWorld: The LLM accesses ALFWorld beliefs through where is <object> queries, while WALL-E validates contemplated actions and returns natural-language validity feedback.Valid actions execute; invalid actions trigger replanning conditioned on the feedback.
- ScienceWorld: ScienceWorld applies analogous state and belief modeling to 24 test task types in a fixed ten-room house, with queries returning ranked rooms or a known container.Its object-location uncertainty is generally less varied than ALFWorld’s.
- ScienceWorld: Because WALL-E lacks a public ScienceWorld implementation, the study uses an oracle validator that always predicts action validity correctly but provides generic rather than failure-specific feedback.Invalid actions are retried within the same retry budget used for WALL-E.
5 Experiments
Experiments on ALFWorld and ScienceWorld evaluate belief queries, simulation, and their combination across LLM agents and frameworks. BB-WM generally improves performance and efficiency, while ablations and qualitative cases indicate that explicit beliefs contribute beyond deterministic memory and complement simulation.
- Experimental Setup: Experiments evaluate Llama-3.1-8B-Instruct, Qwen3-14B, and Sonnet 4.6 with ReAct and ReflAct on ALFWorld and ScienceWorld.ALFWorld uses SR@1 and steps; ScienceWorld uses average reward and average reward per step.
- Methods Evaluated: Belief provides current-state queries, WALL-E provides action-conditioned simulation, and BB-WM combines both interfaces.ScienceWorld uses an oracle WALL-E implementation because no public WALL-E implementation was available.
- Results: On ScienceWorld, BB-WM improves performance and efficiency for Llama and Qwen, whereas Sonnet shows no meaningful performance change but modest efficiency gains.Sonnet appears to saturate both benchmarks, limiting the benefits of additional world modeling on easier tasks.
- Results: BB-WM improves over the base agent in performance and efficiency on ALFWorld, while belief and simulation queries address different shortcomings.For Llama+ReAct on ALFWorld, gains from each method individually compound when combined; in one Qwen3-14B+ReflAct case, WALL-E alone hurts performance but BB-WM surpasses the Belief variant.
- Memory vs. Belief: Removing probabilistic beliefs reduces performance and efficiency, and sometimes makes the Memory variant resemble the base agent.The Memory-only ablation retains observed-object memory but drops beliefs over unobserved object locations.
- Qualitative Analysis: Qualitative cases show BB-WM guiding agents toward valid or sensible locations before targeted actions in both benchmarks.In ALFWorld, belief queries help locate an unusually initialized egg; in ScienceWorld, they help select the greenhouse flower pot for growing an orange plant.
6 Conclusion
The paper concludes that LLM agents benefit from access to accurate beliefs about the current state. BB-WMs expose and maintain these beliefs while supporting active sensing and continual learning.
- LLM agents acting as policies benefit from access to an accurate belief over the current state.
- BB-WMs model, maintain, and expose current-state beliefs to LLM agents through a natural-language interface.
- Benchmark-specific BB-WM instantiations motivate future work on scalable, domain-agnostic methods for learning accurate BB-WMs.
- Belief updating and continual learning address complementary adaptations: new observations reduce state uncertainty, while surprising observations can signal that the world model needs revision.
- BB-WMs can support active sensing by revealing when more information is needed and continual learning by revealing when predictions fail to explain new observations.
A BB-WM Instantiation for ALFWorld
The ALFWorld instantiation combines a simulation-based world model with a belief-based query interface. Agents can query state information and likely object locations without consuming environment steps.
- Agent Prompt: The agent is instructed to reflect on its state and task before selecting an action from the available household-action set.The prompt requires the format “Reflection: ... Action: ...” and lists navigation, manipulation, and object-transformation actions.
- WALL-E World Model: The world model checks proposed actions before execution and returns corrective observations when an action is infeasible.These checks do not execute the action or consume a step.
- Belief-Based World Model: The belief variant tracks the environment state and likely object locations, which agents can consult through free queries.Queries are intended to help agents reorient instead of guessing when uncertain.
- Belief-Based World Model: Agents can query likely locations, observed receptacle contents, search status, and a compact state summary.The interface distinguishes probabilistic location information from deterministic state information.
A.3 In-Context Examples
The ALFWorld world-model state separates deterministic task and environment information from probabilistic beliefs about target-object locations. The target belief is initialized from placement constraints and updated as objects are searched or found.
- Deterministic State Space: The deterministic state records the task goal, receptacles, observed objects, agent location, and inventory.Receptacle and object entries include structural and attribute fields such as search status, contents, location, and transformation states.
- Deterministic State Space: The goal specifies an action type, with optional transformation and destination fields depending on the task.Supported action types include pick_and_place, clean, heat, cool, look, and picktwo.
- Probabilistic State Space: The probabilistic state contains categorical distributions over receptacle locations for task-relevant target objects.The target belief includes a distribution and, after observation, an optional concrete instance identifier and known location.
- Probabilistic State Space: Before observation, the target belief is distributed over compatible receptacles; after observation, it becomes concentrated at the observed receptacle.The representation supports separate target beliefs when multiple same-type objects are required.
A.5 Belief-Query Interface
The belief-query interface exposes world-model information through the same action channel as environment actions. Queries provide ranked location beliefs, deterministic state information, and search status without changing the environment or consuming steps.
- ALFWorld Interface: Queries return world-model answers as observations while leaving the environment unchanged and consuming no environment step.The interface supports four query types.
- ALFWorld Interface: The “where is” query ranks unsearched receptacles for a target object by positive belief probability.Locations are returned in decreasing probability order until cumulative probability reaches at least 0.9, with the remainder summarized separately.
- ALFWorld Interface: For non-target objects, location responses use placement priors or deterministic object state rather than a persistent probabilistic belief.Observed instances are reported as held, located at a receptacle, or not yet located.
- ALFWorld Interface: The “what is in” query reports observed contents or indicates that a receptacle has not been searched.The “searched” query separately reports whether a receptacle has been searched.
- ScienceWorld Interface: In ScienceWorld, the belief interface queries likely rooms, observed room contents, searched rooms, and the agent’s current state.ScienceWorld queries likewise do not consume environment steps.
B.4.2 Probabilistic State Space
The probabilistic state space maintains categorical beliefs over the rooms where tracked object types may be located, alongside a resolved location field updated after observation. Priors are derived from placement constraints and task-provided location information.
- Each tracked object type has a categorical belief over the ten ScienceWorld rooms, representing the probability of being in each room.The belief assigns probability to room-level locations for an object type.
- An object belief also stores a resolved location that is initially undefined and later records its room or inventory status after observation.Observed container contents can be represented through the resolved location state.
- Prior: Without additional task information, the prior is uniform over rooms where the object type may occur, or over all ten rooms when no placement information exists.Candidate-room constraints determine the support of the uniform prior.
- Prior: When a task specifies the target object's room, its belief is initialized deterministically at that room.This task-specific initialization replaces uncertainty for the target location.
- Prior: For other objects with stated locations, the prior assigns 0.9 probability to the stated room and distributes 0.1 across the remaining candidate rooms.The residual probability is distributed uniformly among the object's other candidate rooms.
B.5 Belief-Query Interface
The belief-query interface lets the policy inspect world-model state through the same action channel used for environment actions. It exposes deterministic contents and agent state, as well as probabilistic object locations that are progressively resolved by observations.
- Queries are issued as actions but return world-model observations without modifying the environment or consuming an environment step.The interface provides four query types.
- Belief queries: The `where is` query ranks unsearched rooms by positive belief mass and reports locations until cumulative probability reaches at least 0.9.Remaining probability mass is summarized as other rooms.
- Belief queries: After observation, `where is` can return a deterministic room and container location, or report inventory with probability 1 when the agent holds the object.If all candidate rooms have been searched without finding the object, no candidate rooms remain.
- Deterministic state queries: The `what is in` query returns known contents for searched rooms or observed containers, while unobserved locations are reported as not searched.Container contents can also appear in the inventory summary when the agent holds the container.
- Deterministic state queries: The `searched` query reports whether a room has been observed, and the `state` query summarizes the current room, inventory, focused object, and searched rooms.These queries expose compact deterministic state information to the policy.
C Example ALFWorld Trajectories
The ALFWorld examples contrast trajectories with and without belief queries. In the successful trajectory, querying the egg location narrows the search and supports task completion, whereas the comparison trajectory exhausts its action budget after missing the garbagecan.
- WALL-E trajectory: The comparison trajectory searches many cabinets, drawers, and stoveburners but does not search the garbagecan before reaching the action budget.Its initial observation lists the garbagecan among many possible search targets.
- WALL-E trajectory: The comparison trajectory fails after 30 environment steps, and its termination is attributed to the task's action budget.
- BB-WM trajectory: The BB-WM trajectory queries the egg location before searching and initially receives seven equally likely candidate locations.After checking several locations, the belief narrows to the garbagecan and sinkbasin.
- BB-WM trajectory: The agent finds the egg in the garbagecan, heats it with the microwave, and returns it to the garbagecan.The trajectory includes one rejected action before successfully heating the egg after closing and reopening the microwave.
- BB-WM trajectory: The belief-query trajectory succeeds in 20 environment steps with 2 belief-query steps and 1 WALL-E rejection.
D ScienceWorld Example Trajectories
The ScienceWorld examples show belief queries supporting object search and planning, while also illustrating that successful early localization does not ensure task completion. One trajectory finds the seed and flower pot efficiently, whereas another fails after extended growth actions.
- Task context: The queried trajectory's task description requires growing an orange plant from seed to the reproduction life stage.The task also states that seeds can be found in the bathroom.
- WALL-E Oracle trajectory: The non-query trajectory repeatedly waters and waits for the plant, then explores the greenhouse and outside before returning to the bathroom.Its actions include repeated sink activation, filling, pouring, and waiting.
- WALL-E Oracle trajectory: The non-query trajectory fails with reward 39 after 30 environment steps.
- BB-WM trajectory: The BB-WM trajectory queries the orange seed location, then searches the bathroom specified by the task and finds a seed jar containing orange seeds.The query assigns equal probability across six rooms before the task description guides the initial search.
- BB-WM trajectory: After finding the seed jar, the agent queries the flower pot location and receives greenhouse with probability 1 before teleporting there.