Agima: Teaching AI Agents to "Imagine" Before Acting

AI agents executing actions in real-world environments often face the dilemma of "one wrong step, and the whole game is lost." In the game of Sokoban, a seemingly reasonable push can make a puzzle permanently unsolvable. When environmental feedback is sparse and actions are irreversible, agents must anticipate consequences before acting. A training framework called Agentic Imagination (Agima) proposes an intuitive yet powerful idea: let agents imagine first, then act.

Core Design: One Model Doing Three Things at Once

At the heart of Agima is a unified training framework that trains a single "omni model" to accomplish three tasks: propose candidate actions, generate the future frames these actions might lead to, and then reason over these "imagined futures" to ultimately decide which action to execute in the real environment.

Unlike traditional single-path planning, such agents can simulate a sequence of actions, explore multiple alternatives, and compare the merits of different future scenarios. This transforms "lookahead" from an expensive extra computation into an intrinsic part of the model's decision-making process.

Two-Stage Training: First Learn to Imagine, Then Learn to Use Imagination

Agima's training is divided into two stages.

Stage One: Supervised Fine-Tuning to Unify Policy and World Model

The research team first uses supervised fine-tuning (SFT) to unify the policy and world modeling into a single omni model. This step equips the agent with the ability to predict future outcomes and master basic forms of "imagination"—including forward thinking and counterfactual branching.

World Model is a core concept in reinforcement learning and model predictive control, referring to an agent's internal representation of environment dynamics: given the current state and action, it can predict the next state (or observation). In pixel-based settings, the world model is typically a generative model responsible for synthesizing the future frame after an action is executed. Unlike a policy that only learns a "state-to-action" mapping, a world model allows the agent to "rehearse" consequences without interacting with the environment, thereby supporting planning and counterfactual reasoning. Agima has the same omni model serve both as policy and world model, so that imagination and decision-making share representations—this is also the key to its ability to solve four-box Sokoban with sparse real feedback.

Stage Two: Agentic Reinforcement Learning to Optimize Decision-Making

Imagination alone is not enough; the key is how to use imagination to guide decisions. The second stage employs agentic reinforcement learning (RL) to teach the model to reason over imagined futures and effectively use this reasoning for planning.

Counterfactual branching refers to an agent, in imagination space, starting from the same state and concurrently rolling out future trajectories resulting from multiple different actions, then comparing these "what if..." alternative futures. It differs from single-path forward rollout: forward thinking only predicts what follows along one path, whereas counterfactual branching requires the model to maintain multiple parallel hypotheses and evaluate their consequences—essentially internalizing the idea of tree search into the model. This capability is especially important in environments with irreversible actions—the agent must distinguish, before acting, which branches would lead to unrecoverable failure, in order to avoid "one wrong step, and the whole game is lost."

Cross-Environment Generalization: Leaping from Two Boxes to Four Boxes

Agima was validated in four environments—Maze2D, Sokoban, Pacman, and PushT—and can generalize to harder settings unseen during training.

The most striking data comes from Sokoban: after SFT training only on the "two-box" version, Agima solved 40% of "four-box" puzzles—while seeing the real board only once every five steps. By comparison, an action-only planner solved only 2%. This means the model relies heavily on its imagined internal world to compensate for sparse real feedback.

Supervised Fine-Tuning (SFT) refers to continuing to train a pretrained model on labeled or expert-demonstration data using the standard next-token (or next-frame) prediction loss, so that it acquires the behavior distribution of a specific task. In Agima, the SFT stage provides both action demonstrations and corresponding future observations, enabling the model to learn "given a state and action, generate the resulting frame," thereby unifying policy and world model into one model. SFT provides an imitation-style baseline capability; at this stage the model has not yet learned to weigh the merits of different imaginations—which is precisely why subsequent agentic reinforcement learning is needed to optimize decision-making.

Emergent Behaviors After Reinforcement Learning

After the RL stage, Agima exhibited behaviors in using imagination that were not explicitly taught during training:

• Spontaneous formation of a unified format: the model learned a new unified format never seen in the SFT stage;
• On-demand lookahead: it learned to perform forward thinking only when necessary, rather than indiscriminately rolling out all possibilities;
• Branching at key points: at critical decision points, the model branches more frequently.

These behaviors point to a more interesting conclusion: an agent that can "see where actions lead" will learn on its own when and where it needs to look ahead—that is, it learns to plan.

Agentic RL refers to applying reinforcement learning to agent models with multi-step reasoning and tool-use capabilities, optimizing their decision sequences—not just single-step outputs—through rewards obtained from interaction with the environment. Unlike SFT based on human demonstrations, RL allows the model to explore behaviors beyond demonstrations and discover better strategies through trial and error. Agima's second stage uses RL to train the model to reason over imagined futures and plan accordingly, which allows meta-decisions such as "when to imagine and at which nodes to branch" to be shaped by reward signals, thereby explaining the source of emergent behaviors such as on-demand lookahead and frequent branching at key points.

Why This Matters

Agima's value lies in turning "imagination" from a prompt engineering trick or an external planner into a model capability that can be acquired through training. For robot manipulation, game AI, and any scenario with sparse feedback and irreversible actions, this "simulate first, then execute" paradigm could reduce the cost of trial and error. The project homepage is now public at agima.net, serving as an entry point for further exploration.

One-line Summary

Agima lets agents imagine future frames before acting, using a unified model to achieve planning under sparse feedback.

Paragraph Summary

Agima is a training framework that lets agents "imagine first, then act." At its core, it trains a single omni model to simultaneously propose candidate actions, generate the future frames those actions would lead to, and reason over the imagined futures to make decisions. Training has two stages: first, supervised fine-tuning unifies policy and world model, equipping the model with forward thinking and counterfactual branching; then, agentic reinforcement learning teaches the model to use imagination for planning. Validated in four environments—Maze2D, Sokoban, Pacman, and PushT—the model generalizes to harder settings. Most notably, trained only on the two-box Sokoban version, it solves 40% of four-box tasks, compared with just 2% for an action-only planner. After RL, emergent behaviors such as on-demand lookahead, frequent branching at key points, and a new unified format appear without explicit teaching, showing that imagination can be trained as an intrinsic model capability.