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

Here are the core ideas, introduced step by step using the Sandwich pattern so they’re easy to digest.

1. Reinforcement Learning (RL) 🍞 Top Bread (Hook): Imagine training a puppy. Each time it does the right trick, it gets a treat, so it repeats the good behavior. 🥬 Filling (The Concept): RL is a way for an AI to learn by trying actions and getting rewards or penalties.

• How it works:
  1. The agent sees a situation.
  2. It picks an action.
  3. The world reacts and gives a reward.
  4. The agent updates its strategy to earn more reward next time.
• Why it matters: Without RL, agents can’t improve from experience; they’d keep repeating the same mistakes. 🍞 Bottom Bread (Anchor): A robot vacuum learns that bumping into walls wastes time (low reward), while following the edge cleans faster (higher reward).

1. World Models 🍞 Top Bread (Hook): You know how you can “play a scene in your head” before you do something risky, like stacking books high? 🥬 Filling (The Concept): A world model is an AI’s internal simulator that predicts what will likely happen next if it takes a certain action.

• How it works:
  1. Read the current situation.
  2. Predict the next situation for each possible action.
  3. Chain several predictions to simulate a short future.
• Why it matters: Without a world model, agents must try actions in the real world to learn consequences, which can be slow or unsafe. 🍞 Bottom Bread (Anchor): Before telling a robot to put a hot pan in the fridge, it imagines the fridge reaction—bad idea—so it tries cooling on the counter instead.

1. Partially Observable Markov Decision Process (POMDP) 🍞 Top Bread (Hook): Imagine walking through fog: you can’t see everything, only what’s near you. 🥬 Filling (The Concept): A POMDP describes decision-making when the agent can’t see the whole world state, only partial clues.

• How it works:
  1. The agent gets a limited observation.
  2. It keeps a memory of past observations.
  3. It chooses an action that should help, despite uncertainty.
• Why it matters: Without a POMDP view, agents pretend they see everything and make overconfident, brittle plans. 🍞 Bottom Bread (Anchor): In a text game, the agent only reads the current room description, not the whole house layout, but still must choose where to go.

The World Before This Paper:

• LLM agents were good at reacting to what they just saw and at recalling history, but they lacked deep foresight—the ability to think through consequences several steps ahead. This is called shallow grounding.
• People tried single-step checks (verify the next action) or fixed-horizon rollouts (always imagine, say, 3 steps). These helped a little, but they either missed long-term dependencies (too short) or wasted compute and amplified errors (too long for easy moves).
• Some methods (like MCTS-style planning with LLMs) pushed harder but often used lots of expensive rollouts indiscriminately.

The Problem:

• Complex tasks (like multi-step household chores or science experiments) need different amounts of “thinking ahead” at different moments. A trivial move doesn’t need deep imagination; a risky pivot does. Fixed horizons can’t adapt.

What Failed Before and Why:

• Single-step checks: catch obvious mistakes, but miss chain reactions.
• Fixed-horizon lookahead: either too shallow for tough spots or too deep (slow, error-prone) for easy spots.
• Implicit world modeling inside a policy: helps, but the agent can’t explicitly consult futures when it matters most.

The Gap:

• Agents lacked an adaptive lookahead that flexes based on goal difficulty and current progress, and a formal way to blend what is seen now with what is imagined next.

The Stakes (Why You Should Care):

• In homes, a helper agent should avoid risky actions (like turning on appliances at the wrong time) before causing trouble.
• In education and science games, an agent that plans experiments can save time and materials by previewing likely outcomes.
• In digital tools (email, web tasks), foresight avoids dead ends and reduces back-and-forth retries.

This paper proposes an Imagine-then-Plan (ITP) framework that explicitly imagines futures using a world model and then plans the next action, with an adaptive mechanism for deciding how far to look ahead each time. It also introduces a new formalism—POIMDP—that treats imagined futures as part of the agent’s decision input, reducing shallow grounding and cutting down costly real-world trial-and-error.

