Agentic Imagination, or Agima, trains an agent to generate possible consequences before acting. The research starts from situations such as Sokoban, where a plausible-looking push can make the puzzle impossible to finish.
How imagination enters the policy
A single omni model proposes actions, generates images of the predicted outcomes and reasons over those imagined futures. It can explore alternatives rather than committing immediately to its first proposed action.
The training has two stages. Supervised fine-tuning combines policy and world modelling, giving the model basic forward prediction and branching behaviour. Agentic reinforcement learning then trains it to use those predictions when making decisions.
What the experiments show
The paper tests Maze2D, Sokoban, Pacman and PushT, including harder settings not seen during training. In one supervised-only comparison, a model trained on two-box Sokoban solved 40% of four-box puzzles while observing the real board once every five moves. The action-only comparison solved 2% under the reported setup.
Those figures belong to that specific experiment. They should not be presented as the result of the later reinforcement-learning stage or as a success rate for real-world planning.
After reinforcement learning, the authors report that the model learned to look ahead selectively and branch more often at important decision points. This makes when to simulate part of the learned behaviour.
What to look for beyond the headline
Franklin’s interpretation is that imagined futures are useful only to the extent that they help choose correct actions. A convincing generated image can still depict a transition that the environment will not actually permit.
A practical evaluation should therefore inspect prediction errors, the cost of generating alternatives and whether the agent checks reality before a consequential move. The reported puzzle results motivate that analysis, but do not establish reliability for irreversible actions outside the tested environments.
Read Train An Imaginative Agent for the training design and experimental conditions.
Comments