Imagine-then-Plan: Agent Learning from Adaptive Lookahead with World Models

AI summary

Overview

Research area: LLM-based autonomous agents, model-based planning, and reinforcement learning from world models.

Technical level: Advanced. The paper assumes familiarity with Markov decision processes, reinforcement learning (actor-critic, TD learning), supervised fine-tuning of LLMs, and the agent-planning literature (ReAct, RAP, world-model rollouts).

Scope: The paper introduces Imagine-then-Plan (ITP), a framework that lets an LLM agent "rehearse" future action trajectories inside a learned textual world model and decide, on a per-step basis, how far into the future it should look before acting.

What This Paper Is About

Most LLM-based agents make decisions from the current observation plus a history of past interactions, which leaves them with what the authors call "shallow grounding": they see the environment but cannot predict how today's action will reshape it, so mistakes surface only after they are irreversible. Existing world-model approaches partially address this, but they almost always imagine a single step ahead or a fixed number of steps, which is too shallow for long-horizon tasks and wasteful for trivial ones. ITP's goal is to make agents deliberative by letting them simulate multi-step futures and to make that foresight adaptive, deepening it only for pivotal decisions where the ultimate goal is still far off.

Key Contributions

1. A new decision-process formalism. The paper extends the standard Partially Observable MDP (POMDP) to a Partially Observable and Imaginable MDP (POIMDP), in which the agent's action policy conditions on both the observed present state and an internally imagined future trajectory. When the lookahead horizon is zero, POIMDP reduces exactly to a standard POMDP, so nothing about the environment is changed — only the agent's information state is augmented.

2. The ITP framework with adaptive lookahead. A learned LLM world model and the agent policy interact inside a "mental sandbox" for K steps, producing an imagined trajectory that is fused with the current observation to select the real action. K is chosen dynamically at each step by trading off the ultimate goal against estimated task progress, avoiding both the brittleness of fixed horizons and the cost of always rolling out deeply.

3. Two instantiations spanning the training spectrum. ITP_I is training-free: at inference the agent selects a horizon, imagines, then reflects on the imagined trajectory (checking goal progress, conflicts, bottlenecks) rather than blindly executing the first imagined action. ITP_R is reinforcement-trained: a lightweight K-head predictor on top of the LLM backbone learns how many steps to imagine, trained through a three-stage pipeline of pseudo-labeling horizons, warm-up supervised training, and online A2C optimization with explicit penalties for lookahead and step count.

4. Empirical validation and mechanism analysis. Across four benchmarks and three backbone families, both variants beat prompting and training baselines, and targeted analyses show why: adaptive horizons beat fixed and random ones on both success rate and token budget, world-model accuracy decays with horizon length, and the learned horizons scale with task complexity.

Main Findings

• Both variants beat strong baselines consistently. Without any training, ITP_I substantially improves zero-shot success over CoT, ReAct, and RAP. With training, ITP_R achieves the best overall success rate in every backbone group — for example 88.57% overall on ALFWorld with Qwen3-8B, and 63.91% on the ScienceWorld seen split with Llama-3.1-8B.

• Tool use shows the largest gains. On StableToolBench with Qwen3-8B, ITP_R reaches 68% solvable pass rate versus 44% for IWM and 42% for SFT (a +24 point improvement), and 54% solvable win rate versus 36% for both baselines.

• Online reinforcement learning is the core ingredient, not a detail. Removing reinforced training (w/o RT) drops ALFWorld from 88.57% to 71.42% and ScienceWorld from 59.70% to 46.00%. The authors read this as supervised learning providing only a basic capability, while RL teaches when it is worth imagining.

• Adaptive horizons dominate fixed horizons on both axes. Fixed-k lookahead shows a bell-shaped success curve — performance peaks at a moderate k and then declines while token cost keeps climbing. ITP's adaptive policy attains higher success at a substantially lower normalized budget, so no global horizon tuning is needed.

• Adaptivity, not variability, is what helps. Compared against a random-lookahead strategy (sampling K independently at each step), ITP yields higher success rates at lower and more stable cost, isolating the benefit as state-conditioned allocation.

• Deeper imagination is not automatically better. Using Fact-F1 to measure world-model reliability, prediction quality degrades monotonically with horizon due to compounding errors, while task success peaks around k=3. The design implication is that effective planning means anchoring lookahead inside the world model's reliable range rather than maximizing foresight.

• Learned horizons track task difficulty. Longer-horizon environments such as ScienceWorld induce larger average lookahead values than ALFWorld, matching the intuition that complex tasks need more foresight.

• The method is model-agnostic but world-model choice still matters. In the training-free setting, a larger world model (DeepSeek-V3.2) underperforms Qwen and Llama backbones on ALFWorld, likely from weaker domain alignment with the environment's transition dynamics. ITP_R narrows that gap considerably, showing the adaptive training can distill useful dynamics from varied architectures.

• Per-subtask results are more mixed than the headline numbers. On several ALFWorld sub-tasks (e.g., heat and look), ITP_R trails SFT or IWM for certain backbones, so the aggregate advantage comes from broad consistency rather than uniform dominance.

Methodology in Plain English

The pipeline proceeds in layers:

1. Warm up a policy. Fine-tune a base LLM on expert demonstrations so it can at least emit executable actions. This is the starting point for exploration.

