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

arXiv:2601.08955 · cs.CL, cs.AI, cs.LG · Submitted 2026-01-13 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Imagine-then-Plan: Agent Learning from Adaptive Lookahead with World Models".

Jane: The paper was written by Youwei Liu, Jian Wang, Hanlin Wang, Beichen Guo and Wenjie Li from The Hong Kong Polytechnic University and Central South University.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title: Tom: So, we’ve established that this research moves AI beyond simple reaction and towards internal foresight. When you look at the title, "Imagine-then-Plan: Agent Learning from Adaptive Lookahead with World Models," several concepts jump out as highly significant.

Jane: It's a very dense title, but it actually maps out the entire cognitive process they are trying to build into an AI system. The core idea is that imagination—the ability to simulate—precedes and informs planning.

Lu: For me, the phrase "Imagine-then-Plan" immediately suggests a hierarchical or sequential process, which is exactly what we need for general intelligence. It implies that the agent has a dedicated phase just for simulating possibilities before committing to an action sequence.

Meng: And when you break down "Adaptive Lookahead," it tells us that this isn't a rigid, set-and-forget planning system. The system is designed to be flexible in how far into the future it looks, which is crucial for real-world deployment.

Lalam: Furthermore, coupling this entire mechanism with "World Models" suggests that the imagination isn't just random guessing; it’s constrained by a learned representation of reality—of physics, causality, and environment dynamics.

Tom: Jane, do you think the authors are making a claim that these components are necessary for each other to function effectively?

Jane: I think so. They aren't just listing features; they are describing an integrated architecture. If you took out the "World Model," the "Imagine-then-Plan" process would lack grounding and quickly descend into hallucination, wouldn't it?

Lu: Exactly. The world model provides the structural integrity for the imagination. It gives the agent a consistent set of rules to simulate against, ensuring that its potential futures are plausible rather than arbitrary.

Meng: And this concept of plausibility is key. It means that when the agent imagines failure, it's imagining a failure within the known laws of physics or the operational constraints of the task environment.

Lalam: This ability to constrain imagination based on reality is what elevates this from theoretical modeling to something potentially useful in complex, physical tasks, like robotics or advanced logistics.

Tom: It’s clear that understanding these individual components—imagination, adaptive lookahead, and world models—is the first step toward appreciating the full scope of this agent design. Next, we're going to dig into how they actually build this system using a new formal structure: the POIMDP.

The Core Mechanism: Tom: Building on that idea of foresight, let's zero in on the core mechanism described in "Imagine-then-Plan: Agent Learning from Adaptive Lookahead with World Models." Jane, how does this dual approach—imagining and planning—actually work together to get the agent to plan effectively?

Jane: The paper explains that they are formalizing a new conceptual space, which they call a Partially Observable and Imaginable MDP, or POIMDP. This is the technical cornerstone; it provides the mathematical framework for their entire planning process.

Lu: It’s really about creating an accountability loop, Tom. By using the POIMDP structure, they are forcing the agents to account for every simulated consequence of their potential actions within that formalized environment.

Meng: I find the process of fusing current observations with future imaginations to be incredibly practical from a machine learning standpoint. It means we are combining real-time data streams with predictive modeling to make optimal operational decisions immediately.

Lalam: This suggests that the optimization goal is not merely generating a sequence of good actions, but rather optimizing the *decision-making process itself*. That is a fundamentally holistic improvement for AI design.

Tom: So, if I understand correctly, the agent doesn't just look at what’s happening right now; it sees its current reality alongside a plausible map of its potential futures before it commits to making any final move in those two maps.

Jane: That's precisely right. We are essentially giving the agent both present awareness and predictive power simultaneously, allowing us to anticipate bottlenecks or points of failure long before they actually occur during task execution.

Lu: The world model acts as a mental sandbox here, Tom. It allows us to test our plans against its simulated dynamics without any risk, which is vastly more powerful than simply relying on iterative prompting or chain-of-thought reasoning alone.

Tom: Understanding this POIMDP framework really helps us appreciate how deep the architectural changes are required to achieve this level of planning capability.

Jane: It moves beyond standard Markov Decision Processes (MDPs) because it explicitly incorporates the possibility that the agent doesn't know everything about its current state—it’s partially observable.

Meng: This partial observability is critical because most real-world tasks are messy; you don't have perfect information. The POIMDP allows the planning to account for that inherent uncertainty in the environment.

Lalam: Considering this complexity, the framework really formalizes the difference between predicting *what will happen* and figuring out *what should happen* given all possibilities.

Tom: It's a sophisticated mechanism, but it grounds our understanding of how advanced agents must be designed to operate reliably in complex settings. This leads us to look at the tangible results: how much better is this system compared to older planning methods?

Performance and Improvements: Tom: The research shows significant performance gains, particularly with the adaptive lookahead strategy. Jane, can you elaborate on why this is so much better than a fixed planning approach?

Jane: The paper demonstrates that fixed-k strategies are inherently brittle. They either fail to capture critical long-term dependencies if the lookahead window is too small, or they become computationally intractable if we force them to look too far ahead. Our method adjusts its depth based on the task's intrinsic complexity.

