PlanU: Large Language Model Decision Making through Planning under Uncertainty

AI summary

Overview

Research area: LLM-based decision making (LDM), Monte Carlo Tree Search, and distributional reinforcement learning, applied to planning in stochastic environments.

Technical level: Intermediate. The paper assumes familiarity with MCTS, Markov Decision Processes, and value distributions, but its core ideas are explained through a simple two-stock example.

Scope: The paper introduces PlanU, an LLM-driven MCTS planner that represents each action node's return as a quantile distribution and scores nodes with an Upper Confidence Bounds with Curiosity (UCC) term, and evaluates it on five decision-making benchmarks.

What This Paper Is About

LLMs are increasingly used as agents that choose actions, but they perform poorly in stochastic environments where the same action can lead to different outcomes. Existing LLM decision-making methods mostly attack LLM uncertainty (randomness in text generation) by sampling many reasoning chains or search trees, while ignoring environmental uncertainty (stochastic state transitions), and the few methods that do model uncertainty cannot handle multi-step tasks that require interaction with the environment. PlanU's goal is to handle both kinds of uncertainty inside a single tree search by modeling the return of MCTS action nodes as a distribution rather than a single average value.

Key Contributions

1. Quantile distributions inside MCTS. PlanU augments the standard MCTS tree with explicit action nodes and models the return of each state-action pair, Z(s,a), as a quantile distribution built from n_q Dirac functions rather than the expected value Q(s,a) used in vanilla MCTS. A skewed quantile distribution signals high uncertainty; a uniform one signals low uncertainty.

2. The Upper Confidence Bounds with Curiosity (UCC) score. PlanU selects actions by maximizing UCC(s_t, a_t) = ψ[Z(s_t, a_t)] + c_1 · r_i(s_t)/N(s_t, a_t), combining the uncertainty-aware value term with a novelty reward r_i(s_t) = ||f̂(e(s_t)) − f(e(s_t))||² computed from the difference between a trained predictor network and a fixed randomly initialized target network, inspired by prior work on random network novelty estimation.

3. A text-encoder solution to LLM uncertainty for state matching. Because an LLM may describe the same state differently ("the person is right to the table" versus "the table is left to the person"), PlanU maps text states through a text encoder e(·) and treats states as the same when their embeddings are similar.

4. A distributional back-propagation rule. Node distributions are initialized from the LLM's own generation probability π(s_t, a_t) = ∏ p(t_i|c) and updated with Quantile Regression using the Quantile Huber loss toward the target r + γZ(s_{t+1}, a).

Main Findings

• PlanU outperforms baselines across all benchmarks tested. On Blocksworld with Mistral-7B, PlanU reaches success rates of 1.000, 0.803 ± 0.04, 0.559 ± 0.04, and 0.217 ± 0.03 on the 2-, 4-, 6-, and 8-step tiers, versus RAP's 0.892 ± 0.03, 0.514 ± 0.06, 0.166 ± 0.04, and 0.000 ± 0.00.

• The advantage grows with task depth and model choice. With Llama3.1-8B, PlanU scores 1.000, 0.842 ± 0.04, 0.524 ± 0.04, and 0.238 ± 0.03, compared with RAP-D's 0.973 ± 0.02, 0.750 ± 0.04, 0.428 ± 0.05, and 0.070 ± 0.02; with DeepSeek-R1-Distill-Llama-8B, PlanU scores 1.000, 0.816 ± 0.04, 0.455 ± 0.05, and 0.196 ± 0.02.

• Simple uncertainty prompting does not work. In the stock investment task (stock A yields a fixed 0.9 profit; stock B yields 1 with 60% probability and 0 otherwise), all methods tested — CoT, CoT+U, DeLLMa, RAP, RAP+U, and Reflexion — performed poorly. PlanU learns the correct expected return E[Z(s_0, b)] = 0.6 and takes the optimal action yielding 0.9, which RAP fails to do because vanilla MCTS forces a deterministic transition by picking the most frequent sampled outcome.

• PlanU is the only method to complete the harder Overcooked task. In Overcooked (with a 20% failure rate added to the Chop action), PlanU performs best on both Tomato Salad and Tomato Lettuce Salad; on the more complex Tomato Lettuce Salad, PlanU is the only method that successfully completes the task.

• PlanU succeeds in VirtualHome where MCTS and RL baselines fail. With the failure rate for opening appliances raised to 50%, PlanU reaches the optimal path on the Entertainment task faster than others; RAP (MCTS-based) and QR-DQN (an RL method) fail to succeed.

• Large gains on TravelPlanner. Across 45 travel planning tasks with injected flight and train delays, PlanU achieves a Task Completion Rate of 0.378 ± 0.020 and Constraint Satisfaction Rate of 0.222 ± 0.015, versus LATS at 0.234 ± 0.039 / 0.089 ± 0.017, RAP at 0.222 ± 0.025 / 0.044 ± 0.012, and CoT at 0.156 ± 0.030 / 0.022 ± 0.010. The injected delays used an 80.76% on-time rate for flights and a 2% delay rate for trains.

• Strong WebShop results. On 10 shopping tasks with long-tail log-normal network latency (mean 2s) and action failure above 10s latency, PlanU attains an average reward of 0.73 ± 0.07 and success rate of 0.5, compared with LATS (0.57 ± 0.07, 0.3), CoT (0.46 ± 0.04, 0.1), and RAP (0.41 ± 0.02, 0.2).