2. Collect rollouts and learn a world model. Let the warm-started agent act in the environment, gather those trajectories, and train a separate LLM to predict the next state given the current state and an action. Training data combines expert demonstrations with agent rollouts so the world model is robust to actions the agent actually takes, including imperfect ones.

3. Imagine before acting (POIMDP). At each step, alternate the policy and the world model for K rounds: the policy proposes an imagined action, the world model predicts the resulting imagined state, and so on. The agent then conditions its real action on the current observation plus this imagined trajectory.

4. Choose K adaptively.

   • In ITP_I, the agent itself analyzes the instruction and current state to pick a horizon, imagines, and then reflects on the imagined outcome — checking for goal progress and looming conflicts — before committing to an action. Nothing is trained.
   • In ITP_R, a small linear K-head on the LLM backbone predicts a distribution over horizons. Since expert data has no labels for the "right" horizon, the authors create pseudo-labels: roll out the frozen world model with teacher-forced expert actions, score each candidate horizon by how well it makes the expert action likely under the initial policy, and subtract a penalty proportional to k. The agent is then warm-up trained to both imitate expert actions conditioned on the chosen horizon and predict that horizon. Finally, online A2C jointly optimizes policy and K-head using a reward that is the environment reward minus a lookahead penalty (−λ_K·K_t) and a step penalty (−λ_step), with entropy regularization to keep horizon exploration alive.

The evaluation covers ALFWorld (household embodied tasks), ScienceWorld (procedural science experiments), WebShop (long-horizon web shopping), and StableToolBench (multi-turn tool use), using Qwen2.5-7B, Qwen3-8B, and Llama-3.1-8B-Instruct as matched backbones for all methods.

Why This Matters

The work reframes how LLM agents should use world models: not as a fixed planning module but as an on-demand deliberative resource whose depth is a learned, state-dependent decision. That is a meaningful shift for research on agent grounding, because it treats the horizon itself as part of the policy rather than a hyperparameter.

Real-world applications:

• Robotic and embodied instruction following — household-style manipulation where an irreversible wrong action (dropping, overheating) is costly, and where foresight should be concentrated on high-stakes steps.
• Customer-facing web and e-commerce agents — multi-step navigation and product search where premature purchases or wrong filters waste user time; the paper evaluates exactly this via WebShop.
• Enterprise tool-calling and API orchestration — multi-turn agents that chain external services, where a simulated dry run can catch a conflict before a live call fires, as measured on StableToolBench.
• Scientific and industrial procedure execution — protocol-driven lab or process workflows where the agent must combine information gathering, tool use, and sequenced experimentation.

Industry relevance: The adaptive mechanism is a cost-control story as much as an accuracy story. Deploying deep rollouts everywhere is expensive; ITP's learned budget allocation and explicit inference-cost penalties (λ_K, λ_step) map directly onto the token-cost and latency constraints that production agent deployments face. The training-free variant is especially attractive commercially, since it improves a deployed agent without any retraining.

Future Directions

• Extending POIMDP beyond text. All evaluations are text-based; the authors flag that visual or sensorimotor observations introduce noise that may destabilize adaptive horizon selection, so multimodal and physical-world validation is open.
• Reducing inference overhead. World-model imagination is inherently more expensive than reactive policies. Speculative decoding and better exploration–exploitation balancing are suggested as routes toward real-time viability.
• Mitigating compounding prediction error. Since reliability decays with horizon and success peaks around k=3 on ALFWorld, better world models or error-aware horizon selection could unlock the benefits of deeper foresight than currently pays off.
• Understanding why adaptivity works. The paper shows state-conditioned horizons beat random ones, but a finer characterization of which state features drive large K — and whether those features transfer across environments — remains unexplored.

Target Audience

Researchers and graduate students working on LLM agents, model-based reinforcement learning, and planning with learned dynamics models will get the most from this paper, particularly those interested in horizon selection as a learnable policy component rather than a tuned constant. Practitioners building or fine-tuning multi-step tool-use and web agents will find the training-free ITP_I variant and the cost-penalized reward design directly applicable, while those on the reinforcement-learning side will value the pseudo-labeling scheme for horizons where no expert supervision exists. Readers without background in MDPs and policy optimization will need to work through the preliminaries carefully before the method sections become readable.

Authors’ abstract

Recent advances in world models have shown promise for modeling future dynamics of environmental states, enabling agents to reason and act without accessing real environments. Current methods mainly perform single-step or fixed-horizon rollouts, leaving their potential for complex task planning under-exploited. We propose Imagine-then-Plan (\texttt{ITP}), a unified framework for agent learning via lookahead imagination, where an agent's policy model interacts with the learned world model, yielding multi-step ``imagined'' trajectories. Since the imagination horizon may vary by tasks and stages, we introduce a novel adaptive lookahead mechanism by trading off the ultimate goal and task progress. The resulting imagined trajectories provide rich signals about future consequences, such as achieved progress and potential conflicts, which are fused with current observations, formulating a partially \textit{observable} and \textit{imaginable} Markov decision process to guide policy learning. We instantiate \texttt{ITP} with both training-free and reinforcement-trained variants. Extensive experiments across representative agent benchmarks demonstrate that \texttt{ITP} significantly outperforms competitive baselines. Further analyses validate that our adaptive lookahead largely enhances agents' reasoning capability, providing valuable insights into addressing broader, complex tasks. Our code and data will be publicly available at https://github.com/loyiv/ITP.