Let’s introduce the core innovations with the Sandwich pattern and then connect them.

1. Lookahead Imagination 🍞 Top Bread (Hook): Think of a chess player who quietly plays out future moves in their head before touching a piece. 🥬 Filling (The Concept): Lookahead imagination means the agent simulates several possible next steps to see likely consequences before acting.

• How it works:
  1. Use the world model to predict the next state for a candidate action.
  2. Repeat this for multiple steps to form an imagined mini-trajectory.
  3. Inspect these futures for progress and conflicts.
• Why it matters: Without imagining, agents discover failures only after they happen, wasting time and risking cascades of mistakes. 🍞 Bottom Bread (Anchor): Before picking up a cup in a dark room, the agent imagines turning on the light first to avoid dropping it.

1. Adaptive Lookahead Mechanism 🍞 Top Bread (Hook): When biking on a straight, safe road, you glance a little ahead; when entering busy traffic, you look much farther. 🥬 Filling (The Concept): Adaptive lookahead lets the agent choose how many steps to imagine now, balancing goal progress and compute cost.

• How it works:
  1. Estimate task difficulty and current progress.
  2. Pick a horizon K (maybe 0, maybe 5+) based on need.
  3. Imagine K steps, then decide the next action.
• Why it matters: Without adaptivity, the agent either overthinks easy moves or underthinks hard ones. 🍞 Bottom Bread (Anchor): In ALFWorld, moving from kitchen to hallway may need K=0–1, but planning a wash-then-place sequence may need K=4–5.

1. Imagine-then-Plan (ITP) 🍞 Top Bread (Hook): You know how you might daydream a plan (“If I do A, then B will happen”) and only then choose what to do now? 🥬 Filling (The Concept): ITP is a framework where the agent first imagines multiple future steps with a world model, then plans the immediate action using both what it sees and what it imagined.

• How it works:
  1. Observe the current state.
  2. Choose an adaptive horizon K.
  3. Roll out K imagined steps (actions and resulting states) using the world model.
  4. Reflect on the imagined futures to detect progress or conflicts.
  5. Pick the next action.
• Why it matters: Without this two-stream view (present + imagined), decisions can’t easily account for long-term consequences. 🍞 Bottom Bread (Anchor): The agent imagines that heating food before washing will fail the goal sequence, so it plans to wash first, then heat.

1. Partially Observable and Imaginable MDP (POIMDP) 🍞 Top Bread (Hook): It’s like navigating with both your flashlight (what you see now) and a sketch you just drew of what’s probably ahead (what you imagine). 🥬 Filling (The Concept): POIMDP is a decision model where the agent conditions its policy on two streams: the observed present and the imagined future trajectory.

• How it works:
  1. Keep the usual partial observation and history.
  2. Add imagined future steps from the world model.
  3. Choose the next action using both inputs.
• Why it matters: Without explicitly including imagined futures, the agent can’t properly weigh likely consequences. 🍞 Bottom Bread (Anchor): The policy reads the current room text plus an imagined “foresight” snippet and selects an action that avoids predicted bottlenecks.

The “Aha!” Moment (One sentence): If agents plan with what they can see now and what they just imagined ahead—adjusting how far they imagine each time—they make smarter, safer, and more efficient choices.

Multiple Analogies (3 ways):

• Chess: Don’t just look at the board; simulate a few lines ahead, and simulate farther only in critical positions.
• Driving: Glance near for routine cruising, look far ahead at intersections or in storms.
• Cooking: For a sandwich, no need to pre-visualize every step; for a fancy cake, mentally run through mixing, baking, cooling, and frosting.

Before vs After:

• Before: Agents reacted to observations, or imagined a fixed number of steps regardless of context.
• After: Agents imagine just enough steps when needed, blend foresight with the present, and avoid many preventable mistakes.

