PABU: Progress-Aware Belief Update for Efficient LLM Agents

Technical Summary: PABU (Progress-Aware Belief Update)

PABU: Progress-Aware Belief Update for Efficient LLM Agents addresses a fundamental inefficiency in current Large Language Model (LLM) agents: the reliance on full action-observation histories to manage partial observability.

---

1. The Problem: The "Full-History" Bottleneck

In most agentic tasks (modeled as Partially Observable Markov Decision Processes or POMDPs), the agent cannot see the true state of the environment and must rely on its interaction history as a proxy. Current state-of-the-art (SoTA) methods typically condition actions on the entire history of actions and observations. This leads to three major issues:

• Information Noise: Histories often contain task-irrelevant or redundant information (e.g., failed attempts), which can lead to "representation collapse" and suboptimal decision-making.
• Computational Overhead: As the episode progresses, the context window grows, increasing inference latency and token costs.
• Inefficiency: Agents often repeat failed actions because the "signal" of progress is buried under the "noise" of the full trajectory.

2. Methodology: The PABU Framework

The authors propose a belief-state framework that replaces the monolithic history with a compact, structured representation. Instead of summarizing text (which can lose critical details), PABU explicitly models task progress and uses a selective retention policy.

A. Belief State Composition
The belief state $b_n$ is a compact, text-based representation consisting of:

1. Task Progress ($p_n$): A high-level, semantic estimate of how much of the task has been completed (e.g., "find the knife" $\rightarrow$ "take the knife"). This acts as a task-agnostic anchor.
2. Progress-Conditioned Action Memory ($A_{att}$): A record of actions attempted within the current progress stage to prevent redundant retries.
3. Available Actions ($A_{available}$): The set of permissible actions extracted from the current observation.
4. Selective Observation Memory ($O_{saved}$): A sparse set of past observations deemed critical for future steps, determined by a learned policy.
5. User Query ($q$): The immutable original instruction.

B. Training Objective (Progress-Aware Learning)
The model is trained end-to-end to perform three tasks in a single call: update retention $\rightarrow$ estimate progress $\rightarrow$ select action.
The authors introduce a unique training loss ($L_{PABU}$) that uses Trajectory Augmentation. They identify "critical actions" (those essential for success) and train the model to:

• Predict the next critical action.
• Predict "augmented actions" ($\tilde{a}_i$)—the next progress-consistent action—for non-critical steps. This teaches the model to "skip" redundant or non-advancing steps during training.

3. Key Contributions

• Formalization of Belief-State Abstraction: Moving away from raw history toward a structured, progress-oriented belief state.
• Learned Retention Policy: An end-to-end mechanism that decides which observations are worth keeping, avoiding the information loss common in generic summarization.
• Progress-Aware Training: A novel objective that partitions actions into progress-consistent and non-consistent groups to encourage efficient, goal-oriented behavior.

4. Results and Empirical Evidence

The framework was evaluated on the AgentGym benchmark across eight diverse environments.

• Superior Performance: PABU-8B achieved an 81.0% task completion rate, outperforming previous SoTA models by 23.9%.
• Increased Efficiency: The average number of interaction steps was reduced to 9.5, a 26.9% reduction compared to baselines.
• Robustness: Ablation studies showed that PABU generalizes across different model families (Llama, Qwen) and scales, proving that explicit progress modeling is a robust design principle.
• Component Importance: Ablations revealed that the necessity of specific components (action history vs. observation history) depends on the task structure (e.g., Wordle requires observation history, while ALFWorld requires action history).

5. Significance

PABU represents a shift in LLM agent design from "memory-based" (storing everything) to "belief-based" (storing what matters). By treating task progress as the backbone of state tracking, the authors demonstrate that agents can be both more accurate and significantly more efficient. This approach provides a scalable path for deploying LLM agents in long-horizon, real-world tasks where context window limits and computational costs are primary constraints.