Lu: I see this adaptability as the key to unlocking general intelligence because it suggests that true planning isn't a uniform process; it involves constantly running internal simulations just deep enough to prune impossible paths before they even happen in reality, saving massive amounts of compute power.

Meng: From an engineering standpoint, we aren't wasting massive amounts of compute resources on trivial actions that only require a quick decision. The adaptive nature ensures that resources are focused exactly where the difficulty lies, which is a huge win for scalability and real-world deployment efficiency.

Lalam: This enables AI to be truly deliberative, Tom. Instead of just following pre-programmed commands or simple statistical patterns, it actively calculates the consequences of its actions in a thoughtful, multi-step way that matches the task's demands.

Tom: The results are quite striking; even the training-free ITPI substantially improves zero-shot performance compared to prompting baselines like ReAct, which is genuinely impressive from a generalizability standpoint.

Jane: It’s not just about generating the next action token; it’s about modeling how much further we need to

Conclusion: Tom: So, we’ve really covered a lot of ground today, from the initial concept to how they’ve tested it against traditional methods. To wrap up, I think what we all agree on is that this truly changing the trajectory of how AI systems operate.

Jane: It definitely feels like that shift; instead of just reacting to immediate stimuli, these models are building a genuine internal map of what might happen next before they even take an action, which is huge for making them feel more reliable and intuitive overall.

Lu: I think the most exciting part is how this allows us to model the world dynamically. We're not just running a static script; we're allowing the agents to explore a vast space of possible outcomes and refine their plans based on that entire simulated landscape, which is incredibly powerful for future complex tasks.

Meng: From my perspective, I’m really impressed by the engineering efficiency here. The fact that this system is adaptive means it doesn's just as computationally expensive whether the task is simple or complex; it dedicates exactly the right amount of compute power to solve a problem, which will be crucial for real-world deployment on various hardware setups.

Lalam: I think we should all appreciate how much more deliberative this makes AI. It elevates the machine from being a mere response engine to something that actively contributes to optimizing human goals across diverse domains, truly embodying the potential of thoughtful intelligence in service of culture and science.

Tom: You’re right, Lalam, it’s a massive leap. I want to make sure we highlight the full scope of this achievement one last time—it’s a significant step forward for agent design with "Imagine-then-Plan: Agent Learning from Adaptive Lookahead with World Models."

Jane: It proves that giving an AI the ability to "rehearse" its actions is just as important as having the data, which is a huge lesson for us.

Lu: I hope this paves the way for even more sophisticated, self-corrective agents in research and allows them to see potential failure modes before they are irreversible.

Meng: My team is looking at how to implement this architecture in our next few projects; it' offers a clear path toward scalable planning.

Lalam: It represents a beautiful marriage of prediction and decision-making, giving us a powerful tool for the future.

Tom: Absolutely, and that brings us perfectly to where we need to be before we transition—to the next groundbreaking paper on our agenda.

The Hong Kong Polytechnic University · Central South University

cs.CL, cs.AI, cs.LG

Submitted: 2026-01-13

Updated: 2026-09-03

Code: https://github.com/loyiv/ITP

Importance score: 94/100

The gist: The paper introduces "Imagine-then-Plan," a novel framework designed to enhance agent decision-making by integrating world models for explicit lookahead planning.

Key concepts

Imagine-then-Plan
This core idea suggests that an AI agent first simulates potential outcomes (imagination) before committing to an action sequence (planning). It implies a dedicated, sequential phase for simulating possibilities necessary for general intelligence.
World Models
These models provide the structural integrity for the agent's imagination. They are learned representations of reality—including physics and causality—that constrain simulations, ensuring potential futures are plausible rather than arbitrary.
POIMDP
The Partially Observable and Imaginable MDP is the technical framework used in the paper. It provides a mathematical structure that forces agents to account for every simulated consequence of their actions while acknowledging inherent environmental uncertainty.
Adaptive Lookahead
This strategy allows the planning system to be flexible, adjusting how far into the future it simulates. This prevents computational intractability while ensuring enough depth is analyzed to prune impossible paths effectively.

Terminology

Summary

The paper introduces Imagine-then-Plan, a novel framework designed to enhance agent decision-making by integrating world models for explicit lookahead planning. This methodology addresses limitations in standard sequential interaction by allowing an agent to simulate potential futures before committing to an action. By making the planning process adaptive—determining how far into the future to look and what that future entails—the system aims to achieve more robust and goal-directed behavior across complex environments like household tasks.

Base Agent Policy for Interaction

The foundational interaction mechanism utilizes a ReAct paradigm, requiring the agent to structure its output through explicit reasoning steps. The agent receives the task goal, current state (observation and optional inventory), previous action, and latest environment feedback. The policy must follow a strict sequence:

  1. Reason: Interpreting the current state and feedback to identify progress made versus goals remaining.

  2. Thought: Formulating a plan for subsequent steps based on the reasoning.

  3. Action: Outputting EXACTLY ONE valid action line.

Adaptive Selection of Lookahead Horizon K

