cs.AIOct 19, 2026

Train An Imaginative Agent

Authors: Jiaxin Ge, Junyi Zhang, Tianning Chai, Hanjun Yoo, Chenyang Wang, Xiuyu Li, Weijia Shi, XuDong Wang, 

Organizations: University of California, Berkeley · Impossible Research · University of Washington · Duke University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

A seemingly reasonable push in Sokoban can make the puzzle unsolvable. To avoid such mistakes, an agent has to anticipate the consequences of its actions before execution, especially when feedback from the real environment is sparse or actions are irreversible. We introduce Agentic Imagination (Agima), a training framework that trains an agent to imagine before acting, i.e. to be imaginative. Agima trains a single omni model to propose actions, generate images of their outcomes, and then reason over these imagined futures to decide what to execute in the real environment. Rather than committing to a single plan, the agent can simulate a sequence of actions, explore alternatives, and compare possible futures. We first use supervised fine-tuning (SFT) to unify policy and world modeling within a single omni model. This enables the agent to predict future outcomes and perform basic forms of imagination, such as forward thinking and counterfactual branching. We then use agentic reinforcement learning (RL) to teach the agent to reason with its imagined futures and use them effectively for planning. Across four environments spanning Maze2D, Sokoban, Pacman, and PushT, Agima generalizes to harder settings unseen during training. Trained with SFT alone on two-box Sokoban, Agima solves 40% of four-box puzzles while seeing the real board only once every five moves, compared with an action-only planner that only solves 2%. After RL, new behaviors emerge in how Agima uses its imagination: it learns a new unified format that was unseen during the SFT stage, learns to conduct forward thinking only when necessary, and branches more frequently at critical decision points. An agent that can picture where its actions lead can learn when and where to look ahead, that is, to plan. Project page: https://agima.net/.

Figures & tables

Figure 1: Planning with agentic imagination. (a) A conventional agent receives feedback only after acting in the environment. (b) Agima uses a single model to propose actions and predict the consequences of them. Each prediction serves as an internal observation for the next decision. (c) Planning with Agima in Sokoban: counterfactual branching rejects a deadlocking action, and forward thinking compares candidate plans before execution. The executed actions reach ot+Ko_{t+K} , which grounds the next round of imagination. In (c), o^j\hat{o}_{j} is the predicted outcome of a^j\hat{a}_{j} .
Figure 2: Evaluation environments and imagined visual rollouts. (a) Training and test states in four environments. (b) Agima interleaves actions and predicted states, then executes the actions.
KKGroupMethodPacmanMazeSokobanPushT
Acc.Final Cov.
  Gemini 3.1 Pro0%38%3%0%0.2%
 ClosedGPT-5.6 Sol0%48%32%0%1.8%
