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

1. Problem Statement

Current Large Language Model (LLM) agents often suffer from "shallow grounding." While they can perceive the current environment and react to immediate observations, they lack a deep, causal understanding of how their actions will reshape the environment in the long term.

• Limitations of Existing Methods: Most agents rely on immediate observations or fixed-horizon rollouts (planning a fixed number of steps ahead). Fixed horizons are suboptimal: they may miss long-term dependencies in complex tasks or incur unnecessary computational costs for trivial actions.
• The Gap: Agents need a mechanism to proactively "rehearse" future scenarios to detect potential conflicts, bottlenecks, or catastrophic failures before executing actions in the real environment. However, existing world model approaches often use rigid, single-step, or fixed-horizon strategies that do not adapt to task complexity.

2. Methodology: Imagine-then-Plan (ITP)

The authors propose Imagine-then-Plan (ITP), a unified framework that equips agents with an adaptive lookahead mechanism powered by a learned world model.

Core Concept: POIMDP

The paper redefines the agent's decision-making process from a standard Partially Observable Markov Decision Process (POMDP) to a Partially Observable and Imaginable Markov Decision Process (POIMDP).

• Dual-Stream Representation: The agent's policy is optimized over two streams:
  1. The Concrete Present: Observable state ($s_t$).
  2. The Foresighted Future: Imagined trajectories ($\hat{\tau}$) generated by interacting with a world model.
• Adaptive Lookahead: Instead of a fixed horizon $K$, the agent dynamically selects the imagination horizon $K_t$ based on the estimated task progress and the ultimate goal. This balances foresight reliability against computational cost.

Framework Variants

The authors instantiate ITP in two distinct variants:

A. ITPI (In-Imagination Learning): Training-Free

• Mechanism: Designed for inference-time enhancement of existing LLM agents. The policy and world model remain frozen.
• Process:
  1. Adaptive Horizon Selection: The agent analyzes the current state and instruction to select an optimal $K_t$.
  2. World-Model Imagination: The agent interacts with the world model for $K_t$ steps to generate a future trajectory.
  3. Reflective Policy Generation: The agent uses the imagined trajectory as implicit feedback to reflect on progress and conflicts, refining its decision before acting in the real environment.
• Goal: Plug-and-play improvement without additional training.

B. ITPR (Reinforced Training): Training-Based

• Mechanism: Jointly optimizes the agent's action policy and the lookahead horizon selection mechanism.
• Three-Stage Training Pipeline:
  1. Pseudo-Labeling: Uses the world model and an initial policy to generate "optimal" lookahead horizons ($\tilde{K}_t$) for expert trajectories by scoring how well different horizons predict expert actions.
  2. Warm-Up Training: Fine-tunes the agent policy to imitate expert actions conditioned on these pseudo-trajectories and trains a lightweight K-head predictor to estimate the required horizon $K_t$.
  3. Online Optimization: Uses Advantage Actor-Critic (A2C) reinforcement learning. The reward function includes the environment reward minus penalties for computational overhead (lookahead steps) and reasoning redundancy. This teaches the agent when and how long to imagine.

3. Key Contributions

1. POIMDP Formulation: Conceptualizes a new decision-making framework that integrates imagined futures and historical interactions, moving beyond standard POMDPs.
2. Adaptive Lookahead Mechanism: Introduces a dynamic method to scale the imagination horizon based on task progress, avoiding the pitfalls of rigid fixed-horizon planning.
3. Dual-Variant Framework: Provides both a training-free (ITPI) and a reinforcement-trained (ITPR) solution, demonstrating flexibility across different resource constraints.
4. Empirical Validation: Extensive experiments show that ITP significantly outperforms strong baselines (including CoT, ReAct, RAP, and other world-model methods) across diverse benchmarks.

4. Experimental Results

The authors evaluated ITP on four benchmarks: ALFWorld (embodied household tasks), ScienceWorld (scientific experiments), WebShop (web navigation), and StableToolBench (tool use).

• Performance Gains:
  • ITPI (Zero-Shot): Significantly outperforms prompting-based baselines (e.g., ReAct, RAP) without any additional training. For example, on ALFWorld with Qwen3-8B, ITPI achieved 82.86% success rate vs. 42.86% for RAP.
  • ITPR (Trained): Consistently achieves the highest success rates across all backbone models. On ALFWorld with Qwen3-8B, ITPR reached 97.14% success rate, surpassing the best training-based baseline (IWM) by a large margin.
• Ablation Studies:
  • Removing the online reinforcement training (w/o RT) caused a significant drop in performance (e.g., from 88.57% to 71.42% on ALFWorld), proving that learning when to imagine is crucial.
  • Adaptive vs. Fixed/Random: Adaptive lookahead achieved higher success rates with lower computational budgets compared to fixed-horizon or random lookahead strategies.
• World Model Robustness: The framework is model-agnostic. While the choice of world model backbone affects performance (especially in the training-free setting), ITPR effectively narrows the gap, distilling capabilities from various architectures (Qwen, Llama, DeepSeek).

5. Significance and Impact

• Bridging the Grounding Gap: ITP addresses the critical limitation of "shallow grounding" by enabling agents to simulate long-term consequences, thereby reducing catastrophic failures and erroneous actions.
• Efficiency in Planning: The adaptive mechanism ensures that computational resources are allocated only where necessary (high-stakes decisions), making complex long-horizon planning more feasible.
• Generalizability: The framework is applicable to various domains (embodied AI, web agents, tool use) and works with different LLM backbones.
• Future Direction: The paper suggests that integrating world models with adaptive control is a principled path toward more deliberative and robust autonomous agents, though challenges remain in scaling to multimodal environments and reducing inference latency.

In summary, Imagine-then-Plan represents a significant step forward in agent learning by shifting from reactive decision-making to proactive, adaptive deliberation using world models, resulting in superior performance and efficiency in complex, long-horizon tasks.