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

AI summary

Overview

Research area: Reinforcement learning for vision-language model (VLM) agents in multi-turn, partially observable visual environments.

Technical level: Advanced. The paper assumes familiarity with reinforcement learning (PPO, GAE), partially observable Markov decision processes, and multimodal model post-training.

Scope: The paper proposes and evaluates VAGEN, a framework that trains VLM agents to explicitly reason about visual world state (current and predicted) and rewards that reasoning, across five agentic tasks, using a 3B-parameter backbone.

What This Paper Is About

VLM agents perceive the world through images rather than text, which makes their observations partial and noisy and shifts the problem from a standard MDP to a Partially Observable Markov Decision Process. The paper asks whether VLM agents can build internal world models by explicitly verbalizing two things during their reasoning: what the current visual state is (StateEstimation) and what the next state will be (TransitionModeling). It then trains agents with reinforcement learning over whole multi-turn trajectories, adding a dense reward for the accuracy of those state descriptions and predictions, plus a credit-assignment scheme that operates at both the turn and token level.

Key Contributions

1. A world-model reasoning structure for VLM agents. The paper formulates multi-turn VLM tasks as a POMDP and decomposes the agent's reasoning tokens into StateEstimation (<observation> tokens, approximating the hidden state $s_t$) and TransitionModeling (<prediction> tokens, approximating $s_{t+1}$). It compares five reasoning strategies under controlled RL training: NoThink, FreeThink, StateEstimation, TransitionModeling, and the combined WorldModeling.

2. A systematic study of how internal beliefs should be represented. Across Sokoban, FrozenLake, and PrimitiveSkill, the paper compares Natural Language, Symbolic, and Structured state representations and finds the optimal choice is task-dependent rather than universal.

3. WorldModeling Reward and Bi-Level GAE. The paper introduces a turn-level dense reward derived from an LLM-as-a-Judge (and rule-based comparison) that scores the agent's <observation> and <prediction> fields against ground-truth environment states, and a Bi-Level General Advantage Estimation scheme that computes turn-level advantages first and then propagates them to tokens, using separate γ_turn and γ_token discount factors.

4. The VAGEN framework and empirical results. VAGEN is described as a scalable system that decouples environment setup from model training across diverse visual environments. With it, a 3B model reaches 0.82 overall across five tasks, nearly 3 times its untrained counterpart (0.21), and above the reported figures for GPT-5 (0.75), Gemini 2.5 Pro (0.67), and Claude 4.5 (0.62).

Main Findings

• WorldModeling is the strongest of the five reasoning strategies. It achieves an overall 0.76, ahead of FreeThink (0.67) and well ahead of NoThink (0.28). The paper reports a +0.55 improvement for WorldModeling over the untrained Qwen2.5-VL-3B counterpart.

• Off-the-shelf VLMs struggle on these multi-turn tasks. Among the models benchmarked, the best-performing is GPT-5 at 0.75 out of 1 in overall score, and no model succeeds on the PrimitiveSkill Drawer task. VLM-R1-3B, trained on non-agentic tasks, shows an overall 0.23, below Qwen2.5-VL-72B's 0.51.

• Individual reasoning components have task-specific strengths. StateEstimation performs well on Navigation (0.74 average), where reading the current observation is key, while TransitionModeling is strong on PrimitiveSkill (0.82 average), where predicting future states matters for manipulation. Each alone can underperform on tasks where the model's prior is less aligned with the task structure.

• State representation is not universal. Natural Language beats Symbolic and Structured on Sokoban (0.61 vs 0.49 vs 0.28) and on FrozenLake (0.71 vs 0.49 vs 0.63). On PrimitiveSkill, Structured is slightly better (0.94 average vs 0.91 for Natural Language). Symbolic representation is the least effective and was not run on PrimitiveSkill.

• Existing RL baselines are inadequate for multi-turn VLM agents. Vanilla-PPO scores 0.26 overall, GRPO with masking 0.54, and Turn-PPO with masking 0.55. The paper attributes Vanilla-PPO's failure to the lack of observation-token masking, GRPO's to high trajectory diversity from scene change, and Turn-PPO's to uniform within-turn advantages that cannot capture individual token contributions.