11SFTAgima-A7B84%62%47%42%45.9%
33 Gemini 3.1 Pro0%21%4%0%2.2%
ClosedGPT-5.6 Sol0%39%26%0%3.0%
Table 1: Performance comparison across Pacman, Maze, Sokoban, and PushT environments. KK indicates the number of actions generated before feedback. Reactive ( K=1K=1 ) receives ground-truth feedback after every action. Pacman, Maze, and Sokoban use the hard splits. For PushT, we report success accuracy and final coverage on the far split. A7B denotes 7B active parameters (14B total); No Imagine is trained to predict actions without intermediate visual predictions. Bold indicates the best result at each KK (including ties); gains are relative to Agima-A7B (No Imagine) .
ModelConditionSokobanPacmanMaze2DPushT
No Imagine—2%1%23%1%
Agima-A7BOwn predictions40%37%40%18%
Blank1%0%8%1%
Stale0%0%8%0%
Oracle GT43%81%69%47%
Table 2: Replacing imagined observations at inference. No Imagine is a separately trained action-only baseline; the four Agima-A7B conditions share checkpoint, episodes, and decoding within each environment. Blank replaces predictions with gray images, Stale with the latest real observation, and Oracle GT with ground-truth frames for the proposed actions.
Figure 3: Forward thinking compares alternative futures before execution. A Pacman example (Table 3 (b)): (a) Agima-A7B is caught after executing a single rollout; (b) Agima-A7B (Forward) compares two imagined futures and selects Candidate 1, which succeeds.
(a) SokobanDeadlockStructural Fork
Method2-box3-box4-box2-box3-box4-box
Gemini 3.1 Pro74%42%19%84%34%7%
GPT-5.6 Sol44%23%17%60%20%12%
Qwen2.5-VL-7B1%2%0%12%2%0%
Qwen3-VL-8B6%1%1%19%2%0%
Agima-A7B (No Imagine)12%5%1%15%2%1%
Table 3: Success rates on Sokoban and Pacman reasoning tasks. (a) Sokoban Deadlock and Structural Fork configurations with 2, 3, and 4 boxes. (b) Pacman route-order traps, where G denotes grid size, and F denotes the number of food items. Bold indicates the best result in each column; changes in parentheses are absolute percentage-point differences relative to Agima-A7B.
Figure 4: RL improves training reward and held-out planning performance. (a) Training reward during RL. (b–c) Accuracy on the Deadlock and Structural Fork evaluation sets. Solid and dashed lines denote three-box and four-box tasks, respectively.
Figure 5: RL induces a unified, adaptive planning format. (a) SFT produces a linear rollout or a fixed two-candidate comparison. (b) After RL, the model evaluates each imagined plan and generates alternatives as needed, here rejecting two failed plans before finding a successful third.
Figure 6: RL induces adaptive use of imagination. Forward thinking becomes increasingly failure-triggered, while branching explores more alternatives and shifts toward critical decision points.
Appendix figures & tables105 assets

Supplementary material from the paper’s appendix.

Appendix

 Maze2DSokobanPacmanPushT
