Source-linked AI summary
Reasoning with Language Model is Planning with World Model
Shibo Hao, Yi Gu, Haodi Ma, Joshua Jiahua Hong, Zhen Wang, Daisy Zhe Wang, Zhiting Hu
TL;DR
LLMs struggle with multi-step planning and reasoning because they lack an internal world model for predicting states and simulating outcomes. RAP repurposes an LLM as both world model and reasoning agent and uses MCTS for strategic exploration. Across planning, mathematical, and logical reasoning tasks, RAP improves over strong baselines, including CoT and least-to-most prompting with self-consistency.
Problem
LLMs struggle with action planning and complex mathematical, logical, and commonsense reasoning, while lacking an internal world model for predicting states and long-term outcomes.
Method
RAP repurposes the LLM as both a world model and reasoning agent, using Monte Carlo Tree Search to build and explore reasoning trees with task-specific rewards.
Results
RAP consistently improves over CoT, least-to-most prompting, and self-consistency variants across plan generation, mathematical reasoning, and logical inference.
Takeaways & Limitations
RAP provides a general framework for strategic LLM reasoning across planning, mathematical, and logical tasks.
Takeaways & Limitations
The study focuses mainly on frozen LLMs, whose reasoning and world-model abilities may be bounded by pre-training.
Abstract
from arXiv · showhide
Large language models (LLMs) have shown remarkable reasoning capabilities, especially when prompted to generate intermediate reasoning steps (e.g., Chain-of-Thought, CoT). However, LLMs can still struggle with problems that are easy for humans, such as generating action plans for executing tasks in a given environment, or performing complex math, logical, and commonsense reasoning. The deficiency stems from the key fact that LLMs lack an internal $\textit{world model}$ to predict the world $\textit{state}$ (e.g., environment status, intermediate variable values) and simulate long-term outcomes of actions. This prevents LLMs from performing deliberate planning akin to human brains, which involves exploring alternative reasoning paths, anticipating future states and rewards, and iteratively refining existing reasoning steps. To overcome the limitations, we propose a new LLM reasoning framework, $\underline{R}$easoning vi$\underline{a}$ $\underline{P}$lanning $\textbf{(RAP)}$. RAP repurposes the LLM as both a world model and a reasoning agent, and incorporates a principled planning algorithm (based on Monto Carlo Tree Search) for strategic exploration in the vast reasoning space. During reasoning, the LLM (as agent) incrementally builds a reasoning tree under the guidance of the LLM (as world model) and task-specific rewards, and obtains a high-reward reasoning path efficiently with a proper balance between exploration $\textit{vs.}$ exploitation. We apply RAP to a variety of challenging reasoning problems including plan generation, math reasoning, and logical inference. Empirical results on these tasks demonstrate the superiority of RAP over various strong baselines, including CoT and least-to-most prompting with self-consistency. RAP on LLAMA-33B surpasses CoT on GPT-4 with 33% relative improvement in a plan generation setting.
1 Introduction
LLMs struggle with multi-step reasoning partly because they lack an internal world model and mechanisms for evaluating and exploring alternative reasoning paths. RAP addresses these limitations by combining an LLM-based world model with MCTS-guided reasoning, improving performance across planning, math, and logical tasks.
- GPT-3 achieves only 1% success on Blocksworld action planning, compared with 78% for humans.
- Current LLM reasoning lacks an internal world model for simulating world states and a reward mechanism for guiding reasoning toward desired states.
- RAP repurposes the LLM as both a world model and reasoning agent, using MCTS to explore reasoning trees and balance exploration with exploitation.
- 64% average success rate is achieved by RAP on 2/4/6-step Blocksworld problems, while CoT nearly fails completely.
- RAP consistently improves over CoT, least-to-most prompting, and self-consistency variants on mathematical and logical reasoning tasks.
2 Related Work
Prior LLM reasoning methods generate or search reasoning chains, while classical and code-based approaches address planning through domain-specific procedures. RAP aims to unify world models, rewards, and advanced planning in a general reasoning framework.
- CoT generates intermediate reasoning chains, self-consistency selects among sampled chains, and least-to-most prompting decomposes questions into simpler subquestions.
- Search-guided reasoning methods use procedures such as depth-first or breadth-first search to find better reasoning paths.
- Prior methods do not formally combine a world model, reward, and state within one unified framework.
3 Reasoning via Planning (RAP)
RAP combines an LLM-based world model, task-specific rewards, and MCTS to explore reasoning trees and select high-reward reasoning traces. Its planning process iteratively expands, simulates, evaluates, and updates candidate reasoning paths.
- World model and rewards: The framework uses the same LLM as both reasoning agent and world model, while rewards assess the feasibility and desirability of reasoning steps.Rewards can incorporate the agent’s action log probability and confidence in the predicted state.
- World model: RAP represents reasoning as states and actions, with the world model predicting the next state after each action.State and action definitions vary by task, such as block configurations and block movements, intermediate mathematical values and subquestions, or focused facts and deduction rules.
- MCTS planning: MCTS builds a tree whose nodes are states and edges are actions, using Q values to estimate expected future rewards and guide exploration.UCT balances exploitation of high-value actions with exploration of less-visited nodes, and its exploration weight controls that trade-off.
- MCTS phases: Each MCTS iteration selects a node, expands it with sampled actions and predicted states, simulates future transitions, and back-propagates rewards to update Q values.Expansion can add multiple child nodes, simulation continues to a terminal state, and back-propagation aggregates future rewards along the selected path.
- Trace selection and aggregation: After a fixed computational budget, RAP selects a final reasoning trace from the constructed tree, with the highest-reward iteration path often performing best in practice.For tasks requiring only a final answer, RAP-Aggregation combines multiple traces and answers from different MCTS iterations; complete-trace tasks do not use it.
4 Experiments
Experiments apply RAP across plan generation, mathematical reasoning, and logical inference, using task-specific states, actions, rewards, and MCTS-based exploration. RAP improves reliability over CoT and related baselines, including under limited sampling or iteration budgets.
- RAP is evaluated across embodied plan generation, mathematical reasoning, and logical inference, with task-specific formulations of states, actions, and rewards.The experiments cover Blocksworld, GSM8K, and PrOntoQA.
- 4.1 Plan Generation: 64% average success rate: RAP substantially improves Blocksworld planning over CoT and finds 42% of 6-step plans within 20 iterations.CoT succeeds on only a few 2-step cases and fails on harder problems.
- 4.1 Plan Generation: 33% relative gain: LLaMA-33B with RAP outperforms GPT-4 with CoT on Blocksworld plan generation.
- 4.1 Plan Generation: RAP improves Blocksworld reasoning by maintaining world states, avoiding illegal actions, backtracking after failed choices, and exploring alternative paths.The reasoning trace includes predicted states alongside actions, supporting more grounded inference.
- 4.2 Math Reasoning: RAP-Aggregation consistently outperforms baselines across GSM8K sampling and iteration budgets, especially when only a few are available.Reward guidance helps identify reliable reasoning paths under limited computation.
- 4.3 Logical Reasoning: 94.2% correct answers and 78.8% proof accuracy: RAP surpasses CoT and self-consistency CoT on PrOntoQA.The reported gains are 14% proof accuracy over CoT and 4.4% prediction accuracy over self-consistency CoT.
5 Analysis
Additional experiments examine RAP with a stronger LLM on full Blocksworld and analyze how reward choices affect reasoning performance. RAP substantially outperforms CoT on longer problems, while reward usefulness depends on task type.
- Blocksworld analysis: Experiments evaluate RAP with Llama-2 70B on the full 602-case Blocksworld test set in Easy and Hard settings.The Easy setting uses prior knowledge of each case’s minimum action count and matching demonstrations.
- Blocksworld analysis: RAP outperforms CoT substantially in both Easy and Hard settings, especially when solving cases requires six or more steps.CoT’s success rate drops severely on longer cases, whereas RAP maintains a relatively high success rate.
- Reward analysis: The analysis supplements the main experiments with reward-choice studies for plan generation and mathematical reasoning.These studies investigate how different reward combinations affect LLM reasoning performance.
- Reward analysis: Combining multiple rewards generally improves performance, but individual rewards have task-dependent effects.Action likelihood is essential for plan generation but not very helpful for mathematical reasoning.
6 Conclusion
RAP equips LLMs with a world model and planning-based reasoning process that supports strategic exploration. Across challenging reasoning problems, it outperforms contemporary CoT-based approaches and exceeds GPT-4 with CoT in some settings.
- Conclusion: RAP repurposes one LLM as both a world model and reasoning agent.This lets the model simulate world states and anticipate action outcomes during reasoning.
- Conclusion: Monte Carlo Tree Search helps RAP balance exploration and exploitation while constructing reasoning paths.The framework is designed to produce high-reward reasoning traces through strategic planning.
- Conclusion: RAP outperforms several contemporary CoT-based reasoning approaches and GPT-4 in certain settings.The conclusion reports this pattern across a variety of challenging reasoning problems.
Limitations
The study focuses on frozen LLMs, whose reasoning and world-modeling abilities may be bounded by their pre-training. Future work could investigate fine-tuning and external tools.
- Scope boundary: The experiments mainly use frozen LLMs, whose abilities might be bounded by pre-training.The authors identify fine-tuning as a direction for improving reasoning and world-modeling capabilities.
- Future directions: Combining RAP with external tools is proposed for solving more complex real-world problems.This is presented as future work rather than an evaluated capability in the study.
- Future directions: The paper leaves open how fine-tuning could improve LLMs when they serve as world models.The authors specifically suggest exploring how to fine-tune LLMs for this role.
Ethics Statement
The paper focuses on plan generation, mathematical reasoning, and logical reasoning, and reports no significant ethical concerns for these applications. It warns that broader applications may create misuse risks requiring careful safety consideration.
- Scope: The paper primarily studies plan generation, mathematical reasoning, and logical reasoning.Its ethical discussion is scoped to these application areas.
- Ethical assessment: The authors report no significant ethical concerns for the applications examined.This statement applies to the paper’s primary application scope.
- Safety boundary: Broader future applications of LLM reasoning may pose misuse risks and require careful safety consideration before real-world deployment.The authors recommend considering all aspects of safety before applying relevant techniques to the real world.
B Experiment Settings
The experiments use randomized sampling with a temperature of 0.8 and truncate generation at 2048 tokens or a newline. All experiments run on four NVIDIA A5000 GPUs with 24GB memory each.
- Sampling uses temperature 0.8, with generation capped at 2048 tokens or a newline token.
- Experiments run on 4 × NVIDIA A5000 GPUs with 24GB memory.
C Prompt
The prompt section specifies how RAP and baselines represent block-world tasks, generate actions, predict successor states, and reason through mathematical subquestions. It also gives the MCTS procedure used to search over candidate reasoning steps.
- C.1 Plan Generation: RAP and CoT receive instantiated initial-state and goal prompts, after which the model produces a plan in a designated plan format.Examples represent states and goals as textual statements and list actions such as unstacking, picking up, putting down, and stacking blocks.
- C.1 Plan Generation: The block-world action rules require one block to be picked up or unstacked at a time and update whether the hand is empty or holding a block.A block must be clear for pickup or unstacking, and stacking requires holding the block and a clear destination.
- C.1 Plan Generation: The plan-generation search initializes states, actions, transitions, rewards, value estimates, visit counts, depth, roll-outs, and an exploration weight before iterative MCTS search.The procedure performs selection, expansion, simulation, and backpropagation of returns through the reasoning tree.
- C.1 Plan Generation: RAP’s world model predicts a new state from the current state and selected action using prompts conditioned on the last action.The prompt includes state-transition examples showing how pickup changes hand occupancy, block location, and clearance.
- C.1 Plan Generation: Block-world prompts provide an initial state, a goal, available actions, and action restrictions governing legal manipulation of blocks.The restrictions cover picking up, unstacking, putting down, and stacking blocks, including hand, clearance, and support conditions.
- C.2 Math Reasoning: For math reasoning, the same prompt supports action proposal and next-state prediction by sampling subquestions and answering them as successive reasoning states.After the question is instantiated, previous subquestions and subanswers are appended to guide future actions and state predictions.
D Related work: world model and planning
The paper connects RAP to planning-based world-model methods while distinguishing its use of general LLMs during inference from approaches that train models with external embodied experience.
- Planning algorithms such as AlphaZero and MuZero use tree-structured search to balance exploration and exploitation.
- Model-based reinforcement learning learns world models to support planning or policy learning, including through offline trajectories and latent imagination.World models have also been applied to long-horizon tasks and physical robot learning.
- RAP uses LLMs as world models and applies planning to search for reasoning paths across open-domain reasoning tasks.The framework is described as similar in spirit to model predictive control.
- Unlike work that trains LLMs with an external world model for embodied experience, RAP focuses on inference and remains compatible with any training method.
- RAP’s predicted states reformulate the remaining task as a new test case, creating a demonstration–test difficulty discrepancy because fewer actions remain.The paper identifies this discrepancy from preliminary experiments as an issue requiring mitigation.
F Reward Choice
Reward experiments compare action likelihood, task-specific or transition-confidence rewards, and self-evaluation. Their usefulness depends on the reasoning domain and the kind of errors the model can recognize.
- The paper notes that candidate-action quality requires multi-step reasoning, which reduces the accuracy and usefulness of self-evaluation rewards.This limitation is especially relevant when likelihood already provides a useful search signal.
- Combining action likelihood with task-specific reward significantly outperforms single-reward baselines for Blocksworld.Adding self-evaluation further improves performance slightly in that setting.
- Self-evaluation produces the best overall performance on the reported GSM8k reward comparison.Action likelihood improves one-iteration performance but reduces accuracy with more iterations in that comparison.
- Action likelihood is most indicative when reasoning steps are short and structured, but can be noisy when steps use freer natural-language sentences.Blocksworld actions follow recognizable patterns, whereas mathematical steps permit greater linguistic variation.
- Self-evaluation helps when errors are easier to detect after generation than to avoid during generation.The paper gives mathematical calculation or logic errors as an example of this condition.