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

TL;DR

This paper introduces Imagine-then-Plan (ITP), a framework enabling agents to perform multi-step lookahead planning with adaptive horizons using learned world models, significantly improving reasoning and performance in complex tasks.

Contribution

The paper proposes a novel adaptive lookahead mechanism within a unified framework for agent learning, enhancing multi-step planning with dynamic horizon adjustment.

Findings

ITP outperforms baseline methods on various benchmarks.

Adaptive lookahead improves agents' reasoning capabilities.

The framework is effective for complex, real-world tasks.

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.