• Both components are necessary. Ablations show that replacing the quantile distribution ("PlanU w/o dist") or the novelty reward term ("PlanU w/o ucc") causes failure to find the optimal path on Tomato Lettuce Salad, while on Entertainment the full PlanU learns the optimal path faster than either variant.

• Robustness to LLM uncertainty. Introducing Prompt Shuffling (shuffling sentences while preserving semantics) and Prompt Injection (adding task-irrelevant but environment-related information) only slightly affected PlanU's convergence speed.

Methodology in Plain English

PlanU keeps the familiar four phases of MCTS — selection, expansion, simulation, and back-propagation — but changes what a node stores and how a node is chosen.

During selection, the agent starts at the root state and repeatedly picks the child action node with the highest UCC score. It then executes that action in the real environment; if the resulting state is new, it is added under the action node, otherwise the search moves to the existing matching node.

During expansion, each available action under a leaf state becomes an action node. Its quantile distribution is initialized with n_q identical values, all equal to the probability the LLM assigned the action's token sequence. This lets the LLM's common-sense prior seed the search.

During simulation, trajectories are rolled out from the new node to a terminal state or a depth limit, retrieving the actual next state from the environment at each step because the same action may not always produce the same outcome.

During back-propagation, each action node on the path receives an environment reward and its quantile distribution is updated using Quantile Regression with the Quantile Huber loss, whose target is the reward plus a discounted version of the downstream distribution.

The UCC score replaces the usual UCT term. It combines a scalar summary of the quantile distribution — by default the expectation, though the paper notes alternatives such as adding the spread θ(τ_0.9) − θ(τ_0.1) — with a curiosity bonus divided by the node's visit count. The curiosity bonus comes from comparing the output of a predictor network against a fixed, randomly initialized target network on text-encoder embeddings of the state; the predictor is trained on a buffer of previously visited states across iterations.

Experiments were run on five benchmarks — Blocksworld, Overcooked, VirtualHome, TravelPlanner, and WebShop — against CoT, ToT, Reflexion, DeLLMa, RAP, RAP-D, RAP-E, LATS, and QR-DQN, with every experiment repeated 5 times using different seeds.

Why This Matters

Impact on research. The paper argues that most LLM decision-making research treats uncertainty as a text-generation problem and solves it purely by spending more computation on reasoning chains, leaving environmental stochasticity unaddressed. By importing distributional reinforcement learning ideas into LLM-driven MCTS, it offers a concrete alternative: represent the whole return distribution at each node rather than averaging randomness away, and select actions with an exploration term that measures state novelty. The ablation results give direct evidence that both the quantile distribution and the UCC term matter.

Real-world applications:

• Travel planning agents: itineraries can break when flights or trains are delayed, exactly the scenario simulated here with an 80.76% flight on-time rate and a 2% train delay rate.
• Robotic and embodied assistants: VirtualHome models a simulated household agent, where appliance actions fail a substantial fraction of the time.
• Web agents and shopping assistants: WebShop exercises clicking and searching against a million-item store with variable latency and action failure.
• Multi-agent coordination and scheduling: Overcooked tests whether agents can still deliver dishes when a key action (chopping) fails 20% of the time.

Industry relevance. Any deployed LLM agent that acts in the world — booking, browsing, controlling devices, or coordinating with other agents — faces stochastic outcomes that a purely deterministic planner will mis-handle. PlanU's demonstration that a simple prompt-level request to "consider uncertainty" performs poorly, while a structural change to the search does not, is directly useful guidance for teams building agent frameworks.

Future Directions

• Extending to tasks beyond the reported benchmarks. The evaluation covers five environments; the paper does not report how PlanU scales to longer horizons or to domains with continuous state and action spaces.

• Reducing the cost of the added machinery. PlanU maintains quantile distributions per node, a state buffer, and trains a predictor network during planning; the paper does not report wall-clock or token-cost comparisons against the baselines.

• Better handling of heavy-tailed uncertainty. The paper notes a skew-versus-uniform reading of the quantile distribution and offers alternative ψ operators (such as one using the θ(τ_0.9) − θ(τ_0.1) spread), but only the expectation operator is used by default, leaving the choice of ψ as an open design question.

• Deeper analysis of LLM uncertainty. Prompt shuffling and prompt injection only slightly slowed convergence in the reported tests; the paper states that limitations are discussed in Appendix F, which is not included in the content reviewed here. Robustness to other forms of LLM uncertainty, such as hallucinated world-model transitions, is not reported.

Target Audience

Researchers and graduate students working on LLM agents, planning, and decision making under uncertainty; reinforcement learning practitioners interested in distributional value representations; and engineers building LLM-based agents that must act reliably in environments with stochastic transitions or unreliable action outcomes. Readers unfamiliar with MCTS or quantile regression will need background reading, but the stock investment example and the ablation figures make the central argument accessible.

Authors’ abstract

Large Language Models (LLMs) are increasingly being explored across a range of decision-making tasks. However, LLMs sometimes struggle with decision-making tasks under uncertainty that are relatively easy for humans, such as planning actions in stochastic environments. The adoption of LLMs for decision-making 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 decision-making 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 decision-making tasks under uncertainty.