VAGEN: RL That Teaches VLM Agents World Models

VAGEN is a reinforcement learning method for vision-language agents, the models that act over many steps from screenshots or camera views. The Stanford SAIL blog explains its starting point: each image shows only part of the environment, so an agent has to piece together the real situation across turns. VAGEN makes the model write that understanding out at every step. Before each action it must describe what it thinks the current state is and predict what its action will change. In a grid puzzle, for example, that might be the player's position and where the player will be after moving right. The paper, accepted to NeurIPS 2025, comes from a team led at Northwestern with Stanford coauthors. It reports that a 3B-parameter model trained this way averaged 0.82 across five tasks, up from 0.21 before training. On the same tasks GPT-5 scored 0.75.

The written descriptions and predictions are not just a reasoning format. They are graded during training. For each environment the researchers pull a text version of the true state from the simulator: positions of the player, boxes and targets in Sokoban, object coordinates for the robot-arm tasks, and distances and directions to objects for navigation. A judge model, GPT-4.1 nano in the paper's experiments, then checks the agent's statements against that state, and an accurate statement earns reward. Without this, the agent would mostly be rewarded for eventually finishing the task. Grading every step gives it feedback along the way. The second component, Bi-Level GAE, decides which parts of a long episode deserve credit. It first scores each turn and then spreads that turn's score across the words the model generated within it. GAE, short for generalized advantage estimation, is a standard way to estimate how much better an action turned out than expected.

The comparisons behind the headline are narrower than they first look. The blog's own table compares reasoning formats for the same 3B model trained on these environments. Acting with no reasoning averaged 0.28 and free-form reasoning 0.67. Describing the state and predicting the next state together reached 0.76, and adding the step-by-step reward and the credit method brought it to 0.82. The proprietary models were evaluated off the shelf with a prompt, not trained on these tasks, so the result shows what targeted training on five environments can do. It is not evidence that the small model is generally more capable. The paper also finds that the best way to write a state depends on the task: plain language worked better for the grid puzzles, and a structured list of coordinates worked better for manipulation. Some agents learned to game the judge with generic descriptions that satisfied it without reflecting the actual state. The method also relies on simulators that can report the true state. The authors themselves list the limited set of model architectures and test environments as limitations.