PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Youwei Liu, Jian Wang, Hanlin Wang, Beichen Guo, Wenjie Li

arXiv:2601.08955v2cs.CLcs.AIcs.LG

TL;DR

LLM agents lack deep causal foresight, while single-step or fixed-horizon world-model rollouts can miss long-term dependencies or incur unnecessary computation. ITP combines learned-world-model imagination with adaptive lookahead, and its inference-time and reinforcement-trained variants significantly improve task success and robustness across domains. The evaluation mainly covers interactive text-based benchmarks, and world-model use adds inference-time overhead.

  • Problem

    LLM agents often have shallow grounding, while single-step or fixed-horizon world-model rollouts inadequately address long-term dependencies and computational efficiency in complex tasks.

  • Method

    ITP extends decision making with multi-step imagined trajectories from a learned world model and adaptively selects the imagination horizon using task goals and progress.

  • Results

    ITP's inference-time and reinforcement-trained variants significantly improve task success and robustness across domains, with analyses validating adaptive control of the imagination horizon.

  • Takeaways & Limitations

    ITP provides a framework for reasoning over both present observations and foresighted trajectories in long-horizon agent decision making.

  • Takeaways & Limitations

    Evaluation primarily targets interactive text-based benchmarks, while world models introduce higher inference-time overhead than purely reactive agents.

Abstract

from arXiv · show

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.

1 Introduction

LLM agents often rely on shallow grounding and fixed or single-step world-model rollouts, limiting long-horizon planning. ITP addresses this with adaptive multi-step imagination and demonstrates improved task success through training-free and reinforcement-trained variants.

  • LLM agents can discover erroneous actions or state conflicts only after irreversible execution because they lack deep causal understanding of future consequences.
  • World-model methods commonly use single-step verification or fixed-horizon rollouts, which can miss long-term dependencies and waste computation on trivial actions.
  • ITP conditions decisions on present observations and imagined multi-step trajectories through an adaptive lookahead that trades off the ultimate goal against estimated task progress.
  • ITP extends the POMDP formulation to a Partially Observable and Imaginable MDP, combining concrete present information with foresighted future trajectories.
  • ITP includes ITPI, a training-free inference-time variant, and ITPR, a reinforcement-trained variant that uses imagined futures to optimize policies more effectively and efficiently.
  • Both ITP variants significantly improve task success rates, while further analyses validate the important role of adaptive lookahead.

2 Preliminaries

The paper frames LLM-agent reasoning as a partially observable decision process and interprets LLMs as world models that predict future textual states. ITP builds on this setting by adding imagined futures to agent learning.

  • A POMDP models agent-environment interaction through states, actions, observations, transitions, rewards, and interaction history.
  • In text-based environments, textual observations are treated as state representations, allowing LLMs to model transition regularities by predicting next states from interaction context.
  • ITP equips LLM-based agents with adaptive lookahead through learned world models and conditions decisions on observable present states and imagined future trajectories.

3 Method

ITP equips LLM agents with adaptive world-model lookahead, combining observable states with imagined trajectories for deliberative planning. It provides training-free and reinforcement-trained variants that adapt imagination depth to task progress and costs.

  • Core framework: ITP conditions agent decisions on the observable present and imagined future trajectories generated through adaptive lookahead.The framework extends decision-making with a learned world model and an imagination operator.
  • World model training: The world model is trained on expert demonstrations and agent rollouts to approximate environment dynamics and improve robustness to out-of-distribution actions.An initial policy is first obtained from expert demonstrations, after which expert and rollout data are combined for world-model training.
  • Adaptive lookahead: ITP adaptively selects the imagination horizon by trading off task progress against the risks of missed dependencies, model error, and unnecessary computation.Short horizons can miss long-term dependencies, whereas excessive rollouts increase computation and may amplify errors.
  • ITPI: ITPI uses frozen policy and world models to choose a horizon, imagine future trajectories, and reflect on progress, conflicts, bottlenecks, or failures before acting.The imagined trajectory serves as implicit feedback for self-refinement rather than being executed directly.
  • ITPR: ITPR jointly learns action selection and lookahead prediction through pseudo-labeling, warm-up training, and online reinforcement optimization.A K-head predicts the imagination horizon, while reinforcement learning balances task reward with computational and interaction overhead.
  • ITPR: ITPR’s reinforcement objective penalizes lookahead and redundant reasoning while jointly optimizing policy, value estimation, and horizon exploration.The online stage samples imagined trajectories from a frozen world model and uses A2C to update the policy and K-head parameters.

