Train An Imaginative Agent
Jiaxin Ge, Junyi Zhang, Tianning Chai, Hanjun Yoo, Chenyang Wang, Xiuyu Li, Weijia Shi, XuDong Wang, Sewon Min, Luke Zettlemoyer, Trevor Darrell, Haiwen Feng
cs.AI
2026-10-20
Agima trains one omni model to imagine visual outcomes and plan from them, solving 40% of four-box Sokoban with a real look every five moves versus 2% without imagination.
One legal push in Sokoban can pin a box and leave the puzzle unsolvable. The move cannot be undone, and the board does not reveal that until after the push. A standard agent acts, then sees the next frame. If real feedback shows up only every few steps, the deadlock is already locked in.
Rolling a learned dynamics model forward is an old idea. Imagination-Augmented Agents (I2A) pass imagined rollouts to a policy, but a separate rollout policy still picks which actions get imagined. Agima trains one set of weights to propose an action, draw the frame that follows, and choose what to execute from those frames.
The backbone is BAGEL, an omni model: one set of weights that both reads and generates images and text. It has 14B parameters and 7B active, written A7B. Those weights propose the action and generate the next frame. Predicted frames and real observations share an interface, so a missing environment step can be filled by a picture. After execution, the real frame anchors the next round.
Supervised fine-tuning teacher-forces interleaved actions and frames. Text uses next-token cross-entropy. Images use the generator's mean squared error. The episodes stay fixed. A linear rollout follows one plan. Counterfactual branching tries another action on that plan, draws the immediate result, and rejects it on failure. Forward thinking rolls several complete plans and then picks one. The three traces practice looking ahead, ruling out a local dead end, and choosing among full plans. In Sokoban the rejected push is one a solver has certified as unsolvable, with at most three probes per episode. In Pacman the plans differ in whether the nearer pellet is eaten first.
GRPO, group-relative policy optimization from task rewards, then learns when to use these traces. Groups of 8 and batches of 16 are sampled with no real observations during deliberation. Only the final actions run in the environment. The reward is binary success. Gradients update language tokens. Generated images stay in context and get no image loss, so the world model stays at its supervised quality. Two-box levels are nearly saturated after that stage, so reinforcement learning runs 200 steps on ordinary 8×8 three-box puzzles at learning rate 2e-6, with no new expert traces. Forward thinking and branching start from a checkpoint that is about half linear rollouts and half the matching reasoning traces.
Training uses easier boards. Tests use harder boards held out of training: mazes larger than the 11×11 training size, four-box Sokoban after two-box training on 7×7, Pacman on 12×12 with 12 pellets and a closer ghost after 9×9 with 8 pellets, and PushT poses farther from the center. PushT counts as success when final coverage exceeds 0.80. K is the number of actions executed between real observations. No Imagine is the same BAGEL backbone, fine-tuned on the same data, emitting action chunks with no frames in between.
At K=5 on the hard splits:
| Environment | No Imagine | Agima |
| Pacman | 1% | 37% |
| Maze2D | 23% | 40% |
| Four-box Sokoban | 2% | 40% |
| PushT success | 1% | 18% |
| PushT coverage | 6.7% | 28.5% |
The 40% comes from supervised training on two-box puzzles alone, with the real board shown once every five moves. With a real frame every move (K=1), the same model scores 47% on four-box Sokoban, 84% on Pacman, and 62% on Maze2D. At K=10, four-box Sokoban falls to 15% and Pacman to 7%. At K=∞, with no intermediate real frames, four-box Sokoban is 19%, against 0% for No Imagine and 21% for GPT-5.6 Sol. Maze2D falls to 3%, below Gemini 3.1 Pro at 9% and GPT-5.6 Sol at 36%.
The pixels are in the loop. Replacing K=5 predictions with gray frames or the latest real frame drops four-box Sokoban from 40% to 1% or 0%, and Pacman from 37% to 0%. Substituting the ground-truth next frame raises Pacman to 81%, Maze2D to 69%, and PushT to 47%. Four-box Sokoban only moves from 40% to 43%. That board barely benefits from a perfect picture. The gap in the other three tracks prediction error.
Branching and forward thinking use the same training episodes. Only the trace changes, and the test is fully open-loop. Against the linear model, branching lifts Structural Fork from 34% to 39% with three boxes and from 22% to 32% with four. Forward thinking lifts Deadlock from 22% to 60% with two boxes and from 39% to 48% with three, then drops four-box Deadlock from 27% to 24%. On Pacman route-order traps, where the near-first route has no winning continuation, forward thinking moves 8×8 with four pellets from 22% to 54%, and 10×10 with eight pellets from 5% to 12%. A single rollout takes the near pellet while the ghost is still far, then cannot escape on the way back.
Reinforcement learning rewards success and does not supervise the format. No Imagine's training reward rises to about 15%, but held-out Deadlock and Structural Fork stay near zero. Runs that start with imagination keep improving on four-box puzzles that this stage never trains on. Around step 100 a format missing from supervised data shows up: roll one plan, commit if it looks successful, and open another only if it looks failed. On three-box boards it will open a third plan after two failures. Once a second candidate appears, the model calls the first a failure 100% of the time, against about 32% before RL. Executing that first plan alone fails on 95.5% to 100% of those cases. The supervised habit of always comparing two candidates from the start essentially disappears. Branching gets more selective. Mean branches per episode rise from 1.70 to 1.95 at step 150, more often at states where a wrong legal action makes the puzzle unsolvable, without a longer imagination budget. Pooled over 400 episodes, forward thinking goes from 30.3% to 41.0% success while imagined actions fall from 27.82 to 22.18. Branching goes from 30.0% to 50.3%, with imagined actions falling from 23.62 to 21.59.
Switching plans and probing another action are language tokens in the same trace as the pictures. No external tree search is attached. With one real observation every five moves, those pictures are the gap between 2% and 40% on four-box Sokoban. A binary success reward is enough to hold forward thinking until the first plan looks doomed, and the imagined horizon shortens as success rises.
The evidence is small grids and a 2D pushing task. When feedback arrives every step, the paper does not include the matched no-imagination control on the same backbone.
Visual error is a hard cap. Ground-truth next frames lift Pacman to 81% and PushT to 47%, but reinforcement learning never updates image tokens. The appendix shows selected rollouts and does not report aggregate state-prediction accuracy. Fully open-loop Maze2D leaves Agima at 3%, while GPT-5.6 Sol scores 36%.
The new traces are an uneven upgrade. Forward thinking loses 3 points to a linear rollout on four-box Deadlock. On two-box Deadlock, Gemini 3.1 Pro scores 74%, above forward thinking at 60%. The fine-tuned brancher passes Gemini only on the harder board, 32% versus 19% on four-box Deadlock. Pacman forward-thinking data barely contains the failure the test asks about: both candidates succeed in 96.4% of training traces, and a death occurs in 3.6%, yet the test set is built from near-first traps. Branch and forward labels come from solver certificates. Reinforcement learning is run on Sokoban only.
On K=3 mazes, action-only fine-tuning of Qwen3-VL-8B scores 85%, and Agima scores 59%. For plain navigation, drawing the future can lose to emitting the move. The strongest branching checkpoint still imagines 21.59 actions per episode on average. The paper measures neither latency nor FLOPs.