Why It Works (Intuition):

• Imagined futures reveal hidden dependencies (e.g., “wash before place,” “turn on light before examine”).
• Adaptive K controls error and cost: deep lookahead only when it helps, shallow otherwise.
• Treating imagined futures as inputs (POIMDP) gives the policy a direct way to consider consequences when acting.

Building Blocks:

• World model to simulate textual next states.
• Adaptive horizon chooser (simple reflection in ITP-I; learned predictor in ITP-R).
• Policy that conditions on both observation and imagined trajectory.
• Optional reinforcement learning to jointly optimize actions and horizon choices.

High-level recipe: Input (task + current state) → [Decide horizon K] → [Imagine K-step future with world model] → [Reflect and Plan] → Output (next action).

We’ll introduce two extra helper concepts used in the trained version (ITP-R):

1. K-head Predictor 🍞 Top Bread (Hook): Think of a small dial that sets how far to look ahead based on how tricky things look right now. 🥬 Filling (The Concept): The K-head predictor is a light module on top of the agent that outputs a distribution over horizons (K = 0..Kmax).

• How it works:
  1. Read the current state/history.
  2. Score each possible K.
  3. Sample or pick the K that balances predicted success and cost.
• Why it matters: Without a horizon chooser, the agent can’t automate when to imagine more or less. 🍞 Bottom Bread (Anchor): In a simple hallway step, it outputs K=0; at a multi-appliance step (sink + microwave), it sets K=5.

1. Advantage Actor-Critic (A2C) 🍞 Top Bread (Hook): Imagine a coach (critic) telling a player (actor) how much better a move was than average, so the player learns faster. 🥬 Filling (The Concept): A2C is a reinforcement learning method where the actor proposes actions and the critic estimates how good the situation is to guide learning.

• How it works:
  1. Actor picks horizon K and action a.
  2. Critic estimates value; compute advantage (how much better than expected).
  3. Update actor to favor better-than-expected choices.
• Why it matters: Without advantage guidance, learning which horizons/actions truly help would be slower and noisier. 🍞 Bottom Bread (Anchor): If imagining 4 steps leads to a big reward bump vs. imagining 1 step, A2C increases the chance of choosing 4 next time.

Now the full methodology, like a recipe:

Step 0: Warm up the base policy

• What: Fine-tune a language model on expert demonstrations so it can produce valid, sensible actions.
• Why: Without a decent starter policy, imagined futures won’t be grounded in reasonable behavior.
• Example: In ALFWorld, after warm-up the agent reliably outputs actions like “go to sink,” “turn on faucet,” “wash apple.”

Step 1: Train the world model

• What: Train a text world model to predict the next observation given the current observation and action.
• Why: Without a good predictor, imagination drifts and becomes misleading.
• Example: Input: (state: “You’re in the kitchen; holding apple.”, action: “open fridge”). Output: “The fridge door opens; shelves are visible.”

Step 2: Define POIMDP decision-making

• What: For each time t, the agent can imagine K steps using the world model, producing an imagined trajectory (actions and future states). The policy then picks the next action conditioned on both current state and this imagined trajectory.
• Why: Without combining present and imagined futures, the policy can’t properly anticipate conflicts.
• Example: If foresight shows that heating before washing blocks placement success, the policy chooses to wash first.

Two Variants of ITP:

A) ITP-I (training-free, inference-time learning)

• What happens:
  1. Adaptive horizon selection by reflection: read the goal and current state; decide K (0..Kmax) based on difficulty.
  2. World-model imagination: roll out K steps to get a short foresight text.
  3. Reflective policy generation: use the foresight as feedback; revise plan; choose next action.
• Why this step exists: It upgrades any LLM agent without retraining, letting it “peek ahead” before acting.
• Example: In ScienceWorld, the agent imagines heating water before measuring temperature and realizes the order matters for the goal rubric.

