Train An Imaginative Agent

Jiaxin Ge^1,* Junyi Zhang^1,* Tianning Chai^1,* Hanjun Yoo¹ Chenyang Wang² Xiuyu Li² Weijia Shi³ XuDong Wang^1,4 Sewon Min¹ Luke Zettlemoyer³ Trevor Darrell^1,† Haiwen Feng^1,2,†

¹University of California, Berkeley ²Impossible Research ³University of Washington ⁴Duke University

* Equal contribution† Equal advising

Paper arXiv Code

Agima (Agentic Imagination) trains an agent to imagine before it acts: a single omni model generates images of where its actions lead and plans over these imagined futures.

How Agima imagines

Agima generates actions and imagined frames in one sequence. Each imagined frame is the observation for the next action.

**Figure 1.** Planning with agentic imagination: (a) feedback after acting, (b) feedback through imagination with Agima, (c) counterfactual branching and forward thinking in Sokoban.

Linear rollout

Imagines the frame after each action and plans the next action from it.

Counterfactual branching

Imagines an alternative action and rejects it when its outcome fails.

Forward thinking

Imagines complete candidate plans and commits to the one that succeeds.

Figure 2. The three ways Agima plans with imagined frames. Select one to mark it in Figure 1. Real frames are solid and imagined frames dashed blue. Green marks the plan that works and red a dead end.

Environments

We train Agima on an easier split of four environments (left clip of each pair) and test it on a harder split it never sees in training (right clip).

Maze2D

Sokoban

Pacman

PushT

Figure 3. Ground-truth rollouts from the training and test splits.

Easy-to-hard generalization

K is the number of actions a model plans and executes per real observation. Between observations, Agima plans from its own imagined frames.

Planning traps

Planning traps are hard test sets built specifically to measure planning. In each, a move that looks reasonable fails several steps later. We compare the planning modes of Figure 2, each trained on the same data, at K = ∞.

Deadlock

One push puts the nearer box on the goal beside it. That goal is the only one the farther box can ever reach, so the farther box is left with nowhere to go.

Nearer box first: one push puts it on the goal below it, in the corridor; the farther box, pushed toward the other goal, is stopped by it. — **Nearer box first**

Farther box first: it is pushed along the corridor to the far goal, then the nearer box onto its goal. Solved in 18 moves. — **Nearer box first**

Irreversible push (Structural Fork set)

Two pushes are legal here. Pushing up pins a box to the wall in a row with no goal, and no later move brings it back. Pushing right keeps the puzzle solvable.

Push up: the box is pinned to the top wall, in a row with no goal; nothing can push it back down. — **Push up**

Push right: the puzzle stays solvable; all three boxes reach their goals in 13 moves. — **Push up**

Pacman route order

Going for the nearer food first leaves Pacman no escape from the ghost. Going for the farther food first clears the level.

Nearer food first: Pacman eats the food to its left, heads up for the other one and is caught by the ghost. — **Nearer food first** Agima-A7B

Farther food first: Pacman eats the upper food, comes back down for the food on the left and clears the level. — **Nearer food first** Agima-A7B

Learning when and where to imagine

We run RL from the fine-tuned Agima-A7B models on three-box Sokoban with only a binary success reward, and test on the Deadlock and Fork sets.

Paper Figure 4: training reward and held-out accuracy on the deadlock and fork sets over 200 RL steps, 3 and 4 boxes.

Paper Figure 4. Training reward and held-out accuracy on the Deadlock and Fork sets over 200 RL steps. Solid lines are 3 boxes, dashed lines 4 boxes.

Paper Figure 6. Failure-triggered lookahead, unified-format share, branches per episode and critical-step alignment over RL steps.

RL improves held-out planning

Every imagination variant improves, including on four-box puzzles never seen in RL, while the action-only model (No Imagine) stays near zero.

Figure 6. Held-out accuracy over 200 RL steps: (a) Agima-A7B (Forward) on the Deadlock set and (b) Agima-A7B (Branching) on the Fork set. Solid lines are 3 boxes, dashed lines 4 boxes (paper Fig. 4b–c).

A new unified format emerges

After RL, Agima-A7B (Forward) rolls out one plan and drafts another only if it predicts failure. Second plans drop from 79% to 26–31% of episodes.

Paper Figure 5, the same Sokoban puzzle before and after RL. (a) Before RL: Format 1 imagines one plan step by step; Format 2 imagines Candidate 1, which succeeds, and Candidate 2, which reaches deadlock, and commits to Candidate 1. (b) After RL, the unified format: a step-by-step rollout reaches deadlock, Candidate 2 reaches deadlock, Candidate 3 succeeds and is committed. — **Figure 7.** The same Sokoban puzzle before and after RL (paper Fig. 5).

Figure 8. Agima-A7B (Forward) on the Deadlock set over RL steps: (a) episodes whose first plan is predicted to fail, among those with a second plan, and (b) episodes in the unified format (paper Fig. 6a–b).

Branching moves to critical decisions

After RL, Agima-A7B (Branching) opens more branches per episode, and more of them land where one push would make the puzzle unsolvable.

Figure 9. Agima-A7B (Branching) on the Fork set over RL steps: (a) mean branches per episode and (b) the share of branches placed at a critical step (paper Fig. 6c–d).

Figure 10. An illustration, not a recorded episode. Both models branch, but after RL more branches land on critical steps like this fork (63% to 73%, 3 boxes).

Citation

@misc{ge2026trainimaginativeagent,
  title         = {Train An Imaginative Agent},
  author        = {Ge, Jiaxin and Zhang, Junyi and Chai, Tianning and Yoo, Hanjun and Wang, Chenyang and Li, Xiuyu and Shi, Weijia and Wang, XuDong and Min, Sewon and Zettlemoyer, Luke and Darrell, Trevor and Feng, Haiwen},
  year          = {2026},
  eprint        = {2610.26845},
  archivePrefix = {arXiv},
  primaryClass  = {cs.AI},
  url           = {https://arxiv.org/abs/2610.26845}
}