Grid size11×1111\times 117×77\times 79×99\times 9N/A
Training settingmaze navigation2 boxes8 fooddiverse T start
Expertshortest-path oracleSokoban solveroracle plannerPPO policy
Action space4 moves4 moves4 movescontinuous (x,y)(x,y)
Action budget40208030
Plan length9 (9.9) [5–33]10 (9.9) [2–19]23 (22.5) [9–27]20 (20.0) [4–29]
Appendix
Table 4: Linear-rollout training data per task. Plan length counts the expert’s actions excluding the final STOP and is reported as median (mean) [range]. Grid sizes include the outer wall.
Sokoban Forward Thinking
Both solve79%
At least one candidate deadlocks21%
Sokoban Counterfactual Branching
Branches per episode: 0 / 1 / 2 / 327.1 / 35.3 / 20.6 / 17.1%
Pacman Forward Thinking
Both candidates succeed96.4%
Appendix
Table 5: Forward-Thinking and Counterfactual-Branching trace statistics. Forward-thinking outcome percentages are over two-candidate traces.
Figure 12
Figure 13
Figure 14
Figure 7: A complete expert trajectory and its action chunks. The same 17-action training trajectory, including STOP , is segmented with K∈{1,3,5,10,∞}K\in\{1,3,5,10,\infty\} . Numbers in the lower panel refer to the actions above. Images are saved environment observations used as supervision.
Figure 16
Figure 17
Figure 8: Complete counterfactual-branching training example. The main line contains 12 actions including STOP . The two alternative pushes, inserted after main-line actions 3 and 5, lead to certified deadlocks and are rejected.
Figure 19
Figure 9: Complete forward-thinking training example in Sokoban. Candidate 1 deadlocks after two moves. Candidate 2 solves the puzzle in four moves and ends with STOP .
Figure 21
Figure 22
Figure 23
Figure 24
Figure 25
Figure 26
Figure 10: Complete forward-thinking training example in Pacman. After a shared five-move prefix, Candidate 1 is caught by the ghost after seven further moves. Candidate 2 collects all food after 19 further moves; its executable target ends with STOP . All saved training images are shown.
SettingValue
Standard SFT steps5,000
SFT batch size8
Training-example presentations40,000
Released Sokoban checkpointsBAGEL-7B-MoT, exponential-moving-average (EMA) weights
OptimizerAdamW ( β1=0.9\beta_{1}=0.9 , β2=0.95\beta_{2}=0.95 )
Learning rate2×10−52\times 10^{-5}
Appendix
Table 6: SFT settings.
SetMetricPre-RLRL50RL100RL150RL200
3-box DeadlockGenerates C283%57%18%20%22%
Fixed comparison79%51%0%0%0%
Hybrid4%6%18%20%22%
Predicts C1 failure31.7%33.9%100%100%100%
Actual C1 failure61.0%64.3%94.4%100%95.5%
4-box DeadlockGenerates C276%56%34%36%37%
Appendix
Table 7: Forward-Thinking formats and first-plan failures. C1 and C2 denote the first and second candidate plans. Generates C2, Fixed comparison, and Hybrid use all 100 episodes per set as their denominator. The last two rows use comparisons with readable candidate actions and verdicts: Predicts C1 failure records the model’s judgment; Actual C1 failure records whether C1 fails when executed independently. All values are percentages.
SetCheckpointMean012≥3\mathbf{\geq 3}
3-box DeadlockPre-RL1.1834%31%18%17%
RL501.1333%36%16%15%
RL1001.2034%30%18%18%
RL1501.2535%25%20%20%
RL2001.2535%25%20%20%
4-box DeadlockPre-RL1.6722%23%21%34%
Appendix
Table 8: Number of counterfactual branches per episode. Mean counts branches across all 100 episodes in each row, including episodes with no branch. Columns 0, 1, 2, and ≥3\geq 3 give the percentage of episodes with that many complete branches. Each branch contains one imagined action.
FormatCheckpointSuccessImaginedExecuted
ForwardPre-RL30.3%27.8217.88
 RL5035.0%24.9517.36
 RL10037.5%20.8716.96
 RL15038.5%21.4016.97
 RL20041.0%22.1817.50
