PaperlayerView exact revision on arXiv

Source-linked AI summary

Agent-BRACE: Decoupling Beliefs from Actions in Long-Horizon Tasks via Verbalized State Uncertainty

Joykirat Singh, Zaid Khan, Archiki Prasad, Justin Chih-Yao Chen, Akshay Nambi, Hyunji Lee, Elias Stengel-Eskin, Mohit Bansal

arXiv:2605.11436v1cs.CLcs.AI

TL;DR

Long-horizon partially observable LLM agents need to track uncertain environment states without relying on ever-growing interaction histories. Agent-BRACE jointly learns a verbalized belief state and policy, improving average performance over strong RL baselines while maintaining a near-constant context window.

  • Problem

    LLM agents lack a compact belief-state representation for uncertainty in open-ended partially observable environments, forcing inefficient raw-history conditioning as episodes lengthen.

  • Method

    Agent-BRACE jointly trains belief-state and policy models via reinforcement learning, representing beliefs as atomic claims with ordinal verbalized uncertainty labels.

  • Results

    +14.5% and +5.3% average absolute improvement over the strongest RL-trained baseline on three TextWorld environments for Qwen2.5-3B-Instruct and Qwen3-4B-Instruct, respectively.

  • Takeaways & Limitations

    Agent-BRACE supports high-performing action selection from compact learned beliefs while maintaining a near-constant context window in long-horizon partially observable tasks.

  • Takeaways & Limitations

    The structured belief slots are specified in advance, while automatically discovering belief dimensions in fully open-ended environments remains future work.

Abstract

from arXiv · show

Large language models (LLMs) are increasingly deployed on long-horizon tasks in partially observable environments, where they must act while inferring and tracking a complex environment state over many steps. This leads to two challenges: partial observability requires maintaining uncertainty over unobserved world attributes, and long interaction history causes context to grow without bound, diluting task-relevant information. A principled solution to both challenges is a belief state: a posterior distribution over environment states given past observations and actions, which compactly encodes history for decision making regardless of episode length. In LLM agents, however, the open-ended nature of text makes it unclear how to represent such a distribution. Therefore, we introduce Agent-BRACE: Agent Belief state Representation via Abstraction and Confidence Estimation, a method that decouples an LLM agent into a belief state model and a policy model, jointly optimized via reinforcement learning. The belief state model produces a structured approximation of the belief distribution: a set of atomic natural language claims about the environment, each annotated with an ordinal verbalized certainty label ranging from certain to unknown. The policy model conditions on this compact, structured approximate belief rather than the full history, learning to select actions under explicit uncertainty. Across long-horizon, partially observable embodied language environments, Agent-BRACE achieves an average absolute improvement of +14.5% (Qwen2.5-3B-Instruct) and +5.3% (Qwen3-4B-Instruct), outperforming strong RL baselines while maintaining a near-constant context window independent of episode length. Further analysis shows that the learned belief becomes increasingly calibrated over the course of an episode as evidence accumulates.

1 Introduction

Agent-BRACE addresses uncertainty and unbounded history in long-horizon, partially observable LLM-agent tasks by separating belief-state modeling from action selection. It jointly trains both modules to use compact, uncertainty-aware textual beliefs instead of full interaction histories, improving performance while keeping context nearly constant.

  • Motivation: Raw-history conditioning grows linearly with episode length, making belief-state representation difficult for text-based LLM agents.Current agents generally represent actions and observations in text, complicating explicit belief encoding.
  • Method: Agent-BRACE decouples an LLM agent into a belief-state model and policy model jointly trained with reinforcement learning.The belief represents environment state using textual statements annotated with verbalized certainty estimates.
  • Training: The policy maximizes binary environment reward, while the belief model combines rewards for tracking, correctness, and other belief-quality objectives.PPO jointly trains the modules, with success receiving +1 and failure 0 for the environment reward.
  • Results: +14.5% (Qwen2.5-3B-Instruct) and +5.3% (Qwen3-4B-Instruct) absolute improvement over strong RL baselines demonstrate Agent-BRACE’s effectiveness.It maintains a near-constant context window and generalizes to Treasure and Cooking after training only on Quest.

2 Methodology: Agent-BRACE