Before generating a full plan, the system employs an adaptive selector to determine the necessary depth of foresight, represented by K. This component functions as a planning assistant that maps contextual information—specifically the task instruction and the accumulated dialogue/action history—to a single integer K within the permissible range [0, K]. This mechanism ensures that computational resources are allocated efficiently by only simulating as many steps ahead as is necessary for optimal decision-making.

World Model Foresight Generation

The core imaginative capability resides in the world model. This model is tasked with predicting potential future states based on the current state (st) and a given lookahead depth K. The system prompts the world model to Predict the next k step(s), which results in a concise, structured plan enclosed by .... This output provides the agent with an imagined trajectory that conditions its subsequent planning, allowing it to anticipate outcomes rather than merely reacting to immediate feedback.

Foresight-Conditioned Action Synthesis

The final decision-making step integrates all prior components: the task goal, the current state, and the generated K-step foresight trajectory. The agent’s role shifts to that of a reflective planner. It must first perform a Reflection on whether the foresight indicates progress or contradictions. Subsequently, it generates a Thought detailing its plan based on this imagined future. Critically, when outputting the final action, the agent is constrained by a hard rule: it MUST choose the Action by copying EXACTLY one line from the provided admissible actions list. This structured approach ensures that planning is grounded in both simulation and immediate environmental constraints.

Improvements for AI systems

Based on a rigorous review of this framework, while highly advanced, several critical areas exist for improvement to enhance robustness, generalization, and efficiency. The following improvements are designed to elevate the system from a sophisticated planning agent into a truly adaptive and reliable decision-making architecture.

The Flaw: Figure 10's World Model (WorldModel.imagine) currently outputs a deterministic, concise trajectory (Foresight). In real-world environments (like ALFWorld), predicting the exact future state is inherently uncertain. Relying on a single predicted path makes the agent brittle when encountering novel or stochastic situations.

The Improvement: Modify the World Model to output not just a single foresight, but a distribution of plausible futures. This requires training the World Model to predict parameters for a multimodal distribution (e.g., mean mu and covariance) over key state variables or object locations, rather than just the next state vector.

What the Improved System Can Do:

  • Risk-Aware Planning: The agent can calculate an Expected Utility Function (EUF) that factors in predicted risk. Instead of choosing the action leading to the most likely success, it chooses the action that maximizes success while minimizing variance (i.e., avoiding high-variance outcomes).

  • Early Failure Detection: If multiple plausible futures diverge significantly from the current state's prediction, the system flags high uncertainty and automatically triggers a fallback mechanism (e.g., requesting human intervention or simplifying the task goal).

The Flaw: The agent's reflection step (Figure 11) currently assesses progress based on comparing the current state to the desired goal/foresight. It does not assess how reliable the foresight itself is. If the World Model was highly uncertain, but the agent proceeds as if it were certain, catastrophic failure can occur.

The Improvement: Implement a Confidence Score Module that takes two inputs: (1) The World Model's predicted uncertainty and (2) the historical deviation of real-world observations from the predicted mean (mu). This module generates a single, explicit Confidence Score C WM.

What the Improved System Can Do:

  • Adaptive Planning Depth: If C WM drops below a critical threshold, the system dynamically overrides the adaptive horizon selector (K) and forces a reduction in planning depth (e.g., setting K=1 or K=2), forcing local, safer decisions until stability is regained.

  • Self-Correction Protocol: The agent can explicitly state: My plan is based on a low-confidence prediction regarding object X's location; I must first verify this assumption by performing action Y.

The Flaw: The entire framework relies heavily on complex, multi-step prompt engineering (Figures 8, 9, and 11). While effective, this is fragile. Minor changes in the prompt structure or format expectations can cause catastrophic failure in deployment.

The Improvement: Refactor the system to use structured output APIs (e.g., JSON schema generation) for all intermediate steps: K selection, Foresight Generation, and Action Reflection. The LLM's role should be relegated to reasoning and interpreting the state/feedback, while the underlying planning logic (like calculating the EUF or selecting K) is handled by deterministic code modules receiving structured inputs.

What the Improved System Can Do:

  • Guaranteed Interoperability: The system becomes far more robust to prompt modifications and language shifts. The output of one module can be reliably consumed as a programmatic input for the next, ensuring that the Action taken is always derived from a machine-readable structure, not just free text.

  • Accelerated Development Cycle: New behaviors or environments require updating deterministic code logic rather than complex prompt re-engineering.

The Flaw: The current training pipeline is batch-based and separated: Pseudo-labeling to Warm-up to Online Optimization. The World Model itself is trained separately and its knowledge is static during the online optimization stage.

The Improvement: Establish a Closed Causal Learning Loop. Every time the agent executes an action A t in the real environment, and observes O t+1, this tuple (S t, A t, O t+1) must be immediately used to calculate a prediction error for the World Model. This error signal is then used to perform low-rank adapter updates (LoRA) on the World Model's weights in real-time during the online optimization phase.

What the Improved System Can Do:

  • Rapid Domain Adaptation: The agent does not need massive, pre-collected datasets for every minor environmental change or new object. It learns and adapts its world model knowledge incrementally from its first few interactions in a novel setting, drastically reducing deployment time and cost.

Sources

Related papers