B) ITP-R (reinforcement-trained, adaptive)

• What happens (three stages): Stage 1: Pseudo-label horizons

  • Use expert actions and the frozen world model to simulate several candidate horizons (0..Kmax).
  • Score which K makes the expert action most likely (with a small penalty for larger K) and assign that K as a pseudo-label.
  • Why: Experts don’t tell us how far to imagine; this builds training targets for K.
  • Example: If K=3 best explains an expert’s “turn on light then examine” step, label K=3.

  Stage 2: Warm-up with lookahead

  • Train the policy to act given the imagined trajectory for the labeled K, and train the K-head predictor to predict that K.
  • Why: Without this, the K-head starts from scratch and picks poor horizons; the policy won’t learn to read foresight well.
  • Example: The policy learns to parse a foresight snippet like “(1) go to sink; (2) turn on faucet; (3) wash cup” and then choose “wash cup.”

  Stage 3: Online A2C optimization

  • At run-time, sample K from the K-head, generate foresight, then choose an action.
  • Reward = environment reward − (penalty for larger K) − (small per-step cost).
  • Update both the K-head and the policy using A2C with an entropy term to keep exploring horizons.
  • Why: This teaches the agent not only what to do, but also when deeper imagination is worth it.
  • Example: If K=5 keeps paying off at tricky junctions but is wasteful on easy steps, the learner assigns high K only at tricky spots.

Secret Sauce (What’s clever?):

• Treat imagined futures as first-class inputs (POIMDP), not just side notes.
• Learn to choose K adaptively with real rewards and small penalties so imagination is deep where it counts and shallow where it doesn’t.
• Provide a plug-and-play option (ITP-I) and a stronger trained option (ITP-R) to fit different deployment needs.

The Test: What’s measured and why

• Metric: Success Rate (SR) — percent of tasks fully solved. This captures whether the agent’s whole plan, not just a step or two, works end-to-end.
• Compute: They also track token budgets to measure how costly the planning is. This matters because longer lookahead costs more tokens.

Benchmarks and Setup

• ALFWorld: Text-based household tasks (Pick, Clean, Heat, Cool, Look, Pick2). Average episode ~8 steps.
• ScienceWorld: Text-based science experiments (seen and unseen splits). Average episode ~15 steps.
• Backbones: Qwen2.5-7B, Qwen3-8B, Llama-3.1-8B. Same backbones used across methods for fairness.
• Baselines: Prompting methods (CoT, ReAct, RAP) and training-based methods (SFT, WKM, IWM).

The Competition: Who they beat

• Prompting baselines are strong at reasoning but don’t adapt lookahead.
• Training baselines either clone expert behavior (SFT) or internalize dynamics/knowledge (IWM, WKM) without explicit adaptive foresight.

The Scoreboard (with context)

• Training-free ITP-I: With no extra training, simply adding imagination lifts performance well above prompting baselines. For example, on Qwen2.5-7B in ALFWorld, ITP-I roughly doubles ReAct’s overall SR (about 35.7% vs. 17.1%), like jumping from a D to a solid C+/B- just by peeking ahead.
• Trained ITP-R: Consistently the top performer among trained methods across backbones. On ALFWorld with Qwen3-8B, ITP-R reaches around 88.6% SR overall, which is like scoring an A when others hover at B levels. On ScienceWorld with Llama-3.1-8B, ITP-R hits about 63.9% in seen splits, ahead of other training-based competitors.
• Per-task improvements: ITP-R often leads or ties for best across Pick, Clean, Heat, Cool, Look, Pick2, showing robustness to different long-horizon recipes.

Surprising/Insightful Findings

1. Adaptive lookahead beats fixed horizons on success and cost

• Fixed k: Performance peaks at a medium k and then drops, while compute rises steeply as k grows.
• Adaptive: Keeps high SR without tuning a global k and uses far fewer tokens on average—like getting higher grades while studying fewer wasted hours.

