Training 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 the easy split of four environments and test it on a harder split it never sees in training.

Maze2D

Sokoban

Pacman

PushT

Figure 3. Ground-truth rollouts in each environment. The training split is on the left of each pair and the test split on the right.

Easy-to-hard generalization

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

Planning traps

Test sets where a locally reasonable move fails several steps later. 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.

A 2-box level from Agima's Sokoban generator. Test sets: 2, 3 and 4 boxes, 100 puzzles each.

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.

Test sets: 2, 3 and 4 boxes, 100 puzzles each.

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.

The episode from Figure 3 of the paper. Test grids are 8×8 with 2, 4, 6 or 8 food, and 10×10 with 8.

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 SFT models on three-box Sokoban, with a binary success reward and nothing else.

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 it never sees in RL, while the action-only model 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, each with Agima-A7B and No Imagine for reference. Solid lines are 3 boxes, dashed lines 4 boxes (paper Fig. 4b–c).

A new unified format emerges

After RL, Agima rolls out one plan step by step and drafts the next only once it foresees a 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 branches more often, and more of its branches 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

@article{agima2026,
  title   = {Training 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},
  journal = {arXiv preprint},
  year    = {2026}
}