Align While Search: Belief-Guided Exploratory Inference for World-Grounded Embodied Agents

TL;DR

Align While Search (AWS) tackles exploratory decision making under partial observability by maintaining a structured, test-time belief over environment structure and object locations. It uses a Bayesian-inspired, amortized belief updater implemented via prompting an LLM and selects actions by maximizing expected information gain in belief space, all without gradient updates or additional training. Across ALFWorld, VirtualHome, and BabyAI, AWS achieves improved success–cost trade-offs and strong generalization to larger or multimodal settings, outperforming inference-time and training-time baselines while keeping overhead manageable. The approach demonstrates that posterior-guided exploration can robustly align with latent world states, paving the way for scalable, adaptable embodied agents.

Abstract

In this paper, we propose a test-time adaptive agent that performs exploratory inference through posterior-guided belief refinement without relying on gradient-based updates or additional training for LLM agent operating under partial observability. Our agent maintains an external structured belief over the environment state, iteratively updates it via action-conditioned observations, and selects actions by maximizing predicted information gain over the belief space. We estimate information gain using a lightweight LLM-based surrogate and assess world alignment through a novel reward that quantifies the consistency between posterior belief and ground-truth environment configuration. Experiments show that our method outperforms inference-time scaling baselines such as prompt-augmented or retrieval-enhanced LLMs, in aligning with latent world states with significantly lower integration overhead.

Figures (5)

• Figure 1: Success rate vs. token usage on ALFWorld subtasks. AWS achieves higher success with 2-5$\times$ fewer tokens compared to strong inference-time baselines, highlighting efficient belief-guided exploration under partial observability. Full results in Appendix 20.
• Figure 2: Exploration Failure of Base and SFT Models. (Left) Trajectory diversity measured by entropy and the distinct trajectory ratio. (Right) A breakdown of SFT agent failures reveals that most errors occur due to rigid replay of train-time search patterns.
• Figure 3: Object Counts per House Reveal Latent Diversity. (Left) PCA shows clustered object usage patterns across houses. (Right) Each cluster exhibits distinct object preferences, indicating latent user behavior differences.
• Figure 4: Framework of Align While Search. Given the current trajectory, AWS maps language to a numeric belief via update$\rightarrow$projection: observations update the global textual belief $B^{\mathcal{G}}$, which is then projected to a categorical action posterior $b^{\mathcal{S}}(a)$ (Eq. \ref{['eq:belief_update']}). Candidate actions are scored by expected IG (Eq. \ref{['eq:ig']}) under simulated observations $\hat{o}$, and the top-ranked action is returned to the base agent.
• Figure 5: Evaluating IG-Based Action Ranking and Belief Accuracy.(Left) Success rate under different IG-based selection strategies. (Right) Belief scores at visited locations, grouped by whether the target object was present (hit) or absent (miss).