Have you ever stood in front of a supermarket shelf holding two bags of chips, mentally simulating: "If I buy the original flavor, will it be too bland with cola? If I buy the BBQ flavor, will it clash with hotpot?"
You aren't hesitating. You are imagining.
You're using your brain's "world model" to simulate the future, playing back scenes that haven't happened yet, and making a decision based on what you see. This is one of humanity's core abilities — rehearsing in your head before acting.
Now AI has learned to do it too.
One Wrong Step Ruins Everything
Imagine playing Sokoban. The rules are simple: push boxes to target positions. But there's a fatal constraint — boxes can only be pushed, never pulled.
You have two boxes. Pushing one left looks reasonable. But three moves later, you find the other box is stuck and can never be moved again. The whole level becomes unsolvable.
A seemingly reasonable action can make the entire problem impossible.
This is the cruelty of sparse-feedback environments: actions are irreversible and feedback is delayed. By the time you realize the mistake, it's too late.
The traditional AI agent approach is: act first, observe after. It has no capacity to "imagine." It pushes the box, sees the environment's feedback, and decides the next step. In environments with timely feedback that works fine, but in a "one wrong step ruins everything" world like Sokoban, it's like a blindfolded person navigating a maze.
Agima: An Agent That "Dreams"
The paper *Train An Imaginative Agent* proposes a training framework called Agima (Agentic Imagination), whose core idea fits in one sentence:
> Before acting, generate an image of "what the future will look like," then decide whether to make the move while looking at that picture.
This is not a metaphor. Agima really generates images.
It is built on an "omni model" — the same model plays three roles at once: 1. Policy model: proposes actions ("push left one step") 2. World model: generates the post-action image (imagining what things look like after the push) 3. Reward model: evaluates the imagined image ("does this state look good?")
Three roles, one model. Like a person who can plan, rehearse the outcome, and evaluate themselves.
Three Ways to "Imagine"
Agima doesn't just "think one step ahead." It has three increasingly sophisticated imagination strategies:
Linear Rollout: Follow One Path All the Way
The most basic approach. Imagine: push one step → generate image → based on the image, push another step → generate another image... rolling out K steps ahead. Like mentally playing out five moves in chess.
Without any real intermediate feedback (K=5), Agima reaches a 40% success rate on four-box Sokoban, versus only 2% for a model that just outputs actions. The value of imagination: a 20x improvement.
Counterfactual Branching: Explore Multiple Paths at the Fork
At critical decision points, instead of following one path, imagine several different paths simultaneously. "What if I push left? What if I push right?" Then compare the imagined outcomes.
On Sokoban's "structural fork" test (levels specifically designed to require critical choices), branching imagination lifts four-box success from 22% to 32%.
Lookahead Thinking: Compare Two Full Routes Before Acting
The most complex strategy. Before executing, fully imagine two different multi-step plans, compare them, and choose the better one.
In Pacman's "route-order trap" (eat the near food first or the far one first? The near one looks safer, but after eating it you might get caught by the ghost), lookahead thinking raises success from 22% to 54%.
Why is this remarkable? Because these imaginations come not from an external world model, but from images generated by the same model itself. It previews the future with its own "world model" and then decides based on the preview.
RL Phase: Learning "When to Imagine"
SFT (supervised fine-tuning) taught Agima *how* to imagine, but the format was preset — either linear rollout or a fixed comparison of two paths.
Not enough. Humans don't consider every possibility for every decision. We act directly in simple situations and only pause to imagine at critical points.
The paper uses GRPO (a reinforcement learning algorithm) to let Agima learn "when to imagine and how many paths to imagine" on its own.
The result: a new reasoning format never seen in the SFT stage:
> The agent first imagines one plan and evaluates the predicted outcome. If the prediction succeeds, it executes directly; only if the prediction fails does it generate a second backup plan.
Nobody taught this. It emerged spontaneously from the binary reward of task success.
Meanwhile, counterfactual branching became smarter — the agent learned to explore more alternatives at "critical decision points" (where one wrong step means a dead end) and to imagine less in safe areas.
It is learning to budget its own imagination.
A Key Validation: Is It Really Looking at Its Own Images?
You might ask: is Agima really using its generated images to make decisions, or just going through the motions?
The paper ran an elegant ablation: replace Agima's imagined images with blank images or outdated old ones.
Result: success rate collapses.
Conversely, if imagined images are swapped for real future images (cheating mode), success improves by 29–44 percentage points.
This proves Agima genuinely "looks at" its own imagined images and extracts information from them to decide. Imagination isn't decoration — it's a core information source.
Why This Matters More Than It Seems
1. Imagination Bridges Policy and World Model
In traditional architectures, the policy model handles "what to do" and the world model handles "what will happen" — they're separate. Agima puts both in one model, making imagination the bridge: the policy proposes an action, the world model renders the outcome, and the policy adjusts based on the rendered image.
2. Imagination and Real Observation Share an Interface
Agima's most elegant design: imagined images and images returned by the real environment are indistinguishable to the model. They share the same interface and can be used interchangeably.
This means the agent can substitute imagination when real feedback is unavailable, and real observation when available. Like humans — you can close your eyes and picture a room's layout, open them and see the real room; both kinds of information come in the same format in your brain.
3. RL Makes Imagination Strategies Emerge
SFT teaches "how to imagine"; RL teaches "when to imagine what." The key finding — using only binary task-success rewards, with no supervision about "how to think," the agent spontaneously learned:
- Try one path first; imagine a second only if it fails (saving imagination cost)
- Explore more alternatives at critical forks (concentrating imagination)
- Imagine less in safe areas (not wasting it)
Engineering Takeaways
Add an "Imagination Layer" Rather Than a Bigger Policy Model
When facing sparse feedback and irreversible actions, instead of making the policy model bigger, add an imagination layer — let it rehearse a few steps before acting. The same backbone went from 2% to 40% on four-box Sokoban with imagination added.
Keep the Imagination Format Consistent with Real Observations
Agima can interchange imagination and reality because they share the same visual interface. If your agent's imagined outputs differ in format from real observations, seamless substitution becomes impossible.
Let RL Freely Explore Imagination Strategies
Don't preset "when to imagine what." Give the agent task rewards and let it learn to allocate imagination itself. Agima's most beautiful finding — adaptive imagination strategy — emerged from RL, not design.
My Thoughts
This paper raises a deeper philosophical question: what is imagination?
Before Agima, AI "prediction" was abstract — a vector, a probability distribution, a number. Agima's imagination is concrete — an image that can be "looked at."
Human imagination is concrete too. When you imagine "hiking tomorrow," your mind doesn't hold the text "hiking tomorrow," but a picture — the wind at the summit, the trail underfoot, the pack on your back.
Imagination isn't abstract reasoning; it's rendering a future sensory experience inside your head. Agima achieves this with image generation — it isn't "reasoning" about the future, it's "rendering" it.
What fascinates me even more is the metacognition that emerged from RL. Agima learned "try one path first, imagine another if it fails" — not a rule written by an engineer, but a strategy grown from task rewards.
This suggests: imagination and metacognition may be two sides of the same ability. When a system gains the ability to imagine (to rehearse the future), it naturally develops the ability to manage imagination (deciding when and how much to imagine). The former is taught by SFT; the latter emerges from RL.
Human infant development follows the same path — first learning to simulate scenarios in the mind (pretend play), then making decisions based on simulations, finally controlling the frequency and depth of simulation (metacognition).
Agima has completed the first two steps.
---
Paper: Train An Imaginative Agent (arXiv: 2610.26845)
Core contributions: A unified omni model serving as policy, world model, and reward model simultaneously; SFT teaches imagination, RL teaches metacognition; imagined and real observations share an interface and can be used interchangeably.
One-line summary: Give an agent an "in-head renderer," and it will learn to dream before acting — and find the right path inside the dream.