In Sokoban, a box can look almost home and still be doomed. Push it into the wrong corner and the puzzle may become impossible—not because the move was illegal, but because of what it quietly ruled out. A player can pause, picture what happens next, and choose another route. Researchers are now training an AI agent to do something similar: build a visual picture of possible outcomes before committing to a move.
The key shift is from seeing the game as a sequence of real observations to seeing it as a landscape of imagined ones. The agent does not have to push a box—or move through a maze—to inspect every possible next state. It can propose an action, generate an image of the board that would result, and use that predicted image as the starting point for more planning. The hidden world in this experiment is the one inside the model: a succession of possible scenes that may never happen.
In Train an Imaginative Agent, Jiaxin Ge and colleagues introduce a framework called Agentic Imagination, or Agima. One multimodal model does the imagining and the choosing. It proposes actions, predicts the visual consequences, and reasons over those predictions. After selecting a plan, it carries out actions in the actual environment. Real observations then bring its next round of planning back to the board it really faces.
To teach the system this skill, the researchers trained it on expert examples from four simulated environments: Maze2D, Sokoban, Pacman, and PushT, a task where an agent pushes a T-shaped object toward a target. The training examples interleave actions with images of the states that follow them. In other words, the agent learns not only what an expert did, but also what the world looked like after each step. The model used for Agima has 14 billion total parameters, with 7 billion active parameters.
Those examples teach several different ways to look ahead. In a linear rollout, the agent imagines one action sequence, step by step. In counterfactual branching, it pauses along that route to picture a different move and reject it if it looks dangerous. In forward thinking, it imagines longer alternative plans and compares how they unfold. Each predicted image is fed back into the model as an internal observation, so a plan can grow into a chain of imagined scenes.
The researchers tested whether this look-ahead could help when the agent had to act several times before receiving a fresh view of the real environment. In the reactive setting, it gets real feedback after every action. In the less reactive settings, it must rely on its own predictions between observations. The tests also raised the difficulty: Agima trained on simpler versions of tasks and was evaluated on harder or unfamiliar configurations.
One result came from Sokoban. After training on puzzles with two boxes, Agima solved 40% of four-box puzzles when it received a real observation every five moves. An action-only version trained on the same task data solved 2%. That comparison suggests the generated images were doing more than decorating the plan: the agent benefited from picturing intermediate states even when the real board was temporarily out of view.
Planning also meant comparing the order of actions, not just extending a single route. In a Pacman test, taking the nearest food first could leave the player trapped by a ghost before it reached the rest. The forward-thinking version imagined alternative routes and could favor collecting the farther food first. On one unfamiliar 8-by-8 grid with four food items, success rose from 22% for the linear-rollout version to 54% for the forward-thinking version.
The team then asked whether the agent could learn when to imagine more, rather than simply copying a prescribed planning pattern. They used reinforcement learning: the agent tried a plan, then received a binary reward for whether it completed the task. The researchers did not supply additional expert demonstrations or instructions about how to arrange its imagined futures during this stage. After 200 training steps, the branching version’s success across the held-out three- and four-box Sokoban tests rose from 30.0% to 50.3%. Its average number of imagined actions fell from 23.62 to 21.59.
That is a useful twist. Better performance did not require a longer imagined movie. The agent learned to focus its branches on decision points where a bad choice could make the puzzle unsolvable. In forward thinking, it also learned to seek another plan when its first imagined plan appeared likely to fail. The work suggests that planning can improve not just by looking farther ahead, but by choosing where another look is worth taking.
The results come with important boundaries. These were simulated environments with defined tasks, not tests of a robot navigating an unpredictable home or workplace. The researchers also found room for better predictions: replacing Agima’s imagined frames with ground-truth frames improved success by 29–44 percentage points in Pacman, Maze2D, and PushT, and by 3 points in Sokoban. The model’s imagined view helped, but it was not as reliable as seeing what actually happened.
Agima is therefore not evidence that an AI has an inner life or a humanlike imagination. It is evidence for a practical research idea: a model can be trained to generate visual predictions, inspect them, and use them to decide what to do next. For a box-pushing puzzle, that can mean rejecting a move before it creates a dead end. For AI research, it opens a question about scale: if an agent can use its own imagined scenes as temporary evidence, how much farther can it plan when the real world is costly—or impossible—to check after every move?