• Dense turn-level rewards plus hierarchical credit assignment help. VAGEN-Full, combining the WorldModeling Reward with Bi-Level GAE, improves over VAGEN-Base on the same backbone: 0.82 vs 0.76 overall, with the largest single-task gaps on Sokoban (0.79 vs 0.61) and the PrimitiveSkill average (0.97 vs 0.91).

• The critic-free-style CLIP reward attempt failed. The paper reports that an initial CLIP-based image-text similarity reward was insufficiently sensitive to fine-grained spatial and geometric details, which motivated the switch to LLM-as-a-Judge and rule-based comparison.

• Gemini 2.5 Pro results are partial. Due to that model's safety policy, a subset of evaluation responses for PrimitiveSkill and Navigation was blocked, and the reported results are based solely on unblocked test cases.

Methodology in Plain English

The agent is trained inside five visual environments: Sokoban and FrozenLake (2D grid puzzles), Navigation (3D first-person embodied navigation), PrimitiveSkill (a Panda Arm manipulation task with a hybrid action space such as pick(x,y,z) and a third-person 3D scene), and SVG Reconstruction (generating SVG code to match a target image).

At each turn the agent receives a visual observation and produces a block of text. The researchers force this text to follow a fixed structure: an <observation> section where the agent says what the current state is, a <prediction> section where it says what the next state will be, a <reasoning> section, and an <answer> section containing the executable action. Which of those sections are required defines the five reasoning strategies being compared.

Training uses an actor-critic setup with PPO over entire trajectories, where observation (image) tokens are masked out of the loss so only agent-generated tokens are optimized. To guide learning, the reward at each turn combines three parts: a format reward for following the required structure, the environment's task reward (a scaled binary reward such as {0, 10} for the completion-based tasks, or a dense weighted DreamSim-plus-DINO image similarity for SVG Reconstruction), and a WorldModeling Reward for the reasoning itself. Extracting the agent's <observation> and <prediction> text and comparing it against text-based ground truth obtained from the environment scores that reasoning reward.

For credit assignment, the standard approach of propagating one sparse end-of-trajectory reward backwards token by token is replaced. Bi-Level GAE first computes an advantage for each turn using a turn-level discount factor, then runs a second pass inside each turn at the token level, initialized with that turn's advantage. This delivers turn-specific feedback to the exact tokens that produced it.

Why This Matters

The paper's central claim is that explicit world-model reasoning, rewarded and credited correctly, can let a small open 3B model beat much larger proprietary systems on multi-turn visual agent tasks. If that holds, it shifts the field's emphasis from scaling alone toward the structure of agent reasoning and the design of RL reward signals for multimodal agents. It also provides VAGEN as a system that separates environment definition from training, which makes it easier to add new environments and compare algorithms.

Real-world applications suggested by the tasks and framing:

• Robotic manipulation, where the agent must ground objects in a third-person 3D scene into coordinates and predict the result of actions such as pick(x,y,z).
• Embodied navigation, where an agent follows instructions in a first-person view and must interpret what it currently sees.
• Computer-use and GUI agents, a category the paper explicitly cites alongside games and embodied AI as a target for VLM agents.
• Visual content generation, illustrated by the SVG Reconstruction task, where the agent emits code to match a target image and is rewarded by image similarity.

Industry relevance: the recipe targets small, trainable open models rather than API-hosted frontier models, which matters for cost, latency, and deployment control in agentic products. The reward-shaping components (a judge-scored reasoning reward plus a masking-aware advantage estimator) are also directly reusable patterns for anyone doing RL post-training on multimodal models, and the paper's negative results about unmasked PPO, GRPO, and turn-level PPO give practitioners concrete guidance on what not to do.

Future Directions

• Whether the findings extend beyond a 3B backbone and five tasks. The paper's headline result is on Qwen2.5-VL-3B across Sokoban, FrozenLake, Navigation, PrimitiveSkill, and SVG Reconstruction; scaling behavior and transfer to new environments are open.
• How to choose or unify state representations. Because the best representation is task-dependent (Natural Language for grids, Structured for PrimitiveSkill, with Symbolic weakest), an open question is whether an agent can adaptively select or learn a representation rather than having it fixed per task.
• Extending the WorldModeling Reward to domains without easy textual ground truth. The reward relies on deriving text-based state information from the environment (2D positions, object names and coordinates, relative distances and directions). Tasks where such ground truth is unavailable or expensive would need another mechanism.
• Tuning the reward and credit-assignment knobs. The framework introduces

Authors’ 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.