PlanU: Large Language Model Reasoning through Planning under Uncertainty

TL;DR

PlanU tackles reasoning under uncertainty for LLM-based decision making by modeling environment-influenced returns as quantile distributions within Monte Carlo Tree Search. It introduces an Upper Confidence Bounds with Curiosity (UCC) score that blends distributional uncertainty with state novelty to drive planning, and uses text embeddings to mitigate LLN-related variability in state representations. Across diverse benchmarks (Stock, Blocksworld, Overcooked, VirtualHome, TravelPlanner, WebShop), PlanU consistently outperforms baselines and informative ablations demonstrate the necessity of both quantile distributions and UCC. The work underscores the practical potential of distributional planning for robust, multi-step decision making under both model and environment stochasticity, with implications for real-world planning tasks and future integration with external knowledge sources.

Abstract

Large Language Models (LLMs) are increasingly being explored across a range of reasoning tasks. However, LLMs sometimes struggle with reasoning tasks under uncertainty that are relatively easy for humans, such as planning actions in stochastic environments. The adoption of LLMs for reasoning is impeded by uncertainty challenges, such as LLM uncertainty and environmental uncertainty. LLM uncertainty arises from the stochastic sampling process inherent to LLMs. Most LLM-based Decision-Making (LDM) approaches address LLM uncertainty through multiple reasoning chains or search trees. However, these approaches overlook environmental uncertainty, which leads to poor performance in environments with stochastic state transitions. Some recent LDM approaches deal with uncertainty by forecasting the probability of unknown variables. However, they are not designed for multi-step reasoning tasks that require interaction with the environment. To address uncertainty in LLM decision-making, we introduce PlanU, an LLM-based planning method that captures uncertainty within Monte Carlo Tree Search (MCTS). PlanU models the return of each node in the MCTS as a quantile distribution, which uses a set of quantiles to represent the return distribution. To balance exploration and exploitation during tree search, PlanU introduces an Upper Confidence Bounds with Curiosity (UCC) score which estimates the uncertainty of MCTS nodes. Through extensive experiments, we demonstrate the effectiveness of PlanU in LLM-based reasoning tasks under uncertainty.

Figures (25)

• Figure 1: Impact of Environmental Uncertainty on LDM tasks for Block Stacking. The results are grouped according to the least number of steps to finish tasks. "Deterministic" indicates deterministic state transitions. In "Stochastic" environment, with 20% probability, an action will fail.
• Figure 2: Key stages of PlanU's uncertainty-aware tree search: (a) select and backpropagation phases; (b) UCC-based action selection by combining value distribution $Z(s, a)$ and state novelty, where $\psi[Z(s_t, a_t)]$ maps a quantile distribution to a scalar; (c) distributional updates during backpropagation based on received rewards.
• Figure 3: Stock Investment Tasks. The agent faces two actions: buying stock a, which guarantees a profit (reward) of $0.9$, or buying stock b with a 60% probability of earning a profit of $1$ and 40% probability of zero-profit. The right part shows the average profits for different methods.
• Figure 4: The Return on the Overcooked benchmark: (a) Tomato salad, (b) Tomato lettuce salad, and the VirtualHome benchmark: (c) Entertainment.
• Figure 5: The impact of (a) quantile distribution, (b) UCC score, and (c) LLM uncertainty in PlanU.