Skip to content

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

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

VAGEN trains VLM agents via multi-turn reinforcement learning to build explicit internal world models by structuring their reasoning into StateEstimation (grounding current visual observations) and TransitionModeling (predicting next states). A 3B-parameter model trained with VAGEN achieves an overall score of 0.82 across five diverse agentic tasks, surpassing GPT-5 (0.75), Gemini 2.5 Pro (0.67), and Claude 4.5 (0.62).

Problem

Current VLM agents operating in multi-turn environments lack explicit internal world modeling: they process visual observations without systematically estimating latent environment states or predicting state transitions. This is compounded by the POMDP nature of visual agent tasks—states are only partially observable—making standard single-turn RL and token-level credit assignment inadequate for long-horizon sequential decision-making.

Method

VAGEN formulates multi-turn VLM agent training as a POMDP and enforces a structured reasoning format inside each action token: <observation>ŝt</observation><reasoning>ât</reasoning><prediction>ŝt+1</prediction><answer>aet</answer>. Five reasoning strategies are ablated (NoThink, FreeThink, StateEstimation, TransitionModeling, WorldModeling); WorldModeling is the full combination.

Two key components extend the base PPO framework:

  1. WorldModeling Reward: A dense, turn-level reward computed via LLM-as-a-Judge, comparing the agent's explicit state descriptions/predictions (<observation>, <prediction> fields) against environment-extracted ground-truth states. CLIP-based similarity was discarded as insufficiently sensitive to spatial detail. The reward is rtreason = βs·IStateEstimation + βw·ITransitionModeling.

  2. Bi-Level GAE: A two-stage advantage estimator that first computes a turn-level advantage Aturn_t using a turn-level discount Îłturn, then uses each turn's advantage to initialize the inner token-level GAE backward pass, propagating fine-grained credit to individual generation tokens within each action. This replaces standard single-pass token-level GAE, which suffers from sparse long-horizon credit propagation.

Representation of internal beliefs (Natural Language, Structured dict, Symbolic) is shown to be task-dependent: NL is optimal for semantic tasks (FrozenLake, Sokoban), Structured coordinates are superior for high-precision manipulation (PrimitiveSkill).

Key Contributions

  • Formal POMDP framing of multi-turn VLM agent tasks with explicit world model reasoning via structured reasoning tokens.
  • Systematic study of five reasoning strategies, showing WorldModeling (StateEstimation + TransitionModeling) consistently outperforms ablated variants and FreeThink baselines.
  • WorldModeling Reward: dense, turn-level LLM-as-a-Judge reward supervising visual state grounding and transition prediction.
  • Bi-Level GAE: hierarchical credit assignment with separate turn-level and token-level discount factors, resolving long-horizon sparse-reward instability.
  • VAGEN framework: decoupled environment/training infrastructure supporting five diverse environments (Sokoban, FrozenLake, 3D Navigation, PrimitiveSkill manipulation, SVG Reconstruction).
  • Empirical finding that optimal internal belief representation (NL vs. Structured vs. Symbolic) is task-dependent, with symbolic representations consistently underperforming due to grounding failures.

Results

  • VAGEN-Full (3B Qwen2.5-VL-3B backbone) overall score: 0.82 vs. untrained baseline 0.21 (~3Ă— improvement).
  • Surpasses GPT-5 (0.75), o3 (0.73), Gemini 2.5 Pro (0.64), Claude 4.5 (0.54) on the five-task suite.
  • WorldModeling strategy: 0.76 overall vs. FreeThink 0.67, NoThink 0.28 (all trained on same 3B backbone).
  • VAGEN-Full vs. WorldModeling (VAGEN-Base): most pronounced on PrimitiveSkill (0.97 vs. 0.91 average sub-task score), indicating better generalization to unseen scenes.
  • Vanilla-PPO without observation masking: 0.26; GRPO w/ masking: 0.54; Turn-PPO w/ masking: 0.55—all substantially below VAGEN-Full (0.82).
  • PrimitiveSkill Drawer sub-task: no off-the-shelf proprietary or open-source VLM achieves nonzero success; VAGEN-Full reaches 1.00.

Limitations

  • Experiments restricted to a single model family (Qwen2.5-VL); generalizability across architectures is unverified.
  • Evaluated on five hand-selected environments; coverage of real-world agentic diversity is limited.
  • WorldModeling Reward requires environment-specific ground-truth state extraction (e.g., 2D positions, object coordinates), limiting plug-and-play applicability to novel environments.
  • No supervised fine-tuning stage for multi-turn visual understanding; authors cite this as future work.
  • Gemini 2.5 Pro results are partial due to safety policy blocks on PrimitiveSkill and Navigation subsets.

Relevance to Vision-Language Models

VAGEN directly advances VLM agent capabilities by identifying that explicit world model reasoning—structured StateEstimation and TransitionModeling inside the chain-of-thought—is a critical inductive bias for multi-turn visual decision-making, outperforming both implicit (FreeThink) and no-reasoning (NoThink) baselines. The work bridges RL-based VLM post-training (in the tradition of DeepSeek-R1 and VLM-R1) with classic POMDP theory, offering a principled recipe for reward shaping and credit assignment in visual sequential tasks. The task-dependent representation finding (NL vs. Structured) has direct implications for how VLMs should format their internal monologue when deployed in heterogeneous environments. For researchers tracking VLMs, VAGEN demonstrates that a 3B model trained with structured reasoning RL can surpass much larger proprietary models on complex agentic benchmarks, suggesting world model reasoning as a high-leverage training axis.