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+K , which grounds the next round of imagination. In (c), o^j is the predicted outcome of 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.
K
Group
Method
Pacman
Maze
Sokoban
PushT
Acc.
Final Cov.
Gemini 3.1 Pro
0%
38%
3%
0%
0.2%
Closed
GPT-5.6 Sol
0%
48%
32%
0%
1.8%
1
SFT
Agima-A7B
84%
62%
47%
42%
45.9%
3
Gemini 3.1 Pro
0%
21%
4%
0%
2.2%
Closed
GPT-5.6 Sol
0%
39%
26%
0%
3.0%
Table 1: Performance comparison across Pacman, Maze, Sokoban, and PushT environments. K indicates the number of actions generated before feedback. Reactive ( K=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 K (including ties); gains are relative to Agima-A7B (No Imagine) .
Model
Condition
Sokoban
Pacman
Maze2D
PushT
No Imagine
—
2%
1%
23%
1%
Agima-A7B
Own predictions
40%
37%
40%
18%
Blank
1%
0%
8%
1%
Stale
0%
0%
8%
0%
Oracle GT
43%
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) Sokoban
Deadlock
Structural Fork
Method
2-box
3-box
4-box
2-box
3-box
4-box
Gemini 3.1 Pro
74%
42%
19%
84%
34%
7%
GPT-5.6 Sol
44%
23%
17%
60%
20%
12%
Qwen2.5-VL-7B
1%
2%
0%
12%
2%
0%
Qwen3-VL-8B
6%
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
Maze2D
Sokoban
Pacman
PushT
Grid size
11×11
7×7
9×9
N/A
Training setting
maze navigation
2 boxes
8 food
diverse T start
Expert
shortest-path oracle
Sokoban solver
oracle planner
PPO policy
Action space
4 moves
4 moves
4 moves
continuous (x,y)
Action budget
40
20
80
30
Plan length
9 (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 solve
79%
At least one candidate deadlocks
21%
Sokoban Counterfactual Branching
Branches per episode: 0 / 1 / 2 / 3
27.1 / 35.3 / 20.6 / 17.1%
Pacman Forward Thinking
Both candidates succeed
96.4%
Appendix
Table 5: Forward-Thinking and Counterfactual-Branching trace statistics. Forward-thinking outcome percentages are over two-candidate traces.
Figure 12Figure 13Figure 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,∞} . Numbers in the lower panel refer to the actions above. Images are saved environment observations used as supervision.
Figure 16Figure 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 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.
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.
Set
Checkpoint
Mean
0
1
2
≥3
3-box Deadlock
Pre-RL
1.18
34%
31%
18%
17%
RL50
1.13
33%
36%
16%
15%
RL100
1.20
34%
30%
18%
18%
RL150
1.25
35%
25%
20%
20%
RL200
1.25
35%
25%
20%
20%
4-box Deadlock
Pre-RL
1.67
22%
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 give the percentage of episodes with that many complete branches. Each branch contains one imagined action.
Format
Checkpoint
Success
Imagined
Executed
Forward
Pre-RL
30.3%
27.82
17.88
RL50
35.0%
24.95
17.36
RL100
37.5%
20.87
16.96
RL150
38.5%
21.40
16.97
RL200
41.0%
22.18
17.50
Branching
Pre-RL
30.0%
23.62
18.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=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=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=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=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=∞ . 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=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=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=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=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=∞ . 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=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=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=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=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=∞ . 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=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) .
Figure 27: Complete PushT trajectory at K=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) .
Figure 28: Complete PushT trajectory at K=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) .
Figure 55
Figure 29: Complete PushT trajectory at K=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) .
Figure 30: Complete PushT trajectory at K=∞ . 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) .
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 set
Model
RL50
RL100
RL150
RL200
3-box Deadlock
Forward
Fail
Success
Success
Success
4-box Deadlock
Forward
Fail
Fail
Success
Success
3-box Structural Fork
Branching
Fail
Success
Success
Success
4-box Structural Fork
Branching
Fail
Fail
Success
Success
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 62Figure 63Figure 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 76Figure 77
Figure 41: Branching at RL50 on three-box Structural Fork. Complete trace for episode 17; recorded execution ends in failure.
Figure 79Figure 80Figure 81
Figure 42: Branching at RL100 on three-box Structural Fork. Complete trace for episode 17; recorded execution ends in success.
Figure 83Figure 84Figure 85
Figure 43: Branching at RL150 on three-box Structural Fork. Complete trace for episode 17; recorded execution ends in success.
Figure 87Figure 88Figure 89
Figure 44: Branching at RL200 on three-box Structural Fork. Complete trace for episode 17; recorded execution ends in success.
Figure 91Figure 92
Figure 45: Branching at RL50 on four-box Structural Fork. Complete trace for episode 56; recorded execution ends in failure.
Figure 94Figure 95
Figure 46: Branching at RL100 on four-box Structural Fork. Complete trace for episode 56; recorded execution ends in failure.
Figure 97Figure 98
Figure 47: Branching at RL150 on four-box Structural Fork. Complete trace for episode 56; recorded execution ends in success.
Figure 100Figure 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=∞ . The agent issues STOP before reaching the goal.
Figure 105Figure 106
Figure 50: Sokoban failure at K=∞ . The agent stops with 2 of 4 boxes on goals.
Figure 108Figure 109Figure 110
Figure 51: Pacman failure at K=∞ . The ghost catches Pacman at action 17. The remaining seven predicted states are never executed.
Figure 112Figure 113
Figure 52: PushT failure at K=∞ . The agent stops with 32.4% final target coverage, below the paper’s 80% success threshold.
Recent advances in world models have shown promise for modeling future dynamics of environmental states, enabling agents to reason and act without accessing real environments. Current methods mainly perform single-step or fixed-horizon rollouts, leaving their potential for complex task planning under-exploited. We propose Imagine-then-Plan (ITP), a unified framework for agent learning via lookahead imagination, where an agent's policy model interacts with the learned world model, yielding multi-step ``imagined'' trajectories. Since the imagination horizon may vary by tasks and stages, we introduce a novel adaptive lookahead mechanism by trading off the ultimate goal and task progress. The resulting imagined trajectories provide rich signals about future consequences, such as achieved progress and potential conflicts, which are fused with current observations, formulating a partially observable and imaginable Markov decision process to guide policy learning. We instantiate ITP with both training-free and reinforcement-trained variants. Extensive experiments across representative agent benchmarks demonstrate that ITP significantly outperforms competitive baselines. Further analyses validate that our adaptive lookahead largely enhances agents' reasoning capability, providing valuable insights into addressing broader, complex tasks. Our code and data will be publicly available at https://github.com/loyiv/ITP.
Youwei Liu, Jian Wang, Hanlin Wang +2
Central South University · Department of Computing, The Hong Kong Polytechnic University · College of Computer Science, Sichuan University
Building generalizable agents for diverse applications remains a fundamental challenge. While imitation learning-based policies succeed in specific training environments, they often fail to generalize to novel scenes and tasks. In this work, we propose World Action Planner, a robot planning system that leverages the reasoning capabilities of Vision-Language Models (VLMs) and the physical grounding of a multi-task pose-image conditioned world model. Our system enables an agent to propose initial action plans and iteratively refine them via optimization and search, reasoning over imagined world model rollouts. We demonstrate that our approach achieves superior performance across compositional tasks, new layouts, and zero-shot generalization scenarios, significantly outperforming state-of-the-art end-to-end policy models such as VLAs and WAMs. Project website at worldactionplanner.github.io
Post-training VLA policies typically rely on supervised fine-tuning with costly expert demonstrations or reinforcement learning with expensive and potentially unstable real-world exploration. World models offer a promising alternative by evaluating candidate behaviors through imagined futures, yet effective post-training requires more than accurate prediction: imagination must be scheduled where it is useful, bounded within reliable horizons, and translated into trustworthy policy supervision. In robotic manipulation, the value of imagination varies substantially across execution stages, while extended rollouts can accumulate prediction errors and introduce unreliable learning signals. We introduce WISE (World-model-guided Imagination Scheduling for Efficient Post-training of Vision-Language-Action Models), a unified framework that coordinates when and how world-model imagination is used during policy refinement. WISE selectively invokes imagination at interaction-relevant states, performs bounded multi-view rollouts, evaluates candidate futures using progress and completion signals, and uses their relative outcomes to refine actions generated from real interaction contexts. Extensive experiments with both π0 and π0.5 demonstrate consistent improvements across diverse manipulation tasks while reducing GPU computation time by approximately 80% compared with full imagination. Real-world evaluations further show substantial gains in robustness and generalization under diverse real-world distribution shifts.
Chenhao Zhang, Hanyu Zhao, Hang Cheng +2
1Tsinghua University · 2Beijing Academy of Artificial Intelligence (BAAI)