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

TL;DR

This paper introduces a belief-guided exploratory inference method for embodied agents that improves world state alignment without additional training, using a lightweight LLM surrogate for information gain estimation and achieving better performance with lower overhead.

Contribution

The paper presents a novel test-time adaptive inference approach that refines beliefs and guides actions in partial observability without gradient updates or extra training.

Findings

Outperforms inference-time scaling baselines in world alignment accuracy

Uses a lightweight LLM surrogate to estimate information gain effectively

Achieves better alignment with significantly lower integration overhead

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.