Agent-BRACE jointly trains a belief state model and a policy model for partially observable, long-horizon tasks. The belief model maintains a structured, uncertainty-aware approximation of history that the policy uses for action selection.

  • Joint Architecture: Agent-BRACE jointly trains belief and policy models via PPO to address history growth and missing belief representations under partial observability.The belief state model provides a structured uncertainty-aware approximation of history for downstream action selection.
  • Problem Formulation: An optimal policy in a POMDP must condition on full history or an equivalent belief state because the current observation is insufficient.The belief state is defined as bt = P(st|Ht), where Ht contains the goal and prior observations and actions.
  • Belief State Model: The belief model updates a set of statements annotated with WEP-based uncertainty labels from the goal, prior belief, and new observation.Its update is represented as bt+1 = fϕ(G, bt, ot+1).
  • Policy Model: The policy selects actions from the compact updated belief rather than the full interaction history.The policy is represented as πθ(G, bt+1, ot+1).
  • Training Objectives: The method uses supervised finetuning before PPO to enforce a structured belief representation, while jointly optimizing environment and belief rewards.The policy maximizes binary environment reward Renv, and the belief model maximizes composite reward Rbelief.

3 Experimental Setup and Results

Agent-BRACE is evaluated with Qwen instruction-tuned models on three TextWorld tasks, including held-out environments and long-horizon inference. It achieves the strongest reported accuracy, generalizes beyond its training task, and maintains a near-constant context window while outperforming memory-based comparison methods.

  • Models: Evaluation uses Qwen2.5-3B-Instruct and Qwen3-4B-Instruct as base models, with Qwen3-30B-A3B-Instruct serving as the judge.Both belief-state and policy models are initialized from the same base model.
  • Datasets: Training and evaluation use three TextWorld tasks: Quests, Cooking, and Treasure.The tasks involve navigation, object manipulation, cooking, and quest completion through natural language.
  • Main results: 72.8% and 79.3% average accuracy are achieved on Qwen2.5-3B-Instruct and Qwen3-4B-Instruct, improving over Direct-Action (RL) by +14.5% and +5.3%, respectively.Agent-BRACE achieves the highest average accuracy across the reported baselines.
  • Generalization: 81.5% and 81.0% accuracy on Treasure show the strongest transfer, while Cooking still improves over Direct-Action by +7.0% on Qwen2.5-3B-Instruct.Agent-BRACE is trained exclusively on Quest; Cooking has a different structure requiring sequential sub-goal completion.
  • Long-horizon and cross-environment evaluation: Across 100 steps, Agent-BRACE maintains a near-constant context window and achieves a 78.5% solve rate, while outperforming MEM1 by 5% on ALFWorld.ReAct and Direct-Action (RL) grow linearly and eventually exceed the context limit; Agent-BRACE also generalizes to ALFWorld’s different manipulation actions and observation structure.

4 Ablation and Analysis

The ablations show that Agent-BRACE benefits from jointly training belief and policy models, explicit state-grounding rewards, and fine-grained uncertainty labels. Its belief quality also becomes progressively calibrated as evidence accumulates.

  • Calibration analysis: Brier score drops from 0.40 →0.28 while confirmed claims grow from 21% to 52% across agent steps on Quest for Qwen3-4B-Instruct.The concurrent changes confirm progressive calibration as evidence accumulates.
  • Ablation results: 33.4% and 34.0% average performance drops occur with Summary-Belief for Qwen2.5-3B-Instruct and Qwen3-4B-Instruct, respectively.Summarizing past observations without modeling abstraction and uncertainty substantially weakens performance.
  • Ablation results: 21.5% and 22.6% average accuracy drops occur when the belief model is frozen for Qwen2.5-3B-Instruct and Qwen3-4B-Instruct, respectively.The result demonstrates the importance of optimizing belief and policy models jointly rather than treating belief as fixed.
  • Reward design: 1.5% and 18.5% average accuracy degradations follow removal of state-grounding rewards on Qwen2.5-3B-Instruct and Qwen3-4B-Instruct, respectively.The removed rewards cover state correctness, state tracking, diversity, and format.
  • Uncertainty labels: 79.3% →65.3% average accuracy results from collapsing the 7-level WEP scale to binary certain/unknown on Qwen3-4B-Instru.This ablation indicates that label granularity affects belief quality and downstream performance.

