Agentic imagination lets AI picture actions before deciding
Train An Imaginative Agent
Artificial Intelligence
Summary
Some puzzles like Sokoban become impossible to solve if a wrong move is made early on. The authors introduce Agima, a method that trains an AI agent to imagine what will happen after its possible moves before actually making one. This agent can simulate different future scenarios and choose its action based on those imagined outcomes. When tested, Agima solved more complex puzzles better than agents that acted without imagining the future.
What this means in practice
- •For game ai developers: Create AI agents that can plan moves by simulating future game states to improve decision-making in complex puzzle environments.
- •For robotics control teams: Develop controllers that predict the results of physical actions before execution to avoid irreversible mistakes in tasks with sparse feedback.
Authors
Jiaxin Ge, Junyi Zhang, Tianning Chai, Hanjun Yoo, Chenyang Wang, Xiuyu Li, Weijia Shi, XuDong Wang, Sewon Min, Luke Zettlemoyer, Trevor Darrell, Haiwen Feng
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/.