4 Experiments

ITP is evaluated across multiple agent benchmarks, backbones, ablations, and lookahead strategies. The results show consistent performance gains, with adaptive lookahead improving both success and computational efficiency.

  • Experimental Settings: ITP is evaluated on ALFWorld, ScienceWorld, WebShop, and StableToolBench using Qwen2.5-7B, Qwen3-8B, and Llama-3.1-8B backbones.Success rate is used across benchmarks, with SoPR and SoWR additionally reported for StableToolBench.
  • Main Results: ITPR combines policy learning with selective adaptive lookahead, with a strongest reported result of 63.91% on ScienceWorld test-seen using Llama3.1-8B.ITPI uses the same backbone as prompting baselines and adds world-model rollouts at inference time, isolating explicit lookahead as the additional component.
  • Ablation Study: Removing online reinforced training reduces ALFWorld success from 88.57% to 71.42% and ScienceWorld success from 59.70% to 46.00%.The ablation indicates that online reinforcement optimization is used to learn when to invoke the world model and how to use it efficiently.
  • Lookahead Strategies: Adaptive lookahead achieves higher success with substantially lower computational cost than fixed lookahead and higher, more stable success-budget trade-offs than random lookahead.Fixed-lookahead cost rises sharply with horizon, whereas adaptive allocation avoids global horizon tuning and unnecessary deep rollouts.
  • World-Model Backbones: World-model backbone choice matters most for ITPI, while ITPR narrows performance gaps and makes the DeepSeek-V3.2 world model highly competitive after optimization.Qwen and Llama maintain robust zero-shot success rates, suggesting stronger compatibility with ALFWorld state-transition dynamics.
  • World-Model Reliability: Fact-F1 decreases with longer imagination horizons, success peaks at k=3, and longer-horizon environments induce larger learned lookahead values.These findings characterize a trade-off between foresight and compounding prediction errors while linking horizon selection to task complexity.

5 Related Work

Prior work improves LLM agents through trajectory or step-level optimization and reflective updates, while world-model approaches support planning through predicted environment dynamics. However, these methods often leave agents without active foresight into future environmental changes.

  • Trajectory- and step-level methods optimize exploration, revision, correction, or reflective updates from agent behavior.Examples include ETO, IPR, E2CL, Agent-R, STeCa, and AgentRefine.
  • Historical-trace methods primarily leave agents with shallow grounding and limited ability to anticipate future environmental shifts before execution.This limitation motivates active foresight beyond reactive decision-making.
  • World models provide a mental sandbox for model-based decision-making by predicting environment dynamics.Recent work uses LLMs as implicit or explicit world models for task planning.

6 Conclusion

ITP equips LLM-based agents with adaptive world-model lookahead and represents decision-making through imagined future trajectories alongside present observations. Experiments support improved task success and robustness, while the current evaluation and inference costs limit conclusions beyond text-based benchmarks and real-time deployment.

  • Conclusion: ITP extends the POMDP into a Partially Observable and Imaginable MDP for reasoning over present observations and foresighted trajectories.The framework includes inference-time and reinforcement-trained variants.
  • Conclusion: ITP significantly improves task success and robustness across domains, according to extensive experiments.The supplied conclusion does not provide a specific numerical result.
  • Conclusion: Adaptive control of the imagination horizon supports more deliberative use of world models for autonomous agent learning.The conclusion identifies adaptive horizon control as central to ITP.
  • Limitations: The evaluation primarily covers interactive text-based benchmarks rather than multimodal environments or real-world robotic control.Transferring from linguistic descriptions to visual or sensorimotor observations may introduce noise affecting adaptive-lookahead stability.
  • Limitations: World-model lookahead introduces higher inference-time overhead than purely reactive agents, requiring further optimization for real-time applications.The paper mentions speculative decoding and distilled world models as possible future directions.

A.1 Datasets