5 Related Work

Related work spans classical belief-state estimation, memory-based context compression for LLM agents, and natural-language state representations. Existing LLM approaches either sacrifice interpretability or collapse uncertainty, motivating Agent-BRACE’s structured uncertainty representation.

  • Classical belief estimation: Classical POMDP solvers maintain exact beliefs in finite discrete spaces, while particle filters approximate beliefs in continuous spaces.These methods rely on known transition and observation models or use Monte Carlo approximation.
  • LLM agent memory: ReAct conditions actions on a growing interaction history, while memory-based methods address context burden through techniques such as selective summarization.MEM1 is identified as selectively compressing past interactions through summarization.
  • LLM belief representation: LLM belief-representation methods use either uninterpretable latent states or structured natural-language summaries that collapse beliefs to a single maximum-likelihood estimate.The latter discards uncertainty inherent in the belief distribution.

6 Conclusion · A Belief State Structure · B Belief Model SFT on Teacher Trajectories

Agent-BRACE jointly trains belief and policy models so actions condition on compact, structured natural-language beliefs rather than raw histories. The belief structure separates state tracking from planning, while SFT teaches representation format and reinforcement learning develops belief quality.

  • 6 Conclusion: Agent-BRACE jointly trains a belief state model and policy model via reinforcement learning.The belief model approximates the belief distribution with atomic natural-language claims annotated by ordered Words of Estimative Probability labels.
  • 6 Conclusion: The policy conditions on a compact structured belief instead of raw interaction history, addressing long-history and partial-observability challenges.This representation is intended to compactly encode relevant state information for decision making.
  • A Belief State Structure: In TextWorld, the belief tracks location, room connections, observed-object states, inventory contents, and progress toward specific sub-goals.These five dimensions define the critical environment information represented in the belief state.
  • A Belief State Structure: The belief state is prohibited from generating plans, intentions, or hypothetical future actions, preserving its role as an approximation of environment state.This restriction enforces separation between belief modeling and policy decision making.
  • B Belief Model SFT on Teacher Trajectories: The belief model first undergoes supervised fine-tuning on GPT 5.4 mini teacher trajectories generated with Prompt 1 for state tracking.Training optimizes the standard cross-entropy loss for each transition.
  • B Belief Model SFT on Teacher Trajectories: SFT teaches the belief model to use the Likert certainty scale and follow the structured representation.The SFT stage bootstraps formatting rather than task knowledge.
  • B Belief Model SFT on Teacher Trajectories: >0.96 format compliance remains high from the first training step, while state correctness, diversity, and task rewards rise during belief-model training.The reward trends support the conclusion that reinforcement learning, rather than SFT, develops belief quality.