BranchingPre-RL30.0%23.6218.63
Appendix
Table 9: Mean rollout lengths and pooled task success on the four 3/4-box sets. Each row pools 400 episodes. Task success is a percentage; Imagined and Executed are mean action counts per episode.
Figure 11: Complete Maze2D trajectory at K=1K=1 . Selected successful Ordinary Hard episode; all 17 actions and their states are shown. The reactive reference receives a real observation after every action.
Figure 12: Complete Maze2D trajectory at K=3K=3 . Selected successful Ordinary Hard episode; this page shows steps 0–15 of 27. The trajectory continues on the next page. Blue rows show predictions; teal inputs and feedback arrows show the real observations supplied between chunks.
Figure 34
Figure 13: Complete Maze2D trajectory at K=5K=5 . Selected successful Ordinary Hard episode; all 23 actions and their states are shown. Blue rows show predictions; teal inputs and feedback arrows show the real observations supplied between chunks.
Figure 14: Complete Maze2D trajectory at K=10K=10 . Selected successful Ordinary Hard episode; all 21 actions and their states are shown. Blue rows show predictions; teal inputs and feedback arrows show the real observations supplied between chunks.
Figure 15: Complete Maze2D trajectory at K=∞K=\infty . Selected successful Ordinary Hard episode; all 13 actions and their states are shown. The model receives only the initial observation, with no intermediate feedback. Ground-truth states are shown for comparison.
Figure 16: Complete Sokoban trajectory at K=1K=1 . Selected successful Ordinary Hard episode; all 21 actions and their states are shown. The reactive reference receives a real observation after every action.
Figure 17: Complete Sokoban trajectory at K=3K=3 . Selected successful Ordinary Hard episode; all 21 actions and their states are shown. Blue rows show predictions; teal inputs and feedback arrows show the real observations supplied between chunks.
Figure 18: Complete Sokoban trajectory at K=5K=5 . Selected successful Ordinary Hard episode; all 19 actions and their states are shown. Blue rows show predictions; teal inputs and feedback arrows show the real observations supplied between chunks.
Figure 19: Complete Sokoban trajectory at K=10K=10 . Selected successful Ordinary Hard episode; all 19 actions and their states are shown. Blue rows show predictions; teal inputs and feedback arrows show the real observations supplied between chunks.
Figure 20: Complete Sokoban trajectory at K=∞K=\infty . Selected successful Ordinary Hard episode; all 17 actions and their states are shown. The model receives only the initial observation, with no intermediate feedback. Ground-truth states are shown for comparison.
Figure 21: Complete Pacman trajectory at K=1K=1 . Selected successful Ordinary Hard episode; all 33 actions and their states are shown. The reactive reference receives a real observation after every action.
Figure 22: Complete Pacman trajectory at K=3K=3 . Selected successful Ordinary Hard episode; this page shows steps 0–20 of 33. The trajectory continues on the next page. Blue rows show predictions; teal inputs and feedback arrows show the real observations supplied between chunks.
Figure 45
Figure 23: Complete Pacman trajectory at K=5K=5 . Selected successful Ordinary Hard episode; this page shows steps 0–20 of 33. The trajectory continues on the next page. Blue rows show predictions; teal inputs and feedback arrows show the real observations supplied between chunks.
Figure 47
Figure 24: Complete Pacman trajectory at K=10K=10 . Selected successful Ordinary Hard episode; this page shows steps 0–20 of 35. The trajectory continues on the next page. Blue rows show predictions; teal inputs and feedback arrows show the real observations supplied between chunks.
Figure 49
Figure 25: Complete Pacman trajectory at K=∞K=\infty . Selected successful Ordinary Hard episode; this page shows steps 0–15 of 27. The trajectory continues on the next page. The model receives only the initial observation, with no intermediate feedback. Ground-truth states are shown for comparison.
Figure 51
Figure 26: Complete PushT trajectory at K=1K=1 . Selected successful Far episode; all 19 actions and their states are shown. The reactive reference receives a real observation after every action. Actions give normalized target coordinates (x,y)(x,y) .
Figure 27: Complete PushT trajectory at K=3K=3 . Selected successful Far episode; all 18 actions and their states are shown. Blue rows show predictions; teal inputs and feedback arrows show the real observations supplied between chunks. Actions give normalized target coordinates (x,y)(x,y) .
Figure 28: Complete PushT trajectory at K=5K=5 . Selected successful Far episode; this page shows steps 0–15 of 27. The trajectory continues on the next page. Blue rows show predictions; teal inputs and feedback arrows show the real observations supplied between chunks. Actions give normalized target coordinates (x,y)(x,y) .
Figure 55
Figure 29: Complete PushT trajectory at K=10K=10 . Selected successful Far episode; all 20 actions and their states are shown. Blue rows show predictions; teal inputs and feedback arrows show the real observations supplied between chunks. Actions give normalized target coordinates (x,y)(x,y) .
Figure 30: Complete PushT trajectory at K=∞K=\infty . Selected successful Far episode; this page shows steps 0–15 of 28. The trajectory continues on the next page. The model receives only the initial observation, with no intermediate feedback. Ground-truth states are shown for comparison. Actions give normalized target coordinates (x,y)(x,y) .
Figure 58
Figure 31: Counterfactual Branching. Example trace of counterfactual branching in Sokoban.
Figure 32: Forward Thinking. Example trace of Forward thinking in Sokoban.
Test setModelRL50RL100RL150RL200
3-box DeadlockForwardFailSuccessSuccessSuccess
4-box DeadlockForwardFailFailSuccessSuccess
3-box Structural ForkBranchingFailSuccessSuccessSuccess
4-box Structural ForkBranchingFailFailSuccessSuccess
Appendix
Table 10: Outcomes of the illustrated RL episodes. Each row follows one selected initial state across four checkpoints. Aggregate performance is reported in Figure 4 .
Figure 62
Figure 63
Figure 64
Figure 33: Forward at RL50 on three-box Deadlock. Complete trace for episode 84; recorded execution ends in failure.
Figure 66
Figure 34: Forward at RL100 on three-box Deadlock. Complete trace for episode 84; recorded execution ends in success.
Figure 68
Figure 35: Forward at RL150 on three-box Deadlock. Complete trace for episode 84; recorded execution ends in success.
Figure 70
Figure 36: Forward at RL200 on three-box Deadlock. Complete trace for episode 84; recorded execution ends in success.
Figure 72
Figure 39: Forward at RL150 on four-box Deadlock. Complete trace for episode 21; recorded execution ends in success.
Figure 74
Figure 40: Forward at RL200 on four-box Deadlock. Complete trace for episode 21; recorded execution ends in success.
Figure 76
Figure 77
Figure 41: Branching at RL50 on three-box Structural Fork. Complete trace for episode 17; recorded execution ends in failure.
Figure 79
Figure 80
Figure 81
Figure 42: Branching at RL100 on three-box Structural Fork. Complete trace for episode 17; recorded execution ends in success.
Figure 83
Figure 84
Figure 85
Figure 43: Branching at RL150 on three-box Structural Fork. Complete trace for episode 17; recorded execution ends in success.
Figure 87
Figure 88
Figure 89
Figure 44: Branching at RL200 on three-box Structural Fork. Complete trace for episode 17; recorded execution ends in success.
Figure 91
Figure 92
Figure 45: Branching at RL50 on four-box Structural Fork. Complete trace for episode 56; recorded execution ends in failure.
Figure 94
Figure 95
Figure 46: Branching at RL100 on four-box Structural Fork. Complete trace for episode 56; recorded execution ends in failure.
Figure 97
Figure 98
Figure 47: Branching at RL150 on four-box Structural Fork. Complete trace for episode 56; recorded execution ends in success.
Figure 100
Figure 101
Figure 48: Branching at RL200 on four-box Structural Fork. Complete trace for episode 56; recorded execution ends in success.
Figure 103
Figure 49: Maze2D failure at K=∞K=\infty . The agent issues STOP before reaching the goal.
Figure 105
Figure 106
Figure 50: Sokoban failure at K=∞K=\infty . The agent stops with 2 of 4 boxes on goals.
Figure 108
Figure 109
Figure 110
Figure 51: Pacman failure at K=∞K=\infty . The ghost catches Pacman at action 17. The remaining seven predicted states are never executed.
Figure 112
Figure 113
Figure 52: PushT failure at K=∞K=\infty . The agent stops with 32.4% final target coverage, below the paper’s 80%80\% success threshold.
More in this topic

LLM Agent Planning

Large language model (LLM) agent planning studies how agents decompose goals, choose actions and tools, and adapt plans from feedback.

See all 225 papers

Explore similar work

CardsList

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

    Jan 13, 2026Youwei Liu, Jian Wang, Hanlin Wang +2World Model-Based PlanningModel-Based Planning

    Central South University · Department of Computing, The Hong Kong Polytechnic University · College of Computer Science, Sichuan University

  2. World Action Planner: Generalizable Decision-Making with Action-Conditioned World Models

    Jul 30, 2026Xiangcheng Zhang, Yilun DuWorld Models for RoboticsWorld Model-Based Planning

    Harvard University

  3. WISE: World-model-guided Imagination Scheduling for Efficient Post-training of Vision-Language-Action Models

    Sep 3, 2026Chenhao Zhang, Hanyu Zhao, Hang Cheng +2Robustness of VLA ModelsWorld Models for Robotics

    1Tsinghua University · 2Beijing Academy of Artificial Intelligence (BAAI)