1. Adaptive lookahead beats random horizons

• Randomly varying k is not enough. The state-aware choice of K in ITP dominates random choice, achieving higher SR with lower and more stable budgets. This shows that “when” you imagine matters as much as “how much.”

1. Reinforcement fine-tuning for K is critical

• Removing the online A2C stage (w/o RT) causes large drops: e.g., from ~88.6% to ~71.4% SR in ALFWorld and from ~59.7% to ~46.0% in ScienceWorld. This is like going from an A to a C just by skipping that training step, proving that learning the timing of imagination is a big deal.

1. World model choice matters (especially for ITP-I)

• Training-free ITP-I depends heavily on how good the world model is at predicting futures: stronger, well-aligned world models yield better gains.
• After reinforcement training (ITP-R), even a less specialized world model can become competitive, because the policy learns when and how to rely on foresight.

Bottom line: Across two different domains and multiple backbones, Imagine-then-Plan with adaptive lookahead wins on success rate while being smarter about compute, particularly when reinforcement learning teaches the agent how deep to imagine at each moment.

Limitations

• Modality gap: Evaluations are in text environments. Real robots see images and feel forces; visual/sensor noise could challenge imagination quality and horizon choices.
• Overhead: Imagining futures costs extra tokens and time. While adaptive penalties help, fully real-time settings may still need faster or distilled world models.
• World model bias: If the world model hallucinates, the agent can be misled. Bad foresight is often worse than no foresight.
• Generalization: Adaptive K tuned for ALFWorld/ScienceWorld might need retuning for other domains with very different step structures.

Required Resources

• A capable language model for the policy and another (or the same fine-tuned) for the world model.
• Expert demonstrations to warm-start the policy and world model training data from both expert and agent rollouts.
• Compute for training (including A2C) and for inference-time imagination.

When NOT to Use

• Ultra-low-latency scenarios where any added planning cost is unacceptable.
• Environments where the world model cannot be trained well (e.g., extremely chaotic or adversarial dynamics), making imagination unreliable.
• Tasks with strictly one-step dependencies where imagination brings little value.

Open Questions

• Multimodal extension: How to build vision-and-touch world models that are accurate enough for adaptive lookahead in the physical world?
• Safety under uncertainty: How can the agent detect and mitigate when the world model is uncertain or wrong (e.g., confidence-aware horizons)?
• Learning to imagine efficiently: Can we compress or cache foresight, or use speculative decoding, to keep benefits while slashing cost?
• Curriculum for K: How to train horizon selection that scales to very long tasks without exploding compute?
• Tool integration: How to combine learned world models with external simulators or structured planners for even safer, stronger foresight?

3-Sentence Summary: This paper introduces Imagine-then-Plan (ITP), where agents first imagine multiple future steps using a learned world model and then plan the next move, with an adaptive mechanism deciding how far to look ahead each time. By formalizing decisions as a POIMDP that blends what’s observed now with what’s imagined next, the agent avoids shallow grounding and reduces costly mistakes. Two versions—training-free ITP-I and reinforcement-trained ITP-R—show strong gains across ALFWorld and ScienceWorld.

Main Achievement: They show that explicitly using imagined futures as inputs, and learning when to imagine deeply versus shallowly, significantly boosts success rates while keeping computation in check, with ITP-R delivering the largest improvements.

Future Directions:

• Extend to multimodal (vision, sensors) and real robots with robust, efficient world models.
• Add uncertainty-aware imagination and confidence-guided horizons for safety.
• Speed up inference via distillation, caching, and speculative decoding to meet real-time demands.

Why Remember This: ITP turns “thinking ahead” from a nice idea into a practical, adaptive skill: agents don’t just react—they preview consequences and pick smarter actions. This shift from shallow grounding to foresightful planning is a principled step toward trustworthy, efficient agents that handle long, complex tasks with fewer surprises.