The evaluation uses four text-based agent benchmarks spanning household tasks, science reasoning, web shopping, and tool use. Their descriptions emphasize different combinations of long-horizon planning, noisy observations, experimentation, and execution robustness, with SoPR and SoWR used for solvable instances.

  • Benchmark suite: The evaluation covers ALFWorld, ScienceWorld, WebShop, and StableToolBench.These are described as representative agent benchmarks.
  • ALFWorld: ALFWorld evaluates multi-step household planning under partial observability using natural-language observations and admissible text actions.Tasks include finding, manipulating, and placing objects through compositional action sequences.
  • ScienceWorld: ScienceWorld tests procedural and reasoning-intensive tasks involving information gathering, tool use, and multi-step experimentation.Its task horizons are typically longer than those of household tasks.
  • WebShop: WebShop requires multi-step browsing, comparison, and product selection over long, noisy textual observations.The observations include search results, product titles, attributes, and descriptions.
  • StableToolBench: StableToolBench evaluates correctness and robustness in executable tool calling, feedback handling, and solvability-aware interaction.Episodes are typically shorter, but incorrect tool invocations are more sensitive.
  • Metrics: SoPR measures average success on solvable instances, while SoWR measures head-to-head wins against a fixed baseline set.The baselines are ReAct, RAP, SFT, and IWM.

A.2 Data Preprocessing

The preprocessing pipeline converts expert trajectories into transition-level supervision for text world models and a warm-start policy, using benchmark-specific state serialization. For reinforcement-trained ITP, imagined rollouts across candidate horizons produce adaptive-horizon pseudo-labels balancing action likelihood gains against rollout cost.

  • Supervision construction: Expert trajectories warm-start the agent policy and supervise the text world model at the transition level.Each record contains a state string, executed expert action, and next observation, with optional reward and termination signals.
  • Dataset protocol: Dataset statistics are reported using splits that follow the standard benchmark protocol.
  • State serialization: State serialization incorporates benchmark-specific context such as goals, observations, webpage content, tool descriptions, calls, and execution feedback.The unified format supports embodied, web, and tool-use settings.
  • Adaptive-horizon annotation: ITPR precomputes imagined rollouts over candidate horizons and summarizes them into lookahead text snippets.The process runs at each expert step using the trained world model.
  • Adaptive-horizon annotation: The adaptive-horizon pseudo-label balances improved expert-action likelihood against the cost of deeper rollouts.The selected label and all per-horizon scores are stored for training.

B.1 Tree Search and RL-from-World-Model Methods

ITP is compared with tree-search and RL-from-world-model alternatives on ALFWorld. The results indicate that tree search alone is insufficient, while RL-based methods help but remain below ITP.

  • Tree-search-based methods: 41.43 SR / 0.25 NB gives ITPI the best trade-off, exceeding RAP at 28.57 / 0.42.Both ToT and MCTS underperform ITPI and are less budget-efficient.
  • RL-from-world-model methods: 70.71 → 76.28/78.89 shows that D2PO and VaGen improve over Base SFT.These RL-from-world-model baselines still do not reach ITPR’s performance.
  • RL-from-world-model methods: 88.57 SR with only a modest NB increase to 0.21 places ITPR beyond the compared training-based methods.The comparison uses adapted text-only settings with the same state/action serialization and evaluation protocol.

B.2 Compute-Matched Comparisons with No-World-Model Alternatives

Compute-matched no-world-model alternatives spend similar budgets on repeated direct generation or refinement but still lag behind ITP. The comparison attributes the gains to adaptive lookahead over imagined future states rather than extra compute alone.

  • Prompting-based comparisons: 5.01 points is ITPI’s improvement over the strongest prompting-based no-world-model baseline under matched compute.The alternatives include Self-Refine, Multi-Attempt, and K-candidate Selection.
  • Training-based comparisons: 10.43 points is ITPR’s improvement over the strongest training-based no-world-model baseline under matched compute.The comparison uses repeated direct action generation or refinement without future-state imagination.
  • Interpretation: The reported gains are attributed to allocating computation to state-conditioned adaptive lookahead over imagined future states.Normalized Budget measures algorithmic token efficiency, while deployment profiling separately reports runtime characteristics.

C.4 Training Cost of ITPR

ITPR training is organized into pseudo-label generation, supervised warm-up, and online actor–critic optimization. The appendix also specifies benchmark-dependent horizons and deployment-oriented profiling.

  • Training pipeline: Three stages comprise ITPR training: adaptive-horizon pseudo-label generation, supervised warm-up, and online actor–critic optimization.The warm-up trains the policy and K-head with lookahead labels.
  • Resource cost: 11.8 episodes/hour is reported alongside peak VRAM of 31.0 / 38.4 GB for the policy GPU and 28.7 / 35.9 GB for the world-model GPU.These figures summarize the reported training or deployment resource profile.
  • Hyperparameters: Kmax is benchmark-dependent, set to 5 for ALFWorld, 8 for ScienceWorld, and 7 for WebShop.The setup reflects differences in trajectory lengths and planning horizons across benchmarks.
Loading 2601.08955v2…