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

TL;DR

Imagine-then-Plan (ITP) introduces an adaptive lookahead mechanism that uses a learned world model to imagine multi-step futures and guide policy for LLM-based agents. By formulating a POIMDP, ITP dynamically balances foresight depth against efficiency, and it comes in two forms: a training-free variant (ITP$_{\text{I}}$) and a reinforcement-trained variant (ITP$_{\text{R}}$) that learns when and how long to imagine. Empirical results on ALFWorld and ScienceWorld show substantial improvements over strong prompts and existing world-model baselines, with ablations highlighting the necessity of online horizon optimization. The work advances deliberate, model-based planning for long-horizon tasks, while acknowledging limitations in multimodal transfer and real-time constraints, and points to future work on broader environments and efficiency improvements.

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.

Figures (10)

• Figure 1: Comparison between our ITP framework and conventional agent learning frameworks.
• Figure 2: Overview of the proposed Imageine-then-Plan (ITP) framework. It consists of two variants: (a) $\texttt{ITP}_{\text{I}}$, which is training-free and enables LLM agents to learn from the imagination at inference time. (b) $\texttt{ITP}_{\text{R}}$, which leverages imagined futures to optimize the action policy more effectively and more efficiently.
• Figure 3: Ablation results of $\texttt{ITP}_{\text{R}}$ on ALFWorld and ScienceWorld benchmarks.
• Figure 4: Comparison between our adaptive lookahead mechanism and baselines with fixed lookahead steps, with success rate (left) and computational cost (right).
• Figure 5: Comparison between our adaptive lookahead mechanism (both $\texttt{ITP}_{\text{I}}$ and $\texttt{ITP}_{\text{R}}$) and baselines with a random lookahead strategy (ReAct and SFT).