VAGEN: Reinforcing World Model Reasoning for Multi-Turn VLM Agents

TL;DR

VAGEN advances multi-turn Vision-Language Model agents by enforcing internal world modeling through explicit visual state reasoning in a partially observable setting. By decomposing reasoning into StateEstimation and TransitionModeling, and introducing a dense WorldModeling Reward plus Bi-Level GAE, the framework provides turn-wise supervision and fine-grained credit assignment. Representational choices are demonstrated to be task-dependent, with natural language excelling in general semantic tasks and structured formats aiding precise manipulation. Across five diverse benchmarks, VAGEN-Full achieves superior performance, surpassing proprietary models and showing robust generalization, while the accompanying framework and codebases enable scalable research in visual-agent reasoning.

Abstract

A key challenge in training Vision-Language Model (VLM) agents, compared to Language Model (LLM) agents, lies in the shift from textual states to complex visual observations. This transition introduces partial observability and demands robust world modeling. We ask: Can VLM agents construct internal world models through explicit visual state reasoning? To address this question, we architecturally enforce and reward the agent's reasoning process via reinforcement learning (RL), formulating it as a Partially Observable Markov Decision Process (POMDP). We find that decomposing the agent's reasoning into State Estimation ("what is the current state?") and Transition Modeling ("what comes next?") is critical for success, as demonstrated through five reasoning strategies. Our investigation into how agents represent internal beliefs reveals that the optimal representation is task-dependent: Natural Language excels at capturing semantic relationships in general tasks, while Structured formats are indispensable for precise manipulation and control. Building on these insights, we design a World Modeling Reward that provides dense, turn-level supervision for accurate state prediction, and introduce Bi-Level General Advantage Estimation (Bi-Level GAE) for turn-aware credit assignment. Through this form of visual state reasoning, a 3B-parameter model achieves a score of 0.82 across five diverse agent benchmarks, representing a 3$\times$ improvement over its untrained counterpart (0.21) and outperforming proprietary reasoning models such as GPT-5 (0.75), Gemini 2.5 Pro (0.67) and Claude 4.5 (0.62). All experiments are conducted within our VAGEN framework, a scalable system for training and analyzing multi-turn VLM agents in diverse visual environments. Code and data are publicly available at https://vagen-ai.github.io.

Figures (8)

• Figure 1: Overview of the VAGEN Framework. Within a POMDP, VLM Agents build internal world models through explicit visual state reasoning over the StateEstimation$P(\hat{s}_{t}|o_t)$ and TransitionModeling$P(\hat{s}_{t+1}|o_t,\hat{s}_{t},\hat{a}_{t})$, where $\hat{s}_{t}$ denotes the internal belief of the state, $\hat{a}_{t}$ represents the internal belief of the action, and $a^{\text{e}}_{t}$ is the parsed executable action, as detailed in Table \ref{['tab:vagen_notations']}.
• Figure 2: Examples of visual states from five environments used in our study.
• Figure 3: Token-Level GAE and Bi-Level GAE frameworks. Standard Token-Level GAE, where a single, sparse reward at the end of the trajectory is propagated backward token-by-token. Our proposed Bi-Level GAE assigns rewards at each turn. It first computes a turn-level advantage (purple arrows) before propagating that credit to the individual tokens within the turn (orange arrows), allowing for robust, hierarchical estimation for credit assignment.
• Figure 4: Training success rate for VAGEN-Base, VAGEN-Full and ablations. Since SVG reconstruction is a multi-turn reasoning task without world dynamics, we only apply the Bi-Level GAE configuration.
• Figure 5: Comparison of VAGEN-Base and VAGEN-Full across Navigation, Sokoban, and PrimitiveSkill environments, from left to right. Within each environment, the left column shows VAGEN-Base results while the right column displays VAGEN-Full results. The improvement of VAGEN-Full shows enhanced reasoning ability in understanding environments such as their spatial information.