C Belief State without Task-Specific Supervision · D Reward Design · Belief State Model Rewards (Rbelief

Agent-BRACE remains effective without task-specific belief supervision, though performance declines across tasks. Its jointly optimized rewards combine belief quality, structured uncertainty, and environment success to train belief and policy models.

  • C Belief State without Task-Specific Supervision: 69.3% average accuracy without belief supervision, down from 79.3% across all three tasks.The ablation removes supervised warm-start initialization and domain-specific belief prompt structure.
  • C Belief State without Task-Specific Supervision: 69.3% average accuracy still exceeds PABU (62.7%), ReAct RL (54.2%), MEM1 (45.0%), ReAct (47.8%), and Base Model (65.3%), but trails Direct-Action RL (74.0%).These comparisons use the full Qwen3-4B-Instruct baseline suite from Table 2.
  • D Reward Design: Agent-BRACE jointly optimizes distinct rewards for the belief state model and policy model, shaping belief representation around decision-making needs.The belief model receives Rbelief, while the policy model receives Renv.
  • Belief State Model Rewards (Rbelief: Rbelief combines five signals, and belief-state reward ablations test their importance.The composite design includes state tracking, state correctness, format, diversity, and discounted success rewards.
  • Belief State Model Rewards (Rbelief: State tracking rewards belief updates that incorporate new observations while avoiding missing information and stale contradictions.An LLM judge evaluates logical consistency of the transition bt−1 → bt.
  • Belief State Model Rewards (Rbelief: State correctness verifies each claim and uncertainty score against the environment, distinguishing fully correct, partially correct, and incorrect tuples.Claims are parsed as (subject, predicate, certainty) tuples before verification.
  • Belief State Model Rewards (Rbelief: Format reward gates the composite signal, assigning zero reward to structurally invalid outputs, while diversity rewards uncertainty labels spread across seven canonical levels.Diversity uses normalized Shannon entropy over mapped uncertainty-label distributions.
  • Belief State Model Rewards (Rbelief: The policy reward Renv is +1 for task success and 0 for task failure, providing the primary reinforcement for action selection.Belief states also receive a time-step-discounted success reward γt × 1[success].

E Dataset Details … F.1 Models Used

The paper evaluates Agent-BRACE on procedurally generated TextWorld benchmarks spanning Quest, Treasure, and Cooking, alongside ALFWorld, using Qwen instruction-tuned models and a separate judge model. The datasets vary task structure and controlled difficulty, with training and evaluation settings designed around distinct generation schemes.

  • E Dataset Details: TextWorld provides fully observable game graphs and natural-language descriptions for the Quest, Treasure, and Cooking benchmark suites.Each suite contains 1,000 training, 100 validation, and 200 test games generated with non-overlapping base seeds.
  • E.1 Quest: Quest requires ordered navigation, object manipulation, and collection of a designated target across multiple rooms.Quest generation varies rooms, objects, and quest length, with validation and test configurations harder than training.
  • E.2 Treasure: Treasure places the agent in a procedurally generated maze and requires locating a named object hidden in a random room.The generation level jointly controls world size, container nesting depth, and distractor-object count, with harder validation and test levels.
  • E.3 Cooking: Cooking requires finding a cookbook recipe, gathering ingredients, performing preparation steps, and preparing and eating the resulting meal.Generation parameters control ingredients, pickups, rooms, and mechanics including containers, cooking, cutting, and limited inventory.
  • E.4 ALFWorld: The ALFWorld training and testing datasets are directly taken from Wang and Ammanabrolu [2025].This provides an additional embodied-language evaluation source beyond the procedurally generated TextWorld suites.
  • F Implementation Details and Hyperparameters: The codebase builds on Wang and Ammanabrolu [2025], with all training run on a single node using four NVIDIA A100 GPUs.The implementation therefore uses a single-node, four-GPU training setup.
  • F.1 Models Used: Agent-BRACE uses Qwen2.5-3B-Instruct and Qwen3-4B-Instruct as its instruction-tuned base models.These are the models used to train and evaluate the agent.
  • F.1 Models Used: Qwen3-30B-A3B-Instruct-25075 serves as the judge for state-tracking evaluation and correctness reward.The judge model is separate from the instruction-tuned base models used by the agent.

F.2 Belief State Model: Supervised Pre-training (SFT) · F.3 Policy PPO Training · F.4 Joint Belief-State PPO Training

The training pipeline first supervised-fine-tunes the belief-state model, then trains the policy with PPO, and finally updates the belief model jointly with the policy during reinforcement learning. Hyperparameters for each stage are documented in separate tables.

  • F.2 Belief State Model: Supervised Pre-training (SFT): Before joint reinforcement learning, the belief-state language model is fine-tuned with supervised learning.This pre-training uses belief-state trajectories generated from teacher demonstrations.
  • F.2 Belief State Model: Supervised Pre-training (SFT): The supervised belief-state pre-training hyperparameters are summarized in Table 8.
  • F.3 Policy PPO Training: The policy is trained with Proximal Policy Optimization.
  • F.3 Policy PPO Training: The policy-training critic uses Generalized Advantage Estimation.
  • F.3 Policy PPO Training: The policy PPO hyperparameters are listed in Table 9.
  • F.4 Joint Belief-State PPO Training: During reinforcement learning, the belief-state model is updated jointly with the policy after every rollout batch.
  • F.4 Joint Belief-State PPO Training: The joint belief-model training hyperparameters are provided in Table 10.

G Belief States are underconfident but improve over training

The belief state model is initially underconfident, especially when emitting “unknown,” but its uncertainty calibration improves over training. High-confidence labels remain well calibrated, while underconfidence is characterized as safer than overconfidence in partially observable decision making.

  • Calibration over training: 68% of claims labeled unknown are true early in training, falling to 54% late as the model better reserves unknown for genuinely uncertain claims.Unknown corresponds to a nominal probability of approximately 0, so both rates indicate underconfidence, with the later rate showing improved selectivity.
  • Calibration over training: ≥91% of almost certain claims and ≥95% of confirmed claims remain well calibrated throughout training.These labels occupy the high-confidence end of the WEP scale and remain stable across early and late training stages.
  • Calibration over training: Underconfidence is identified as a safer failure mode than overconfidence for sequential decision making under partial observability.The analysis frames conservative confidence estimates as preferable when agents must act despite incomplete state information.

H Statistical Reliability of Main Results

The section assesses Agent-BRACE’s statistical reliability using mean accuracy and standard deviation across three independent runs, covering methods trained on Quest and evaluated in three TextWorld environments. It also reports Qwen3-4B-Instruct performance comparisons and calibration behavior of verbalized belief labels.

  • Statistical reliability: Statistical reliability is assessed with mean accuracy and standard deviation across three independent runs.The evaluation covers all methods trained on Quest and evaluated across three TextWorld environments.
  • Calibration: The belief model is systematically underconfident, assigning conservative labels to beliefs that are more often true than their nominal probabilities imply.Figure 5 compares verified claim frequencies with nominal WEP-label probabilities at early and late episode steps.
  • Performance comparison: Table 11 compares Agent-BRACE with baselines across Quest, Treasure, and Cooking using Qwen3-4B-Instruct.The table reports standard deviation across three independent runs.

I Belief State Model Training … — FORMAT —

Agent-BRACE trains a structured belief state that preserves observed facts, explicit uncertainty, and state updates while separating belief representation from action selection. Training evidence shows format structure is successfully bootstrapped, and qualitative examples demonstrate calibrated uncertainty, contradiction handling, and fine-grained inventory tracking.

  • I Belief State Model Training: Format compliance starts near 0.98 and remains consistently high throughout belief-state training, confirming that SFT bootstrapped the required structure.The observed pattern contradicts the expectation that SFT had already distilled task knowledge, because state correctness and total task reward did not instead begin high and plateau quickly.
  • J Agent-BRACE without state correctness and tracking reward: Removing both LLM-judged state correctness and tracking rewards degrades average accuracy by 4% (79.3% →75.3%), while still outperforming all Table 2 baselines on average.The ablation does not catastrophically degrade Agent-BRACE performance.
  • Prompt 1: Belief State Generation/Update: The belief-state update prompt conditions on the goal, previous belief state, and current observation, and requires output only inside <belief_state> tags.This defines the belief model’s recurrent update interface.
  • — STRICT RULES —: Belief states record observed or directly inferred facts rather than next actions, plans, intentions, or command recommendations.The restriction explicitly prevents the representation from prescribing behavior.
  • — UPDATE RULES —: Updates upgrade contradicted or confirmed claims and replace stale beliefs, using ordinal markers from confirmed and certain through probable, possible, unlikely, doubtful, and unknown.The certainty scale maps labels to evidence strength, while update rules require direct observations to override stale claims.
  • — MANDATORY COVERAGE —: Each belief state uses one bullet per distinct fact, exactly one certainty marker, and covers location, exits, goal-relevant objects, inventory, and sub-goal progress without JSON, percentages, or key-value pairs.The mandatory coverage and format constraints make the representation structured while remaining natural language.
  • Prompt 2: State-Tracking Reward: The state-tracking reward counts correctly captured new facts, missing or wrong facts, stale contradicted priors, and total claims before producing a score in [0.00, 1.00].The evaluator explicitly identifies new, missing, and stale information before computing the reward.
  • Prompt 3: Claim Extraction / Prompt 4: State Correctness: The belief-state evaluators extract claims as subject–predicate–certainty triples and verify them as correct, incorrect, partially correct, or unverifiable against ground truth.Partial correctness captures factual claims whose certainty labels are badly miscalibrated.
Loading 2605.11436v1…